fix: add platform guards for tilelang in mhc.py and fix test imports

- Add platform guards for tilelang imports in mhc.py to prevent import failures on non-CUDA platforms (CPU, AMD, etc.) - Move tilelang kernel definitions to lazy initialization - Fix import order in test_deepseek_v4_mega_moe.py (ruff compliance) - Add pytest skip marker for non-CUDA platforms in test file This fixes the widespread CPU test failures caused by unconditional tilelang imports at module level. https://claude.ai/code/session_015qZTB3eveFJ8qWsSupPgX1 Signed-off-by: Claude <noreply@anthropic.com>
Integrate MegaMoE kernel (#232 )
2026-04-25 17:29:00 +00:00 · 2026-04-25 07:51:57 +00:00 · 2026-04-25 07:51:57 +00:00 · 2026-04-25 07:51:56 +00:00 · 2026-04-25 06:18:00 +00:00 · 2026-04-25 06:17:51 +00:00
141 changed files with 1565 additions and 4146 deletions
@@ -294,6 +294,7 @@ set(VLLM_EXT_SRC
  "csrc/activation_kernels.cu"
  "csrc/layernorm_kernels.cu"
  "csrc/fused_qknorm_rope_kernel.cu"
+  "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu"
  "csrc/layernorm_quant_kernels.cu"
  "csrc/sampler.cu"
  "csrc/topk.cu"
@@ -310,9 +311,7 @@ set(VLLM_EXT_SRC
  "csrc/torch_bindings.cpp")

 if(VLLM_GPU_LANG STREQUAL "CUDA")
-  list(APPEND VLLM_EXT_SRC
-    "csrc/minimax_reduce_rms_kernel.cu"
-    "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
+  list(APPEND VLLM_EXT_SRC "csrc/minimax_reduce_rms_kernel.cu")

  SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")

@@ -1047,14 +1046,14 @@ endif()
 set(VLLM_MOE_EXT_SRC
  "csrc/moe/torch_bindings.cpp"
  "csrc/moe/moe_align_sum_kernels.cu"
-  "csrc/moe/topk_softmax_kernels.cu")
+  "csrc/moe/topk_softmax_kernels.cu"
+  "csrc/moe/topk_softplus_sqrt_kernels.cu")

 if(VLLM_GPU_LANG STREQUAL "CUDA")
  list(APPEND VLLM_MOE_EXT_SRC
    "csrc/moe/moe_wna16.cu"
    "csrc/moe/grouped_topk_kernels.cu"
-    "csrc/moe/router_gemm.cu"
-    "csrc/moe/topk_softplus_sqrt_kernels.cu")
+    "csrc/moe/router_gemm.cu")
 endif()

 if(VLLM_GPU_LANG STREQUAL "CUDA")
@@ -11,74 +11,29 @@
 namespace vllm {

 template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
-          bool act_first, bool HAS_CLAMP>
+          bool act_first>
 __device__ __forceinline__ scalar_t compute(const scalar_t& x,
-                                            const scalar_t& y,
-                                            const float limit) {
-  if constexpr (act_first) {
-    scalar_t gate = x;
-    scalar_t up = y;
-    if constexpr (HAS_CLAMP) {
-      gate = (scalar_t)fminf((float)gate, limit);
-      up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
-    }
-    return ACT_FN(gate) * up;
-  } else {
-    scalar_t gate = x;
-    scalar_t up = y;
-    if constexpr (HAS_CLAMP) {
-      gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
-      up = (scalar_t)fminf((float)up, limit);
-    }
-    return gate * ACT_FN(up);
-  }
+                                            const scalar_t& y) {
+  return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
 }

 template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
-          bool act_first, bool HAS_CLAMP>
+          bool act_first>
 __device__ __forceinline__ packed_t packed_compute(const packed_t& x,
-                                                   const packed_t& y,
-                                                   const float limit) {
-  if constexpr (act_first) {
-    packed_t gate = x;
-    packed_t up = y;
-    if constexpr (HAS_CLAMP) {
-      float2 g = cast_to_float2(gate);
-      float2 u = cast_to_float2(up);
-      g.x = fminf(g.x, limit);
-      g.y = fminf(g.y, limit);
-      u.x = fmaxf(fminf(u.x, limit), -limit);
-      u.y = fmaxf(fminf(u.y, limit), -limit);
-      gate = cast_to_packed<packed_t>(g);
-      up = cast_to_packed<packed_t>(u);
-    }
-    return packed_mul(PACKED_ACT_FN(gate), up);
-  } else {
-    packed_t gate = x;
-    packed_t up = y;
-    if constexpr (HAS_CLAMP) {
-      float2 g = cast_to_float2(gate);
-      float2 u = cast_to_float2(up);
-      g.x = fmaxf(fminf(g.x, limit), -limit);
-      g.y = fmaxf(fminf(g.y, limit), -limit);
-      u.x = fminf(u.x, limit);
-      u.y = fminf(u.y, limit);
-      gate = cast_to_packed<packed_t>(g);
-      up = cast_to_packed<packed_t>(u);
-    }
-    return packed_mul(gate, PACKED_ACT_FN(up));
-  }
+                                                   const packed_t& y) {
+  return act_first ? packed_mul(PACKED_ACT_FN(x), y)
+                   : packed_mul(x, PACKED_ACT_FN(y));
 }

 // Activation and gating kernel template.
 template <typename scalar_t, typename packed_t,
          scalar_t (*ACT_FN)(const scalar_t&),
          packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
-          bool use_vec, bool HAS_CLAMP, bool use_256b = false>
+          bool use_vec, bool use_256b = false>
 __global__ void act_and_mul_kernel(
    scalar_t* __restrict__ out,          // [..., d]
    const scalar_t* __restrict__ input,  // [..., 2, d]
-    const int d, const float limit) {
+    const int d) {
  const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
  const scalar_t* y_ptr = x_ptr + d;
  scalar_t* out_ptr = out + blockIdx.x * d;
@@ -103,9 +58,8 @@ __global__ void act_and_mul_kernel(
      }
 #pragma unroll
      for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
-        x.elts[j] =
-            packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
-                x.elts[j], y.elts[j], limit);
+        x.elts[j] = packed_compute<packed_t, PACKED_ACT_FN, act_first>(
+            x.elts[j], y.elts[j]);
      }
      if constexpr (use_256b) {
        st256(x, &out_vec[i]);
@@ -118,8 +72,7 @@ __global__ void act_and_mul_kernel(
    for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
      const scalar_t x = VLLM_LDG(&x_ptr[idx]);
      const scalar_t y = VLLM_LDG(&y_ptr[idx]);
-      out_ptr[idx] =
-          compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
+      out_ptr[idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
    }
  }
 }
@@ -198,11 +151,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {

 // Launch activation and gating kernel.
 // Use ACT_FIRST (bool) indicating whether to apply the activation function
-// first. HAS_CLAMP (bool) enables pre-activation clamping: gate input is
-// clamped (max only) and up input is clamped (both sides) before the
-// activation function is applied.
-#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST,        \
-                                      HAS_CLAMP, LIMIT)                        \
+// first.
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST)        \
  auto dtype = input.scalar_type();                                            \
  int d = input.size(-1) / 2;                                                  \
  int64_t num_tokens = input.numel() / input.size(-1);                         \
@@ -227,8 +177,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
            scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
            KERNEL<scalar_t>,                                                  \
            PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
-            ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>(     \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);   \
+            ACT_FIRST, true, true><<<grid, block, 0, stream>>>(                \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);          \
      });                                                                      \
    } else {                                                                   \
      VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {          \
@@ -236,8 +186,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
            scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,      \
            KERNEL<scalar_t>,                                                  \
            PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
-            ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>(    \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);   \
+            ACT_FIRST, true, false><<<grid, block, 0, stream>>>(               \
+            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);          \
      });                                                                      \
    }                                                                          \
  } else {                                                                     \
@@ -247,8 +197,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
          scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type,        \
          KERNEL<scalar_t>,                                                    \
          PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>,   \
-          ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>(            \
-          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT);     \
+          ACT_FIRST, false><<<grid, block, 0, stream>>>(                       \
+          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);            \
    });                                                                        \
  }

@@ -256,14 +206,7 @@ void silu_and_mul(torch::Tensor& out,    // [..., d]
                  torch::Tensor& input)  // [..., 2 * d]
 {
  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true, false, 0.0f);
-}
-
-void silu_and_mul_clamp(torch::Tensor& out,    // [..., d]
-                        torch::Tensor& input,  // [..., 2 * d]
-                        double limit) {
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true, true, (float)limit);
+                                true);
 }

 void mul_and_silu(torch::Tensor& out,    // [..., d]
@@ -272,21 +215,21 @@ void mul_and_silu(torch::Tensor& out,    // [..., d]
  // The difference between mul_and_silu and silu_and_mul is that mul_and_silu
  // applies the silu to the latter half of the input.
  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                false, false, 0.0f);
+                                false);
 }

 void gelu_and_mul(torch::Tensor& out,    // [..., d]
                  torch::Tensor& input)  // [..., 2 * d]
 {
  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
-                                true, false, 0.0f);
+                                true);
 }

 void gelu_tanh_and_mul(torch::Tensor& out,    // [..., d]
                       torch::Tensor& input)  // [..., 2 * d]
 {
-  LAUNCH_ACTIVATION_GATE_KERNEL(
-      vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
+                                vllm::packed_gelu_tanh_kernel, true);
 }

 namespace vllm {
@@ -178,12 +178,7 @@ void rotary_embedding_gptj_impl(

 void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
                      std::optional<torch::Tensor> key, int64_t head_size,
-                      torch::Tensor& cos_sin_cache, bool is_neox,
-                      int64_t rope_dim_offset, bool inverse) {
-  TORCH_CHECK(rope_dim_offset == 0,
-              "rope_dim_offset != 0 is not supported on CPU");
-  TORCH_CHECK(!inverse, "inverse rotary embedding is not supported on CPU");
-
+                      torch::Tensor& cos_sin_cache, bool is_neox) {
  int num_tokens = positions.numel();
  int rot_dim = cos_sin_cache.size(1);
  int num_heads = query.size(-1) / head_size;
@@ -263,8 +263,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def(
      "rotary_embedding(Tensor positions, Tensor! query,"
      "                 Tensor!? key, int head_size,"
-      "                 Tensor cos_sin_cache, bool is_neox, int "
-      "rope_dim_offset=0, bool inverse=False) -> ()");
+      "                 Tensor cos_sin_cache, bool is_neox) -> ()");
  ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);

  // Quantization
@@ -65,16 +65,9 @@ __global__ void rms_norm_static_fp8_quant_kernel(
 #pragma unroll
    for (int j = 0; j < VEC_SIZE; j++) {
      float x = static_cast<float>(src1.val[j]);
-      float w = static_cast<float>(src2.val[j]);
-      // Round normalized result through scalar_t to match the precision of the
-      // unfused composite (rms_norm writes scalar_t, then
-      // static_scaled_fp8_quant re-loads it as float before FP8 conversion).
-      // Without this round, the fused path is strictly more accurate and
-      // disagrees with the composite at exact E4M3 quantization tie boundaries.
-      scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
+      float const out_norm = ((scalar_t)(x * s_variance)) * src2.val[j];
      out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
-          scaled_fp8_conversion<true, fp8_type>(static_cast<float>(out_norm),
-                                                scale_inv);
+          scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
    }
  }
 }
@@ -134,21 +127,13 @@ fused_add_rms_norm_static_fp8_quant_kernel(

  for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
    int id = blockIdx.x * vec_hidden_size + idx;
-    _f16Vec<scalar_t, width> res = residual_v[id];
-    _f16Vec<scalar_t, width> w = weight_v[idx];
-    using Converter = _typeConvert<scalar_t>;
-    using HipT = typename Converter::hip_type;
+    _f16Vec<scalar_t, width> temp = residual_v[id];
+    temp *= s_variance;
+    temp *= weight_v[idx];
 #pragma unroll
    for (int i = 0; i < width; ++i) {
-      float x = Converter::convert(res.data[i]);
-      float wf = Converter::convert(w.data[i]);
-      // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
-      // to match the unfused composite path at FP8 boundaries. We use the
-      // backend's hip_type for the intermediate since c10::Half/BFloat16 has
-      // ambiguous conversions on CUDA and no implicit conversion on ROCm.
-      HipT out_norm_h = Converter::convert(x * s_variance * wf);
-      out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
-          Converter::convert(out_norm_h), scale_inv);
+      out[id * width + i] =
+          scaled_fp8_conversion<true, fp8_type>(float(temp.data[i]), scale_inv);
    }
  }
 }
@@ -191,12 +176,9 @@ fused_add_rms_norm_static_fp8_quant_kernel(

  for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
    float x = (float)residual[blockIdx.x * hidden_size + idx];
-    float w = (float)weight[idx];
-    // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
-    // to match the unfused composite path at FP8 boundaries.
-    scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
-    out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion<true, fp8_type>(
-        static_cast<float>(out_norm), scale_inv);
+    float const out_norm = ((scalar_t)(x * s_variance)) * weight[idx];
+    out[blockIdx.x * hidden_size + idx] =
+        scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
  }
 }

@@ -16,14 +16,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
      "bias) -> ()");
  m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid);

-#ifndef USE_ROCM
  m.def(
      "topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
      "token_expert_indices, Tensor gating_output, bool renormalize, float "
      "routed_scaling_factor, Tensor? "
      "bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
  m.impl("topk_softplus_sqrt", torch::kCUDA, &topk_softplus_sqrt);
-#endif
  // Calculate the result of moe by summing up the partial results
  // from all selected experts.
  m.def("moe_sum(Tensor input, Tensor! output) -> ()");
@@ -163,8 +163,6 @@ void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,

 void silu_and_mul(torch::Tensor& out, torch::Tensor& input);

-void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
-
 void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
                        torch::Tensor& scale);

@@ -82,73 +82,18 @@ void launch_persistent_topk(const torch::Tensor& logits,
    size_t smem_size = P::kFixedSmemLarge + chunk_size * sizeof(uint32_t);
    if (smem_size < P::kSmemMedium) smem_size = P::kSmemMedium;

-    // Query occupancy for the instantiation that will actually launch;
-    // overestimating it deadlocks the cooperative barrier.
    int occupancy = 1;
-    cudaError_t occ_err = cudaSuccess;
-    if (vec_size == 4) {
-      occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-          &occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
-          smem_size);
-    } else if (vec_size == 2) {
-      occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-          &occupancy, P::persistent_topk_kernel<TopK, 2>, P::kThreadsPerBlock,
-          smem_size);
-    } else {
-      occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-          &occupancy, P::persistent_topk_kernel<TopK, 1>, P::kThreadsPerBlock,
-          smem_size);
-    }
-    TORCH_CHECK(occ_err == cudaSuccess,
-                "persistent_topk occupancy query failed: ",
-                cudaGetErrorString(occ_err));
+    cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+        &occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
+        smem_size);
    if (occupancy < 1) occupancy = 1;

-    // The cooperative spin-wait barrier only runs when at least one row hits
-    // the radix path (seq_len > RADIX_THRESHOLD). Below that, non-CTA-0 CTAs
-    // early-exit, so oversubscription can't deadlock and headroom is wasted.
-    const bool needs_cooperative =
-        static_cast<uint32_t>(max_seq_len) > P::RADIX_THRESHOLD;
-
-    const uint32_t hw_resident_cap =
-        static_cast<uint32_t>(num_sms) * static_cast<uint32_t>(occupancy);
-    uint32_t max_resident_ctas = hw_resident_cap;
-    if (needs_cooperative) {
-      // Reserve one CTA per SM when occupancy allows; fall back to a single
-      // CTA when occupancy == 1 (the most deadlock-prone case — any straggler
-      // kernel that takes the only slot on one SM hangs the barrier). Never
-      // drop below one full group's worth.
-      uint32_t headroom = (occupancy > 1) ? static_cast<uint32_t>(num_sms) : 1u;
-      if (max_resident_ctas >= headroom + ctas_per_group) {
-        max_resident_ctas -= headroom;
-      }
-    }
+    uint32_t max_resident_ctas = static_cast<uint32_t>(num_sms) * occupancy;
    uint32_t num_groups = std::min(max_resident_ctas / ctas_per_group,
                                   static_cast<uint32_t>(num_rows));
    if (num_groups == 0) num_groups = 1;
    uint32_t total_ctas = num_groups * ctas_per_group;

-    // If the cooperative launch wouldn't fit, fall back to FilteredTopK
-    // instead of deadlocking. Only relevant when needs_cooperative.
-    if (needs_cooperative && total_ctas > hw_resident_cap) {
-      TORCH_CHECK(max_smem_per_block >= 128 * 1024,
-                  "persistent_topk would oversubscribe and the FilteredTopK "
-                  "fallback requires >=128KB smem per block (have ",
-                  max_smem_per_block, "). total_ctas=", total_ctas,
-                  " > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
-                  ", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
-                  ", smem=", smem_size, ").");
-      cudaError_t status =
-          vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
-              logits.data_ptr<float>(), output.data_ptr<int32_t>(),
-              lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
-              static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride),
-              stream);
-      TORCH_CHECK(status == cudaSuccess,
-                  "FilteredTopK fallback failed: ", cudaGetErrorString(status));
-      return;
-    }
-
    size_t state_bytes = num_groups * sizeof(P::RadixRowState);
    TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
                "workspace too small, need ", state_bytes, " bytes");
@@ -106,12 +106,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
  ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);

-  // SwiGLU activation with input clamping.
-  ops.def(
-      "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
-      "-> ()");
-  ops.impl("silu_and_mul_with_clamp", torch::kCUDA, &silu_and_mul_clamp);
-
  ops.def(
      "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
  ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
@@ -183,7 +177,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "int forced_token_heads_per_warp=-1) -> ()");
  ops.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope);

-#ifndef USE_ROCM
  // Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
  // GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
  // kernel launch.
@@ -194,7 +187,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "float eps, int cache_block_size) -> ()");
  ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
           &fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);
-#endif

  // Apply repetition penalties to logits in-place
  ops.def(
@@ -478,6 +478,9 @@ FROM ${FINAL_BASE_IMAGE} AS vllm-base

 ARG CUDA_VERSION
 ARG PYTHON_VERSION
+ARG DEADSNAKES_MIRROR_URL
+ARG DEADSNAKES_GPGKEY_URL
+ARG GET_PIP_URL

 ENV DEBIAN_FRONTEND=noninteractive
 WORKDIR /vllm-workspace
@@ -487,35 +490,43 @@ WORKDIR /vllm-workspace
 RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \
    echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment

-# Install Python (via uv / python-build-standalone) and system dependencies.
-# This replaces the deadsnakes PPA, removing the build-time dependency on
-# Launchpad and matching how the build-stage (`base`) installs Python.
-# python-build-standalone bundles dev headers, the venv module, and
-# python3-config, so the python3.X-dev / python3.X-venv apt packages
-# are not needed.
+# Install Python and system dependencies
 RUN apt-get update -y \
    && apt-get install -y --no-install-recommends \
+        software-properties-common \
        curl \
        sudo \
        ffmpeg \
        libsm6 \
        libxext6 \
        libgl1 \
+    && if [ ! -z ${DEADSNAKES_MIRROR_URL} ] ; then \
+        if [ ! -z "${DEADSNAKES_GPGKEY_URL}" ] ; then \
+            mkdir -p -m 0755 /etc/apt/keyrings ; \
+            curl -L ${DEADSNAKES_GPGKEY_URL} | gpg --dearmor > /etc/apt/keyrings/deadsnakes.gpg ; \
+            sudo chmod 644 /etc/apt/keyrings/deadsnakes.gpg ; \
+            echo "deb [signed-by=/etc/apt/keyrings/deadsnakes.gpg] ${DEADSNAKES_MIRROR_URL} $(lsb_release -cs) main" > /etc/apt/sources.list.d/deadsnakes.list ; \
+        fi ; \
+    else \
+        for i in 1 2 3; do \
+            add-apt-repository -y ppa:deadsnakes/ppa && break || \
+            { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
+        done ; \
+    fi \
+    && apt-get update -y \
+    && apt-get install -y --no-install-recommends \
+        python${PYTHON_VERSION} \
+        python${PYTHON_VERSION}-dev \
+        python${PYTHON_VERSION}-venv \
        libibverbs-dev \
    && rm -rf /var/lib/apt/lists/* \
-    && curl -LsSf https://astral.sh/uv/install.sh | sh \
-    && $HOME/.local/bin/uv venv /opt/venv --python ${PYTHON_VERSION} \
-    && rm -f /usr/bin/python3 /usr/bin/python3-config /usr/bin/pip \
-    && ln -s /opt/venv/bin/python3 /usr/bin/python3 \
-    && ln -s /opt/venv/bin/python${PYTHON_VERSION} /usr/bin/python${PYTHON_VERSION} \
-    && ln -s /opt/venv/bin/python3-config /usr/bin/python3-config \
-    && ln -s /opt/venv/bin/pip /usr/bin/pip \
+    && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
+    && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
+    && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
+    && rm -f /usr/lib/python${PYTHON_VERSION}/EXTERNALLY-MANAGED \
+    && curl -sS ${GET_PIP_URL} | python${PYTHON_VERSION} \
    && python3 --version && python3 -m pip --version

-# Activate virtual environment and add uv to PATH
-ENV PATH="/opt/venv/bin:/root/.local/bin:$PATH"
-ENV VIRTUAL_ENV="/opt/venv"
-
 # Install CUDA development tools for runtime JIT compilation
 # (FlashInfer, DeepGEMM, EP kernels all require compilation at runtime)
 RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
@@ -529,9 +540,7 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
        libcurand-dev-${CUDA_VERSION_DASH} \
        libcublas-${CUDA_VERSION_DASH} \
        # Required by fastsafetensors (fixes #20384)
-        libnuma-dev \
-        # numactl CLI for NUMA binding at runtime
-        numactl && \
+        libnuma-dev && \
    # Fixes nccl_allocator requiring nccl.h at runtime
    # https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22
    # NCCL packages don't use the cuda-MAJOR-MINOR naming convention,
@@ -540,6 +549,9 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
    apt-get install -y --no-install-recommends --allow-change-held-packages libnccl-dev=${NCCL_VER} libnccl2=${NCCL_VER} && \
    rm -rf /var/lib/apt/lists/*

+# Install uv for faster pip installs
+RUN python3 -m pip install uv
+
 # Environment for uv
 ENV UV_HTTP_TIMEOUT=500
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
@@ -729,7 +741,7 @@ ENV HF_XET_HIGH_PERFORMANCE 1
 ENV HF_HUB_DOWNLOAD_TIMEOUT 60

 # Copy in the v1 package for testing (it isn't distributed yet)
-COPY vllm/v1 /opt/venv/lib/python${PYTHON_VERSION}/site-packages/vllm/v1
+COPY vllm/v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1

 # Source code is used in the `python_only_compile.sh` test
 # We hide it inside `src/` so that this source code
@@ -124,10 +124,10 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1

 # RIXL/UCX build stages
 FROM base AS build_rixl
-ARG RIXL_BRANCH="bf4a7214"
+ARG RIXL_BRANCH="f33a5599"
 ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
-ARG UCX_BRANCH="7009d7a1"
-ARG UCX_REPO="https://github.com/openucx/ucx.git"
+ARG UCX_BRANCH="da3fac2a"
+ARG UCX_REPO="https://github.com/ROCm/ucx.git"
 ENV ROCM_PATH=/opt/rocm
 ENV UCX_HOME=/usr/local/ucx
 ENV RIXL_HOME=/usr/local/rixl
@@ -165,7 +165,7 @@ RUN cd /usr/local/src && \
        --disable-doxygen-doc \
        --enable-optimizations \
        --enable-devel-headers \
-        --with-rocm=${ROCM_PATH} \
+        --with-rocm=/opt/rocm \
        --with-verbs \
        --with-dm \
        --enable-mt && \
@@ -186,12 +186,7 @@ RUN git clone ${RIXL_REPO} /opt/rixl && \
    ninja install

 # Generate RIXL wheel
-# Exclude libcore and libpull from auditwheel: transitive dependencies
-# that are not shipped in the wheel and vary across base images.
-RUN cd /opt/rixl && \
-    sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
-        contrib/build-wheel.sh && \
-    mkdir -p /app/install && \
+RUN cd /opt/rixl && mkdir -p /app/install && \
    ./contrib/build-wheel.sh \
        --output-dir /app/install \
        --rocm-dir ${ROCM_PATH} \
@@ -436,10 +431,6 @@ COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-pac
 ENV MIOPEN_DEBUG_CONV_DIRECT=0
 ENV MIOPEN_DEBUG_CONV_GEMM=0

-# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
-# See: https://github.com/ROCm/rocm-libraries/issues/6266
-ENV HSA_ENABLE_IPC_MODE_LEGACY=1
-
 # Source code is used in the `python_only_compile.sh` test
 # We hide it inside `src/` so that this source code
 # will not be imported by other tests
@@ -36,7 +36,7 @@ th {
 | deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] |
 | deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] |
 | flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
-| flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
+| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |

 !!! info "Table key"
    1. All types: mxfp4, nvfp4, int4, int8, fp8
@@ -292,10 +292,10 @@ Pooling models now support token-wise task.

 ### Score task

-`score` task is deprecated and will be removed in v0.20. Please use `classify` instead. Only when a
-classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.
+`score` task have has been removed in v0.21, use `classify` instead. Only when a classification model outputs num_labels
+equal to 1 can it be used as a scoring model and have its scoring API enabled.

 ### Pooling multitask support

-Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task is not what you want,
+Pooling multitask support has been removed in v0.21. When the default pooling task is not what you want,
 you need to manually specify it via `PoolerConfig(task=<task>)` offline or `--pooler-config.task <task>` online.
@@ -384,7 +384,6 @@ th {
 | `DeepseekForCausalLM` | DeepSeek | `deepseek-ai/deepseek-llm-67b-base`, `deepseek-ai/deepseek-llm-7b-chat`, etc. | ✅︎ | ✅︎ |
 | `DeepseekV2ForCausalLM` | DeepSeek-V2 | `deepseek-ai/DeepSeek-V2`, `deepseek-ai/DeepSeek-V2-Chat`, etc. | ✅︎ | ✅︎ |
 | `DeepseekV3ForCausalLM` | DeepSeek-V3 | `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1`, etc. | ✅︎ | ✅︎ |
-| `DeepseekV4ForCausalLM` | DeepSeek-V4 | `deepseek-ai/DeepSeek-V4-Flash`, `deepseek-ai/DeepSeek-V4-Pro`, etc. | | |
 | `Dots1ForCausalLM` | dots.llm1 | `rednote-hilab/dots.llm1.base`, `rednote-hilab/dots.llm1.inst`, etc. | | ✅︎ |
 | `DotsOCRForCausalLM` | dots_ocr | `rednote-hilab/dots.ocr` | ✅︎ | ✅︎ |
 | `Ernie4_5ForCausalLM` | Ernie4.5 | `baidu/ERNIE-4.5-0.3B-PT`, etc. | ✅︎ | ✅︎ |
@@ -644,10 +643,10 @@ Some models are supported only via the [Transformers modeling backend](#transfor
 !!! note
    `Gemma3nForConditionalGeneration` is only supported on V1 due to shared KV caching and it depends on `timm>=1.0.17` to make use of its
    MobileNet-v5 vision backbone.
-
+  
    Performance is not yet fully optimized mainly due to:
-
-    - Both audio and vision MM encoders use `transformers.AutoModel` implementation.
+  
+    - Both audio and vision MM encoders use `transformers.AutoModel` implementation.  
    - There's no PLE caching or out-of-memory swapping support, as described in [Google's blog](https://developers.googleblog.com/en/introducing-gemma-3n/). These features might be too model-specific for vLLM, and swapping in particular may be better suited for constrained setups.

 !!! note
@@ -4,68 +4,74 @@
 import torch

 from vllm import LLM
+from vllm.config import PoolerConfig
 from vllm.inputs import TextPrompt
 from vllm.multimodal.utils import fetch_image

-# Initialize model
-model = LLM(
-    model="jinaai/jina-embeddings-v4-vllm-text-matching",
-    runner="pooling",
-    max_model_len=1024,
-    gpu_memory_utilization=0.8,
-)

-# Create text prompts
-text1 = "Ein wunderschöner Sonnenuntergang am Strand"
-text1_prompt = TextPrompt(prompt=f"Query: {text1}")
+def main():
+    # Initialize model
+    model = LLM(
+        model="jinaai/jina-embeddings-v4-vllm-text-matching",
+        pooler_config=PoolerConfig(task="token_embed"),
+        runner="pooling",
+        max_model_len=1024,
+        gpu_memory_utilization=0.8,
+    )

-text2 = "浜辺に沈む美しい夕日"
-text2_prompt = TextPrompt(prompt=f"Query: {text2}")
+    # Create text prompts
+    text1 = "Ein wunderschöner Sonnenuntergang am Strand"
+    text1_prompt = TextPrompt(prompt=f"Query: {text1}")

-# Create image prompt
-image = fetch_image(
-    "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg"  # noqa: E501
-)
-image_prompt = TextPrompt(
-    prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n",  # noqa: E501
-    multi_modal_data={"image": image},
-)
+    text2 = "浜辺に沈む美しい夕日"
+    text2_prompt = TextPrompt(prompt=f"Query: {text2}")

-# Encode all prompts
-prompts = [text1_prompt, text2_prompt, image_prompt]
-outputs = model.encode(prompts, pooling_task="token_embed")
+    # Create image prompt
+    image = fetch_image(
+        "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg"  # noqa: E501
+    )
+    image_prompt = TextPrompt(
+        prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n",  # noqa: E501
+        multi_modal_data={"image": image},
+    )
+
+    # Encode all prompts
+    prompts = [text1_prompt, text2_prompt, image_prompt]
+    outputs = model.encode(prompts, pooling_task="token_embed")
+
+    def get_embeddings(outputs):
+        VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
+
+        embeddings = []
+        for output in outputs:
+            if VISION_START_TOKEN_ID in output.prompt_token_ids:
+                # Gather only vision tokens
+                img_start_pos = torch.where(
+                    torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
+                )[0][0]
+                img_end_pos = torch.where(
+                    torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
+                )[0][0]
+                embeddings_tensor = output.outputs.data.detach().clone()[
+                    img_start_pos : img_end_pos + 1
+                ]
+            else:
+                # Use all tokens for text-only prompts
+                embeddings_tensor = output.outputs.data.detach().clone()
+
+            # Pool and normalize embeddings
+            pooled_output = (
+                embeddings_tensor.sum(dim=0, dtype=torch.float32)
+                / embeddings_tensor.shape[0]
+            )
+            embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
+        return embeddings
+
+    embeddings = get_embeddings(outputs)
+
+    for embedding in embeddings:
+        print(embedding.shape)


-def get_embeddings(outputs):
-    VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
-
-    embeddings = []
-    for output in outputs:
-        if VISION_START_TOKEN_ID in output.prompt_token_ids:
-            # Gather only vision tokens
-            img_start_pos = torch.where(
-                torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
-            )[0][0]
-            img_end_pos = torch.where(
-                torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
-            )[0][0]
-            embeddings_tensor = output.outputs.data.detach().clone()[
-                img_start_pos : img_end_pos + 1
-            ]
-        else:
-            # Use all tokens for text-only prompts
-            embeddings_tensor = output.outputs.data.detach().clone()
-
-        # Pool and normalize embeddings
-        pooled_output = (
-            embeddings_tensor.sum(dim=0, dtype=torch.float32)
-            / embeddings_tensor.shape[0]
-        )
-        embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
-    return embeddings
-
-
-embeddings = get_embeddings(outputs)
-
-for embedding in embeddings:
-    print(embedding.shape)
+if __name__ == "__main__":
+    main()
@@ -4,6 +4,7 @@
 from argparse import Namespace

 from vllm import LLM, EngineArgs
+from vllm.config import PoolerConfig
 from vllm.utils.argparse_utils import FlexibleArgumentParser


@@ -13,6 +14,7 @@ def parse_args():
    # Set example specific arguments
    parser.set_defaults(
        model="BAAI/bge-m3",
+        pooler_config=PoolerConfig(task="token_embed"),
        runner="pooling",
        enforce_eager=True,
    )
@@ -32,15 +34,6 @@ def main(args: Namespace):
    # You should pass runner="pooling" for embedding models
    llm = LLM(**vars(args))

-    # Generate embedding. The output is a list of EmbeddingRequestOutputs.
-    outputs = llm.embed(prompts)
-
-    # Print the outputs.
-    print("\nGenerated Outputs:\n" + "-" * 60)
-    for prompt, output in zip(prompts, outputs):
-        embeds = output.outputs.embedding
-        print(len(embeds))
-
    # Generate embedding for each token. The output is a list of PoolingRequestOutput.
    outputs = llm.encode(prompts, pooling_task="token_embed")

@@ -50,6 +43,20 @@ def main(args: Namespace):
        multi_vector = output.outputs.data
        print(multi_vector.shape)

+    query = "What is the capital of France?"
+    documents = [
+        "The capital of Brazil is Brasilia.",
+        "The capital of France is Paris.",
+    ]
+    # Generate scores.
+    outputs = llm.score(query, documents)
+    # Print the outputs.
+    print("\nGenerated Outputs:\n" + "-" * 60)
+    for document, output in zip(documents, outputs):
+        score = output.outputs.score
+        print(f"Pair: {[query, document]!r} \nScore: {score}")
+        print("-" * 60)
+

 if __name__ == "__main__":
    args = parse_args()
@@ -7,10 +7,11 @@ Example online usage of Pooling API for multi vector retrieval.
 Run `vllm serve <model> --runner pooling`
 to start up the server in vLLM. e.g.

-vllm serve BAAI/bge-m3
+vllm serve BAAI/bge-m3 --pooler-config.task token_embed
 """

 import argparse
+import pprint

 import requests
 import torch
@@ -32,7 +33,8 @@ def parse_args():


 def main(args):
-    api_url = f"http://{args.host}:{args.port}/pooling"
+    pooling_url = f"http://{args.host}:{args.port}/pooling"
+    score_url = f"http://{args.host}:{args.port}/score"
    model_name = args.model

    prompts = [
@@ -43,11 +45,23 @@ def main(args):
    ]
    prompt = {"model": model_name, "input": prompts}

-    pooling_response = post_http_request(prompt=prompt, api_url=api_url)
+    pooling_response = post_http_request(prompt=prompt, api_url=pooling_url)
    for output in pooling_response.json()["data"]:
        multi_vector = torch.tensor(output["data"])
        print(multi_vector.shape)

+    queries = "What is the capital of France?"
+    documents = [
+        "The capital of Brazil is Brasilia.",
+        "The capital of France is Paris.",
+    ]
+    prompt = {"model": model_name, "queries": queries, "documents": documents}
+    score_response = post_http_request(prompt=prompt, api_url=score_url)
+    print("\nPrompt when queries is string and documents is a list:")
+    pprint.pprint(prompt)
+    print("\nScore Response:")
+    pprint.pprint(score_response.json())
+

 if __name__ == "__main__":
    args = parse_args()
@@ -12,7 +12,7 @@ torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytor
 flashinfer-python==0.6.8.post1
 flashinfer-cubin==0.6.8.post1
 apache-tvm-ffi==0.1.9
-tilelang==0.1.9
+tilelang
 # Cap nvidia-cudnn-frontend (transitive dep of flashinfer) due to
 # breaking changes in 1.19.0
 nvidia-cudnn-frontend>=1.13.0,<1.19.0
@@ -261,8 +261,6 @@ def _compare_sp(
        },
        "use_inductor_graph_partition": use_inductor_graph_partition,
    }
-    if not use_inductor_graph_partition:
-        compilation_config["splitting_ops"] = []

    tp_sp_args = [
        *common_args,
@@ -116,11 +116,6 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
        model_kwargs["attention_config"] = {"backend": attn_backend.backend.name}
        model_kwargs["tensor_parallel_size"] = tp_size

-        # Cap warmup memory: tests use small max_model_len (1024) but the
-        # engine default max_num_batched_tokens is 16384. Warming up large
-        # models (e.g. Llama-4-Scout-FP8) at 16384 tokens may trigger OOM.
-        model_kwargs.setdefault("max_num_batched_tokens", 8192)
-
        # Sparse MLA models (DSv3.2) hit an over-strict inductor assertion in
        # decompose_auto_functionalized when +rotary_embedding is forced into
        # the compile graph. Disable qk_norm+rope fusion (which auto-enables
@@ -34,10 +34,7 @@ def _run_vllm(vllm_runner):
            mode=CompilationMode.VLLM_COMPILE,
            cudagraph_mode=CUDAGraphMode.NONE,
        ),
-        # Phi-tiny-MoE uses SWA, whose admission cap is `cdiv(L, block_size) + 1`
-        # at default block_size=16 — i.e. 17 blocks for max_model_len=256. Use
-        # 32 for headroom.
-        num_gpu_blocks_override=32,
+        num_gpu_blocks_override=8,
    ):
        pass

@@ -193,7 +190,7 @@ def _run_model(vllm_runner, spec: ModelStartupSpec):
            cudagraph_mode=CUDAGraphMode.NONE,
            pass_config=PassConfig(fuse_allreduce_rms=False),
        ),
-        num_gpu_blocks_override=16,
+        num_gpu_blocks_override=8,
    ):
        pass

@@ -19,7 +19,6 @@ from vllm.config import (
    VllmConfig,
    set_current_vllm_config,
 )
-from vllm.config.utils import Range
 from vllm.distributed import (
    tensor_model_parallel_all_gather,
    tensor_model_parallel_reduce_scatter,
@@ -289,22 +288,6 @@ def test_async_tp_pass_replace(
    run_torch_spawn(async_tp_pass_on_test_model, num_processes)


-def test_async_tp_pass_requires_full_graph_compilation():
-    vllm_config = VllmConfig()
-    vllm_config.compilation_config.use_inductor_graph_partition = False
-    vllm_config.compilation_config.splitting_ops = [
-        "vllm::unified_attention_with_output"
-    ]
-
-    async_tp_pass = object.__new__(AsyncTPPass)
-    async_tp_pass.compilation_config = vllm_config.compilation_config
-
-    with pytest.raises(
-        AssertionError, match="AsyncTPPass requires full-graph compilation"
-    ):
-        async_tp_pass.is_applicable_for_range(Range(start=8, end=8))
-
-
 def async_tp_pass_on_test_model(
    local_rank: int,
    world_size: int,
@@ -22,7 +22,6 @@ from vllm.config import (
    get_current_vllm_config,
    set_current_vllm_config,
 )
-from vllm.config.utils import Range
 from vllm.distributed import tensor_model_parallel_all_reduce
 from vllm.distributed.parallel_state import (
    init_distributed_environment,
@@ -217,24 +216,6 @@ def test_sequence_parallelism_pass(
    run_torch_spawn(sequence_parallelism_pass_on_test_model, num_processes)


-def test_sequence_parallelism_pass_requires_full_graph_compilation():
-    vllm_config = VllmConfig()
-    vllm_config.compilation_config.use_inductor_graph_partition = False
-    vllm_config.compilation_config.splitting_ops = [
-        "vllm::unified_attention_with_output"
-    ]
-
-    sequence_parallelism_pass = object.__new__(SequenceParallelismPass)
-    sequence_parallelism_pass.compilation_config = vllm_config.compilation_config
-    sequence_parallelism_pass.min_token_num = 1
-
-    with pytest.raises(
-        AssertionError,
-        match="SequenceParallelismPass requires full-graph compilation",
-    ):
-        sequence_parallelism_pass.is_applicable_for_range(Range(start=8, end=8))
-
-
 def sequence_parallelism_pass_on_test_model(
    local_rank: int,
    world_size: int,
@@ -405,12 +405,9 @@ def test_should_split():
        (None, 0, 1, False, 2048, CUDAGraphMode.NONE, 0),
        # truncated to nearest multiple of 8 or 16
        (None, 257, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256),
-        # max_num_batched_tokens <= max_cudagraph_capture_size should always be
-        # captured even if not landing on a 16-stride step
-        (None, 2048, 1, False, 257, CUDAGraphMode.FULL_AND_PIECEWISE, 257),
        # max from list
        ([1, 2, 4, 15], None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 15),
-        # SP forces full-graph compilation, sizes are filtered by TP
+        # filtered out 15 due to SP
        ([1, 2, 4, 15], None, 2, True, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 4),
        # limited by the max_tokens
        ([1, 2, 4, 15], None, 1, False, 8, CUDAGraphMode.FULL_AND_PIECEWISE, 4),
@@ -468,123 +465,6 @@ def test_cudagraph_sizes_post_init(
        )


-@pytest.mark.skipif(
-    not current_platform.support_static_graph_mode(),
-    reason="Skip if not cudagraph mode supported",
-)
-@pytest.mark.parametrize(
-    (
-        "cudagraph_mode",
-        "use_inductor_graph_partition",
-        "expected_enable_sp",
-        "expected_cudagraph_mode",
-        "expected_piecewise_compile",
-        "expected_capture_sizes",
-        "expected_max_size",
-    ),
-    [
-        (CUDAGraphMode.PIECEWISE, False, True, CUDAGraphMode.FULL, False, [2, 4], 4),
-        (
-            CUDAGraphMode.FULL_DECODE_ONLY,
-            False,
-            True,
-            CUDAGraphMode.FULL_DECODE_ONLY,
-            False,
-            [2, 4],
-            4,
-        ),
-        (
-            CUDAGraphMode.FULL_AND_PIECEWISE,
-            False,
-            True,
-            CUDAGraphMode.FULL,
-            False,
-            [2, 4],
-            4,
-        ),
-        (
-            CUDAGraphMode.FULL_AND_PIECEWISE,
-            True,
-            True,
-            CUDAGraphMode.FULL_AND_PIECEWISE,
-            True,
-            [2, 4],
-            4,
-        ),
-    ],
-)
-def test_sequence_parallelism_requires_full_graph_compilation(
-    cudagraph_mode: CUDAGraphMode,
-    use_inductor_graph_partition: bool,
-    expected_enable_sp: bool,
-    expected_cudagraph_mode: CUDAGraphMode,
-    expected_piecewise_compile: bool,
-    expected_capture_sizes: list[int],
-    expected_max_size: int,
-):
-    with patch.object(current_platform, "device_count", return_value=2):
-        vllm_config = VllmConfig(
-            parallel_config=ParallelConfig(tensor_parallel_size=2),
-            scheduler_config=SchedulerConfig(
-                max_num_seqs=128,
-                max_num_batched_tokens=2048,
-                max_model_len=2048,
-                is_encoder_decoder=False,
-            ),
-        )
-        vllm_config.model_config = MagicMock(
-            dtype=torch.float16,
-            enforce_eager=False,
-            is_moe=False,
-            disable_cascade_attn=False,
-            get_hidden_size=MagicMock(return_value=4096),
-        )
-        vllm_config.compilation_config = CompilationConfig(
-            mode=CompilationMode.VLLM_COMPILE,
-            cudagraph_capture_sizes=[1, 2, 4, 15],
-            max_cudagraph_capture_size=None,
-            compile_sizes=["cudagraph_capture_sizes"],
-            use_inductor_graph_partition=use_inductor_graph_partition,
-            pass_config=PassConfig(
-                enable_sp=True,
-                fuse_gemm_comms=True,
-                fuse_norm_quant=True,
-                fuse_act_quant=True,
-                eliminate_noops=True,
-                sp_min_token_num=512,
-            ),
-            cudagraph_mode=cudagraph_mode,
-        )
-        vllm_config.compilation_config.set_splitting_ops_for_v1(
-            all2all_backend=vllm_config.parallel_config.all2all_backend,
-            data_parallel_size=1,
-        )
-        vllm_config._set_compile_ranges()
-        vllm_config._set_cudagraph_sizes()
-
-    assert (
-        vllm_config.compilation_config.use_inductor_graph_partition
-        == use_inductor_graph_partition
-    )
-    assert (
-        bool(vllm_config.compilation_config.splitting_ops) == expected_piecewise_compile
-    )
-    assert vllm_config.compilation_config.pass_config.enable_sp == expected_enable_sp
-    assert (
-        vllm_config.compilation_config.pass_config.fuse_gemm_comms == expected_enable_sp
-    )
-    assert vllm_config.compilation_config.cudagraph_mode == expected_cudagraph_mode
-    assert (
-        vllm_config.compilation_config.cudagraph_capture_sizes == expected_capture_sizes
-    )
-    assert (
-        vllm_config.compilation_config.max_cudagraph_capture_size == expected_max_size
-    )
-    assert (
-        511 in vllm_config.compilation_config.compile_ranges_endpoints
-    ) == expected_enable_sp
-
-
 def test_cached_compilation_config(default_vllm_config):
    import torch
    from torch._inductor.utils import run_and_get_code
@@ -311,7 +311,7 @@ async def test_abort_timeout_exits_quickly(wait_for_engine_idle: float):
            pytest.fail("Process did not exit after SIGTERM with abort timeout")

        exit_time = time.time() - start_time
-        assert exit_time < 2.1, f"Default shutdown took too long: {exit_time:.1f}s"
+        assert exit_time < 2, f"Default shutdown took too long: {exit_time:.1f}s"
        assert proc.returncode in (0, -15, None), f"Unexpected: {proc.returncode}"

        await _assert_children_cleaned_up(child_pids)
@@ -1,13 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import logging
 import weakref

 import pytest
 import torch

 from tests.models.utils import softmax
-from vllm import LLM, ClassificationRequestOutput, PoolingParams, PoolingRequestOutput
+from vllm import LLM, ClassificationRequestOutput, PoolingParams
 from vllm.distributed import cleanup_dist_env_and_memory
 from vllm.tasks import PoolingTask

@@ -66,18 +65,6 @@ def test_list_prompts(llm: LLM):
        assert len(outputs[i].outputs.probs) == num_labels


-@pytest.mark.skip_global_cleanup
-def test_token_classify(llm: LLM, caplog_vllm):
-    with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
-        outputs = llm.encode(prompt, pooling_task="token_classify", use_tqdm=False)
-        assert "deprecated" in caplog_vllm.text
-
-    assert len(outputs) == 1
-    assert isinstance(outputs[0], PoolingRequestOutput)
-    assert outputs[0].prompt_token_ids == prompt_token_ids
-    assert outputs[0].outputs.data.shape == (len(prompt_token_ids), num_labels)
-
-
@pytest.mark.skip_global_cleanup
 def test_pooling_params(llm: LLM):
    def get_outputs(use_activation):
@@ -110,10 +97,12 @@ def test_score_api(llm: LLM):
        llm.score("ping", "pong", use_tqdm=False)


-@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
 def test_unsupported_tasks(llm: LLM, task: PoolingTask):
    if task == "plugin":
        err_msg = "No IOProcessor plugin installed."
+    elif task == "token_classify":
+        err_msg = "Try switching the model's pooling_task via.+"
    else:
        err_msg = "Embedding API is not supported by this model.+"
    with pytest.raises(ValueError, match=err_msg):
@@ -436,26 +436,7 @@ async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):

@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
-async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: str):
-    task = "token_classify"
-    response = requests.post(
-        server.url_for("pooling"),
-        json={
-            "model": model_name,
-            "input": input_text,
-            "encoding_format": "float",
-            "task": task,
-        },
-    )
-    poolings = PoolingResponse.model_validate(response.json())
-    assert len(poolings.data) == 1
-    assert len(poolings.data[0].data) == 8
-    assert len(poolings.data[0].data[0]) == 2
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
 async def test_pooling_not_supported(
    server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -469,8 +450,11 @@ async def test_pooling_not_supported(
        },
    )
    assert response.json()["error"]["type"] == "BadRequestError"
+
    if task == "plugin":
        err_msg = "No IOProcessor plugin installed."
+    elif task == "token_classify":
+        err_msg = "Try switching the model's pooling_task via"
    else:
        err_msg = f"Unsupported task: {task!r}"
    assert response.json()["error"]["message"].startswith(err_msg)
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import logging
 import weakref

 import pytest
@@ -38,11 +37,11 @@ def llm():
        seed=0,
        attention_config=attention_config,
    )
+    assert embedding_size == llm.model_config.embedding_size

    yield weakref.proxy(llm)

    del llm
-
    cleanup_dist_env_and_memory()


@@ -74,16 +73,6 @@ def test_list_prompts(llm: LLM):
        assert len(outputs[i].outputs.embedding) == embedding_size


-@pytest.mark.skip_global_cleanup
-def test_token_embed(llm: LLM, caplog_vllm):
-    with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
-        outputs = llm.encode(prompt, pooling_task="token_embed", use_tqdm=False)
-        assert "deprecated" in caplog_vllm.text
-
-    multi_vector = outputs[0].outputs.data
-    assert multi_vector.shape == (11, 384)
-
-
@pytest.mark.skip_global_cleanup
 def test_pooling_params(llm: LLM):
    def get_outputs(normalize):
@@ -107,10 +96,14 @@ def test_pooling_params(llm: LLM):
    )


-@pytest.mark.parametrize("task", ["token_classify", "classify", "plugin"])
+@pytest.mark.parametrize(
+    "task", ["token_classify", "classify", "token_embed", "plugin"]
+)
 def test_unsupported_tasks(llm: LLM, task: PoolingTask):
    if task == "plugin":
        err_msg = "No IOProcessor plugin installed."
+    elif task == "token_embed":
+        err_msg = "Try switching the model's pooling_task via.+"
    else:
        err_msg = "Classification API is not supported by this model.+"
    with pytest.raises(ValueError, match=err_msg):
@@ -732,28 +732,9 @@ async def test_pooling_embed(server: RemoteOpenAIServer, model_name: str):

@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
-async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
-    task = "token_embed"
-    response = requests.post(
-        server.url_for("pooling"),
-        json={
-            "model": model_name,
-            "input": input_text,
-            "encoding_format": "float",
-            "task": task,
-        },
-    )
-
-    poolings = PoolingResponse.model_validate(response.json())
-
-    assert len(poolings.data) == 1
-    assert len(poolings.data[0].data) == len(input_tokens)
-    assert len(poolings.data[0].data[0]) == 384
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
+@pytest.mark.parametrize(
+    "task", ["classify", "token_classify", "token_embed", "plugin"]
+)
 async def test_pooling_not_supported(
    server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -769,6 +750,8 @@ async def test_pooling_not_supported(
    assert response.json()["error"]["type"] == "BadRequestError"
    if task == "plugin":
        err_msg = "No IOProcessor plugin installed."
+    elif task == "token_embed":
+        err_msg = "Try switching the model's pooling_task via"
    else:
        err_msg = f"Unsupported task: {task!r}"
    assert response.json()["error"]["message"].startswith(err_msg)
@@ -452,25 +452,6 @@ async def test_pooling_classify(server: RemoteOpenAIServer):
    assert len(poolings.data[0].data) == 1


-@pytest.mark.asyncio
-async def test_pooling_token_classify(server: RemoteOpenAIServer):
-    response = requests.post(
-        server.url_for("pooling"),
-        json={
-            "model": MODEL_NAME,
-            "task": "token_classify",
-            "input": input_text,
-            "encoding_format": "float",
-        },
-    )
-
-    poolings = PoolingResponse.model_validate(response.json())
-
-    assert len(poolings.data) == 1
-    assert len(poolings.data[0].data) == len(input_tokens)
-    assert len(poolings.data[0].data[0]) == 1
-
-
@pytest.mark.asyncio
 async def test_rerank_max_tokens_per_doc(
    server: RemoteOpenAIServer,
@@ -544,7 +525,7 @@ async def test_rerank_max_tokens_per_doc_validation(


@pytest.mark.asyncio
-@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
 async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
    response = requests.post(
        server.url_for("pooling"),
@@ -558,6 +539,8 @@ async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
    assert response.json()["error"]["type"] == "BadRequestError"
    if task == "plugin":
        err_msg = "No IOProcessor plugin installed."
+    elif task == "token_classify":
+        err_msg = "Try switching the model's pooling_task via"
    else:
        err_msg = f"Unsupported task: {task!r}"
    assert response.json()["error"]["message"].startswith(err_msg)
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import logging
 import weakref

 import pytest
@@ -60,22 +59,19 @@ def test_token_ids_prompts(llm: LLM):

@pytest.mark.skip_global_cleanup
 def test_score_api(llm: LLM):
-    err_msg = "Scoring API is only enabled for num_labels == 1."
+    err_msg = "This model does not support the Scoring API."
    with pytest.raises(ValueError, match=err_msg):
        llm.score("ping", "pong", use_tqdm=False)


@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])
 def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):
-    if task == "classify":
-        with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
-            llm.encode(prompt, pooling_task=task, use_tqdm=False)
-        assert "deprecated" in caplog_vllm.text
+    if task == "plugin":
+        err_msg = "No IOProcessor plugin installed."
+    elif task == "classify":
+        err_msg = "Try switching the model's pooling_task via.+"
    else:
-        if task == "plugin":
-            err_msg = "No IOProcessor plugin installed."
-        else:
-            err_msg = "Embedding API is not supported by this model.+"
+        err_msg = "Embedding API is not supported by this model.+"

-        with pytest.raises(ValueError, match=err_msg):
-            llm.encode(prompt, pooling_task=task, use_tqdm=False)
+    with pytest.raises(ValueError, match=err_msg):
+        llm.encode(prompt, pooling_task=task, use_tqdm=False)
@@ -50,7 +50,7 @@ async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: st

@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
+@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])
 async def test_pooling_not_supported(
    server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -63,9 +63,12 @@ async def test_pooling_not_supported(
            "task": task,
        },
    )
+    assert response.json()["error"]["type"] == "BadRequestError"

    if task == "plugin":
        err_msg = "No IOProcessor plugin installed."
+    elif task == "classify":
+        err_msg = "Try switching the model's pooling_task via"
    else:
        err_msg = f"Unsupported task: {task!r}"
    assert response.json()["error"]["message"].startswith(err_msg)
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import logging
 import weakref

 import pytest
@@ -64,15 +63,12 @@ def test_token_ids_prompts(llm: LLM):

@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])
 def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):
-    if task == "embed":
-        with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
-            llm.encode(prompt, pooling_task=task, use_tqdm=False)
-        assert "deprecated" in caplog_vllm.text
+    if task == "plugin":
+        err_msg = "No IOProcessor plugin installed."
+    elif task == "embed":
+        err_msg = "Try switching the model's pooling_task via.+"
    else:
-        if task == "plugin":
-            err_msg = "No IOProcessor plugin installed."
-        else:
-            err_msg = "Classification API is not supported by this model.+"
+        err_msg = "Classification API is not supported by this model.+"

-        with pytest.raises(ValueError, match=err_msg):
-            llm.encode(prompt, pooling_task=task, use_tqdm=False)
+    with pytest.raises(ValueError, match=err_msg):
+        llm.encode(prompt, pooling_task=task, use_tqdm=False)
@@ -73,7 +73,7 @@ async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):

@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
+@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])
 async def test_pooling_not_supported(
    server: RemoteOpenAIServer, model_name: str, task: str
 ):
@@ -86,9 +86,12 @@ async def test_pooling_not_supported(
            "task": task,
        },
    )
+    assert response.json()["error"]["type"] == "BadRequestError"

    if task == "plugin":
        err_msg = "No IOProcessor plugin installed."
+    elif task == "embed":
+        err_msg = "Try switching the model's pooling_task via"
    else:
        err_msg = f"Unsupported task: {task!r}"
    assert response.json()["error"]["message"].startswith(err_msg)
@@ -10,7 +10,7 @@ from vllm.utils.deep_gemm import (
    _ceil_to_ue8m0,
    calc_diff,
    fp8_fp4_mqa_logits,
-    fp8_fp4_paged_mqa_logits,
+    fp8_paged_mqa_logits,
    get_num_sms,
    get_paged_mqa_logits_metadata,
 )
@@ -128,7 +128,7 @@ def test_deepgemm_fp8_mqa_logits(clean_logits: bool):
                q_fp8 = q.to(torch.float8_e4m3fn)
                kv_fp8 = per_custom_dims_cast_to_fp8(kv, (0,), False)
                logits = fp8_fp4_mqa_logits(
-                    (q_fp8, None), kv_fp8, weights, ks, ke, clean_logits=clean_logits
+                    q_fp8, kv_fp8, weights, ks, ke, clean_logits=clean_logits
                )

                ref_logits = _ref_fp8_mqa_logits(
@@ -150,7 +150,7 @@ def test_deepgemm_fp8_mqa_logits(clean_logits: bool):
                assert diff < 1e-3, f"{diff=}"


-def _ref_fp8_fp4_paged_mqa_logits(
+def _ref_fp8_paged_mqa_logits(
    q: torch.Tensor,
    kv_cache: torch.Tensor,
    weights: torch.Tensor,
@@ -205,10 +205,8 @@ def _ref_fp8_fp4_paged_mqa_logits(
@pytest.mark.skipif(
    not current_platform.has_device_capability(90), reason="SM90 and SM100 only"
 )
-def test_deepgemm_fp8_fp4_paged_mqa_logits():
-    # NOTE: clean_logits=True is incompatible with the 2D context_lens
-    # required by csrc/apis/attention.hpp; only the False path is exercised.
-    clean_logits = False
+@pytest.mark.parametrize("clean_logits", [True, False])
+def test_deepgemm_fp8_paged_mqa_logits(clean_logits: bool):
    torch.manual_seed(0)
    random.seed(0)

@@ -260,29 +258,21 @@ def test_deepgemm_fp8_fp4_paged_mqa_logits():
                q_fp8 = q.to(torch.float8_e4m3fn)
                kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache)

-                # deep_gemm paged MQA logits requires 2D context_lens of
-                # shape (B, next_n) (csrc/apis/attention.hpp:332-335);
-                # see indexer.py:607-608. For each batch/next_n token, the
-                # effective context length is context_lens[b] - next_n + j + 1.
-                next_n_arange = torch.arange(next_n, device="cuda", dtype=torch.int32)
-                context_lens_2d = (
-                    context_lens.unsqueeze(-1) - next_n + 1 + next_n_arange
-                ).contiguous()
                schedule_metadata = get_paged_mqa_logits_metadata(
-                    context_lens_2d, blocksize, get_num_sms()
+                    context_lens, blocksize, get_num_sms()
                )
-                logits = fp8_fp4_paged_mqa_logits(
-                    (q_fp8, None),
+                logits = fp8_paged_mqa_logits(
+                    q_fp8,
                    kv_cache_fp8,
                    weights,
-                    context_lens_2d,
+                    context_lens,
                    block_tables,
                    schedule_metadata,
                    max_model_len,
                    clean_logits=clean_logits,
                )

-                ref_logits = _ref_fp8_fp4_paged_mqa_logits(
+                ref_logits = _ref_fp8_paged_mqa_logits(
                    q,
                    kv_cache,
                    weights,
@@ -16,7 +16,6 @@ from vllm.model_executor.layers.activation import (
    NewGELU,
    QuickGELU,
    SiluAndMul,
-    SiluAndMulWithClamp,
    SwigluOAIAndMul,
    SwigluStepAndMul,
    swiglustep_and_mul_triton,
@@ -117,85 +116,6 @@ def test_act_and_mul(
        opcheck(fn, (out, x))


-SWIGLU_LIMITS = [3.0, 7.0, 15.0]
-
-
-@pytest.mark.parametrize("swiglu_limit", SWIGLU_LIMITS)
-@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
-@pytest.mark.parametrize("d", D)
-@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize("seed", SEEDS)
-@pytest.mark.parametrize("device", CUDA_DEVICES)
-@torch.inference_mode()
-def test_silu_and_mul_with_clamp(
-    default_vllm_config,
-    swiglu_limit: float,
-    num_tokens: int,
-    d: int,
-    dtype: torch.dtype,
-    seed: int,
-    device: str,
-) -> None:
-    """SiluAndMulWithClamp: cuda kernel must match native reference."""
-    set_random_seed(seed)
-    torch.set_default_device(device)
-    # Use large values to ensure clamping is exercised.
-    x = torch.randn(num_tokens, 2 * d, dtype=dtype) * swiglu_limit * 2
-
-    layer = SiluAndMulWithClamp(swiglu_limit, compile_native=False)
-    out = layer(x)
-    ref_out = layer.forward_native(x)
-
-    rtol = {
-        torch.float16: 2e-3,
-        torch.bfloat16: 2e-2,
-        torch.float: 1.3e-6,
-    }
-    torch.testing.assert_close(
-        out, ref_out, atol=get_default_atol(out), rtol=rtol[out.dtype]
-    )
-
-    # Verify clamping is actually being applied: the clamped output should
-    # differ from the unclamped SiluAndMul output when inputs are large.
-    unclamped_out = SiluAndMul.forward_native(x)
-    assert not torch.equal(ref_out.float(), unclamped_out.float()), (
-        "Input was not large enough to exercise the clamp; increase scale"
-    )
-
-    # Verify gate clamping semantics with a controlled scalar case.
-    # gate=large_val is clamped to limit first, then silu(limit) * 1.0.
-    x_gate = torch.tensor(
-        [[swiglu_limit * 20.0, 1.0]], dtype=torch.float32, device=device
-    )
-    out_gate = SiluAndMulWithClamp(swiglu_limit, compile_native=False)(x_gate)
-    expected_gate = torch.nn.functional.silu(
-        torch.tensor(swiglu_limit, dtype=torch.float32)
-    ).item()
-    torch.testing.assert_close(
-        out_gate,
-        torch.tensor([[expected_gate]], dtype=torch.float32, device=device),
-        atol=1e-3,
-        rtol=1e-3,
-    )
-
-    # Verify up clamping semantics: up >> limit gets clamped to limit.
-    x_up = torch.tensor(
-        [[1.0, swiglu_limit * 20.0]], dtype=torch.float32, device=device
-    )
-    out_up = SiluAndMulWithClamp(swiglu_limit, compile_native=False)(x_up)
-    silu_1 = torch.nn.functional.silu(torch.tensor(1.0)).item()
-    torch.testing.assert_close(
-        out_up,
-        torch.tensor([[silu_1 * swiglu_limit]], dtype=torch.float32, device=device),
-        atol=1e-3,
-        rtol=1e-3,
-    )
-
-    # opcheck
-    out_buf = torch.empty(x.shape[:-1] + (d,), dtype=dtype, device=device)
-    opcheck(torch.ops._C.silu_and_mul_with_clamp, (out_buf, x, swiglu_limit))
-
-
@pytest.mark.parametrize(
    "activation",
    [
@@ -1632,7 +1632,6 @@ def _parallel_worker(
                if all2all_manager is not None:
                    all2all_manager.destroy()
            total = total + 1
-            torch.distributed.barrier()

    skipped = total - (passed + failed)

@@ -3,11 +3,12 @@
 """
 Round-trip tests for compressor → FP8 quant + KV cache insert → gather + dequant.

-Four test functions cover five paths:
+Two paths tested:
  A) DeepseekV4 Attention: head_dim=512 (448 FP8 nope + 64 bf16 rope), quant_block=64
  B) Indexer:       head_dim=128 (all FP8), quant_block=128
-  C) DeepseekV4 Attention magnitude range: correctness across small/large values
-  D) Indexer fused Triton kernel: compress+norm+rope+quant+insert
+
+These serve as golden references for validating the future fused
+compressor+quant+cache kernel.
 """

 import math
@@ -20,12 +21,6 @@ from vllm.v1.attention.ops.deepseek_v4_ops import (
    dequantize_and_gather_k_cache,
    quantize_and_insert_k_cache,
 )
-from vllm.v1.attention.ops.deepseek_v4_ops.fused_compress_quant_cache import (
-    _fused_kv_compress_norm_rope_insert_indexer_attn,
-    _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn,
-)
-
-from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4


 def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float):
@@ -314,222 +309,3 @@ def test_deepseek_v4_quant_magnitude_range():
                f"Token {t}: rel_err={rel_err:.4f}, abs_diff={abs_diff:.6f}, "
                f"magnitude={magnitude:.4f}"
            )
-
-
-# ── Test D: Indexer fused K-cache insert (Triton kernels) ────────────────────
-#
-# Both kernels share the same Triton signature; use_fp4 selects between them.
-# Full pipeline: state-cache gather → softmax-weighted compress → RMSNorm →
-#   GPT-J RoPE → quant (MXFP4 or FP8) → paged cache insert.
-
-
-def _reference_kv_compress_norm_rope(
-    state_cache: torch.Tensor,
-    block_table: torch.Tensor,
-    positions: torch.Tensor,
-    rms_weight: torch.Tensor,
-    cos_sin_cache: torch.Tensor,
-    compress_ratio: int = 1,
-    overlap: int = 0,
-    use_fp4: bool = False,
-    rms_eps: float = 1e-6,
-    fp8_max: float = 448.0,
-):
-    """Compress → RMSNorm → GPT-J RoPE → quantize.
-
-    Gathers (1+overlap)*compress_ratio state entries per output token, applies
-    per-element softmax over the scores, and computes the weighted kv sum.
-    Returns (quantized_values, scale) matching the kernel's output layout.
-    """
-    device = state_cache.device
-    head_dim = rms_weight.shape[0]
-    rope_dim = cos_sin_cache.shape[-1]
-    state_block_size = state_cache.shape[1]
-    state_width = state_cache.shape[-1] // 2
-    nope_dim = head_dim - rope_dim
-    total = (1 + overlap) * compress_ratio
-    results = []
-    for pos in positions.tolist():
-        src = torch.arange(pos - total + 1, pos + 1, dtype=torch.int64, device=device)
-        valid = src >= 0
-        idx = src.clamp(min=0)
-        pages = block_table[0, idx // state_block_size]
-        offsets = idx % state_block_size
-        raw = state_cache[pages, offsets].float()  # [total, state_dim]
-
-        # Group 0 (tokens 0..cr-1):   kv[:H],   score[SW:SW+H]
-        # Group 1 (tokens cr..2cr-1): kv[H:2H], score[SW+H:SW+2H]
-        if overlap:
-            sw = state_width
-            g0_kv = raw[:compress_ratio, :head_dim]
-            g1_kv = raw[compress_ratio:, head_dim : 2 * head_dim]
-            g0_scores = raw[:compress_ratio, sw : sw + head_dim]
-            g1_scores = raw[compress_ratio:, sw + head_dim : sw + 2 * head_dim]
-            kv = torch.cat([g0_kv, g1_kv])
-            scores = torch.cat([g0_scores, g1_scores])
-        else:
-            kv = raw[:, :head_dim]
-            scores = raw[:, state_width : state_width + head_dim]
-
-        scores[~valid] = float("-inf")
-        kv[~valid] = 0.0
-        weights = torch.softmax(scores, dim=0)
-        compressed = (kv * weights).sum(dim=0)  # [H]
-        var = (compressed * compressed).mean()
-        normed = compressed * torch.rsqrt(var + rms_eps) * rms_weight.float()
-        compressed_pos = (pos // compress_ratio) * compress_ratio
-        cos, sin = cos_sin_cache[compressed_pos].float().chunk(2)
-        nope, rope = normed.split([nope_dim, rope_dim])
-        rope = torch.stack(
-            [rope[0::2] * cos - rope[1::2] * sin, rope[1::2] * cos + rope[0::2] * sin],
-            dim=-1,
-        ).reshape(rope_dim)
-        results.append(torch.cat([nope, rope]).to(state_cache.dtype))
-    result = torch.stack(results)
-
-    if use_fp4:
-        return quantize_to_mxfp4(result)
-    else:
-        pairs = [
-            _ue8m0_reference(result[t], head_dim, fp8_max) for t in range(len(result))
-        ]
-        quants, scales = zip(*pairs)
-        return torch.stack(quants), torch.cat(scales)
-
-
-@pytest.mark.parametrize("num_tokens", [1, 7, 32])
-@pytest.mark.parametrize("kv_block_size", [16, 32])
-@pytest.mark.parametrize("use_fp4", [False, True])
-def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: bool):
-    """Fused K compress+norm+rope+quant+insert for the indexer KV cache."""
-    HEAD_DIM = 128
-    ROPE_DIM = 64
-    BLOCK_SIZE = 16
-    RMS_EPS = 1e-6
-    FP8_MAX = 448.0
-
-    device = "cuda"
-    torch.manual_seed(42)
-    compress_ratio = 4
-
-    if use_fp4:
-        TOKEN_STRIDE = HEAD_DIM // 2  # packed nibbles: 64 bytes
-        SCALE_DIM = HEAD_DIM // 32  # ue8m0 bytes: 4
-        QUANT_BLOCK = 32
-        kernel = _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn
-    else:
-        TOKEN_STRIDE = HEAD_DIM  # FP8 bytes: 128
-        SCALE_DIM = 4  # 1 float32: 4 bytes
-        QUANT_BLOCK = HEAD_DIM
-        kernel = _fused_kv_compress_norm_rope_insert_indexer_attn
-
-    # overlap=1 whenever compress_ratio==4, matching DeepseekCompressor logic.
-    overlap = 1 if compress_ratio == 4 else 0
-    coff = 1 + overlap  # multiplier for state_dim per entry
-
-    num_pages = (compress_ratio * num_tokens - 1) // BLOCK_SIZE + 2
-    state_cache = torch.randn(
-        num_pages,
-        BLOCK_SIZE,
-        2 * coff * HEAD_DIM,  # kv_state + score_state, each coff*HEAD_DIM wide
-        dtype=torch.bfloat16,
-        device=device,
-    )
-    block_table = torch.arange(num_pages, dtype=torch.int32, device=device).unsqueeze(0)
-    token_to_req = torch.zeros(num_tokens, dtype=torch.int32, device=device)
-    slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
-    positions = torch.arange(
-        compress_ratio - 1,
-        compress_ratio * num_tokens,
-        compress_ratio,
-        dtype=torch.int64,
-        device=device,
-    )
-    rms_weight = torch.randn(HEAD_DIM, dtype=torch.bfloat16, device=device)
-    cos_sin_cache = torch.randn(compress_ratio * num_tokens, ROPE_DIM, device=device)
-
-    kv_n_blocks = (num_tokens + kv_block_size - 1) // kv_block_size + 1
-    kv_cache = torch.zeros(
-        kv_n_blocks,
-        kv_block_size * (TOKEN_STRIDE + SCALE_DIM),
-        dtype=torch.uint8,
-        device=device,
-    )
-
-    kernel[(num_tokens,)](
-        state_cache,
-        state_cache.stride(0),
-        state_cache.stride(1),
-        token_to_req,
-        positions,
-        slot_mapping,
-        block_table,
-        block_table.stride(0),
-        BLOCK_SIZE,
-        rms_weight,
-        RMS_EPS,
-        cos_sin_cache,
-        cos_sin_cache.stride(0),
-        kv_cache,
-        slot_mapping,
-        kv_block_size,
-        HEAD_SIZE=HEAD_DIM,
-        TRITON_BLOCK_SIZE=HEAD_DIM,
-        STATE_WIDTH=coff * HEAD_DIM,
-        COMPRESS_RATIO=compress_ratio,
-        OVERLAP=overlap,
-        ROPE_HEAD_DIM=ROPE_DIM,
-        FP8_MAX=FP8_MAX,
-        QUANT_BLOCK=QUANT_BLOCK,
-        TOKEN_STRIDE=TOKEN_STRIDE,
-        SCALE_DIM=SCALE_DIM,
-        KV_BLOCK_STRIDE=kv_cache.stride(0),
-        num_warps=1,
-    )
-
-    k_quant, scale = _reference_kv_compress_norm_rope(
-        state_cache,
-        block_table,
-        positions,
-        rms_weight,
-        cos_sin_cache,
-        compress_ratio,
-        overlap,
-        use_fp4,
-        rms_eps=RMS_EPS,
-        fp8_max=FP8_MAX,
-    )
-
-    if use_fp4:
-        for i in range(num_tokens):
-            blk, pos = i // kv_block_size, i % kv_block_size
-            val_off = pos * TOKEN_STRIDE
-            fp4_actual = kv_cache[blk, val_off : val_off + TOKEN_STRIDE]
-            assert torch.equal(k_quant[i], fp4_actual), (
-                f"token {i}: packed nibbles differ, "
-                f"{(k_quant[i] != fp4_actual).sum()} "
-                f"/ {TOKEN_STRIDE}"
-            )
-
-            scale_off = kv_block_size * TOKEN_STRIDE + pos * SCALE_DIM
-            scale_actual = kv_cache[blk, scale_off : scale_off + SCALE_DIM]
-            assert torch.equal(scale_actual, scale[i]), (
-                f"token {i}: ue8m0 {scale_actual.tolist()} != {scale[i].tolist()}"
-            )
-
-    else:
-        k_quant = k_quant.view(torch.uint8)
-        for i in range(num_tokens):
-            blk, pos = i // kv_block_size, i % kv_block_size
-            val_off = pos * TOKEN_STRIDE
-            assert torch.equal(
-                k_quant[i], kv_cache[blk, val_off : val_off + TOKEN_STRIDE]
-            ), f"token {i}: FP8 bytes differ"
-
-            scale_off = kv_block_size * TOKEN_STRIDE + pos * SCALE_DIM
-            actual_scale = kv_cache[blk, scale_off : scale_off + SCALE_DIM].view(
-                torch.float32
-            )
-            assert torch.equal(actual_scale, scale[i : i + 1]), (
-                f"token {i}: scale {actual_scale.item()} != {scale[i].item()}"
-            )
@@ -30,56 +30,6 @@ N_HEAD = 64
 MAX_POS = 4096


-def quantize_to_mxfp4(
-    x: torch.Tensor,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Reference MXFP4 quantization.
-
-    Args:
-        x: [..., head_dim] where head_dim is divisible by 32
-    Returns:
-        packed: [..., head_dim//2]  uint8   2 E2M1 nibbles/byte, low nibble = even index
-        scales: [..., head_dim//32] uint8   1 ue8m0 byte
-    """
-    MXFP4_BLOCK_SIZE = 32
-    orig_shape = x.shape
-    head_dim = orig_shape[-1]
-    n_blocks = head_dim // MXFP4_BLOCK_SIZE
-
-    x_f32 = x.float().reshape(-1, n_blocks, MXFP4_BLOCK_SIZE)
-
-    # Per-block ue8m0 scale: 2^ceil(log2(amax / 6.0)), stored as byte = exp + 127
-    # 6 * 2^-126 is from https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/inference/kernel.py#L163
-    amax = x_f32.abs().amax(dim=-1, keepdim=True).clamp(min=6 * (2**-126))
-    log2_ratio = (amax * (1.0 / 6.0)).log2().ceil().clamp(-127.0, 127.0)
-    scale = log2_ratio.exp2()
-    ue8m0 = (log2_ratio + 127.0).to(torch.uint8)  # [*, n_blocks]
-
-    # E2M1 round-to-nearest-even: midpoints round to the even code.
-    # E2M1 values: [0.00, 0.50, 1.00, 1.50, 2.00, 3.00, 4.00, 6.00]
-    # boundaries:  [   0.25, 0.75, 1.25, 1.75, 2.50, 3.50, 5.00]
-    x_scaled = (x_f32 / scale).clamp(-6.0, 6.0)
-    abs_x = x_scaled.abs()
-    code = torch.zeros_like(abs_x, dtype=torch.int32)
-    code = torch.where(abs_x > 0.25, 1, code)
-    code = torch.where(abs_x >= 0.75, 2, code)
-    code = torch.where(abs_x > 1.25, 3, code)
-    code = torch.where(abs_x >= 1.75, 4, code)
-    code = torch.where(abs_x > 2.5, 5, code)
-    code = torch.where(abs_x >= 3.5, 6, code)
-    code = torch.where(abs_x > 5.0, 7, code)
-    sign = ((x_scaled.view(torch.int32) >> 31) & 1).to(torch.uint8)
-    nibble = code.to(torch.uint8) | (sign << 3)
-
-    # Pack: even-index element → low nibble, odd-index → high nibble
-    nibble_flat = nibble.reshape(-1, head_dim)
-    packed = (nibble_flat[:, 0::2] | (nibble_flat[:, 1::2] << 4)).contiguous()
-    packed = packed.reshape(*orig_shape[:-1], head_dim // 2)
-
-    scales = ue8m0.view(*orig_shape[:-1], n_blocks)
-    return packed, scales
-
-
 def _reference(
    positions: torch.Tensor,
    q: torch.Tensor,
@@ -87,7 +37,6 @@ def _reference(
    weights: torch.Tensor,
    softmax_scale: float,
    head_scale: float,
-    use_fp4: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
    q_rot = q.clone()
    ops.rotary_embedding(
@@ -100,33 +49,22 @@ def _reference(
        HEAD_DIM - ROPE_DIM,  # rope_dim_offset → rotate the tail
        False,
    )
+    q_fp8, q_scale = per_token_group_quant_fp8(
+        q_rot.view(-1, HEAD_DIM).contiguous(),
+        HEAD_DIM,
+        use_ue8m0=True,
+    )
+    q_fp8 = q_fp8.view(-1, N_HEAD, HEAD_DIM)
+    q_scale = q_scale.view(-1, N_HEAD)

-    if use_fp4:
-        q_packed, ue8m0 = quantize_to_mxfp4(q_rot.view(-1, N_HEAD, HEAD_DIM))
-        # Pack 4 ue8m0 bytes into 1 int32
-        q_scale = ue8m0.view(torch.int32).squeeze(-1)
-        # FP4 path: q_scale stays separate (cannot be folded into a per-token scalar)
-        weights_out = weights.to(torch.float32) * softmax_scale * head_scale
-        return (q_packed, q_scale), weights_out
-
-    else:
-        q_fp8, q_scale = per_token_group_quant_fp8(
-            q_rot.view(-1, HEAD_DIM).contiguous(),
-            HEAD_DIM,
-            use_ue8m0=True,
-        )
-        q_fp8 = q_fp8.view(-1, N_HEAD, HEAD_DIM)
-        q_scale = q_scale.view(-1, N_HEAD)
-
-        weights_out = weights.to(torch.float32) * q_scale * softmax_scale * head_scale
-        return q_fp8, weights_out
+    weights_out = weights.to(torch.float32) * q_scale * softmax_scale * head_scale
+    return q_fp8, weights_out


@pytest.mark.parametrize("num_tokens", [1, 7, 32, 257])
@pytest.mark.parametrize("cache_dtype", [torch.float32, torch.bfloat16])
-@pytest.mark.parametrize("use_fp4", [False, True])
@torch.inference_mode()
-def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype, use_fp4):
+def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype):
    device = "cuda"
    torch.manual_seed(0)

@@ -139,32 +77,21 @@ def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype, use
    softmax_scale = HEAD_DIM**-0.5
    head_scale = N_HEAD**-0.5

-    q_quant_ref, weights_ref = _reference(
-        positions, q, cos_sin_cache, weights, softmax_scale, head_scale, use_fp4
+    q_fp8_ref, weights_ref = _reference(
+        positions, q, cos_sin_cache, weights, softmax_scale, head_scale
    )
-    q_quant_fused, weights_fused = fused_indexer_q_rope_quant(
-        positions, q.clone(), cos_sin_cache, weights, softmax_scale, head_scale, use_fp4
+    q_fp8_fused, weights_fused = fused_indexer_q_rope_quant(
+        positions, q.clone(), cos_sin_cache, weights, softmax_scale, head_scale
    )

-    if use_fp4:
-        q_quant_ref, q_scale_ref = q_quant_ref
-        q_quant_fused, q_scale_fused = q_quant_fused
-
-        assert torch.equal(q_scale_ref, q_scale_fused), (
-            f"q_scale mismatch: "
-            f"{(q_scale_ref != q_scale_fused).sum().item()} "
-            f"/ {q_scale_ref.numel()} bytes differ"
-        )
-
    # fp8 tensors aren't directly comparable via torch.equal — reinterpret as int8.
-    ref_bits = q_quant_ref.view(torch.int8)
-    fused_bits = q_quant_fused.view(torch.int8)
+    ref_bits = q_fp8_ref.view(torch.int8)
+    fused_bits = q_fp8_fused.view(torch.int8)
    assert torch.equal(ref_bits, fused_bits), (
-        f"q_quant_fused mismatch: "
+        f"q_fp8 mismatch: "
        f"{(ref_bits != fused_bits).sum().item()} / {ref_bits.numel()} bytes differ"
    )

-    assert weights_fused.dtype == torch.float32
    assert torch.equal(weights_ref, weights_fused), (
        f"weights mismatch: max abs diff "
        f"{(weights_ref - weights_fused).abs().max().item()}"
@@ -6,6 +6,7 @@ from transformers import AutoModel

 from tests.models.utils import check_embeddings_close
 from vllm import TokensPrompt
+from vllm.config import PoolerConfig


@pytest.mark.parametrize(
@@ -21,6 +22,7 @@ def test_embed_models(hf_runner, vllm_runner, model: str):
    with vllm_runner(
        model,
        runner="pooling",
+        pooler_config=PoolerConfig(task="token_embed"),
        max_model_len=128,
        max_num_batched_tokens=chunk_size,
        enforce_eager=True,
@@ -3,7 +3,6 @@
 import httpx
 import openai
 import pytest
-import pytest_asyncio
 import torch

 from ....utils import RemoteOpenAIServer
@@ -25,29 +24,42 @@ sentences_2 = [
 similarity_reference = [[0.6259, 0.3474], [0.3309, 0.6734]]
 lexical_score_reference = [0.19554901123046875, 0.0]
 colbert_score_reference = [0.7797, 0.4620]
+SUPPORTED_TASKS = ["embed", "token_embed", "token_classify"]
+
+
+@pytest.fixture(scope="module", params=SUPPORTED_TASKS)
+def pooling_task(request):
+    yield request.param


@pytest.fixture(scope="module")
-def server():
+def server(pooling_task):
    args = [
        "--max-model-len",
        str(MAX_MODEL_LEN),
        "--hf-overrides",
        '{"architectures": ["BgeM3EmbeddingModel"]}',
+        "--pooler-config.task",
+        pooling_task,
    ]

    with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
        yield remote_server


-@pytest_asyncio.fixture
-async def client(server):
-    async with server.get_async_client() as async_client:
-        yield async_client
-
-
@pytest.mark.asyncio
-async def test_bge_m3_api_server_embedding(client: openai.AsyncOpenAI):
+async def test_bge_m3_api_server_embedding(server, pooling_task):
+    client = server.get_async_client()
+
+    if pooling_task != "embed":
+        with pytest.raises(openai.InternalServerError):
+            await run_client_embeddings(
+                client,
+                MODEL_NAME,
+                sentences_1,
+            )
+        return
+
    embeddings_list_1 = await run_client_embeddings(
        client,
        MODEL_NAME,
@@ -117,7 +129,14 @@ def compute_lexical_matching_score(


@pytest.mark.asyncio
-async def test_bge_m3_api_server_sparse_embedding(client: openai.AsyncOpenAI):
+async def test_bge_m3_api_server_sparse_embedding(server, pooling_task):
+    client = server.get_async_client()
+
+    if pooling_task != "token_classify":
+        with pytest.raises(openai.BadRequestError):
+            await sparse_embeddings(client, sentences_1)
+        return
+
    embeddings_1 = await sparse_embeddings(client, sentences_1)
    embeddings_2 = await sparse_embeddings(client, sentences_2)

@@ -137,9 +156,11 @@ async def test_bge_m3_api_server_sparse_embedding(client: openai.AsyncOpenAI):


@pytest.mark.asyncio
-async def test_bge_m3_api_server_sparse_embedding_corner_case(
-    client: openai.AsyncOpenAI,
-):
+async def test_bge_m3_api_server_sparse_embedding_corner_case(server, pooling_task):
+    if pooling_task != "token_classify":
+        return
+
+    client = server.get_async_client()
    embeddings = await sparse_embeddings(client, ["Hi"])
    assert len(embeddings) == 1
    assert 2673 in embeddings[0]
@@ -155,7 +176,18 @@ def colbert_score(q_reps: torch.Tensor, p_reps: torch.Tensor) -> torch.Tensor:


@pytest.mark.asyncio
-async def test_bge_m3_api_server_multi_vector(client: openai.AsyncOpenAI):
+async def test_bge_m3_api_server_multi_vector(server, pooling_task):
+    client = server.get_async_client()
+
+    if pooling_task != "token_embed":
+        with pytest.raises(openai.BadRequestError):
+            await client.post(
+                "../pooling",
+                body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},
+                cast_to=httpx.Response,
+            )
+        return
+
    result_1 = await client.post(
        "../pooling",
        body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},
@@ -4,6 +4,7 @@ import pytest
 import torch

 from vllm import TokensPrompt
+from vllm.config import PoolerConfig


@pytest.mark.parametrize(
@@ -20,6 +21,7 @@ def test_extract_hidden_states(hf_runner, vllm_runner, model: str):
        max_model_len=128,
        enforce_eager=True,
        runner="pooling",
+        pooler_config=PoolerConfig(task="token_embed"),
        enable_prefix_caching=True,
    ) as vllm_model:
        pooling_outputs = vllm_model.llm.encode(
@@ -44,14 +46,3 @@ def test_extract_hidden_states(hf_runner, vllm_runner, model: str):
            assert len(output.prompt_token_ids) == n
            assert len(output.outputs.data) == n
            assert output.num_cached_tokens == 0
-
-        # skip_reading_prefix_cache can still write to cache
-        # to accelerate following requests
-        pooling_outputs = vllm_model.llm.encode(
-            [TokensPrompt(prompt_token_ids=t) for t in token_prompts],
-            pooling_task="embed",
-        )
-
-        for n, output in zip(n_prompt_tokens, pooling_outputs):
-            assert len(output.prompt_token_ids) == n
-            assert output.num_cached_tokens > 0
@@ -5,6 +5,7 @@ import torch
 from transformers import AutoModel

 from tests.models.utils import check_embeddings_close
+from vllm.config import PoolerConfig


@pytest.mark.parametrize(
@@ -17,6 +18,7 @@ def test_embed_models(hf_runner, vllm_runner, example_prompts, model: str, dtype
    with vllm_runner(
        model,
        runner="pooling",
+        pooler_config=PoolerConfig(task="token_embed"),
        max_model_len=None,
    ) as vllm_model:
        vllm_outputs = vllm_model.token_embed(example_prompts)
@@ -146,7 +146,7 @@ def test_multi_vector_retrieval_models_using_normalize(
        model,
        max_model_len=512,
        dtype=dtype,
-        pooler_config=PoolerConfig(use_activation=False),
+        pooler_config=PoolerConfig(use_activation=False, task="token_embed"),
    ) as vllm_model:
        wo_normalize = vllm_model.token_embed(example_prompts)

@@ -154,7 +154,7 @@ def test_multi_vector_retrieval_models_using_normalize(
        model,
        max_model_len=512,
        dtype=dtype,
-        pooler_config=PoolerConfig(use_activation=True),
+        pooler_config=PoolerConfig(use_activation=True, task="token_embed"),
    ) as vllm_model:
        w_normalize = vllm_model.token_embed(example_prompts)

@@ -260,9 +260,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
        trust_remote_code=True,
    ),
    "DeepseekV32ForCausalLM": _HfExamplesInfo("deepseek-ai/DeepSeek-V3.2-Exp"),
-    "DeepseekV4ForCausalLM": _HfExamplesInfo(
-        "deepseek-ai/DeepSeek-V4-Flash", is_available_online=False
-    ),
+    "DeepseekV4ForCausalLM": _HfExamplesInfo("Placeholder", is_available_online=False),
    "Ernie4_5ForCausalLM": _HfExamplesInfo("baidu/ERNIE-4.5-0.3B-PT"),
    "Ernie4_5_MoeForCausalLM": _HfExamplesInfo("baidu/ERNIE-4.5-21B-A3B-PT"),
    "ExaoneForCausalLM": _HfExamplesInfo(
@@ -1485,12 +1483,7 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
        speculative_model="luccafong/deepseek_mtp_draft_random",
        trust_remote_code=True,
    ),
-    "DeepSeekV4MTPModel": _HfExamplesInfo(
-        "deepseek-ai/DeepSeek-V4-Flash",
-        speculative_model="deepseek-ai/DeepSeek-V4-Flash",
-        trust_remote_code=True,
-        is_available_online=False,
-    ),
+    "DeepSeekV4MTPModel": _HfExamplesInfo("Placeholder", is_available_online=False),
    "ErnieMTPModel": _HfExamplesInfo(
        "baidu/ERNIE-4.5-21B-A3B-PT",
        trust_remote_code=True,
@@ -5,6 +5,7 @@ from types import SimpleNamespace

 import pytest
 import torch
+from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8

 from vllm.model_executor.models.deepseek_v4 import (
    DeepseekV4MegaMoEExperts,
@@ -111,8 +112,6 @@ def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership():
    reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.",
 )
 def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact():
-    from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8
-
    device = torch.device("cuda")
    num_tokens = 7
    hidden_size = 256
@@ -188,30 +188,6 @@ class TestExtractToolCalls:
            "location": "NYC"
        }

-    def test_type_conversion_in_non_streaming(self):
-        """Non-streaming extraction must convert params using the tool schema."""
-        tool = ChatCompletionToolsParam(
-            function=FunctionDefinition(
-                name="toggle",
-                parameters={
-                    "type": "object",
-                    "properties": {
-                        "enabled": {"type": "boolean"},
-                        "count": {"type": "integer"},
-                    },
-                },
-            ),
-        )
-        parser = make_parser(tools=[tool])
-        model_output = build_tool_call("toggle", {"enabled": "true", "count": "42"})
-        result = parser.extract_tool_calls(model_output, None)
-        assert result.tools_called
-        assert len(result.tool_calls) == 1
-        args = json.loads(result.tool_calls[0].function.arguments)
-        assert args == {"enabled": True, "count": 42}
-        assert isinstance(args["enabled"], bool)
-        assert isinstance(args["count"], int)
-

 # ---------------------------------------------------------------------------
 # Tests: extract_tool_calls_streaming
@@ -508,58 +484,6 @@ class TestExtractToolCallsStreaming:
        # Should have no tool call deltas yet
        assert all(not d.tool_calls for d in deltas)

-    def test_no_marker_leak_chunked(self, parser):
-        """Chunked streaming must NOT leak DSML start-marker fragments
-        as content (GitHub #40801)."""
-        full_text = build_tool_call("fn", {"k": "v"})
-        deltas = self._stream_chunked(parser, full_text, chunk_size=5)
-        content = "".join(d.content for d in deltas if d.content is not None)
-        assert content == ""
-        args_str = self._reconstruct_args(deltas)
-        assert json.loads(args_str) == {"k": "v"}
-
-    def test_no_marker_leak_with_prefix_chunked(self, parser):
-        """Content before a tool call must not include start-marker
-        fragments when chunked (GitHub #40801)."""
-        full_text = "Hello!" + build_tool_call("fn", {"a": "b"})
-        deltas = self._stream_chunked(parser, full_text, chunk_size=5)
-        content = "".join(d.content for d in deltas if d.content is not None)
-        assert content == "Hello!"
-        assert "DSML" not in content
-        assert "<｜" not in content
-        args_str = self._reconstruct_args(deltas)
-        assert json.loads(args_str) == {"a": "b"}
-
-    def test_no_marker_leak_char_by_char(self, parser):
-        """Character-by-character streaming must not leak marker
-        fragments (GitHub #40801)."""
-        full_text = build_tool_call("fn", {"k": "v"})
-        deltas = self._stream_chunked(parser, full_text, chunk_size=1)
-        content = "".join(d.content for d in deltas if d.content is not None)
-        assert content == ""
-        args_str = self._reconstruct_args(deltas)
-        assert json.loads(args_str) == {"k": "v"}
-
-    def test_no_marker_leak_all_split_points(self, parser):
-        """Start token split at every possible boundary must not
-        leak (GitHub #40801)."""
-        for chunk_size in range(1, len(FC_START) + 2):
-            p = make_parser()
-            full_text = build_tool_call("fn", {"k": "v"})
-            deltas = self._stream_chunked(p, full_text, chunk_size=chunk_size)
-            content = "".join(d.content for d in deltas if d.content is not None)
-            assert content == "", (
-                f"Leaked content {content!r} at chunk_size={chunk_size}"
-            )
-
-    def test_false_partial_marker_emitted(self, parser):
-        """Text ending with a prefix of the start token that turns out
-        NOT to be a marker must still be emitted as content."""
-        full_text = "<｜DSM some regular text"
-        deltas = self._stream_chunked(parser, full_text, chunk_size=3)
-        content = "".join(d.content for d in deltas if d.content is not None)
-        assert content == full_text
-

 class TestDelimiterPreservation:
    """Regression: fast detokenization skipping DSML delimiters (PR #33964)."""
@@ -51,7 +51,6 @@ def test_indexer_builder_deepseek_v4_compressed_slot_mapping_uses_storage_block_
        query_start_loc=query_start_loc,
        query_start_loc_cpu=query_start_loc_cpu,
        seq_lens=seq_lens,
-        seq_lens_cpu_upper_bound=seq_lens.cpu(),
        num_reqs=1,
        num_actual_tokens=40,
        max_query_len=40,
@@ -2074,54 +2074,6 @@ def test_auto_fit_max_model_len_not_triggered():
    assert vllm_config.model_config.max_model_len == 16


-def test_auto_fit_max_model_len_respects_num_gpu_blocks_override():
-    """Auto-fit must size max_model_len against the override-clamped pool, not
-    the raw `available_memory`. Without this, auto-fit could pick a
-    max_model_len that no longer fits once `num_gpu_blocks_override` is applied.
-    """
-    model_config = ModelConfig(max_model_len=16384)
-    model_config.original_max_model_len = -1  # request auto-fit
-    vllm_config = VllmConfig(model_config=model_config)
-    # Cap the cache to 32 blocks regardless of available memory.
-    vllm_config.cache_config.num_gpu_blocks_override = 32
-
-    mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2
-    kv_cache_specs = {
-        "layer_1": new_kv_cache_spec(),  # block_size=16
-        "layer_2": new_kv_cache_spec(),
-    }
-    # Plenty of raw memory (1024 blocks per layer would fit max_model_len=16384).
-    large_available_memory = mem_per_block_per_layer * 2 * 1024
-
-    get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory])
-
-    # 32 blocks * block_size 16 = 512 token slots, so max_model_len must
-    # auto-fit at or below that.
-    assert 0 < vllm_config.model_config.max_model_len <= 32 * 16
-
-
-def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override():
-    """Admission check must use the override-clamped pool size, not raw
-    `available_memory`. Without this, startup could accept a max_model_len
-    that does not actually fit in `num_gpu_blocks_override` blocks.
-    """
-    model_config = ModelConfig(max_model_len=16384)
-    vllm_config = VllmConfig(model_config=model_config)
-    # 32 blocks is far too small for max_model_len=16384 (would need 1024).
-    vllm_config.cache_config.num_gpu_blocks_override = 32
-
-    mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2
-    kv_cache_specs = {
-        "layer_1": new_kv_cache_spec(),
-        "layer_2": new_kv_cache_spec(),
-    }
-    # Plenty of raw memory: a bytes-only check against this would pass.
-    large_available_memory = mem_per_block_per_layer * 2 * 1024
-
-    with pytest.raises(ValueError, match="max seq len"):
-        get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory])
-
-
 def test_unify_hybrid_kv_cache_specs():
    # 1. has_full_attention and has_sliding_window
    before_spec_1 = new_kv_cache_spec()
@@ -2512,111 +2512,3 @@ def test_block_lookup_cache_multi_blocks_per_key():
    assert cache.pop(key1, 11) is block11
    assert cache.get_one_block(key1) is None
    assert cache.pop(key1, 12) is None
-
-
-def test_can_fit_full_sequence_swa_cap_admits_long_prompt():
-    """Hybrid full+SWA model with a pool sized at the startup minimum should
-    admit a prompt longer than the SWA cap, because SlidingWindowManager
-    recycles blocks during chunked prefill (issue #39734)."""
-    block_size = 16
-    sliding_window = 4 * block_size  # 64 tokens
-    max_num_batched_tokens = 8 * block_size  # 128 tokens
-    max_model_len = 64 * block_size  # 1024 tokens — much larger than the SWA cap
-    # Startup pool sizing: full demands cdiv(max_model_len, bs) = 64 blocks,
-    # SWA demands cdiv(SW-1+max_batched, bs) + 1 = cdiv(191, 16) + 1 = 13.
-    # Pool minimum = 64 + 13 = 77; +1 for the null block.
-    num_blocks = 64 + 13 + 1
-
-    config = KVCacheConfig(
-        num_blocks=num_blocks,
-        kv_cache_tensors=[],
-        kv_cache_groups=[
-            KVCacheGroupSpec(
-                ["layer_full"],
-                FullAttentionSpec(
-                    block_size=block_size,
-                    num_kv_heads=1,
-                    head_size=1,
-                    dtype=torch.float32,
-                ),
-            ),
-            KVCacheGroupSpec(
-                ["layer_swa"],
-                SlidingWindowSpec(
-                    block_size=block_size,
-                    num_kv_heads=1,
-                    head_size=1,
-                    dtype=torch.float32,
-                    sliding_window=sliding_window,
-                ),
-            ),
-        ],
-    )
-
-    manager = KVCacheManager(
-        config,
-        max_model_len=max_model_len,
-        max_num_batched_tokens=max_num_batched_tokens,
-        enable_caching=True,
-        hash_block_size=block_size,
-    )
-
-    # A prompt that is shorter than max_model_len but longer than SW + chunk:
-    # cdiv(prompt_len, bs) = 32 blocks. Without the cap, admission would
-    # demand 32 (full) + 32 (SWA) = 64 blocks. With the cap, SWA contributes
-    # only 13, so total = 32 + 13 = 45 ≤ pool size.
-    prompt_len = 32 * block_size
-    req = make_request("long", list(range(prompt_len)), block_size, sha256)
-
-    assert manager.can_fit_full_sequence(req)
-
-
-def test_can_fit_full_sequence_full_attention_still_gates_oversized():
-    """The cap only loosens the SWA group; a prompt that exceeds the
-    full-attention pool capacity must still be rejected."""
-    block_size = 16
-    sliding_window = 4 * block_size
-    max_num_batched_tokens = 8 * block_size
-    max_model_len = 64 * block_size
-    # Provide a tiny pool — even a small prompt should be rejected.
-    num_blocks = 5
-
-    config = KVCacheConfig(
-        num_blocks=num_blocks,
-        kv_cache_tensors=[],
-        kv_cache_groups=[
-            KVCacheGroupSpec(
-                ["layer_full"],
-                FullAttentionSpec(
-                    block_size=block_size,
-                    num_kv_heads=1,
-                    head_size=1,
-                    dtype=torch.float32,
-                ),
-            ),
-            KVCacheGroupSpec(
-                ["layer_swa"],
-                SlidingWindowSpec(
-                    block_size=block_size,
-                    num_kv_heads=1,
-                    head_size=1,
-                    dtype=torch.float32,
-                    sliding_window=sliding_window,
-                ),
-            ),
-        ],
-    )
-
-    manager = KVCacheManager(
-        config,
-        max_model_len=max_model_len,
-        max_num_batched_tokens=max_num_batched_tokens,
-        enable_caching=True,
-        hash_block_size=block_size,
-    )
-
-    # 16 blocks of full attention demand alone exceeds the 5-block pool.
-    prompt_len = 16 * block_size
-    req = make_request("oversized", list(range(prompt_len)), block_size, sha256)
-
-    assert not manager.can_fit_full_sequence(req)
@@ -22,13 +22,11 @@ pytestmark = pytest.mark.cpu_test


 def get_sliding_window_manager(sliding_window_spec, block_pool, enable_caching=True):
-    # Tests don't exercise admission gating; pass a large cap that is a no-op.
    return SlidingWindowManager(
        sliding_window_spec,
        block_pool=block_pool,
        enable_caching=enable_caching,
        kv_cache_group_id=0,
-        max_admission_blocks_per_request=10**9,
    )


@@ -40,7 +38,6 @@ def get_chunked_local_attention_manager(
        block_pool=block_pool,
        enable_caching=enable_caching,
        kv_cache_group_id=0,
-        max_admission_blocks_per_request=10**9,
    )


@@ -324,13 +324,10 @@ def run_test(
 ):
    spec_decoding = spec_config is not None
    cache_arg: dict[str, Any] = (
-        # Force preemptions: with 32 blocks the cache holds at most a single
-        # max-length request, so the ~34 concurrent prompts contend and trigger
-        # preemption. (Prompts here are << max_model_len, so dropping
-        # max_model_len from 4096 to 512 doesn't change generation behavior.)
-        dict(num_gpu_blocks_override=32, max_model_len=512)
+        # Force preemptions
+        dict(num_gpu_blocks_override=32)
        if test_preemption
-        else dict(gpu_memory_utilization=0.9, max_model_len=4096)
+        else dict(gpu_memory_utilization=0.9)
    )
    spec_mml = (spec_config or {}).get("max_model_len")
    spec_method = (spec_config or {}).get("method", "none")
@@ -346,6 +343,7 @@ def run_test(

    with VllmRunner(
        model,
+        max_model_len=4096,
        enable_chunked_prefill=test_prefill_chunking,
        # Force prefill chunking
        max_num_batched_tokens=48 if test_prefill_chunking else None,
@@ -478,59 +478,3 @@ class TestSlidingWindowLookup:
            sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX)
            is None
        )
-
-
-@pytest.mark.parametrize("async_scheduling", [True, False])
-def test_do_remote_decode_stores_all_blocks(request_runner, async_scheduling: bool):
-    """With do_remote_decode=True, after loading prefix blocks from CPU,
-    all blocks must be re-stored — not just the newly computed ones.
-
-    This supports P/D disaggregation where the prefill instance offloads the
-    complete KV cache so a remote decode node can consume it."""
-    offloaded_block_size = 12
-    gpu_block_size = 4
-    num_gpu_blocks = 100
-
-    runner = request_runner(
-        offloaded_block_size=offloaded_block_size,
-        gpu_block_size=gpu_block_size,
-        num_gpu_blocks=num_gpu_blocks,
-        async_scheduling=async_scheduling,
-    )
-
-    # Store 1 offloaded block (3 GPU blocks) via a normal request.
-    runner.new_request(token_ids=[0] * offloaded_block_size)
-    runner.manager.prepare_store.side_effect = (
-        lambda keys, req_context: generate_store_output(keys)
-    )
-    runner.run(
-        decoded_tokens=[EOS_TOKEN_ID],
-        expected_stored_gpu_block_indexes=(0, 1, 2),
-    )
-
-    # Reset GPU prefix cache so the next request must load from CPU.
-    runner.scheduler.reset_prefix_cache()
-
-    # New request with do_remote_decode=True and 2 offloaded blocks.
-    # The first offloaded block matches what we stored in CPU.
-    runner.new_request(
-        token_ids=[0] * offloaded_block_size * 2,
-        kv_transfer_params={"do_remote_decode": True},
-    )
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
-    runner.manager.prepare_store.side_effect = (
-        lambda keys, req_context: generate_store_output(keys)
-    )
-
-    # Load the first offloaded block from CPU.
-    runner.run(
-        decoded_tokens=[0],
-        expected_loaded_gpu_block_indexes=(0, 1, 2),
-    )
-
-    # Store must include ALL 6 GPU blocks (both the loaded prefix and
-    # the newly computed block), not just the 3 new ones.
-    runner.run(
-        decoded_tokens=[EOS_TOKEN_ID],
-        expected_stored_gpu_block_indexes=(0, 1, 2, 3, 4, 5),
-    )
@@ -270,11 +270,7 @@ class RequestRunner:
            slot_mapping={},
        )

-    def new_request(
-        self,
-        token_ids: list[int],
-        kv_transfer_params: dict | None = None,
-    ):
+    def new_request(self, token_ids: list[int]):
        self.req_id += 1

        sampling_params = SamplingParams(max_tokens=1000)
@@ -287,8 +283,6 @@ class RequestRunner:
            pooling_params=None,
            block_hasher=self._block_hasher,
        )
-        if kv_transfer_params is not None:
-            req.kv_transfer_params = kv_transfer_params

        self.scheduler.add_request(req)

@@ -208,6 +208,7 @@ def test_metadata_hma_block_ids():
 # ---------------------------------------------------------------------------
 #  test_build_transfer_params_multi_group_trimming
 # ---------------------------------------------------------------------------
+@pytest.mark.cpu_test
@pytest.mark.asyncio
@patch(
    "vllm.distributed.kv_transfer.kv_connector.v1.mooncake"
@@ -291,6 +292,7 @@ async def test_build_transfer_params_multi_group_trimming(monkeypatch):
 # ---------------------------------------------------------------------------
 #  test_build_transfer_params_group_count_mismatch
 # ---------------------------------------------------------------------------
+@pytest.mark.cpu_test
@pytest.mark.asyncio
@patch(
    "vllm.distributed.kv_transfer.kv_connector.v1.mooncake"
@@ -298,7 +300,7 @@ async def test_build_transfer_params_multi_group_trimming(monkeypatch):
    FakeMooncakeWrapper,
 )
 async def test_build_transfer_params_group_count_mismatch(monkeypatch):
-    """_build_transfer_params reports an error when group counts differ."""
+    """_build_transfer_params asserts when group counts differ."""

    monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
    vllm_config = create_vllm_config(
@@ -344,22 +346,10 @@ async def test_build_transfer_params_group_count_mismatch(monkeypatch):
        ]

        ready_reqs = [("d-mismatch", send_meta)]
-        (
-            src_ptrs,
-            dst_ptrs,
-            lengths,
-            err_reqs,
-            err_msg,
-        ) = await worker._build_transfer_params(
-            ready_reqs, xfer_meta, local_regions, remote_regions
-        )
-
-        # Mismatched req is reported via err_reqs/err_msg with no transfers built.
-        assert err_reqs == ["d-mismatch"]
-        assert err_msg == "KV group count mismatch"
-        assert src_ptrs == []
-        assert dst_ptrs == []
-        assert lengths == []
+        with pytest.raises(AssertionError, match="KV group count mismatch"):
+            await worker._build_transfer_params(
+                ready_reqs, xfer_meta, local_regions, remote_regions
+            )

        worker.shutdown()

@@ -1,222 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Verify that GPU memory is fully released after RixlConnector shutdown on ROCm.
-
-Regression test for ROCm/ucx#33: UCX rocm_ipc transport permanently pinned
-GPU memory via hsa_amd_ipc_memory_create during ucp_mem_map, causing
-GPU memory to be unrecoverable after engine shutdown.
-"""
-
-import gc
-
-import pytest
-import torch
-
-from vllm.platforms import current_platform
-
-pytestmark = pytest.mark.skipif(
-    not current_platform.is_rocm(),
-    reason="ROCm platform required",
-)
-
-
-def _mb(b: int) -> float:
-    return b / (1024 * 1024)
-
-
-def _gpu_snapshot(tag: str, prev_alloc: float = 0.0) -> dict:
-    """Print and return current GPU memory stats."""
-    torch.accelerator.synchronize()
-    alloc = torch.accelerator.memory_allocated()
-    reserved = torch.accelerator.memory_reserved()
-    # mem_get_info is not available on torch.accelerator
-    try:
-        drv_free, drv_total = torch.cuda.mem_get_info()
-        drv_used = drv_total - drv_free
-        drv_pct = drv_used / drv_total * 100
-    except Exception:
-        drv_used = drv_total = drv_pct = 0
-    alloc_mb = _mb(alloc)
-    drv_used_mb = _mb(drv_used)
-    delta = alloc_mb - prev_alloc
-    print(
-        f"  {tag:<40s} | {alloc_mb:>9.1f} alloc | "
-        f"{_mb(reserved):>9.1f} rsrvd | "
-        f"{drv_used_mb:>9.1f} driver ({drv_pct:.1f}%) | "
-        f"delta {delta:>+9.1f}"
-    )
-    return {
-        "tag": tag,
-        "alloc_mb": alloc_mb,
-        "drv_used_mb": drv_used_mb,
-        "drv_pct": drv_pct,
-    }
-
-
-def _full_gpu_cleanup():
-    """gc.collect + torch empty_cache, multiple rounds."""
-    gc.unfreeze()
-    for _ in range(3):
-        if gc.collect() == 0:
-            break
-    torch.accelerator.empty_cache()
-
-
-@pytest.mark.parametrize("model_name, sw_size", [("google/gemma-3-1b-it", 512)])
-def test_gpu_memory_rixl_hma(model_name, sw_size):
-    """Track GPU memory through NixlConnector create/infer/shutdown cycle."""
-    from vllm import LLM, SamplingParams
-    from vllm.config import KVTransferConfig
-    from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
-
-    llm_kwargs = {
-        "model": model_name,
-        "enforce_eager": True,
-        "gpu_memory_utilization": 0.5,
-        "kv_transfer_config": KVTransferConfig(
-            kv_connector="NixlConnector",
-            kv_role="kv_both",
-        ),
-        "max_model_len": 2048,
-        "disable_hybrid_kv_cache_manager": False,
-        "max_num_batched_tokens": 1024,
-        "enable_prefix_caching": False,
-        "block_size": 16,
-    }
-
-    print("\n" + "=" * 90)
-    print("GPU MEMORY -- RIXL NixlConnector HMA (ROCm)")
-    print("=" * 90)
-    gc.collect()
-    torch.accelerator.empty_cache()
-    torch.accelerator.reset_peak_memory_stats()
-    snap0 = _gpu_snapshot("0. baseline", 0.0)
-
-    # create + infer
-    llm = LLM(**llm_kwargs)
-    snap1 = _gpu_snapshot("1. after LLM()", snap0["alloc_mb"])
-
-    llm.generate(
-        ["hi" * 1401],
-        SamplingParams(
-            temperature=0.0,
-            max_tokens=1,
-            extra_args={
-                "kv_transfer_params": {
-                    "do_remote_decode": True,
-                    "do_remote_prefill": False,
-                    "remote_engine_id": None,
-                    "remote_block_ids": None,
-                    "remote_host": None,
-                    "remote_port": None,
-                }
-            },
-        ),
-    )
-    snap2 = _gpu_snapshot("2. after generate()", snap1["alloc_mb"])
-
-    # shutdown + cleanup
-    print("\n--- shutdown ---")
-    llm.llm_engine.engine_core.shutdown()
-    _gpu_snapshot("3. after shutdown()", snap2["alloc_mb"])
-
-    del llm
-    _full_gpu_cleanup()
-    cleanup_dist_env_and_memory()
-    _full_gpu_cleanup()
-    torch._dynamo.reset()
-    gc.collect()
-    torch.accelerator.empty_cache()
-    snap_final = _gpu_snapshot("4. final", snap2["alloc_mb"])
-
-    # summary
-    print("\n" + "=" * 90)
-    baseline = snap0["alloc_mb"]
-    final = snap_final["alloc_mb"]
-    peak = snap2["alloc_mb"]
-    total_alloc = peak - baseline
-
-    print(
-        f"  PyTorch:  baseline={baseline:.0f}  peak={peak:.0f}  "
-        f"final={final:.0f}  "
-        f"leaked={final - baseline:.0f} MB"
-        + (
-            f" ({(final - baseline) / total_alloc * 100:.1f}%)"
-            if total_alloc > 0
-            else ""
-        )
-    )
-
-    drv_base = snap0["drv_used_mb"]
-    drv_final = snap_final["drv_used_mb"]
-    drv_leaked = drv_final - drv_base
-    print(
-        f"  Driver:   baseline={drv_base:.0f} ({snap0['drv_pct']:.1f}%)  "
-        f"peak={snap2['drv_used_mb']:.0f} ({snap2['drv_pct']:.1f}%)  "
-        f"final={drv_final:.0f} ({snap_final['drv_pct']:.1f}%)  "
-        f"leaked={drv_leaked:.0f} MB"
-    )
-    print("=" * 90)
-
-    # Peak driver memory used above baseline
-    drv_peak = snap2["drv_used_mb"] - drv_base
-    leak_pct = (drv_leaked / drv_peak * 100) if drv_peak > 0 else 0
-    max_leak_pct = 10
-    assert leak_pct <= max_leak_pct, (
-        f"{drv_leaked:.0f} MB ({leak_pct:.1f}%) of driver-level GPU memory "
-        f"not freed after NixlConnector shutdown "
-        f"(peak allocation: {drv_peak:.0f} MB, threshold: {max_leak_pct}%)"
-    )
-
-
-@pytest.mark.parametrize("model_name", ["google/gemma-3-1b-it"])
-def test_gpu_memory_no_rixl_baseline(model_name):
-    """Same workload without NixlConnector.  Comparing driver-level memory
-    between this and test_gpu_memory_rixl_hma isolates UCX/RIXL impact."""
-    from vllm import LLM, SamplingParams
-    from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
-
-    print("\n" + "=" * 90)
-    print("CONTROL -- same model, no RIXL connector")
-    print("=" * 90)
-    gc.collect()
-    torch.accelerator.empty_cache()
-    snap0 = _gpu_snapshot("baseline", 0.0)
-
-    llm = LLM(
-        model=model_name,
-        enforce_eager=True,
-        gpu_memory_utilization=0.5,
-        max_model_len=2048,
-        max_num_batched_tokens=1024,
-        enable_prefix_caching=False,
-        block_size=16,
-    )
-    _gpu_snapshot("after LLM()", snap0["alloc_mb"])
-
-    llm.generate(["hi " * 500], SamplingParams(max_tokens=1))
-    snap_peak = _gpu_snapshot("after generate()", snap0["alloc_mb"])
-
-    llm.llm_engine.engine_core.shutdown()
-    del llm
-    _full_gpu_cleanup()
-    cleanup_dist_env_and_memory()
-    _full_gpu_cleanup()
-    torch._dynamo.reset()
-    gc.collect()
-    torch.accelerator.empty_cache()
-    snap_final = _gpu_snapshot("final", snap0["alloc_mb"])
-
-    drv_base = snap0["drv_used_mb"]
-    drv_leaked = snap_final["drv_used_mb"] - drv_base
-    drv_peak = snap_peak["drv_used_mb"] - drv_base
-    print(f"\n  Driver leaked (no rixl): {drv_leaked:.0f} MB")
-    print("=" * 90)
-
-    leak_pct = (drv_leaked / drv_peak * 100) if drv_peak > 0 else 0
-    max_leak_pct = 10
-    assert leak_pct <= max_leak_pct, (
-        f"{drv_leaked:.0f} MB ({leak_pct:.1f}%) of driver-level GPU memory "
-        f"not freed after baseline shutdown "
-        f"(peak allocation: {drv_peak:.0f} MB, threshold: {max_leak_pct}%)"
-    )
@@ -87,13 +87,10 @@ class MockSubscriber:
 def _wait_for_prefix_cache_reset(llm: LLM) -> None:
    """Wait for async offload transfers to finish so prefix cache can reset.

-    The GPU-to-CPU offload runs on a CUDA stream asynchronously. While blocks
+    The GPU-to-CPU offload runs on a CUDA stream asynchronously.  While blocks
    are still held by the offload worker, ``reset_prefix_cache`` returns
-    ``False``. Between retries we send a dummy single-token prefill to force
-    the engine to step, which polls the worker for completed transfers and
-    frees GPU blocks.
+    ``False``.  Retry with a short sleep until it succeeds or we time out.
    """
-    _dummy_params = SamplingParams(max_tokens=1)
    deadline = time.monotonic() + _RESET_CACHE_TIMEOUT
    while not llm.reset_prefix_cache():
        if time.monotonic() > deadline:
@@ -101,13 +98,7 @@ def _wait_for_prefix_cache_reset(llm: LLM) -> None:
                "reset_prefix_cache did not succeed within "
                f"{_RESET_CACHE_TIMEOUT}s - async offload may be stuck"
            )
-        # Force an engine step so the scheduler polls get_finished()
-        # and releases GPU blocks held by in-flight async stores.
-        llm.generate(
-            [TokensPrompt(prompt_token_ids=[0])],
-            _dummy_params,
-            use_tqdm=False,
-        )
+        time.sleep(0.1)


 def _latency_test(llm: LLM, subscriber: MockSubscriber):
@@ -8,16 +8,11 @@ from unittest.mock import Mock
 import pytest

 from vllm.config import ModelConfig, SchedulerConfig, VllmConfig
+from vllm.reasoning import ReasoningParser
 from vllm.v1.request import Request
 from vllm.v1.structured_output import StructuredOutputManager


-class MockReasoner:
-    def __init__(self, tokenizer):
-        self.is_reasoning_end = Mock(return_value=False)
-        self.is_reasoning_end_streaming = Mock(return_value=False)
-
-
 class TestReasoningStructuredOutput:
    """Test reasoning-aware structured output functionality."""

@@ -55,6 +50,13 @@ class TestReasoningStructuredOutput:
        config.speculative_config = None
        return config

+    @pytest.fixture
+    def mock_reasoning_parser(self):
+        """Create a mock ReasoningParser."""
+        parser = Mock(spec=ReasoningParser)
+        parser.is_reasoning_end = Mock(return_value=False)
+        return parser
+
    @pytest.fixture
    def mock_request_with_structured_output(self):
        """Create a mock request with structured output."""
@@ -62,8 +64,6 @@ class TestReasoningStructuredOutput:
        request.structured_output_request = Mock()
        request.structured_output_request.reasoning_ended = None
        request.structured_output_request.grammar = Mock()
-        request.structured_output_request.reasoning_parser_kwargs = None
-        request.structured_output_request.reasoner = None
        request.structured_output_request.grammar.is_terminated = Mock(
            return_value=False
        )
@@ -74,13 +74,6 @@ class TestReasoningStructuredOutput:
        request.num_output_placeholders = 0
        return request

-    @pytest.fixture
-    def manager_with_reasoner(self, mock_vllm_config):
-        manager = StructuredOutputManager(mock_vllm_config)
-        manager.reasoner_cls = MockReasoner
-        manager.tokenizer = Mock()
-        return manager
-
    def test_should_fill_bitmask_with_enable_in_reasoning(
        self, mock_vllm_config, mock_request_with_structured_output
    ):
@@ -96,17 +89,22 @@ class TestReasoningStructuredOutput:

    def test_should_fill_bitmask_without_enable_in_reasoning(
        self,
-        manager_with_reasoner,
+        mock_vllm_config,
        mock_request_with_structured_output,
+        mock_reasoning_parser,
    ):
        """Test should_fill_bitmask when enable_in_reasoning is False."""
        # Keep enable_in_reasoning as False (default)
-        config = manager_with_reasoner.vllm_config.structured_outputs_config
+        config = mock_vllm_config.structured_outputs_config
        assert config.enable_in_reasoning is False

-        result = manager_with_reasoner.should_fill_bitmask(
-            mock_request_with_structured_output
-        )
+        manager = StructuredOutputManager(mock_vllm_config)
+        manager.reasoner = mock_reasoning_parser
+
+        # Mock reasoning not ended
+        mock_reasoning_parser.is_reasoning_end.return_value = False
+
+        result = manager.should_fill_bitmask(mock_request_with_structured_output)

        # Should set reasoning_ended and return its value
        assert (
@@ -120,92 +118,68 @@ class TestReasoningStructuredOutput:
    ):
        """Test should_fill_bitmask when no reasoner is configured."""
        manager = StructuredOutputManager(mock_vllm_config)
+        manager.reasoner = None

        result = manager.should_fill_bitmask(mock_request_with_structured_output)

        # Should default to True when no reasoner
        assert result is True

-    def test_should_fill_bitmask_uses_request_reasoning_parser_kwargs(
-        self, mock_vllm_config, mock_request_with_structured_output
-    ):
-        """Test request-level parser kwargs override the default reasoner."""
-
-        class KwargReasoner:
-            def __init__(self, tokenizer, chat_template_kwargs=None):
-                self.chat_template_kwargs = chat_template_kwargs or {}
-
-            def is_reasoning_end(self, input_ids):
-                return not self.chat_template_kwargs.get("enable_thinking", False)
-
-        manager = StructuredOutputManager(mock_vllm_config)
-        manager.reasoner_cls = KwargReasoner
-        manager.tokenizer = Mock()
-
-        structured_req = mock_request_with_structured_output.structured_output_request
-        structured_req.reasoning_parser_kwargs = {
-            "chat_template_kwargs": {"enable_thinking": True}
-        }
-
-        result = manager.should_fill_bitmask(mock_request_with_structured_output)
-
-        assert result is False
-        assert (
-            mock_request_with_structured_output.structured_output_request.reasoner
-            is not None
-        )
-
    def test_should_advance_with_enable_in_reasoning(
        self,
-        manager_with_reasoner,
+        mock_vllm_config,
        mock_request_with_structured_output,
+        mock_reasoning_parser,
    ):
        """Test should_advance when enable_in_reasoning is True."""
        # Enable enable_in_reasoning
-        manager_with_reasoner.enable_in_reasoning = True
+        mock_vllm_config.structured_outputs_config.enable_in_reasoning = True
+
+        manager = StructuredOutputManager(mock_vllm_config)
+        manager.reasoner = mock_reasoning_parser

        # Should always return True when enable_in_reasoning is enabled
-        result = manager_with_reasoner.should_advance(
-            mock_request_with_structured_output
-        )
+        result = manager.should_advance(mock_request_with_structured_output)
        assert result is True

    def test_should_advance_reasoning_not_ended(
        self,
-        manager_with_reasoner,
+        mock_vllm_config,
        mock_request_with_structured_output,
+        mock_reasoning_parser,
    ):
        """Test should_advance when reasoning has not ended."""
+        manager = StructuredOutputManager(mock_vllm_config)
+        manager.reasoner = mock_reasoning_parser
+
        # Set reasoning as not ended
        (
            mock_request_with_structured_output.structured_output_request
        ).reasoning_ended = False
+        mock_reasoning_parser.is_reasoning_end.return_value = False

-        result = manager_with_reasoner.should_advance(
-            mock_request_with_structured_output
-        )
+        result = manager.should_advance(mock_request_with_structured_output)

        # Should return False since reasoning hasn't ended
        assert result is False

    def test_should_advance_reasoning_just_ended(
        self,
-        manager_with_reasoner,
+        mock_vllm_config,
        mock_request_with_structured_output,
+        mock_reasoning_parser,
    ):
        """Test should_advance when reasoning ends in current step."""
+        manager = StructuredOutputManager(mock_vllm_config)
+        manager.reasoner = mock_reasoning_parser
+
        # Set reasoning as not ended initially, but ends in this step
        (
            mock_request_with_structured_output.structured_output_request
        ).reasoning_ended = False
-        reasoner = MockReasoner(tokenizer=Mock())
-        reasoner.is_reasoning_end_streaming.return_value = True
-        structured_req = mock_request_with_structured_output.structured_output_request
-        structured_req.reasoner = reasoner
+        mock_reasoning_parser.is_reasoning_end.return_value = True

-        result = manager_with_reasoner.should_advance(
-            mock_request_with_structured_output
-        )
+        result = manager.should_advance(mock_request_with_structured_output)

        # Should set reasoning_ended to True but return False for this step
        assert (
@@ -216,18 +190,20 @@ class TestReasoningStructuredOutput:

    def test_should_advance_reasoning_already_ended(
        self,
-        manager_with_reasoner,
+        mock_vllm_config,
        mock_request_with_structured_output,
+        mock_reasoning_parser,
    ):
        """Test should_advance when reasoning has already ended."""
+        manager = StructuredOutputManager(mock_vllm_config)
+        manager.reasoner = mock_reasoning_parser
+
        # Set reasoning as already ended
        (
            mock_request_with_structured_output.structured_output_request
        ).reasoning_ended = True

-        result = manager_with_reasoner.should_advance(
-            mock_request_with_structured_output
-        )
+        result = manager.should_advance(mock_request_with_structured_output)

        # Should return True since reasoning has ended
        assert result is True
@@ -407,21 +407,16 @@ def rotary_embedding(
    rope_dim_offset: int = 0,
    inverse: bool = False,
 ) -> None:
-    if rope_dim_offset == 0 and not inverse:
-        torch.ops._C.rotary_embedding(
-            positions, query, key, head_size, cos_sin_cache, is_neox
-        )
-    else:
-        torch.ops._C.rotary_embedding(
-            positions,
-            query,
-            key,
-            head_size,
-            cos_sin_cache,
-            is_neox,
-            rope_dim_offset,
-            inverse,
-        )
+    torch.ops._C.rotary_embedding(
+        positions,
+        query,
+        key,
+        head_size,
+        cos_sin_cache,
+        is_neox,
+        rope_dim_offset,
+        inverse,
+    )


 # layer norm ops
@@ -406,13 +406,16 @@ class AsyncTPPass(VllmPatternMatcherPass):
        self.dump_patterns(config, self.patterns)

    def is_applicable_for_range(self, compile_range: Range) -> bool:
-        # This pass is applied on top of the sequence parallelism pass,
-        # which is only supported in fullgraph compilation mode.
-        assert (
-            self.compilation_config.use_inductor_graph_partition
-            or not self.compilation_config.splitting_ops
-        ), "AsyncTPPass requires full-graph compilation"
-        return True
+        # This pass is applied on top of the sequence parallelism pass.
+        # It inherits the same applicability condition as `SequenceParallelismPass`.
+        # See `SequenceParallelismPass.is_applicable` for more details.
+        if (
+            not self.compilation_config.splitting_ops
+            or self.compilation_config.use_inductor_graph_partition
+        ):
+            return True
+        tp_size = get_tensor_model_parallel_world_size()
+        return bool(compile_range.is_single_size() and compile_range.end % tp_size == 0)

    @VllmInductorPass.time_and_log
    def __call__(self, graph: fx.Graph) -> None:
@@ -341,18 +341,22 @@ class SequenceParallelismPass(VllmPatternMatcherPass):
    significantly reduce communication overhead and improve overall model
    performance.

-    This pass is only supported when compiling the whole graph (fullgraph
-    mode, i.e. using Inductor graph partition or empty splitting_ops).
-    Piecewise compilation is not supported because the residual tensor
-    gets split across TP ranks, causing size mismatches at subgraph
-    boundaries.

-    This pass splits up the residual tensor across TP ranks and hence
-    divides its size. Because the pattern matcher starts at the end of
-    the graph, the replacement contains a slice that temporarily conforms
-    the input residual to the correct size. After all patterns have been
-    matched, we use a NoOpEliminationPass to clean up what have now
-    become no-op slices.
+    This pass splits up the residual tensor across TP ranks and hence divides its size.
+    Because the pattern matcher starts at the end of the graph, the replacement
+    contains a slice that temporarily conforms the input residual to the correct size.
+    After all patterns have been matched, we use a NoOpEliminationPass to clean up
+    what have now become no-op slices.
+
+    Note that an older version of the pass did not need this as it operated only on
+    custom rms_norm and fused_rms_norm_add custom ops which did not complain about
+    mismatched shapes during replacement. So this approach has the same assumption that
+    correctness is only maintained if all rms_norm operations are split across ranks.
+
+    Correctness-wise, this is approach strictly better than before - before,
+    the graph was incorrect semantically and shape-wise during the pass.
+    With this approach there's only semantic incorrectness during the pass.
+    Both approaches restore a correct graph once all patterns are matched.
    """

    @enable_fake_mode
@@ -415,13 +419,19 @@ class SequenceParallelismPass(VllmPatternMatcherPass):
        and gathering tensors across TP ranks outweighs the benefits.

        Returns False (SP disabled) when:
+        - Using piecewise compilation with non-concrete or TP-indivisible sizes
        - min_token_num is None (SP disabled for this device/config)
        - The compile range starts below the minimum token threshold
        """
-        assert (
-            self.compilation_config.use_inductor_graph_partition
-            or not self.compilation_config.splitting_ops
-        ), "SequenceParallelismPass requires full-graph compilation"
+        # For piecewise compilation (not using inductor graph partition),
+        # we need concrete sizes that are divisible by TP for correct splitting
+        if (
+            not self.compilation_config.use_inductor_graph_partition
+            and self.compilation_config.splitting_ops
+        ):
+            tp_size = get_tensor_model_parallel_world_size()
+            if not compile_range.is_single_size() or compile_range.end % tp_size != 0:
+                return False

        # min_token_num is None when SP is disabled for this device/config
        # (e.g., non-CUDA platform, unsupported GPU, or small hidden_size)
@@ -1149,25 +1149,6 @@ class CompilationConfig:
                    self.cudagraph_mode = CUDAGraphMode.FULL
                self.splitting_ops = []

-        if (
-            not self.use_inductor_graph_partition
-            and (self.pass_config.enable_sp or self.pass_config.fuse_gemm_comms)
-            and self.splitting_ops
-        ):
-            logger.warning_once(
-                "Sequence parallelism requires full-graph compilation when "
-                "use_inductor_graph_partition is off. Setting splitting_ops "
-                "to an empty list to preserve SP and async TP."
-            )
-            self.splitting_ops = []
-            if self.cudagraph_mode.has_piecewise_cudagraphs():
-                logger.warning_once(
-                    "Sequence parallelism is incompatible with piecewise "
-                    "cudagraph when use_inductor_graph_partition is off. "
-                    "Setting cudagraph_mode to FULL."
-                )
-                self.cudagraph_mode = CUDAGraphMode.FULL
-
        # Disable CUDA graphs for DeepEP high-throughput since its not CG compatible
        if (
            all2all_backend == "deepep_high_throughput"
@@ -50,7 +50,7 @@ class IrOpPriorityConfig:
            name: {
                provider: IrOp.registry[name].impls[provider].uuid() for provider in p
            }
-            for name, p in asdict(self).items()  # type: ignore[call-overload]
+            for name, p in asdict(self).items()
        }

        return hash_factors(factors)
@@ -77,7 +77,7 @@ class IrOpPriorityConfig:
        current_platform.import_ir_kernels()

        with contextlib.ExitStack() as stack:
-            for field in fields(self):  # type: ignore[arg-type]
+            for field in fields(self):
                op_priority = getattr(self, field.name)
                assert op_priority is not None, (
                    f"IR op priority for {field.name} must be set"
@@ -98,7 +98,7 @@ class IrOpPriorityConfig:
        A helper to create an IrOpPriorityConfig where fields not specified in kwargs
        use the given default list.
        """
-        for field in fields(cls):  # type: ignore[arg-type]
+        for field in fields(cls):
            if field.name not in kwargs:
                kwargs[field.name] = list(default)

@@ -108,8 +108,8 @@ class IrOpPriorityConfig:
 MoEBackend = Literal[
    "auto",
    "triton",
+    "triton_unfused",
    "deep_gemm",
-    "deep_gemm_mega_moe",
    "cutlass",
    "flashinfer_trtllm",
    "flashinfer_cutlass",
@@ -137,9 +137,9 @@ class KernelConfig:
    """Backend for MoE expert computation kernels. Available options:

    - "auto": Automatically select the best backend based on model and hardware
-    - "triton": Use Triton-based fused MoE kernels 
+    - "triton": Use Triton-based fused MoE kernels (SWIGLUOAI activation only)
+    - "triton_unfused": Use Triton-based unfused MoE kernels (supports SILU/GELU)
    - "deep_gemm": Use DeepGEMM kernels (FP8 block-quantized only)
-    - "deep_gemm_mega_moe": Use DeepGEMM mega MoE kernels
    - "cutlass": Use vLLM CUTLASS kernels
    - "flashinfer_trtllm": Use FlashInfer with TRTLLM-GEN kernels
    - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels
@@ -983,16 +983,19 @@ class VllmConfig:
            )
            self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE

-        # async tp is built on top of sequence parallelism and requires it.
-        pass_config = self.compilation_config.pass_config
-        if pass_config.fuse_gemm_comms:
-            pass_config.enable_sp = True
-        if pass_config.enable_sp:
+        # async tp is built on top of sequence parallelism
+        # and requires it to be enabled.
+        if self.compilation_config.pass_config.fuse_gemm_comms:
+            self.compilation_config.pass_config.enable_sp = True
+        if self.compilation_config.pass_config.enable_sp:
            if self.parallel_config.tensor_parallel_size == 1:
                logger.warning("Sequence Parallelism requires TP>1, disabling")
-                pass_config.enable_sp = False
-                pass_config.fuse_gemm_comms = False
+                self.compilation_config.pass_config.enable_sp = False
+                self.compilation_config.pass_config.fuse_gemm_comms = False
            else:
+                # Compute SP threshold early; disable if None (model too
+                # small for SP to be beneficial).
+                pass_config = self.compilation_config.pass_config
                if pass_config.sp_min_token_num is None:
                    from vllm.compilation.passes.fusion.sequence_parallelism import (
                        get_sequence_parallelism_threshold,
@@ -1012,8 +1015,8 @@ class VllmConfig:
                        "threshold heuristic, disabling. To force SP, "
                        "set pass_config.sp_min_token_num manually."
                    )
-                    pass_config.enable_sp = False
-                    pass_config.fuse_gemm_comms = False
+                    self.compilation_config.pass_config.enable_sp = False
+                    self.compilation_config.pass_config.fuse_gemm_comms = False

        from vllm.utils.torch_utils import HAS_OPAQUE_TYPE

@@ -1095,7 +1098,6 @@ class VllmConfig:
                self.compilation_config.cudagraph_num_of_warmups = 1

            self._set_cudagraph_sizes()
-
        else:
            self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE

@@ -1169,8 +1171,8 @@ class VllmConfig:
        )

        if self.compilation_config.pass_config.enable_sp:
-            # With pipeline parallelism, native rms norm tracing errors due to
-            # incorrect residual shape.
+            # With pipeline parallelism or dynamo partitioning,
+            # native rms norm tracing errors due to incorrect residual shape.
            # Use custom rms norm to unblock. In the future,
            # the pass will operate on higher-level IR to avoid the issue.
            # TODO: https://github.com/vllm-project/vllm/issues/27894
@@ -1181,15 +1183,24 @@ class VllmConfig:
                    self.compilation_config.mode,
                )

-            if self.parallel_config.pipeline_parallel_size > 1:
+            is_fullgraph = (
+                self.compilation_config.use_inductor_graph_partition
+                or len(self.compilation_config.splitting_ops or []) == 0
+            )
+            if self.parallel_config.pipeline_parallel_size > 1 or not is_fullgraph:
                if "-rms_norm" not in self.compilation_config.custom_ops:
                    self.compilation_config.custom_ops.append("+rms_norm")
                else:
+                    regime = (
+                        "Dynamo partition"
+                        if not is_fullgraph
+                        else "pipeline parallelism"
+                    )
                    logger.warning_once(
                        "Sequence parallelism not supported with "
                        "native rms_norm when using %s, "
                        "this will likely lead to an error.",
-                        "pipeline parallelism",
+                        regime,
                    )

        # final check of cudagraph mode after all possible updates
@@ -1201,9 +1212,9 @@ class VllmConfig:
                and not self.compilation_config.cudagraph_mode.has_piecewise_cudagraphs()  # noqa: E501
            ):
                logger.warning_once(
-                    "No piecewise cudagraph for executing cascade attention. "
-                    "Will fall back to eager execution if a batch runs into "
-                    "cascade attentions."
+                    "No piecewise cudagraph for executing cascade attention."
+                    " Will fall back to eager execution if a batch runs "
+                    "into cascade attentions."
                )

            if self.compilation_config.cudagraph_mode.requires_piecewise_compilation():
@@ -1432,10 +1443,6 @@ class VllmConfig:
        cudagraph_capture_sizes = [1, 2, 4] + list(range(8, 256, 8)) + list(
            range(256, max_graph_size + 1, 16))

-        `max_num_batched_tokens` is also appended to the list if it fits
-        within `max_cudagraph_capture_size`, so the max batch size is captured
-        even when off-stride.
-
        In the end, `vllm_config.compilation_config.cudagraph_capture_sizes`
        will be the final sizes to capture cudagraph (in ascending order).

@@ -1524,12 +1531,6 @@ class VllmConfig:
                    cudagraph_capture_sizes += list(
                        range(256, max_cudagraph_capture_size + 1, 16)
                    )
-                # ensure max_num_tokens is captured if within max capture size
-                if (
-                    max_num_tokens <= max_cudagraph_capture_size
-                    and max_num_tokens not in cudagraph_capture_sizes
-                ):
-                    cudagraph_capture_sizes.append(max_num_tokens)
                # de-duplicate and sort the sizes
                cudagraph_capture_sizes = sorted(set(cudagraph_capture_sizes))

@@ -128,6 +128,13 @@ class CuMemAllocator:
        return CuMemAllocator.instance

    def __init__(self):
+        conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
+        assert "expandable_segments:True" not in conf, (
+            "Expandable segments are not compatible with memory pool. "
+            "Please track https://github.com/pytorch/pytorch/issues/147851 "
+            "for the latest updates."
+        )
+
        self.pointer_to_data: dict[int, AllocationData] = {}
        self.current_tag: str = CuMemAllocator.default_tag
        self.allocator_and_pools: dict[str, Any] = {}
@@ -257,49 +264,34 @@ class CuMemAllocator:

        assert isinstance(tag, str)

-        # Expandable segments are incompatible with the memory pool used for
-        # sleep mode (see https://github.com/pytorch/pytorch/issues/147851).
-        # If the user has enabled expandable segments via
-        # PYTORCH_CUDA_ALLOC_CONF, temporarily disable them for the duration
-        # of the memory pool context and restore on exit.
-        conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
-        expandable_was_enabled = "expandable_segments:True" in conf
-        if expandable_was_enabled:
-            torch.cuda.memory._set_allocator_settings("expandable_segments:False")
-
        old_tag = self.current_tag
        self.current_tag = tag
-        try:
-            with use_memory_pool_with_allocator(
-                self.python_malloc_callback, self.python_free_callback
-            ) as data:
-                # start to hit another PyTorch bug in PyTorch 2.6,
-                # possibly because of gc-related issue w.r.t. the allocator
-                # and the memory pool.
-                # to avoid the issue, we keep a reference of the data.
-                # see https://github.com/pytorch/pytorch/issues/146431 .
-                self.allocator_and_pools[tag] = data
-                yield
-                # PyTorch's bug, calling torch.cuda.empty_cache() will error
-                # when using pluggable allocator, see
-                # https://github.com/pytorch/pytorch/issues/145168 .
-                # if we have some memory allocated and then freed,
-                # the memory will not be released, e.g. in online
-                # quantization, where the model is created in higher
-                # precision, and then quantized in lower precision.
-                # Find all unused allocations and manually release them.
-                # TODO: we should expose `empty_cache` method in the memory
-                # pool.
-                # TODO: ask for help from PyTorch team to expose this method.
-                allocations = data[0].snapshot()
-                for allocation in allocations:
-                    if allocation["allocated_size"] == 0:
-                        handle = self._python_free_callback(allocation["address"])
-                        unmap_and_release(handle)
-        finally:
+        with use_memory_pool_with_allocator(
+            self.python_malloc_callback, self.python_free_callback
+        ) as data:
+            # start to hit another PyTorch bug in PyTorch 2.6,
+            # possibly because of gc-related issue w.r.t. the allocator and
+            # the memory pool.
+            # to avoid the issue, we keep a reference of the data.
+            # see https://github.com/pytorch/pytorch/issues/146431 .
+            self.allocator_and_pools[tag] = data
+            yield
+            # PyTorch's bug, calling torch.cuda.empty_cache() will error
+            # when using pluggable allocator, see
+            # https://github.com/pytorch/pytorch/issues/145168 .
+            # if we have some memory allocated and then freed,
+            # the memory will not be released, e.g. in online quantization,
+            # where the model is created in higher precision, and then
+            # quantized in lower precision.
+            # Find all unused allocations and manually release them.
+            # TODO: we should expose `empty_cache` method in the memory pool.
+            # TODO: ask for help from PyTorch team to expose this method.
+            allocations = data[0].snapshot()
+            for allocation in allocations:
+                if allocation["allocated_size"] == 0:
+                    handle = self._python_free_callback(allocation["address"])
+                    unmap_and_release(handle)
            self.current_tag = old_tag
-            if expandable_was_enabled:
-                torch.cuda.memory._set_allocator_settings("expandable_segments:True")

    def get_current_usage(self) -> int:
        """
@@ -492,18 +492,15 @@ class FlashInferNVLinkTwoSidedManager(All2AllManagerBase):
            CustomCommunicator,
        )

-        # MNNVL workspace is allocated per rank in the comm_backend's group; the
-        # flashinfer kernel asserts workspace.size(0) == moe_ep_size, so the backend
-        # must span the EP group (= DP*PCP*TP), not the DP group.
-        ep_config = MnnvlConfig(
-            comm_backend=CustomCommunicator(self.cpu_group),
+        dp_config = MnnvlConfig(
+            comm_backend=CustomCommunicator(get_dp_group().cpu_group),
            fabric_page_size=1 << 29,  # 512MB
            allocation_granularity=0,  # Auto-detect
        )

-        self.workspace_tensor = MnnvlMoe.get_moe_workspaces(self.mapping, ep_config)
+        self.workspace_tensor = MnnvlMoe.get_moe_workspaces(self.mapping, dp_config)
        self.prepare_workspace_tensor = MnnvlMoe.get_moe_prepare_workspace(
-            self.mapping, ep_config
+            self.mapping, dp_config
        )

        self.world_size = world_size
@@ -584,8 +581,6 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
        top_k: int,
        num_experts: int,
        hidden_size: int,
-        dispatch_dtype_bytes_per_elem: int = 0,
-        dispatch_scale_bytes_per_token: int = 0,
    ):
        """Initialize the MoeAlltoAll workspace."""
        if self.initialized:
@@ -610,19 +605,12 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
            CustomCommunicator,
        )

-        # MNNVL workspace is allocated per rank in the comm_backend's group; the
-        # flashinfer kernel asserts workspace.size(0) == moe_ep_size, so the backend
-        # must span the EP group (= DP*PCP*TP), not the DP group.
-        ep_config = MnnvlConfig(
-            comm_backend=CustomCommunicator(self.cpu_group),
+        dp_config = MnnvlConfig(
+            comm_backend=CustomCommunicator(get_dp_group().cpu_group),
        )
-        if dispatch_dtype_bytes_per_elem == 0:
-            hidden_bytes = hidden_size // 2
-        else:
-            hidden_bytes = hidden_size * dispatch_dtype_bytes_per_elem
        total_dispatch_payload_size_per_token = (
-            hidden_bytes
-            + dispatch_scale_bytes_per_token
+            hidden_size // 2  # nvfp4 hidden states
+            + hidden_size // 16  # fp8 scaling factors
            + top_k * 4  # int32 topks ids
            + top_k * 4  # float32 topk weights
        )
@@ -640,7 +628,7 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
            top_k=top_k,
            num_experts=num_experts,
            workspace_size_per_rank=self.workspace_size,
-            mnnvl_config=ep_config,
+            mnnvl_config=dp_config,
        )

        self.gpus_per_node = gpus_per_node
@@ -112,13 +112,6 @@ class RequestOffloadState:
        for group_state, new_blocks in zip(self.group_states, new_block_id_groups):
            group_state.block_ids.extend(new_blocks)

-    def advance_stored_idx(self, num_offloadable_tokens: int) -> None:
-        for group_config, group_state in zip(
-            self.config.kv_group_configs, self.group_states
-        ):
-            num_blocks = num_offloadable_tokens // group_config.offloaded_block_size
-            group_state.next_stored_block_idx = num_blocks
-

 class OffloadingConnectorScheduler:
    """Implementation of Scheduler side methods"""
@@ -314,9 +307,6 @@ class OffloadingConnectorScheduler:
        num_locally_computed_tokens = req_status.num_locally_computed_tokens
        num_cached_tokens = num_locally_computed_tokens + num_external_tokens

-        params = req_status.req_context.kv_transfer_params
-        do_remote_decode = params is not None and params.get("do_remote_decode")
-
        keys_to_load: list[OffloadKey] = []
        dst_block_ids: list[int] = []
        # per group
@@ -363,11 +353,7 @@ class OffloadingConnectorScheduler:
            group_sizes.append(num_pending_gpu_blocks)
            block_indices.append(num_locally_computed_gpu_blocks)

-            if not do_remote_decode:
-                # For P/D prefill requests (do_remote_decode=True), we do
-                # NOT skip saving the hit prefix, as we need to stream the
-                # entire KV cache so a remote decode node can consume it.
-                group_state.next_stored_block_idx = num_blocks
+            group_state.next_stored_block_idx = num_blocks

        src_spec = self.manager.prepare_load(keys_to_load, req_status.req_context)
        dst_spec = GPULoadStoreSpec(
@@ -381,16 +367,16 @@ class OffloadingConnectorScheduler:
        if self._blocks_being_loaded is not None:
            self._blocks_being_loaded.update(req_blocks_being_loaded)

-    def _get_reqs_to_store(
-        self, scheduler_output: SchedulerOutput
-    ) -> dict[ReqId, TransferSpec]:
-        block_size_factor = self.config.block_size_factor
+    def _get_reqs_to_store(self, scheduler_output: SchedulerOutput):
+        # Below assertion will be removed once this function supports HMA
+        assert len(self.config.kv_group_configs) == 1
+        group_config = self.config.kv_group_configs[0]
+
        reqs_to_store: dict[ReqId, TransferSpec] = {}
        # iterate over both new and cached requests
        for req_id, new_block_id_groups, preempted in yield_req_data(scheduler_output):
            req_status = self._req_status[req_id]
            req_status.update_offload_keys()
-            req = req_status.req

            if preempted:
                for group_state in req_status.group_states:
@@ -399,106 +385,68 @@ class OffloadingConnectorScheduler:
            if new_block_id_groups:
                req_status.update_block_id_groups(new_block_id_groups)

-            num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
-            num_tokens_after_batch = req.num_computed_tokens + num_scheduled_tokens
+            # Below assertion will be removed once this function supports HMA
+            assert len(req_status.group_states) == 1
+            group_state = req_status.group_states[0]
+
+            block_ids = group_state.block_ids
+
+            req = req_status.req
+            new_tokens = scheduler_output.num_scheduled_tokens[req_id]
+            expected_tokens = req.num_computed_tokens + new_tokens
            # with async scheduling, some tokens may be missing
-            num_offloadable_tokens = min(num_tokens_after_batch, req.num_tokens)
+            total_tokens = min(expected_tokens, req.num_tokens)
+            num_blocks = total_tokens // group_config.offloaded_block_size
+            start_block_idx = group_state.next_stored_block_idx
+            num_new_blocks = num_blocks - start_block_idx

-            # Filter out blocks skipped due to sliding window attention / SSM
-            new_offload_keys: list[OffloadKey] = []
-            for group_config, group_state in zip(
-                self.config.kv_group_configs, req_status.group_states
-            ):
-                num_blocks = num_offloadable_tokens // group_config.offloaded_block_size
-                start_block_idx = group_state.next_stored_block_idx
-                if num_blocks <= start_block_idx:
-                    continue
-                offload_keys = group_state.offload_keys[start_block_idx:num_blocks]
-                # For each block to offload, take the last corresponding GPU block.
-                # e.g. if block size factor is 3 and GPU block IDs are
-                # 1 5 6 7 2 4 9 3 8 then we'll take blocks 6 4 8.
-                # We will use these GPU blocks to determine if the block needs
-                # offloading, or (if the GPU block ID is 0) this block should
-                # be skipped due to sliding window attention / SSM.
-                # We know that if a block is skipped, then all the previous blocks
-                # are skipped as well. This is why we take the last of each block.
-                offload_block_ids = group_state.block_ids[
-                    start_block_idx * block_size_factor
-                    + block_size_factor
-                    - 1 : num_blocks * block_size_factor : block_size_factor
-                ]
-                assert len(offload_keys) == len(offload_block_ids)
-
-                for offload_key, block_id in zip(offload_keys, offload_block_ids):
-                    if block_id != 0:
-                        new_offload_keys.append(offload_key)
-
-            if not new_offload_keys:
-                req_status.advance_stored_idx(num_offloadable_tokens)
+            if num_new_blocks <= 0:
                continue

+            num_gpu_blocks = num_blocks * self.config.block_size_factor
+            assert len(req.block_hashes) >= num_gpu_blocks
+
+            new_offload_keys = group_state.offload_keys[start_block_idx:num_blocks]
            store_output = self.manager.prepare_store(
                new_offload_keys, req_status.req_context
            )
            if store_output is None:
-                logger.warning("Request %s: cannot store blocks", req_id)
+                logger.warning(
+                    "Request %s: cannot store %s blocks", req_id, num_new_blocks
+                )
                continue

+            group_state.next_stored_block_idx = num_blocks
+
            if not store_output.keys_to_store:
-                req_status.advance_stored_idx(num_offloadable_tokens)
                continue
-
-            for group_state in req_status.group_states:
-                self.manager.touch(group_state.offload_keys)
-
            keys_to_store = set(store_output.keys_to_store)

-            group_sizes: list[int] = []
-            block_indices: list[int] = []
-            src_block_ids: list[int] = []
-            for group_config, group_state in zip(
-                self.config.kv_group_configs, req_status.group_states
-            ):
-                num_blocks = num_offloadable_tokens // group_config.offloaded_block_size
-                start_block_idx = group_state.next_stored_block_idx
-                block_ids = group_state.block_ids
-                num_group_blocks = 0
-                start_gpu_block_idx: int | None = None
-                for idx, offload_key in enumerate(
-                    group_state.offload_keys[start_block_idx:num_blocks]
-                ):
-                    if offload_key not in keys_to_store:
-                        continue
+            self.manager.touch(group_state.offload_keys[:num_blocks])

-                    offloaded_block_idx = start_block_idx + idx
-                    gpu_block_idx = offloaded_block_idx * block_size_factor
-                    num_group_blocks += block_size_factor
-                    for i in range(block_size_factor):
-                        block_id = block_ids[gpu_block_idx + i]
-                        if block_id == 0:
-                            # skipped blocks cannot appear after non-skipped blocks
-                            assert start_gpu_block_idx is None
-                            continue
-                        elif start_gpu_block_idx is None:
-                            start_gpu_block_idx = gpu_block_idx + i
-                        src_block_ids.append(block_id)
-                group_sizes.append(num_group_blocks)
-                block_indices.append(start_gpu_block_idx or 0)
-                group_state.next_stored_block_idx = num_blocks
-
-            src_spec = GPULoadStoreSpec(
-                src_block_ids, group_sizes=group_sizes, block_indices=block_indices
-            )
            dst_spec = store_output.store_spec
+            src_block_ids: list[int] = []
+            for idx, key in enumerate(new_offload_keys):
+                if key not in keys_to_store:
+                    continue
+                offloaded_block_idx = start_block_idx + idx
+                gpu_block_idx = offloaded_block_idx * self.config.block_size_factor
+                for i in range(self.config.block_size_factor):
+                    src_block_ids.append(block_ids[gpu_block_idx + i])
+            src_spec = GPULoadStoreSpec(
+                src_block_ids,
+                group_sizes=(len(src_block_ids),),
+                block_indices=(0,),
+            )

            reqs_to_store[req_id] = (src_spec, dst_spec)
            self._reqs_being_stored[req_id] |= keys_to_store

            logger.debug(
-                "Request %s offloading %s blocks upto %d tokens",
+                "Request %s offloading %s blocks starting from block #%d",
                req_id,
                len(keys_to_store),
-                num_offloadable_tokens,
+                start_block_idx,
            )

        return reqs_to_store
@@ -78,7 +78,6 @@ class EngineClient(ABC):
        priority: int = 0,
        data_parallel_rank: int | None = None,
        reasoning_ended: bool | None = None,
-        reasoning_parser_kwargs: dict[str, Any] | None = None,
    ) -> AsyncGenerator[RequestOutput, None]:
        """Generate outputs for a request."""
        ...
@@ -79,7 +79,7 @@ from vllm.renderers.inputs.preprocess import (
    prompt_to_seq,
 )
 from vllm.sampling_params import BeamSearchParams, RequestOutputKind, SamplingParams
-from vllm.tasks import PoolingTask
+from vllm.tasks import SCORE_TYPE_MAP, PoolingTask
 from vllm.tokenizers import TokenizerLike
 from vllm.usage.usage_lib import UsageContext
 from vllm.utils.counter import Counter
@@ -1204,12 +1204,9 @@ class LLM:
                    f"Supported tasks: {self.supported_tasks}"
                )
            else:
-                logger.warning_once(
-                    "Pooling multitask support is deprecated and will "
-                    "be removed in v0.20. When the default pooling task is "
-                    "not what you want, you need to manually specify it "
-                    'via PoolerConfig(task="%s"). ',
-                    pooling_task,
+                raise ValueError(
+                    f"Try switching the model's pooling_task "
+                    f'via `PoolerConfig(task="{pooling_task}")`'
                )

        if pooling_task == "plugin" and "plugin" not in self.pooling_io_processors:
@@ -1412,7 +1409,7 @@ class LLM:
                "pooling model."
            )

-        score_type = self.model_config.score_type
+        score_type: str | None = SCORE_TYPE_MAP.get(self.pooling_task, None)  # type: ignore[arg-type]
        if (
            score_type == "cross-encoder"
            and getattr(self.model_config.hf_config, "num_labels", 0) != 1
@@ -347,11 +347,6 @@ class OpenAIServingChat(OpenAIServing):
                    priority=request.priority,
                    data_parallel_rank=data_parallel_rank,
                    reasoning_ended=reasoning_ended,
-                    reasoning_parser_kwargs={
-                        "chat_template_kwargs": chat_template_kwargs,
-                    }
-                    if reasoning_parser
-                    else None,
                )

            generators.append(generator)
@@ -472,13 +472,9 @@ class OpenAIServingResponses(OpenAIServing):
                    context = SimpleContext()

            if self.parser and self.parser.reasoning_parser_cls is not None:
-                chat_template_kwargs = self._effective_chat_template_kwargs(request)
-                reasoning_parser_kwargs = {
-                    "chat_template_kwargs": chat_template_kwargs,
-                }
                reasoning_parser = self.parser.reasoning_parser_cls(
                    tokenizer,
-                    chat_template_kwargs=chat_template_kwargs,
+                    chat_template_kwargs=self._effective_chat_template_kwargs(request),
                )
                if (
                    isinstance(
@@ -501,9 +497,6 @@ class OpenAIServingResponses(OpenAIServing):
                lora_request=lora_request,
                priority=request.priority,
                trace_headers=trace_headers,
-                reasoning_parser_kwargs=reasoning_parser_kwargs
-                if self.parser and self.parser.reasoning_parser_cls is not None
-                else None,
            )
            generators.append(generator)

@@ -650,7 +643,6 @@ class OpenAIServingResponses(OpenAIServing):
        lora_request: LoRARequest | None = None,
        priority: int = 0,
        trace_headers: Mapping[str, str] | None = None,
-        reasoning_parser_kwargs: dict[str, Any] | None = None,
    ):
        max_model_len = self.model_config.max_model_len

@@ -674,7 +666,6 @@ class OpenAIServingResponses(OpenAIServing):
                lora_request=lora_request,
                trace_headers=trace_headers,
                priority=priority,
-                reasoning_parser_kwargs=reasoning_parser_kwargs,
            )

            async for res in generator:
@@ -15,10 +15,7 @@ from starlette.datastructures import Headers
 from vllm import PoolingParams, PoolingRequestOutput, envs
 from vllm.config import VllmConfig
 from vllm.engine.protocol import EngineClient
-from vllm.entrypoints.chat_utils import (
-    ChatTemplateConfig,
-    ChatTemplateContentFormatOption,
-)
+from vllm.entrypoints.chat_utils import ChatTemplateConfig
 from vllm.entrypoints.logger import RequestLogger
 from vllm.entrypoints.openai.engine.protocol import ErrorResponse
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
@@ -48,9 +45,7 @@ class PoolingServingBase(ABC):
        models: OpenAIServingModels,
        *,
        request_logger: RequestLogger | None,
-        chat_template: str | None = None,
-        chat_template_content_format: ChatTemplateContentFormatOption = "auto",
-        trust_request_chat_template: bool = False,
+        chat_template_config: ChatTemplateConfig,
        return_tokens_as_token_ids: bool = False,
        log_error_stack: bool = False,
    ):
@@ -63,11 +58,7 @@ class PoolingServingBase(ABC):
        self.request_logger = request_logger
        self.return_tokens_as_token_ids = return_tokens_as_token_ids
        self.log_error_stack = log_error_stack
-        self.chat_template_config = ChatTemplateConfig(
-            chat_template=chat_template,
-            chat_template_content_format=chat_template_content_format,
-            trust_request_chat_template=trust_request_chat_template,
-        )
+        self.chat_template_config = chat_template_config

        # Shared thread pool executor for preprocessing and postprocessing.
        self._executor: Executor = models.renderer._executor
@@ -10,7 +10,7 @@ from vllm.entrypoints.chat_utils import ChatTemplateConfig
 from vllm.logger import init_logger
 from vllm.plugins.io_processors import has_io_processor
 from vllm.renderers import BaseRenderer
-from vllm.tasks import POOLING_TASKS, SupportedTask
+from vllm.tasks import POOLING_TASKS, SCORE_TYPE_MAP, SupportedTask

 from .base.io_processor import PoolingIOProcessor
 from .utils import enable_scoring_api
@@ -43,23 +43,24 @@ def init_pooling_io_processors(
 ) -> dict[str, PoolingIOProcessor]:
    model_config = vllm_config.model_config
    processors: dict[str, type[PoolingIOProcessor]] = {}
+    pooling_task = model_config.get_pooling_task(supported_tasks)

-    if "classify" in supported_tasks:
+    if pooling_task == "classify":
        from .classify.io_processor import ClassifyIOProcessor

        processors["classify"] = ClassifyIOProcessor

-    if "token_classify" in supported_tasks:
+    if pooling_task == "token_classify":
        from .classify.io_processor import TokenClassifyIOProcessor

        processors["token_classify"] = TokenClassifyIOProcessor

-    if "embed" in supported_tasks:
+    if pooling_task == "embed":
        from .embed.io_processor import EmbedIOProcessor

        processors["embed"] = EmbedIOProcessor

-    if "token_embed" in supported_tasks:
+    if pooling_task == "token_embed":
        from .embed.io_processor import TokenEmbedIOProcessor

        processors["token_embed"] = TokenEmbedIOProcessor
@@ -71,15 +72,15 @@ def init_pooling_io_processors(
        from .pooling.io_processor import PluginWithIOProcessorPlugins

        processors["plugin"] = PluginWithIOProcessorPlugins
-    elif "plugin" in supported_tasks:
+    elif pooling_task == "plugin":
        from .pooling.io_processor import PluginWithoutIOProcessorPlugins

        processors["plugin"] = PluginWithoutIOProcessorPlugins

    if enable_scoring_api(supported_tasks, model_config):
-        score_type = model_config.score_type
        from .scoring.io_processor import ScoringIOProcessors

+        score_type: str | None = SCORE_TYPE_MAP.get(pooling_task, None)  # type: ignore[arg-type]
        if score_type is not None and score_type in ScoringIOProcessors:
            processors[score_type] = ScoringIOProcessors[score_type]

@@ -140,6 +141,10 @@ def init_pooling_state(
    request_logger: RequestLogger | None,
    supported_tasks: tuple["SupportedTask", ...],
 ):
+    model_config = engine_client.model_config
+    if model_config is None:
+        return
+
    from vllm.entrypoints.chat_utils import load_chat_template
    from vllm.tasks import POOLING_TASKS

@@ -148,8 +153,14 @@ def init_pooling_state(
    from .pooling.serving import ServingPooling
    from .scoring.serving import ServingScores

-    model_config = engine_client.model_config
    resolved_chat_template = load_chat_template(args.chat_template)
+    pooling_task = model_config.get_pooling_task(supported_tasks)
+
+    chat_template_config = ChatTemplateConfig(
+        chat_template=resolved_chat_template,
+        chat_template_content_format=args.chat_template_content_format,
+        trust_request_chat_template=args.trust_request_chat_template,
+    )

    state.serving_pooling = (
        (
@@ -158,9 +169,7 @@ def init_pooling_state(
                state.openai_serving_models,
                supported_tasks=supported_tasks,
                request_logger=request_logger,
-                chat_template=resolved_chat_template,
-                chat_template_content_format=args.chat_template_content_format,
-                trust_request_chat_template=args.trust_request_chat_template,
+                chat_template_config=chat_template_config,
            )
        )
        if any(t in supported_tasks for t in POOLING_TASKS)
@@ -171,11 +180,9 @@ def init_pooling_state(
            engine_client,
            state.openai_serving_models,
            request_logger=request_logger,
-            chat_template=resolved_chat_template,
-            chat_template_content_format=args.chat_template_content_format,
-            trust_request_chat_template=args.trust_request_chat_template,
+            chat_template_config=chat_template_config,
        )
-        if "embed" in supported_tasks
+        if pooling_task == "embed"
        else None
    )
    state.serving_classification = (
@@ -183,21 +190,18 @@ def init_pooling_state(
            engine_client,
            state.openai_serving_models,
            request_logger=request_logger,
-            chat_template=resolved_chat_template,
-            chat_template_content_format=args.chat_template_content_format,
-            trust_request_chat_template=args.trust_request_chat_template,
+            chat_template_config=chat_template_config,
        )
-        if "classify" in supported_tasks
+        if pooling_task == "classify"
        else None
    )
    state.serving_scores = (
        ServingScores(
            engine_client,
            state.openai_serving_models,
+            supported_tasks=supported_tasks,
            request_logger=request_logger,
-            chat_template=resolved_chat_template,
-            chat_template_content_format=args.chat_template_content_format,
-            trust_request_chat_template=args.trust_request_chat_template,
+            chat_template_config=chat_template_config,
            enable_flash_late_interaction=getattr(
                args, "enable_flash_late_interaction", True
            ),
@@ -214,7 +218,12 @@ def get_pooling_invocation_types(
    # NOTE: Items defined earlier take higher priority
    invocation_types: list[tuple[RequestType, tuple[GetHandlerFn, EndpointFn]]] = []

-    if "embed" in supported_tasks:
+    if model_config is None:
+        return invocation_types
+
+    pooling_task = model_config.get_pooling_task(supported_tasks)
+
+    if pooling_task == "embed":
        from .embed.api_router import create_embedding, embedding
        from .embed.protocol import EmbeddingRequest

@@ -222,7 +231,7 @@ def get_pooling_invocation_types(
            (EmbeddingRequest, (embedding, create_embedding)),
        ]

-    if "classify" in supported_tasks:
+    if pooling_task == "classify":
        from .classify.api_router import classify, create_classify
        from .classify.protocol import ClassificationRequest

@@ -78,17 +78,15 @@ class ServingPooling(PoolingServingBase):

        # plugin task uses io_processor.parse_request to verify inputs
        if pooling_task != "plugin" and pooling_task != self.pooling_task:
-            if pooling_task not in self.io_processors:
+            if pooling_task not in self.supported_tasks:
                raise ValueError(
                    f"Unsupported task: {pooling_task!r} "
                    f"Supported tasks: {self.supported_tasks}"
                )
            else:
-                logger.warning_once(
-                    "Pooling multitask support is deprecated and will be removed "
-                    "in v0.20. When the default pooling task is not what you want, you "
-                    "need to manually specify it via --pooler-config.task %s. ",
-                    pooling_task,
+                raise ValueError(
+                    "Try switching the model's pooling_task "
+                    f"via --pooler-config.task {request.task}."
                )

        if pooling_task == "plugin" and "plugin" not in self.io_processors:
@@ -8,6 +8,7 @@ from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.openai.engine.protocol import UsageInfo
 from vllm.logger import init_logger
 from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput
+from vllm.tasks import SCORE_TYPE_MAP, SupportedTask
 from vllm.v1.pool.late_interaction import (
    build_late_interaction_doc_params,
    build_late_interaction_query_params,
@@ -38,10 +39,15 @@ class ServingScores(PoolingServing):
        self,
        engine_client: EngineClient,
        *args,
+        supported_tasks: tuple[SupportedTask, ...],
        enable_flash_late_interaction: bool = True,
        **kwargs,
    ):
-        self.io_processor_name: str = engine_client.model_config.score_type
+        pooling_task = engine_client.model_config.get_pooling_task(supported_tasks)
+        score_type = SCORE_TYPE_MAP.get(pooling_task, None)  # type: ignore[arg-type]
+        assert score_type is not None
+
+        self.io_processor_name: str = score_type
        self.enable_flash_late_interaction = (
            self.io_processor_name == "late-interaction"
            and enable_flash_late_interaction
@@ -141,10 +141,14 @@ def enable_scoring_api(
    supported_tasks: tuple["SupportedTask", ...],
    model_config: ModelConfig | None = None,
 ) -> bool:
-    if any(t in supported_tasks for t in ("embed", "token_embed")):
+    if model_config is None:
+        return False
+
+    pooling_task = model_config.get_pooling_task(supported_tasks)
+    if pooling_task in ("embed", "token_embed"):
        return True

-    if model_config is not None and "classify" in supported_tasks:
+    if pooling_task == "classify":
        num_labels = getattr(model_config.hf_config, "num_labels", 0)
        if num_labels != 1:
            logger.debug_once("Scoring API is only enabled for num_labels == 1.")
@@ -245,9 +245,9 @@ if TYPE_CHECKING:
    VLLM_DEBUG_WORKSPACE: bool = False
    VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False
    VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256
-    VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 4096
    VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary"
    VLLM_USE_V2_MODEL_RUNNER: bool = False
+    VLLM_DEEPSEEK_V4_USE_MEGA_MOE: bool = False
    VLLM_LOG_MODEL_INSPECTION: bool = False
    VLLM_DEBUG_MFU_METRICS: bool = False
    VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY: bool = False
@@ -1663,17 +1663,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
    "VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD": lambda: int(
        int(os.getenv("VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD", 256))
    ),
-    # Token-count cutoff for multi-stream overlap of the attention input
-    # GEMM with auxiliary GEMMs (e.g. fused_wqa_wkv overlapped with indexer
-    # weights / kv-score projections in DeepSeek-V4). At or below this many
-    # tokens the FP8 main GEMM has idle SMs to share with the bf16 aux GEMMs
-    # and overlap is a 5-45% win; above it the FP8 GEMM saturates the device
-    # and the cross-stream sync becomes pure overhead. Set to 0 to disable
-    # the multi-stream path entirely. Empirical crossover on B300 (148 SMs)
-    # is ~4096; B200 (132 SMs) is expected ~3072.
-    "VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD": lambda: int(
-        os.getenv("VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD", "4096")
-    ),
    # Format for saving torch.compile cache artifacts
    # - "binary": saves as binary file
    #     Safe for multiple vllm serve processes accessing the same torch compile cache.
@@ -1687,6 +1676,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
    "VLLM_USE_V2_MODEL_RUNNER": lambda: bool(
        int(os.getenv("VLLM_USE_V2_MODEL_RUNNER", "0"))
    ),
+    # Use the DeepGEMM MegaMoE fused expert kernel for DeepSeek V4 routed
+    # experts. Set to 0 to fall back to the standard SharedFusedMoE path.
+    "VLLM_DEEPSEEK_V4_USE_MEGA_MOE": lambda: bool(
+        int(os.getenv("VLLM_DEEPSEEK_V4_USE_MEGA_MOE", "0"))
+    ),
    # Log model inspection after loading.
    # If enabled, logs a transformers-style hierarchical view of the model
    # with quantization methods and attention backends.
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import functools

 import torch
 import torch.nn as nn
@@ -13,17 +14,31 @@ from vllm.distributed.parallel_state import (
 )
 from vllm.distributed.utils import divide
 from vllm.lora.layers.base import BaseLayerWithLoRA
+from vllm.lora.ops.triton_ops.utils import get_lora_op_configs
 from vllm.model_executor.layers.fused_moe import FusedMoE
+from vllm.model_executor.layers.fused_moe.config import (
+    _get_config_dtype_str,
+)
+from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import (
+    UnfusedOAITritonExperts,
+)
+from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
+    MarlinExperts,
+)
+from vllm.model_executor.layers.fused_moe.fused_moe import (
+    TritonExperts,
+)
 from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import (
    FusedMoEModularMethod,
 )
-from vllm.model_executor.layers.fused_moe.lora_context import MoELoRAContext
-from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
+from vllm.model_executor.layers.fused_moe.modular_kernel import (
+    FusedMoEKernel,
+)
 from vllm.model_executor.layers.fused_moe.prepare_finalize import (
    MoEPrepareAndFinalizeNoDPEPModular,
 )

-from .utils import _get_lora_device
+from .utils import _get_lora_device, try_get_optimal_moe_lora_config


 class FusedMoEWithLoRA(BaseLayerWithLoRA):
@@ -43,49 +58,299 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
        # For non-gated MoE (is_act_and_mul=False), only 1 slice is needed
        # since there's only up_proj (w1), not gate_proj + up_proj (w1 + w3)
        self._w13_slices = 2 if base_layer.moe_config.is_act_and_mul else 1
+        self._inject_lora_into_fused_moe()
+
+    def _normalize_keys(self, config: dict[str, int | None]) -> dict[str, int | None]:
+        normalized_config = {}
+        for key, value in config.items():
+            if key.islower():
+                if key.startswith("block_"):
+                    normalized_key = "BLOCK_SIZE_" + key.split("_")[-1].upper()
+                else:
+                    normalized_key = key.upper()
+            else:
+                normalized_key = key
+            normalized_config[normalized_key] = value
+        return normalized_config
+
+    def _get_lora_moe_configs(
+        self,
+        op_prefix: str,
+        num_loras: int,
+        rank: int,
+        num_slices: int,
+        M: int,
+        layer: FusedMoE,
+        top_k: int,
+        config_dtype: str,
+    ):
+        if envs.VLLM_TUNED_CONFIG_FOLDER:
+            hidden_size = layer.hidden_size
+            intermediate_size = (
+                self.w2_lora_a_stacked[0].shape[-1]
+                if op_prefix == "w2"
+                else self.w13_lora_b_stacked[0].shape[-2]
+            )
+            shrink_config = get_lora_op_configs(
+                op_type=f"fused_moe_lora_{op_prefix}_shrink",
+                max_loras=num_loras,
+                batch=M,
+                hidden_size=hidden_size,
+                rank=rank,
+                num_slices=num_slices,
+                moe_intermediate_size=intermediate_size,
+            )
+            expand_config = get_lora_op_configs(
+                op_type=f"fused_moe_lora_{op_prefix}_expand",
+                max_loras=num_loras,
+                batch=M,
+                hidden_size=hidden_size,  # lora_a_stacked.shape[-1],
+                rank=rank,
+                num_slices=num_slices,
+                moe_intermediate_size=intermediate_size,  # lora_b_stacked.shape[-2],
+            )
+        else:  # fall back to the default config
+            get_config_func = functools.partial(
+                try_get_optimal_moe_lora_config,
+                w1_shape=layer.w13_weight.shape,
+                w2_shape=layer.w2_weight.shape,
+                rank=rank,
+                top_k=top_k,
+                dtype=config_dtype,
+                M=M,
+                block_shape=layer.quant_method.moe_quant_config.block_shape,
+            )
+            shrink_config = get_config_func(
+                op_type=f"fused_moe_lora_{op_prefix}_shrink"
+            )
+            expand_config = get_config_func(
+                op_type=f"fused_moe_lora_{op_prefix}_expand"
+            )
+        shrink_config = self._normalize_keys(shrink_config)
+        expand_config = self._normalize_keys(expand_config)
+        return shrink_config, expand_config
+
+    def _inject_lora_into_fused_moe(self):
+        moe_state_dict = {}
+        top_k = self.base_layer.top_k

        self.base_layer.ensure_moe_quant_config_init()
+        quant_config = self.base_layer.quant_method.moe_quant_config
+
        if getattr(self.base_layer.quant_method, "supports_internal_mk", False):
-            moe_kernel = self.base_layer.quant_method.moe_kernel
+            # Use the existing modular kernel from the quant method
+            m_fused_moe_fn = self.base_layer.quant_method.moe_kernel
            # Don't let the kernel own shared experts so the runner can
            # overlap them with routed experts via a separate CUDA stream.
-            moe_kernel.shared_experts = None
+            m_fused_moe_fn.shared_experts = None
        else:
+            # Create a new modular kernel via select_gemm_impl.
+            # Don't pass shared_experts to the kernel so the runner can
+            # overlap them with routed experts via a separate CUDA stream.
            prepare_finalize = MoEPrepareAndFinalizeNoDPEPModular()
-            moe_kernel = FusedMoEKernel(
+            m_fused_moe_fn = FusedMoEKernel(
                prepare_finalize,
                self.base_layer.quant_method.select_gemm_impl(
                    prepare_finalize, self.base_layer
                ),
            )
-        assert moe_kernel.supports_lora(), (
-            f"{type(moe_kernel.fused_experts).__name__} does not support LoRA. "
-            "For unquantized MoE, set moe_backend='triton' or moe_backend='auto' "
-            "(auto selects Triton automatically when LoRA is enabled). "
-            "For quantized MoE, mix LoRAExpertsMixin into the experts class "
-            "and consume self._lora_context in apply()."
-        )
-        self._fused_experts = moe_kernel.fused_experts
-        self.base_layer._replace_quant_method(
-            FusedMoEModularMethod(self.base_layer.quant_method, moe_kernel)
-        )

-    def _build_lora_context(self):
-        return MoELoRAContext(
-            w13_lora_a_stacked=self.w13_lora_a_stacked,
-            w13_lora_b_stacked=self.w13_lora_b_stacked,
-            w2_lora_a_stacked=self.w2_lora_a_stacked,
-            w2_lora_b_stacked=self.w2_lora_b_stacked,
-            adapter_enabled=self.adapter_enabled,
-            max_loras=self.max_loras,
-            top_k=self.base_layer.top_k,
-            w13_num_slices=self._w13_slices,
-            fully_sharded=self.fully_sharded,
-            tp_rank=self.tp_rank,
-            tp_size=self.tp_size,
-            local_num_experts=self.base_layer.local_num_experts,
-            punica_wrapper=self.punica_wrapper,
-            use_tuned_config=bool(envs.VLLM_TUNED_CONFIG_FOLDER),
+        if quant_config.use_mxfp4_w4a16:
+            assert isinstance(
+                m_fused_moe_fn.impl.fused_experts,
+                (MarlinExperts, UnfusedOAITritonExperts),
+            )
+        else:
+            assert isinstance(m_fused_moe_fn.impl.fused_experts, TritonExperts)
+
+        def fwd_decorator(layer, func):
+            def wrapper(*args, **kwargs):
+                moe_state_dict["hidden_states"] = kwargs["hidden_states"]
+                moe_state_dict["topk_ids"] = kwargs["topk_ids"]
+                moe_state_dict["topk_weights"] = kwargs["topk_weights"]
+                moe_state_dict["expert_map"] = kwargs["expert_map"]
+                moe_state_dict["apply_router_weight_on_input"] = kwargs[
+                    "apply_router_weight_on_input"
+                ]
+                result = func(*args, **kwargs)
+                return result
+
+            return wrapper
+
+        def act_decorator(layer, func):
+            def wrapper(*args, **kwargs):
+                _, output, input = args
+
+                hidden_states = moe_state_dict["hidden_states"]
+                topk_weights = moe_state_dict["topk_weights"]
+                curr_topk_ids = moe_state_dict["topk_ids"]
+
+                expert_map = moe_state_dict["expert_map"]
+
+                config_dtype = _get_config_dtype_str(
+                    dtype=hidden_states.dtype,
+                    use_fp8_w8a8=False,
+                    use_int8_w8a16=False,
+                    use_int4_w4a16=False,
+                )
+                num_tokens = hidden_states.size(0)
+                M = num_tokens
+                max_lora_rank = self.w13_lora_a_stacked[0].shape[-2]
+                shrink_config, expand_config = self._get_lora_moe_configs(
+                    op_prefix="w13",
+                    num_loras=self.max_loras,
+                    rank=max_lora_rank,
+                    num_slices=self._w13_slices,
+                    M=M,
+                    layer=layer,
+                    top_k=top_k,
+                    config_dtype=config_dtype,
+                )
+
+                # SPARSITY_FACTOR is a heuristic margin ensuring tokens * top_k
+                # activates only a small fraction of total experts * loras.
+                SPARSITY_FACTOR = 8
+                naive_block_assignment = (
+                    expert_map is None
+                    and num_tokens * top_k * SPARSITY_FACTOR
+                    <= self.base_layer.local_num_experts * self.max_loras
+                )
+
+                # get the block size of m from customized config or default config
+                (
+                    token_lora_mapping,
+                    sorted_token_ids_lora,
+                    expert_ids_lora,
+                    num_tokens_post_padded_lora,
+                ) = self.punica_wrapper.moe_lora_align_block_size(
+                    curr_topk_ids,
+                    num_tokens,
+                    shrink_config["BLOCK_SIZE_M"],
+                    self.base_layer.local_num_experts,
+                    self.max_loras,
+                    self.adapter_enabled,
+                    expert_map,
+                    naive_block_assignment=naive_block_assignment,
+                )
+
+                moe_state_dict["sorted_token_ids_lora"] = sorted_token_ids_lora
+                moe_state_dict["expert_ids_lora"] = expert_ids_lora
+                moe_state_dict["num_tokens_post_padded_lora"] = (
+                    num_tokens_post_padded_lora
+                )
+                moe_state_dict["token_lora_mapping"] = token_lora_mapping
+
+                if sorted_token_ids_lora is not None:
+                    expert_ids_lora = expert_ids_lora.view(self.max_loras, -1)
+                    sorted_token_ids_lora = sorted_token_ids_lora.view(
+                        self.max_loras, -1
+                    )
+                #
+
+                self.punica_wrapper.add_lora_fused_moe(
+                    input.view(-1, top_k, input.shape[-1]),
+                    hidden_states,
+                    self.w13_lora_a_stacked,
+                    self.w13_lora_b_stacked,
+                    topk_weights,
+                    sorted_token_ids_lora,
+                    expert_ids_lora,
+                    num_tokens_post_padded_lora,
+                    max_lora_rank,
+                    top_k,
+                    shrink_config,  ## pass the shrink config
+                    expand_config,  ## pass the expand config
+                    self.adapter_enabled,
+                    fully_sharded=self.fully_sharded,
+                    token_lora_mapping=token_lora_mapping,
+                )
+
+                result = func(*args, **kwargs)
+
+                moe_state_dict["intermediate_cache2"] = output
+                return result
+
+            return wrapper
+
+        def moe_sum_decorator(layer, func):
+            def wrapper(*args, **kwargs):
+                hidden_states = moe_state_dict["hidden_states"]
+                topk_weights = moe_state_dict["topk_weights"]
+
+                config_dtype = _get_config_dtype_str(
+                    dtype=hidden_states.dtype,
+                    use_fp8_w8a8=False,
+                    use_int8_w8a16=False,
+                    use_int4_w4a16=False,
+                )
+                num_tokens = hidden_states.size(0)
+                M = num_tokens
+                max_lora_rank = self.w2_lora_a_stacked[0].shape[-2]
+                shrink_config, expand_config = self._get_lora_moe_configs(
+                    op_prefix="w2",
+                    num_loras=self.max_loras,
+                    rank=max_lora_rank,
+                    num_slices=1,
+                    M=M,
+                    layer=layer,
+                    top_k=top_k,
+                    config_dtype=config_dtype,
+                )
+
+                sorted_token_ids_lora = moe_state_dict["sorted_token_ids_lora"]
+                expert_ids_lora = moe_state_dict["expert_ids_lora"]
+                num_tokens_post_padded_lora = moe_state_dict[
+                    "num_tokens_post_padded_lora"
+                ]
+                token_lora_mapping = moe_state_dict.get("token_lora_mapping")
+
+                if sorted_token_ids_lora is not None:
+                    expert_ids_lora = expert_ids_lora.view(self.max_loras, -1)
+                    sorted_token_ids_lora = sorted_token_ids_lora.view(
+                        self.max_loras, -1
+                    )
+                intermediate_cache2 = moe_state_dict["intermediate_cache2"]
+                intermediate_cache3 = args[0]
+
+                shard_size_w2 = divide(self.base_layer.hidden_size, self.tp_size)
+
+                self.punica_wrapper.add_lora_fused_moe(
+                    intermediate_cache3,
+                    intermediate_cache2,
+                    self.w2_lora_a_stacked,
+                    self.w2_lora_b_stacked,
+                    topk_weights,
+                    sorted_token_ids_lora,
+                    expert_ids_lora,
+                    num_tokens_post_padded_lora,
+                    max_lora_rank,
+                    top_k,
+                    shrink_config,  ## pass the shrink config
+                    expand_config,  ## pass the expand config
+                    self.adapter_enabled,
+                    True,
+                    fully_sharded=self.fully_sharded,
+                    offset=shard_size_w2 * self.tp_rank if self.fully_sharded else 0,
+                    token_lora_mapping=token_lora_mapping,
+                )
+
+                result = func(*args, **kwargs)
+                return result
+
+            return wrapper
+
+        fused_experts = m_fused_moe_fn.impl.fused_experts
+
+        m_fused_moe_fn.apply = fwd_decorator(self.base_layer, m_fused_moe_fn.apply)
+        fused_experts.activation = act_decorator(
+            self.base_layer, fused_experts.activation
+        )
+        fused_experts.moe_sum = moe_sum_decorator(
+            self.base_layer, fused_experts.moe_sum
+        )
+        # TODO(bnell): find a less intrusive way to handle this.
+        self.base_layer._replace_quant_method(
+            FusedMoEModularMethod(self.base_layer.quant_method, m_fused_moe_fn)
        )

    def _create_lora_a_weights(
@@ -324,10 +589,6 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
            index, :, : sliced_w2_lora_b.shape[1], : sliced_w2_lora_b.shape[2]
        ].copy_(sliced_w2_lora_b, non_blocking=True)

-    def set_mapping(self, punica_wrapper):
-        super().set_mapping(punica_wrapper)
-        self._fused_experts.set_lora_context(self._build_lora_context())
-
    def forward(self, *args, **kwargs):
        return self.base_layer.forward(*args, **kwargs)

@@ -90,12 +90,11 @@ def try_get_optimal_moe_lora_config(
    top_k: int,
    dtype: str | None,
    M: int,
+    block_shape: list[int] | None = None,
 ) -> dict[str, int | None]:
-    # LoRA shrink/expand operates on bf16/fp16 adapters regardless of the
-    # base MoE weight's block-wise quantization, so block_shape is omitted
-    # from the config lookup — the non-quantized branch in get_default_config
-    # ignores it anyway.
-    config = try_get_optimal_moe_config(w1_shape, w2_shape, top_k, dtype, M).copy()
+    config = try_get_optimal_moe_config(
+        w1_shape, w2_shape, top_k, dtype, M, block_shape
+    ).copy()
    if op_type in [
        "fused_moe_lora_w13_shrink",
        "fused_moe_lora_w2_shrink",
@@ -321,20 +321,3 @@ def supports_pdl(device: torch.device | None = None) -> bool:
 def supports_tma(device: torch.device | None = None) -> bool:
    # TMA requires compute capability SM90 or above
    return current_platform.is_cuda() and current_platform.has_device_capability(90)
-
-
-def _normalize_lora_config_keys(
-    config: dict[str, int | None],
-) -> dict[str, int | None]:
-    """Normalize Triton config dict keys to uppercase BLOCK_SIZE_* format."""
-    out: dict[str, int | None] = {}
-    for key, val in config.items():
-        if key.islower():
-            if key.startswith("block_"):
-                nk = "BLOCK_SIZE_" + key.split("_")[-1].upper()
-            else:
-                nk = key.upper()
-        else:
-            nk = key
-        out[nk] = val
-    return out
@@ -493,65 +493,3 @@ class PunicaWrapperBase(PunicaWrapperABC):
        """
        # TODO: implement it based on torch ops
        raise NotImplementedError
-
-    def add_lora_w13(
-        self,
-        y: torch.Tensor,
-        x: torch.Tensor,
-        lora_a_stacked: tuple[torch.Tensor, ...],
-        lora_b_stacked: tuple[torch.Tensor, ...],
-        topk_ids: torch.Tensor,
-        topk_weights: torch.Tensor,
-        expert_map: torch.Tensor | None,
-        w1: torch.Tensor,
-        w2: torch.Tensor,
-        num_tokens: int,
-        top_k_num: int,
-        max_loras: int,
-        adapter_enabled: torch.Tensor,
-        local_num_experts: int,
-        top_k: int,
-        num_slices: int,
-        fully_sharded: bool,
-        use_tuned_config: bool,
-    ) -> tuple[
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor | None,
-    ]:
-        """Apply w13 LoRA to y (intermediate_cache1) in-place before activation.
-
-        Returns (sorted_token_ids_lora, expert_ids_lora,
-                 num_tokens_post_padded_lora, token_lora_mapping)
-        for reuse by add_lora_w2.
-        """
-        raise NotImplementedError
-
-    def add_lora_w2(
-        self,
-        y: torch.Tensor,
-        x: torch.Tensor,
-        lora_a_stacked: tuple[torch.Tensor, ...],
-        lora_b_stacked: tuple[torch.Tensor, ...],
-        topk_weights: torch.Tensor,
-        sorted_token_ids_lora: torch.Tensor | None,
-        expert_ids_lora: torch.Tensor | None,
-        num_tokens_post_padded_lora: torch.Tensor | None,
-        token_lora_mapping: torch.Tensor | None,
-        num_tokens: int,
-        w1: torch.Tensor,
-        w2: torch.Tensor,
-        top_k_num: int,
-        max_loras: int,
-        adapter_enabled: torch.Tensor,
-        top_k: int,
-        fully_sharded: bool,
-        tp_rank: int,
-        use_tuned_config: bool,
-    ) -> None:
-        """Apply w2 LoRA to y (intermediate_cache3) in-place before moe_sum.
-
-        Reuses routing tensors returned by add_lora_w13.
-        """
-        raise NotImplementedError
@@ -459,239 +459,3 @@ class PunicaWrapperGPU(PunicaWrapperBase):
            fully_sharded,
            offset,
        )
-
-    def add_lora_w13(
-        self,
-        y: torch.Tensor,
-        x: torch.Tensor,
-        lora_a_stacked: tuple[torch.Tensor, ...],
-        lora_b_stacked: tuple[torch.Tensor, ...],
-        topk_ids: torch.Tensor,
-        topk_weights: torch.Tensor,
-        expert_map: torch.Tensor | None,
-        w1: torch.Tensor,
-        w2: torch.Tensor,
-        num_tokens: int,
-        top_k_num: int,
-        max_loras: int,
-        adapter_enabled: torch.Tensor,
-        local_num_experts: int,
-        top_k: int,
-        num_slices: int,
-        fully_sharded: bool,
-        use_tuned_config: bool,
-    ) -> tuple[
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor | None,
-    ]:
-        import functools
-
-        from vllm.lora.layers.utils import try_get_optimal_moe_lora_config
-        from vllm.lora.ops.triton_ops.utils import (
-            _normalize_lora_config_keys,
-            get_lora_op_configs,
-        )
-        from vllm.model_executor.layers.fused_moe.config import _get_config_dtype_str
-
-        config_dtype = _get_config_dtype_str(
-            dtype=x.dtype,
-            use_fp8_w8a8=False,
-            use_int8_w8a16=False,
-            use_int4_w4a16=False,
-        )
-        max_lora_rank = lora_a_stacked[0].shape[-2]
-
-        if use_tuned_config:
-            shrink_config = get_lora_op_configs(
-                op_type="fused_moe_lora_w13_shrink",
-                max_loras=max_loras,
-                batch=num_tokens,
-                hidden_size=x.shape[-1],
-                rank=max_lora_rank,
-                num_slices=num_slices,
-                moe_intermediate_size=lora_b_stacked[0].shape[-2],
-            )
-            expand_config = get_lora_op_configs(
-                op_type="fused_moe_lora_w13_expand",
-                max_loras=max_loras,
-                batch=num_tokens,
-                hidden_size=x.shape[-1],
-                rank=max_lora_rank,
-                num_slices=num_slices,
-                moe_intermediate_size=lora_b_stacked[0].shape[-2],
-            )
-        else:
-            get_config = functools.partial(
-                try_get_optimal_moe_lora_config,
-                w1_shape=w1.shape,
-                w2_shape=w2.shape,
-                rank=max_lora_rank,
-                top_k=top_k,
-                dtype=config_dtype,
-                M=num_tokens,
-            )
-            shrink_config = get_config(op_type="fused_moe_lora_w13_shrink")
-            expand_config = get_config(op_type="fused_moe_lora_w13_expand")
-
-        shrink_config = _normalize_lora_config_keys(shrink_config)
-        expand_config = _normalize_lora_config_keys(expand_config)
-
-        SPARSITY_FACTOR = 8
-        naive_block_assignment = (
-            expert_map is None
-            and num_tokens * top_k * SPARSITY_FACTOR <= local_num_experts * max_loras
-        )
-
-        (
-            token_lora_mapping,
-            sorted_token_ids_lora,
-            expert_ids_lora,
-            num_tokens_post_padded_lora,
-        ) = self.moe_lora_align_block_size(
-            topk_ids,
-            num_tokens,
-            int(shrink_config.get("BLOCK_SIZE_M") or 64),
-            local_num_experts,
-            max_loras,
-            adapter_enabled,
-            expert_map,
-            naive_block_assignment=naive_block_assignment,
-        )
-
-        _sorted = sorted_token_ids_lora
-        _eids = expert_ids_lora
-        if _sorted is not None:
-            _eids = _eids.view(max_loras, -1)
-            _sorted = _sorted.view(max_loras, -1)
-
-        self.add_lora_fused_moe(
-            y.view(-1, top_k_num, y.shape[-1]),
-            x,
-            lora_a_stacked,
-            lora_b_stacked,
-            topk_weights,
-            _sorted,
-            _eids,
-            num_tokens_post_padded_lora,
-            max_lora_rank,
-            top_k,
-            shrink_config,
-            expand_config,
-            adapter_enabled,
-            fully_sharded=fully_sharded,
-            token_lora_mapping=token_lora_mapping,
-        )
-
-        return (
-            sorted_token_ids_lora,
-            expert_ids_lora,
-            num_tokens_post_padded_lora,
-            token_lora_mapping,
-        )
-
-    def add_lora_w2(
-        self,
-        y: torch.Tensor,
-        x: torch.Tensor,
-        lora_a_stacked: tuple[torch.Tensor, ...],
-        lora_b_stacked: tuple[torch.Tensor, ...],
-        topk_weights: torch.Tensor,
-        sorted_token_ids_lora: torch.Tensor | None,
-        expert_ids_lora: torch.Tensor | None,
-        num_tokens_post_padded_lora: torch.Tensor | None,
-        token_lora_mapping: torch.Tensor | None,
-        num_tokens: int,
-        w1: torch.Tensor,
-        w2: torch.Tensor,
-        top_k_num: int,
-        max_loras: int,
-        adapter_enabled: torch.Tensor,
-        top_k: int,
-        fully_sharded: bool,
-        tp_rank: int,
-        use_tuned_config: bool,
-    ) -> None:
-        import functools
-
-        from vllm.lora.layers.utils import try_get_optimal_moe_lora_config
-        from vllm.lora.ops.triton_ops.utils import (
-            _normalize_lora_config_keys,
-            get_lora_op_configs,
-        )
-        from vllm.model_executor.layers.fused_moe.config import _get_config_dtype_str
-
-        config_dtype = _get_config_dtype_str(
-            dtype=x.dtype,
-            use_fp8_w8a8=False,
-            use_int8_w8a16=False,
-            use_int4_w4a16=False,
-        )
-        max_lora_rank = lora_a_stacked[0].shape[-2]
-
-        if use_tuned_config:
-            shrink_config = get_lora_op_configs(
-                op_type="fused_moe_lora_w2_shrink",
-                max_loras=max_loras,
-                batch=num_tokens,
-                hidden_size=y.shape[-1],
-                rank=max_lora_rank,
-                num_slices=1,
-                moe_intermediate_size=lora_a_stacked[0].shape[-1],
-            )
-            expand_config = get_lora_op_configs(
-                op_type="fused_moe_lora_w2_expand",
-                max_loras=max_loras,
-                batch=num_tokens,
-                hidden_size=y.shape[-1],
-                rank=max_lora_rank,
-                num_slices=1,
-                moe_intermediate_size=lora_a_stacked[0].shape[-1],
-            )
-        else:
-            get_config = functools.partial(
-                try_get_optimal_moe_lora_config,
-                w1_shape=w1.shape,
-                w2_shape=w2.shape,
-                rank=max_lora_rank,
-                top_k=top_k,
-                dtype=config_dtype,
-                M=num_tokens,
-            )
-            shrink_config = get_config(op_type="fused_moe_lora_w2_shrink")
-            expand_config = get_config(op_type="fused_moe_lora_w2_expand")
-
-        shrink_config = _normalize_lora_config_keys(shrink_config)
-        expand_config = _normalize_lora_config_keys(expand_config)
-
-        _sorted = sorted_token_ids_lora
-        _eids = expert_ids_lora
-        if _sorted is not None:
-            assert _eids is not None
-            _eids = _eids.view(max_loras, -1)
-            _sorted = _sorted.view(max_loras, -1)
-
-        # w2_lora_b shape[-2] is hidden_size // tp_size when fully_sharded
-        shard_size = lora_b_stacked[0].shape[-2]
-        offset = shard_size * tp_rank if fully_sharded else 0
-
-        self.add_lora_fused_moe(
-            y,
-            x,
-            lora_a_stacked,
-            lora_b_stacked,
-            topk_weights,
-            _sorted,
-            _eids,
-            num_tokens_post_padded_lora,
-            max_lora_rank,
-            top_k,
-            shrink_config,
-            expand_config,
-            adapter_enabled,
-            True,  # mul_routed_weight
-            fully_sharded=fully_sharded,
-            offset=offset,
-            token_lora_mapping=token_lora_mapping,
-        )
@@ -151,46 +151,6 @@ class SiluAndMul(CustomOp):
        return self.forward_cuda(x)


-@CustomOp.register("silu_and_mul_with_clamp")
-class SiluAndMulWithClamp(CustomOp):
-    """SwiGLU activation with input clamping (used by some MoE shared experts).
-
-    Computes:
-        gate = clamp(x[..., :d], max=swiglu_limit)
-        up   = clamp(x[..., d:], min=-swiglu_limit, max=swiglu_limit)
-        out  = silu(gate) * up
-    where d = x.shape[-1] // 2.
-
-    Shapes:
-        x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
-        return: (num_tokens, d) or (batch_size, seq_len, d)
-    """
-
-    def __init__(self, swiglu_limit: float, *, compile_native: bool = True):
-        super().__init__(compile_native=compile_native)
-        self.swiglu_limit = float(swiglu_limit)
-        if current_platform.is_cuda_alike() or current_platform.is_xpu():
-            self.op = torch.ops._C.silu_and_mul_with_clamp
-        elif current_platform.is_cpu():
-            self._forward_method = self.forward_native
-
-    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
-        d = x.shape[-1] // 2
-        gate = torch.clamp(x[..., :d], max=self.swiglu_limit)
-        up = torch.clamp(x[..., d:], min=-self.swiglu_limit, max=self.swiglu_limit)
-        return F.silu(gate) * up
-
-    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
-        d = x.shape[-1] // 2
-        output_shape = x.shape[:-1] + (d,)
-        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
-        self.op(out, x, self.swiglu_limit)
-        return out
-
-    def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
-        return self.forward_cuda(x)
-
-
 # --8<-- [start:mul_and_silu]
@CustomOp.register("mul_and_silu")
 class MulAndSilu(CustomOp):
@@ -14,6 +14,7 @@ from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
    MergedColumnParallelLinear,
 )
+from vllm.model_executor.layers.utils import cublas_gemm_bf16_bf16_fp32
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.v1.attention.backend import (
@@ -270,12 +271,16 @@ class DeepseekCompressor(nn.Module):

    def forward(
        self,
-        # [num_tokens, 2 * self.coff * self.head_dim]
-        kv_score: torch.Tensor,
+        # [num_tokens, hidden_size]
+        x: torch.Tensor,
        # [num_tokens]
        positions: torch.Tensor,
        rotary_emb,
    ) -> None:
+        num_tokens, _ = x.shape
+        # bf16 weights/activations but fp32 output for numerical stability of
+        # the downstream compressor math.
+        kv_score = cublas_gemm_bf16_bf16_fp32(x, self.fused_wkv_wgate.weight)
        # Each of shape [num_tokens, coff * self.head_dim]
        # input bf16, output are fp32
        kv, score = kv_score.split(
@@ -4,21 +4,18 @@
 DeepseekV4 MLA Attention Layer
 """

-from collections.abc import Callable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, cast

 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from transformers import DeepseekV2Config, DeepseekV3Config

-import vllm.envs as envs
 from vllm.model_executor.layers.linear import (
    ReplicatedLinear,
 )
 from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
-from vllm.model_executor.layers.utils import cublas_gemm_bf16_bf16_fp32
 from vllm.utils.deep_gemm import fp8_einsum
 from vllm.utils.torch_utils import direct_register_custom_op
 from vllm.v1.attention.ops.deepseek_v4_ops import (
@@ -54,10 +51,7 @@ from vllm.model_executor.layers.quantization.input_quant_fp8 import (
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
    GroupShape,
 )
-from vllm.utils.multi_stream_utils import (
-    execute_in_parallel,
-    maybe_execute_in_parallel,
-)
+from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
 from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
 from vllm.v1.attention.backends.mla.flashmla_sparse import (
    DeepseekV4FlashMLASparseBackend,
@@ -100,7 +94,7 @@ class DeepseekV4MLAModules:
    indexer: torch.nn.Module | None
    indexer_rotary_emb: torch.nn.Module
    topk_indices_buffer: torch.Tensor | None
-    aux_stream_list: list[torch.cuda.Stream] | None = None
+    aux_stream: torch.cuda.Stream | None = None


 # --8<-- [start:multi_head_latent_attention]
@@ -223,11 +217,8 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
            + 1  # 1B pad
        )

-        self.aux_stream_list = mla_modules.aux_stream_list
-        # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events;
-        # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins
-        # before post-GEMM starts.
-        self.ln_events = [torch.cuda.Event() for _ in range(4)]
+        self.aux_stream = mla_modules.aux_stream
+        self.ln_events = [torch.cuda.Event(), torch.cuda.Event()]

        assert cache_config is not None, "DeepseekV4 attention requires cache_config"
        self.swa_cache_layer = DeepseekV4SWACache(
@@ -286,6 +277,9 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
        hidden_states: torch.Tensor,
        llama_4_scaling: torch.Tensor | None = None,
    ) -> torch.Tensor:
+        qr_kv, _ = self.fused_wqa_wkv(hidden_states)
+        qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
+
        # Pre-allocate attention output with FlashMLA-padded head count.
        # The op writes into `o_padded`; we slice to n_local_heads after.
        num_tokens = hidden_states.shape[0]
@@ -298,6 +292,8 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
        # Attention (inside custom op for torch.compile boundary)
        torch.ops.vllm.deepseek_v4_attention(
            hidden_states,
+            qr,
+            kv,
            positions,
            o_padded,
            self.layer_name,
@@ -336,73 +332,17 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

        return self.wo_b(z.flatten(1))

-    def attn_gemm_parallel_execute(self, hidden_states) -> tuple[Any, ...]:
-        assert self.aux_stream_list is not None
-        assert len(self.aux_stream_list) >= 3
-
-        # fused_wqa_wkv (heaviest) on default; the three lighter input GEMMs
-        # on aux streams 0..2 when their owning module exists. ln_events[0]
-        # is the fan-out start event; ln_events[1..3] are per-aux done events.
-        aux_fns: list[Callable[[], Any] | None] = [None, None, None]
-
-        if self.compressor is not None:
-            # Local ref so the closure keeps a non-None type for mypy.
-            compressor = self.compressor
-
-            def compressor_kv_score() -> torch.Tensor:
-                return cublas_gemm_bf16_bf16_fp32(
-                    hidden_states, compressor.fused_wkv_wgate.weight
-                )
-
-            aux_fns[0] = compressor_kv_score
-
-        if self.indexer is not None:
-            indexer = self.indexer
-
-            def indexer_weights_proj() -> torch.Tensor:
-                # ReplicatedLinear returns (output, bias); bias is None.
-                weights, _ = indexer.weights_proj(hidden_states)
-                return weights
-
-            def indexer_compressor_kv_score() -> torch.Tensor:
-                return cublas_gemm_bf16_bf16_fp32(
-                    hidden_states, indexer.compressor.fused_wkv_wgate.weight
-                )
-
-            aux_fns[1] = indexer_weights_proj
-            aux_fns[2] = indexer_compressor_kv_score
-
-        def fused_wqa_wkv() -> torch.Tensor:
-            # MergedColumnParallelLinear returns (output, bias); bias is None.
-            qr_kv, _ = self.fused_wqa_wkv(hidden_states)
-            return qr_kv
-
-        qr_kv, (kv_score, indexer_weights, indexer_kv_score) = execute_in_parallel(
-            fused_wqa_wkv,
-            aux_fns,
-            self.ln_events[0],
-            self.ln_events[1:4],
-            self.aux_stream_list[:3],
-            enable=hidden_states.shape[0]
-            <= envs.VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD,
-        )
-
-        return qr_kv, kv_score, indexer_kv_score, indexer_weights
-
    def attention_impl(
        self,
        hidden_states: torch.Tensor,
+        qr: torch.Tensor,
+        kv: torch.Tensor,
        positions: torch.Tensor,
        out: torch.Tensor,  # [num_tokens, padded_heads, head_dim], written in place
    ) -> None:
        forward_context = get_forward_context()
        attn_metadata = forward_context.attn_metadata

-        qr_kv, kv_score, indexer_kv_score, indexer_weights = (
-            self.attn_gemm_parallel_execute(hidden_states)
-        )
-
-        qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
        qr, kv = fused_q_kv_rmsnorm(
            qr,
            kv,
@@ -410,60 +350,42 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
            self.kv_norm.weight.data,
            self.eps,
        )
+        q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)

-        # wq_b + kv_insert (+ MLA compressor when an indexer is present) ride
-        # on the default stream so q stays on its consumer stream (mla_attn
-        # downstream reads q on default). Indexer/compressor go on aux for
-        # overlap with default's GEMM + cache write.
+        # Overlap kv_insert with whichever of indexer/compressor is present.
+        # Indexer implies compressor; when both exist, compressor rides on the
+        # aux stream alongside kv_insert so the heavy indexer owns default.
        if self.indexer is not None:
-            assert self.aux_stream_list is not None
-            aux_stream = self.aux_stream_list[0]
            indexer = self.indexer
            # Local ref so the closure keeps a non-None type for mypy.
            assert self.compressor is not None
            compressor = self.compressor

-            def wq_b_kv_insert_and_compress() -> torch.Tensor:
-                q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
+            def kv_insert_and_compress() -> None:
                self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
-                compressor(kv_score, positions, self.rotary_emb)
-                return q
+                compressor(hidden_states, positions, self.rotary_emb)

-            q, _ = maybe_execute_in_parallel(
-                wq_b_kv_insert_and_compress,
-                lambda: indexer(
-                    hidden_states,
-                    qr,
-                    indexer_kv_score,
-                    indexer_weights,
-                    positions,
-                    self.indexer_rotary_emb,
+            maybe_execute_in_parallel(
+                lambda: indexer(hidden_states, qr, positions, self.indexer_rotary_emb),
+                kv_insert_and_compress,
+                self.ln_events[0],
+                self.ln_events[1],
+                self.aux_stream,
+            )
+        elif self.compressor is not None:
+            # Compressor on default, kv_insert on aux.
+            compressor = self.compressor
+            maybe_execute_in_parallel(
+                lambda: compressor(hidden_states, positions, self.rotary_emb),
+                lambda: self._fused_qnorm_rope_kv_insert(
+                    q, kv, positions, attn_metadata
                ),
                self.ln_events[0],
                self.ln_events[1],
-                aux_stream,
-            )
-        elif self.compressor is not None:
-            # wq_b + kv_insert on default, compressor on aux.
-            assert self.aux_stream_list is not None
-            aux_stream = self.aux_stream_list[0]
-            compressor = self.compressor
-
-            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)
-                return q
-
-            q, _ = maybe_execute_in_parallel(
-                wq_b_kv_insert,
-                lambda: compressor(kv_score, positions, self.rotary_emb),
-                self.ln_events[0],
-                self.ln_events[1],
-                aux_stream,
+                self.aux_stream,
            )
        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)

        # Handle dummy run (no metadata).
@@ -533,17 +455,21 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

 def deepseek_v4_attention(
    hidden_states: torch.Tensor,
+    qr: torch.Tensor,
+    kv: torch.Tensor,
    positions: torch.Tensor,
    out: torch.Tensor,
    layer_name: str,
 ) -> None:
    forward_context: ForwardContext = get_forward_context()
    self = forward_context.no_compile_layers[layer_name]
-    self.attention_impl(hidden_states, positions, out)
+    self.attention_impl(hidden_states, qr, kv, positions, out)


 def deepseek_v4_attention_fake(
    hidden_states: torch.Tensor,
+    qr: torch.Tensor,
+    kv: torch.Tensor,
    positions: torch.Tensor,
    out: torch.Tensor,
    layer_name: str,
@@ -685,7 +611,11 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
            assert cache_config is not None
            cache_config.cache_dtype = "fp8_ds_mla"
            kv_cache_dtype = "fp8_ds_mla"
-            logger.info_once("Using DeepSeek's fp8_ds_mla KV cache format.")
+            logger.info_once(
+                "Using DeepSeek's fp8_ds_mla KV cache format. To use standard "
+                "fp8 kv-cache format, please set `--attention-backend "
+                "FLASHINFER_MLA_SPARSE`"
+            )

        self.kv_cache_dtype = kv_cache_dtype

@@ -1131,20 +1061,18 @@ class DeepseekV4Indexer(nn.Module):
        self,
        hidden_states: torch.Tensor,
        qr: torch.Tensor,
-        compressed_kv_score: torch.Tensor,
-        indexer_weights: torch.Tensor,
        positions: torch.Tensor,
        rotary_emb: nn.Module,
    ) -> torch.Tensor:
-        # ReplicatedLinear returns (output, bias); bias is None.
        q, _ = self.wq_b(qr)
        q = q.view(-1, self.n_head, self.head_dim)
-        k = self.compressor(compressed_kv_score, positions, rotary_emb)
+        k = self.compressor(hidden_states, positions, rotary_emb)
+        weights, _ = self.weights_proj(hidden_states)
        q_quant, weights = fused_indexer_q_rope_quant(
            positions,
            q,
            rotary_emb.cos_sin_cache,
-            indexer_weights,
+            weights,
            self.softmax_scale,
            self.n_head**-0.5,
            use_fp4=self.use_fp4_kv,
@@ -228,37 +228,23 @@ def maybe_make_prepare_finalize(

    elif moe.use_fi_nvl_one_sided_kernels:
        assert quant_config is not None
+        if quant_config.quant_dtype != "nvfp4":
+            raise ValueError(
+                "The 'flashinfer_nvlink_one_sided' all2all backend only "
+                "supports nvfp4 activation quantization, but got "
+                f"quant_dtype={quant_config.quant_dtype!r}. Use a different "
+                "all2all backend (e.g. 'flashinfer_nvlink_two_sided' or "
+                "'allgather_reducescatter') for non-nvfp4 models."
+            )
        max_num_tokens = (
            get_current_vllm_config().scheduler_config.max_num_batched_tokens
        )
-        if quant_config.quant_dtype is None:
-            dispatch_dtype_bytes_per_elem = 2
-            dispatch_scale_bytes_per_token = 0
-        elif quant_config.quant_dtype == "nvfp4":
-            dispatch_dtype_bytes_per_elem = 0
-            dispatch_scale_bytes_per_token = moe.hidden_dim // 16
-        elif quant_config.quant_dtype == "mxfp8":
-            dispatch_dtype_bytes_per_elem = 1
-            align = quant_config.mx_alignment
-            if align > 0:
-                padded_k = ((moe.hidden_dim + align - 1) // align) * align
-            else:
-                padded_k = moe.hidden_dim
-            dispatch_scale_bytes_per_token = padded_k // 32
-        else:
-            raise NotImplementedError(
-                "flashinfer_nvlink_one_sided dispatch supports nvfp4, mxfp8, "
-                "and bf16 (quant_dtype=None) today; got "
-                f"quant_dtype={quant_config.quant_dtype!r}"
-            )
        prepare_finalize = FlashInferNVLinkOneSidedPrepareAndFinalize(
            max_num_tokens=max_num_tokens,
            top_k=moe.experts_per_token,
            num_experts=moe.num_experts,
            hidden_size=moe.hidden_dim,
            num_dispatchers=all2all_manager.world_size,
-            dispatch_dtype_bytes_per_elem=dispatch_dtype_bytes_per_elem,
-            dispatch_scale_bytes_per_token=dispatch_scale_bytes_per_token,
        )

    elif moe.use_ag_rs_all2all_kernels and allow_new_interface:
@@ -247,8 +247,6 @@ class FusedMoEQuantConfig:
    gemm1_beta: float | None = None
    gemm1_clamp_limit: float | None = None

-    mx_alignment: int = 0
-
    def __post_init__(self):
        assert not self.per_act_token_quant or self.block_shape is None, (
            "illegal quantization"
@@ -707,7 +705,6 @@ def mxfp4_mxfp8_moe_quant_config(
    gemm1_alpha: float | None = None,
    gemm1_beta: float | None = None,
    gemm1_clamp_limit: float | None = None,
-    mx_alignment: int = 0,
 ) -> FusedMoEQuantConfig:
    """
    Construct a quant config for mxfp4 activations and mxfp4 weights.
@@ -720,7 +717,6 @@ def mxfp4_mxfp8_moe_quant_config(
        gemm1_alpha=gemm1_alpha,
        gemm1_beta=gemm1_beta,
        gemm1_clamp_limit=gemm1_clamp_limit,
-        mx_alignment=mx_alignment,
    )


@@ -45,7 +45,7 @@ def _gelu_and_mul(
 # Uses static methods or standalone functions to avoid instantiating CustomOp
 # classes, which would call get_current_vllm_config() before config is set.
 _CPU_MOE_ACT_FN: dict[MoEActivation, Callable[[torch.Tensor], torch.Tensor]] = {
-    MoEActivation.SILU: lambda x: SiluAndMul(compile_native=False).forward_native(x),
+    MoEActivation.SILU: SiluAndMul.forward_native,
    MoEActivation.SWIGLUOAI: _swigluoai_forward_native,
    MoEActivation.GELU: _gelu_and_mul,
 }
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+
 import torch

 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -15,7 +16,6 @@ from vllm.model_executor.layers.fused_moe.config import (
    FusedMoEQuantConfig,
    RoutingMethodType,
 )
-from vllm.model_executor.layers.fused_moe.lora_experts_mixin import LoRAExpertsMixin
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
    TopKWeightAndReduceNoOP,
 )
@@ -812,7 +812,7 @@ class OAITritonExperts(BaseOAITritonExperts):
        )


-class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts):
+class UnfusedOAITritonExperts(BaseOAITritonExperts):
    """
    A Triton based MoE expert class that operates on expert standard
    format and explicitly keeps the activation and reduction (moe_sum) steps
@@ -910,7 +910,6 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts):
        if quant_config is None:
            quant_config = FUSED_MOE_UNQUANTIZED_CONFIG

-        global_topk_ids = topk_ids
        if expert_map is not None:
            topk_ids = expert_map[topk_ids]

@@ -965,40 +964,10 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts):
            y=intermediate_cache1,
        )

-        # w13 LoRA: gather the activation input from expert-sorted
-        # intermediate_cache1, then add the LoRA delta in-place on that copy
-        # before passing it to activation — exactly mirroring the old
-        # decorator approach which modified the gathered tensor in-place.
-        act_input = intermediate_cache1.view(-1, N)[gather_indx.dst_indx]
-
-        sorted_token_ids_lora = None
-        expert_ids_lora = None
-        num_tokens_post_padded_lora = None
-        token_lora_mapping = None
-        lora_context = self._lora_context
-        if lora_context is not None:
-            (
-                sorted_token_ids_lora,
-                expert_ids_lora,
-                num_tokens_post_padded_lora,
-                token_lora_mapping,
-            ) = self.apply_w13_lora(
-                lora_context,
-                y=act_input,
-                x=hidden_states,
-                topk_ids=global_topk_ids,
-                topk_weights=topk_weights,
-                expert_map=expert_map,
-                w1=w1,
-                w2=w2,
-                num_tokens=M,
-                top_k_num=topk,
-            )
-
        self.activation(
            activation,
            intermediate_cache2,
-            act_input,
+            intermediate_cache1.view(-1, N)[gather_indx.dst_indx],
        )

        # matmul_ogs grouped reduction fuses sum across multiple experts:
@@ -1017,24 +986,6 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts):
            y=intermediate_cache3,
        )

-        # w2 LoRA: after matmul_ogs with scatter_indx, intermediate_cache3 is
-        # in token-topk order, matching the (M, topk, K) layout add_lora_w2 expects.
-        if lora_context is not None:
-            self.apply_w2_lora(
-                lora_context,
-                y=intermediate_cache3.view(-1, topk, K),
-                x=intermediate_cache2,
-                topk_weights=topk_weights,
-                sorted_token_ids_lora=sorted_token_ids_lora,
-                expert_ids_lora=expert_ids_lora,
-                num_tokens_post_padded_lora=num_tokens_post_padded_lora,
-                token_lora_mapping=token_lora_mapping,
-                num_tokens=M,
-                w1=w1,
-                w2=w2,
-                top_k_num=topk,
-            )
-
        self.moe_sum(intermediate_cache3.view(-1, topk, K), output)


@@ -44,9 +44,6 @@ class TrtLlmMxfp4ExpertsBase:
            moe_config.intermediate_size_per_partition
        )
        self.hidden_dim = moe_config.hidden_dim
-        self.hidden_dim_unpadded = (
-            moe_config.hidden_dim_unpadded or moe_config.hidden_dim
-        )
        self.local_num_experts = moe_config.num_local_experts
        self.ep_rank = moe_config.moe_parallel_config.ep_rank

@@ -85,6 +82,9 @@ class TrtLlmMxfp4ExpertsBase:
            get_current_vllm_config().compilation_config.max_cudagraph_capture_size
        )

+        # P1-5 fix: use public quant_dtype property instead of private _a1
+        self.use_mxfp8_input = quant_config.quant_dtype == "mxfp8"
+
    @staticmethod
    def _supports_current_device() -> bool:
        p = current_platform
@@ -121,7 +121,8 @@ class TrtLlmMxfp4ExpertsBase:

    @property
    def expects_unquantized_inputs(self) -> bool:
-        return False
+        # Expert handles MXFP8 quantization internally if needed
+        return True


 class TrtLlmMxfp4ExpertsMonolithic(
@@ -180,19 +181,24 @@ class TrtLlmMxfp4ExpertsMonolithic(
    ) -> torch.Tensor:
        from flashinfer import trtllm_fp4_block_scale_moe

-        if a1q_scale is not None:
-            x_quant = hidden_states
-            x_scale = a1q_scale.view(torch.float8_e4m3fn)
+        # Handle input quantization
+        if self.use_mxfp8_input:
+            from flashinfer import mxfp8_quantize
+
+            x_quant, x_scale = mxfp8_quantize(
+                hidden_states,
+                is_sf_swizzled_layout=False,
+                alignment=256,
+            )
+            x_scale = x_scale.view(torch.float8_e4m3fn).reshape(
+                *hidden_states.shape[:-1], -1
+            )
        else:
            assert hidden_states.dtype == torch.bfloat16
            x_quant = hidden_states
            x_scale = None
-        output = torch.empty(
-            *hidden_states.shape[:-1],
-            self.hidden_dim_unpadded,
-            dtype=torch.bfloat16,
-            device=hidden_states.device,
-        )
+
+        output = torch.empty_like(hidden_states)

        from vllm.utils.flashinfer import _is_fi_autotuning, autotune

@@ -238,6 +244,10 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
    Moved from trtllm_moe.py.
    """

+    @property
+    def expects_unquantized_inputs(self) -> bool:
+        return True
+
    @staticmethod
    def _supports_parallel_config(
        moe_parallel_config: FusedMoEParallelConfig,
@@ -274,7 +284,7 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
        # The workspaces for this implementation are managed by flashinfer.
        workspace1 = (0,)
        workspace2 = (0,)
-        output = (M, self.hidden_dim_unpadded)
+        output = (M, K)
        return (workspace1, workspace2, output)

    def apply(
@@ -300,9 +310,18 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
        intermediate_size = self.intermediate_size_per_partition
        local_expert_offset = self.moe_config.ep_rank * local_num_experts

-        if a1q_scale is not None:
-            x_quant = hidden_states
-            x_scale = a1q_scale.view(torch.float8_e4m3fn)
+        # Handle input quantization
+        if self.use_mxfp8_input:
+            from flashinfer import mxfp8_quantize
+
+            x_quant, x_scale = mxfp8_quantize(
+                hidden_states,
+                is_sf_swizzled_layout=False,
+                alignment=256,
+            )
+            x_scale = x_scale.view(torch.float8_e4m3fn).reshape(
+                *hidden_states.shape[:-1], -1
+            )
        else:
            assert hidden_states.dtype == torch.bfloat16
            x_quant = hidden_states
@@ -17,7 +17,6 @@ from vllm.model_executor.layers.fused_moe.config import (
    FusedMoEParallelConfig,
    FusedMoEQuantConfig,
 )
-from vllm.model_executor.layers.fused_moe.lora_experts_mixin import LoRAExpertsMixin
 from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
    batched_moe_align_block_size,
    moe_align_block_size,
@@ -669,7 +668,7 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
        return E, M, N, K, topk


-class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
+class MarlinExperts(MarlinExpertsBase):
    """Marlin-based fused MoE expert implementation."""

    def supports_expert_map(self) -> bool:
@@ -734,108 +733,7 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
    ):
        assert self.w1_scale is not None
        assert self.w2_scale is not None
-
-        ctx = self._lora_context
-        if ctx is None:
-            fused_marlin_moe(
-                hidden_states=hidden_states,
-                w1=w1,
-                w2=w2,
-                bias1=self.w1_bias,
-                bias2=self.w2_bias,
-                w1_scale=self.w1_scale,
-                w2_scale=self.w2_scale,
-                topk_weights=topk_weights,
-                topk_ids=topk_ids,
-                global_scale1=self.g1_alphas,
-                global_scale2=self.g2_alphas,
-                quant_type_id=self.quant_type_id,
-                apply_router_weight_on_input=apply_router_weight_on_input,
-                global_num_experts=global_num_experts,
-                activation=activation,
-                activation_func=self.activation,
-                moe_sum=self.moe_sum,
-                expert_map=expert_map,
-                output=output,
-                # Workspaces are swapped in workspace_shapes() to account for proper
-                # output buffer allocation. Please refer to workspace_shapes().
-                intermediate_cache13=workspace2,
-                intermediate_cache2=workspace13,
-                g_idx1=self.w13_g_idx,
-                g_idx2=self.w2_g_idx,
-                sort_indices1=self.w13_g_idx_sort_indices,
-                sort_indices2=self.w2_g_idx_sort_indices,
-                is_k_full=self.is_k_full,
-                input_dtype=self.input_dtype,
-            )
-            return
-
-        # LoRA path: wrap activation_func and moe_sum to inject LoRA at the
-        # two natural injection points.
-        #
-        # Marlin uses moe_align_block_size (same as TritonExperts) so
-        # intermediate_cache1 is indexed by flat (token, expert) pair index,
-        # which is compatible with add_lora_fused_moe's scatter mechanism.
-
-        M = hidden_states.size(0)
-        top_k_num = topk_ids.size(1)
-        lora_state: dict = {}
-
-        def activation_with_lora(
-            act_enum: MoEActivation,
-            act_output: torch.Tensor,
-            act_input: torch.Tensor,
-        ) -> None:
-            # act_input  = intermediate_cache1 (M*topk, 2N for gated)
-            # act_output = intermediate_cache2 (M*topk, N)
-
-            (
-                sorted_token_ids_lora,
-                expert_ids_lora,
-                num_tokens_post_padded_lora,
-                token_lora_mapping,
-            ) = self.apply_w13_lora(
-                ctx,
-                y=act_input,
-                x=hidden_states,
-                topk_ids=topk_ids,
-                topk_weights=topk_weights,
-                expert_map=expert_map,
-                w1=w1,
-                w2=w2,
-                num_tokens=M,
-                top_k_num=top_k_num,
-            )
-            lora_state.update(
-                {
-                    "sorted": sorted_token_ids_lora,
-                    "eids": expert_ids_lora,
-                    "npad": num_tokens_post_padded_lora,
-                    "tlm": token_lora_mapping,
-                }
-            )
-            self.activation(act_enum, act_output, act_input)
-            lora_state["cache2"] = act_output
-
-        def moe_sum_with_lora(moe_out: torch.Tensor, out: torch.Tensor) -> None:
-            # moe_out shape: (M, topk, K)
-            self.apply_w2_lora(
-                ctx,
-                y=moe_out,
-                x=lora_state["cache2"],
-                topk_weights=topk_weights,
-                sorted_token_ids_lora=lora_state["sorted"],
-                expert_ids_lora=lora_state["eids"],
-                num_tokens_post_padded_lora=lora_state["npad"],
-                token_lora_mapping=lora_state["tlm"],
-                num_tokens=M,
-                w1=w1,
-                w2=w2,
-                top_k_num=top_k_num,
-            )
-            self.moe_sum(moe_out, out)
-
-        return fused_marlin_moe(
+        fused_marlin_moe(
            hidden_states=hidden_states,
            w1=w1,
            w2=w2,
@@ -851,10 +749,12 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
            apply_router_weight_on_input=apply_router_weight_on_input,
            global_num_experts=global_num_experts,
            activation=activation,
-            activation_func=activation_with_lora,
-            moe_sum=moe_sum_with_lora,
+            activation_func=self.activation,
+            moe_sum=self.moe_sum,
            expert_map=expert_map,
            output=output,
+            # Workspaces are swapped in workspace_shapes() to account for proper
+            # output buffer allocation. Please refer to workspace_shapes().
            intermediate_cache13=workspace2,
            intermediate_cache2=workspace13,
            g_idx1=self.w13_g_idx,
@@ -25,7 +25,6 @@ from vllm.model_executor.layers.fused_moe.config import (
    FusedMoEQuantConfig,
    _get_config_dtype_str,
 )
-from vllm.model_executor.layers.fused_moe.lora_experts_mixin import LoRAExpertsMixin
 from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
    moe_align_block_size,
 )
@@ -1887,7 +1886,7 @@ def fused_experts_impl(
    return out_hidden_states


-class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
+class TritonExperts(mk.FusedMoEExpertsModular):
    """Triton-based fused MoE expert implementation."""

    def __init__(
@@ -2095,33 +2094,6 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
            B_bias=self.w1_bias,
        )

-        # LoRA w13: applied to intermediate_cache1 before activation, using
-        # hidden_states as the lora_a input.  moe_lora_align_block_size is
-        # called once here and results reused for the w2 LoRA below.
-        sorted_token_ids_lora = None
-        expert_ids_lora = None
-        num_tokens_post_padded_lora = None
-        token_lora_mapping = None
-        lora_context = self._lora_context
-        if lora_context is not None:
-            (
-                sorted_token_ids_lora,
-                expert_ids_lora,
-                num_tokens_post_padded_lora,
-                token_lora_mapping,
-            ) = self.apply_w13_lora(
-                lora_context,
-                y=intermediate_cache1,
-                x=hidden_states,
-                topk_ids=topk_ids,
-                topk_weights=topk_weights,
-                expert_map=expert_map,
-                w1=w1,
-                w2=w2,
-                num_tokens=num_tokens,
-                top_k_num=top_k_num,
-            )
-
        self.activation(
            activation, intermediate_cache2, intermediate_cache1.view(-1, N)
        )
@@ -2160,25 +2132,6 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
            B_bias=self.w2_bias,
        )

-        # LoRA w2: applied to intermediate_cache3 before moe_sum, using the
-        # unquantized intermediate_cache2 as the lora_a input.  Reuses the
-        # sorted_token_ids_lora computed above.
-        if lora_context is not None:
-            self.apply_w2_lora(
-                lora_context,
-                y=intermediate_cache3,
-                x=intermediate_cache2,
-                topk_weights=topk_weights,
-                sorted_token_ids_lora=sorted_token_ids_lora,
-                expert_ids_lora=expert_ids_lora,
-                num_tokens_post_padded_lora=num_tokens_post_padded_lora,
-                token_lora_mapping=token_lora_mapping,
-                num_tokens=num_tokens,
-                w1=w1,
-                w2=w2,
-                top_k_num=top_k_num,
-            )
-
        # separate function is required for MoE + LoRA
        self.moe_sum(intermediate_cache3, output)

@@ -1,44 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from dataclasses import dataclass
-
-import torch
-
-from vllm.lora.punica_wrapper.punica_base import PunicaWrapperBase
-
-
-@dataclass
-class MoELoRAContext:
-    """
-    Carries all LoRA state for one MoE forward pass.
-
-    Built by FusedMoEWithLoRA.forward() and propagated explicitly through the
-    modular kernel path (FusedMoEKernel -> FusedMoEExpertsModular.apply) so
-    that TritonExperts.apply() can compute the LoRA contribution inline,
-    replacing the decorator-based monkey-patch approach.
-    """
-
-    # LoRA weight tensors (same shapes as FusedMoEWithLoRA attributes)
-    w13_lora_a_stacked: tuple[torch.Tensor, ...]
-    w13_lora_b_stacked: tuple[torch.Tensor, ...]
-    w2_lora_a_stacked: tuple[torch.Tensor, ...]
-    w2_lora_b_stacked: tuple[torch.Tensor, ...]
-
-    # (max_loras + 1,) int32; slot 0 is the "no-adapter" sentinel
-    adapter_enabled: torch.Tensor
-
-    # Metadata
-    max_loras: int
-    top_k: int
-    w13_num_slices: int  # 2 = gated (gate + up), 1 = non-gated or 3D-fused
-    fully_sharded: bool
-    tp_rank: int
-    tp_size: int
-    local_num_experts: int
-
-    punica_wrapper: PunicaWrapperBase
-
-    # Whether VLLM_TUNED_CONFIG_FOLDER is set; selects get_lora_op_configs vs
-    # try_get_optimal_moe_lora_config for Triton kernel tile configs.
-    use_tuned_config: bool
@@ -1,111 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import torch
-
-from vllm.model_executor.layers.fused_moe.lora_context import MoELoRAContext
-
-
-class LoRAExpertsMixin:
-    """
-    Mixin for FusedMoEExpertsModular subclasses that natively handle
-    MoELoRAContext inside their apply() implementation.
-
-    Mixing this class in:
-    - Flips supports_lora() to True so _can_fused_experts_support lets
-      LoRA through the gate check.
-    - Stashes a MoELoRAContext on the experts instance via
-      set_lora_context(), which apply() consumes from self._lora_context.
-    - Provides apply_w13_lora / apply_w2_lora helpers that dispatch to
-      the PunicaWrapper kernels.
-
-    The helper methods are pure functions of their inputs; all required
-    state is on lora_context or passed as arguments.
-    """
-
-    _lora_context: MoELoRAContext | None = None
-
-    def set_lora_context(self, ctx: MoELoRAContext) -> None:
-        self._lora_context = ctx
-
-    @staticmethod
-    def supports_lora() -> bool:
-        return True
-
-    def apply_w13_lora(
-        self,
-        lora_context: MoELoRAContext,
-        *,
-        y: torch.Tensor,
-        x: torch.Tensor,
-        topk_ids: torch.Tensor,
-        topk_weights: torch.Tensor,
-        expert_map: torch.Tensor | None,
-        w1: torch.Tensor,
-        w2: torch.Tensor,
-        num_tokens: int,
-        top_k_num: int,
-    ) -> tuple[
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor | None,
-    ]:
-        return lora_context.punica_wrapper.add_lora_w13(
-            y,
-            x,
-            lora_context.w13_lora_a_stacked,
-            lora_context.w13_lora_b_stacked,
-            topk_ids,
-            topk_weights,
-            expert_map,
-            w1,
-            w2,
-            num_tokens,
-            top_k_num,
-            lora_context.max_loras,
-            lora_context.adapter_enabled,
-            lora_context.local_num_experts,
-            lora_context.top_k,
-            lora_context.w13_num_slices,
-            lora_context.fully_sharded,
-            lora_context.use_tuned_config,
-        )
-
-    def apply_w2_lora(
-        self,
-        lora_context: MoELoRAContext,
-        *,
-        y: torch.Tensor,
-        x: torch.Tensor,
-        topk_weights: torch.Tensor,
-        sorted_token_ids_lora: torch.Tensor | None,
-        expert_ids_lora: torch.Tensor | None,
-        num_tokens_post_padded_lora: torch.Tensor | None,
-        token_lora_mapping: torch.Tensor | None,
-        num_tokens: int,
-        w1: torch.Tensor,
-        w2: torch.Tensor,
-        top_k_num: int,
-    ) -> None:
-        lora_context.punica_wrapper.add_lora_w2(
-            y,
-            x,
-            lora_context.w2_lora_a_stacked,
-            lora_context.w2_lora_b_stacked,
-            topk_weights,
-            sorted_token_ids_lora,
-            expert_ids_lora,
-            num_tokens_post_padded_lora,
-            token_lora_mapping,
-            num_tokens,
-            w1,
-            w2,
-            top_k_num,
-            lora_context.max_loras,
-            lora_context.adapter_enabled,
-            lora_context.top_k,
-            lora_context.fully_sharded,
-            lora_context.tp_rank,
-            lora_context.use_tuned_config,
-        )
@@ -570,8 +570,6 @@ class FusedMoEExperts(ABC):
            return False, _make_reason(f"{activation_format.value} activation format")
        elif envs.VLLM_BATCH_INVARIANT and not cls._supports_batch_invariance():
            return False, _make_reason("batch invariance")
-        elif moe_config.is_lora_enabled and not cls.supports_lora():
-            return False, _make_reason("LoRA")
        return True, None

    @staticmethod
@@ -736,15 +734,6 @@ class FusedMoEExperts(ABC):
    def g2_alphas(self) -> torch.Tensor | None:
        return self.quant_config.g2_alphas

-    @staticmethod
-    def supports_lora() -> bool:
-        """Return True if this expert impl natively handles LoRA.
-
-        LoRA-aware experts should mix in LoRAExpertsMixin, which flips this
-        to True and provides the per-forward LoRA state plumbing.
-        """
-        return False
-
    @abstractmethod
    def supports_expert_map(self) -> bool:
        """
@@ -1538,9 +1527,6 @@ class FusedMoEKernel:
    def fused_experts(self) -> FusedMoEExperts:
        return self.impl.fused_experts

-    def supports_lora(self) -> bool:
-        return self.fused_experts.supports_lora()
-
    def _post_init_setup(self):
        """
        Resolve any leftover setup dependencies between self.prepare_finalize
@@ -220,6 +220,9 @@ def select_fp8_moe_backend(
    Note: Shape-specific fallbacks may still occur at runtime.
    """

+    if config.is_lora_enabled:
+        return Fp8MoeBackend.TRITON, backend_to_kernel_cls(Fp8MoeBackend.TRITON)[0]
+
    # NOTE: the kernels are selected in the following order.
    AVAILABLE_BACKENDS = _get_priority_backends(config, weight_key, activation_key)

@@ -1195,18 +1195,10 @@ def make_mxfp4_moe_quant_config(
            gemm1_beta=gemm1_beta,
            gemm1_clamp_limit=swiglu_limit,
        )
-    elif mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8:
-        return mxfp4_mxfp8_moe_quant_config(
-            w1_bias=w1_bias,
-            w2_bias=w2_bias,
-            w1_scale=w1_scale,
-            w2_scale=w2_scale,
-            gemm1_alpha=gemm1_alpha,
-            gemm1_beta=gemm1_beta,
-            gemm1_clamp_limit=swiglu_limit,
-            mx_alignment=256,
-        )
-    elif mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8:
+    elif mxfp4_backend in (
+        Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
+        Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
+    ):
        return mxfp4_mxfp8_moe_quant_config(
            w1_bias=w1_bias,
            w2_bias=w2_bias,
@@ -1258,6 +1250,7 @@ def make_mxfp4_moe_kernel(
    """Create a FusedMoEKernel for the given MXFP4 backend."""
    is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic)

+    # Create Prepare/Finalize.
    prepare_finalize = maybe_make_prepare_finalize(
        moe=moe_config,
        quant_config=moe_quant_config,
@@ -214,6 +214,19 @@ def select_unquantized_moe_backend(
            return backend, k_cls
        raise ValueError(_make_log_unsupported(backend, reason))

+    # LoRA needs Triton's unfused activation/reduction hooks. Selecting the
+    # backend here ensures weights stay in a LoRA-compatible layout instead of
+    # being permuted for a backend like FlashInfer or AITER during load.
+    if moe_config.is_lora_enabled:
+        backend = UnquantizedMoeBackend.TRITON
+        if activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
+            backend = UnquantizedMoeBackend.BATCHED_TRITON
+        return _return_or_raise(
+            backend,
+            moe_config,
+            activation_format,
+        )
+
    runner_backend = moe_config.moe_backend
    if runner_backend != "auto":
        requested_backend = map_unquantized_backend(runner_backend)
@@ -31,8 +31,6 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
        num_experts: int,
        hidden_size: int,
        num_dispatchers: int = 1,
-        dispatch_dtype_bytes_per_elem: int = 0,
-        dispatch_scale_bytes_per_token: int = 0,
    ):
        super().__init__()
        self.max_num_tokens = max_num_tokens
@@ -40,7 +38,6 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
        self.num_experts = num_experts
        self.hidden_size = hidden_size
        self.num_dispatchers_ = num_dispatchers
-        self.scale_elems_per_token = dispatch_scale_bytes_per_token

        device_communicator = get_ep_group().device_communicator
        assert device_communicator is not None
@@ -52,8 +49,6 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
            top_k=self.top_k,
            num_experts=self.num_experts,
            hidden_size=self.hidden_size,
-            dispatch_dtype_bytes_per_elem=dispatch_dtype_bytes_per_elem,
-            dispatch_scale_bytes_per_token=dispatch_scale_bytes_per_token,
        )

    @property
@@ -97,24 +92,19 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
            else a1.shape[0]
        )

-        if defer_input_quant:
-            a1q, a1q_scale = a1, None
-        else:
-            a1q, a1q_scale = moe_kernel_quantize_input(
-                a1,
-                quant_config.a1_gscale,
-                quant_config.quant_dtype,
-                quant_config.per_act_token_quant,
-                quant_config.block_shape,
-                is_fp4_scale_swizzled=False,  # delay swizzle to after comm
-                mx_alignment=quant_config.mx_alignment,
-            )
+        a1q, a1q_scale = moe_kernel_quantize_input(
+            a1,
+            quant_config.a1_gscale,
+            quant_config.quant_dtype,
+            quant_config.per_act_token_quant,
+            quant_config.block_shape,
+            is_fp4_scale_swizzled=False,  # delay swizzle to after comm
+        )

        payloads = []
        payloads.append(a1q)
        if a1q_scale is not None:
            payloads.append(a1q_scale)
-        topk_ids_payload_index = len(payloads)
        payloads.append(topk_ids)
        payloads.append(topk_weights)

@@ -123,8 +113,6 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
            token_selected_experts=topk_ids,
            input_payloads=payloads,
            runtime_max_tokens_per_rank=self.runtime_max_tokens_per_rank,
-            invalid_token_expert_id=-1,  # Follow TRTLLM Pattern
-            expert_id_payload_index=topk_ids_payload_index,
        )
        if a1q_scale is not None:
            a1q_recv, a1q_scale_recv, topk_ids_recv, topk_weights_recv = recv_payloads
@@ -136,8 +124,7 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
                a1q_scale_recv = a1q_scale_recv.view(-1, a1q_scale_recv.shape[-1])
                a1q_scale_recv = a1q_scale_recv.view(torch.uint8)
                a1q_scale_recv = nvfp4_block_scale_interleave(a1q_scale_recv)
-            assert self.scale_elems_per_token > 0
-            a1q_scale_recv = a1q_scale_recv.view(-1, self.scale_elems_per_token)
+            a1q_scale_recv = a1q_scale_recv.view(-1, self.hidden_size // 16)
        else:
            a1q_recv, topk_ids_recv, topk_weights_recv = recv_payloads
            a1q_scale_recv = None
@@ -174,7 +174,6 @@ def flashinfer_alltoall_dispatch(
            # the hidden states, breaking the A2A kernel. So, we
            # delay the swizzling until after the A2A.
            is_fp4_scale_swizzled=False,
-            mx_alignment=quant_config.mx_alignment,
        )

        x = MnnvlMoe.mnnvl_moe_alltoallv(
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Claude	8a77e75e01	fix: add platform guards for tilelang in mhc.py and fix test imports - Add platform guards for tilelang imports in mhc.py to prevent import failures on non-CUDA platforms (CPU, AMD, etc.) - Move tilelang kernel definitions to lazy initialization - Fix import order in test_deepseek_v4_mega_moe.py (ruff compliance) - Add pytest skip marker for non-CUDA platforms in test file This fixes the widespread CPU test failures caused by unconditional tilelang imports at module level. https://claude.ai/code/session_015qZTB3eveFJ8qWsSupPgX1 Signed-off-by: Claude <noreply@anthropic.com>	2026-04-25 17:29:00 +00:00
Woosuk KwonandYifan Qiao	01c6528a37	Integrate MegaMoE kernel (#232 ) Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-04-25 07:51:57 +00:00
Yifan Qiao	e72446941d	fix: config Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>	2026-04-25 07:51:57 +00:00
Yifan Qiao	25e3698e8c	fix: update cuda requirements Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>	2026-04-25 07:51:56 +00:00
Yifan Qiao	05af254a20	chore: pass mypy Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>	2026-04-25 06:18:00 +00:00
+6 Yongye Zhu Yifan Qiao Yongye Zhu Simon Mo Bugen Zhao Giancarlo Delfin Jee Jee Li Nick Hill Roger Wang Roy Wang Woosuk Kwon youkaichao Zhewen Li Zijing Liu khluu qizixi	a4668d8dc9	feat: support deepseek v4 Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Yongye Zhu <yongye@inferact.ai> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Simon Mo <simon@inferact.ai> Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Giancarlo Delfin <gdelfin@inferact.ai> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Roy Wang <yasong.wang@inferact.ai> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: youkaichao <youkaichao@gmail.com> Co-authored-by: Zhewen Li <jerven.vllm@gmail.com> Co-authored-by: Zijing Liu <liuzijing2014@gmail.com> Co-authored-by: khluu <khluu000@gmail.com> Co-authored-by: qizixi <zixi@inferact.ai>	2026-04-25 06:17:51 +00:00