Sparse MLA on Hopper: Use SGLang's kernel for the sparse mla low latency runs

Signed-off-by: Alexander Matveev <amatveev@redhat.com>
2026-04-17 16:51:32 +00:00
123 changed files with 5564 additions and 4501 deletions
@@ -92,8 +92,8 @@ check_and_skip_if_image_exists() {
 }

 ecr_login() {
-    aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
-    aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true
+    aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
+    aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com
 }

 prepare_cache_tags() {
@@ -11,7 +11,7 @@ REPO=$2
 BUILDKITE_COMMIT=$3

 # authenticate with AWS ECR
-aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
+aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"

 # skip build if image already exists
 if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu) ]]; then
@@ -11,7 +11,7 @@ REPO=$2
 BUILDKITE_COMMIT=$3

 # authenticate with AWS ECR
-aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
+aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"

 # skip build if image already exists
 if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64-cpu) ]]; then
@@ -2613,7 +2613,6 @@ steps:
  - vllm/platforms/rocm.py
  commands:
  - export TORCH_NCCL_BLOCKING_WAIT=1
-  - VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
  - pytest -v -s tests/v1/distributed/test_dbo.py


@@ -3602,6 +3601,7 @@ steps:
  commands:
  - export TORCH_NCCL_BLOCKING_WAIT=1
  - pytest -v -s tests/distributed/test_context_parallel.py
+  - VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
  - pytest -v -s tests/v1/distributed/test_dbo.py


@@ -141,7 +141,6 @@ steps:
    - pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py
    - pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py
    - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
-    - pytest -v -s tests/kernels/moe/test_mxfp4_moe.py
    - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
    - pytest -v -s tests/kernels/moe/test_flashinfer.py
    - pytest -v -s tests/kernels/moe/test_flashinfer_moe.py
@@ -952,9 +952,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
      "csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
      "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
      "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
-      "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
-      "csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu"
-      "csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu")
+      "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
    set_gencode_flags_for_srcs(
      SRCS "${SRCS}"
      CUDA_ARCHS "${FP4_ARCHS}")
@@ -1242,6 +1240,7 @@ endif()
 if (VLLM_GPU_LANG STREQUAL "CUDA")
    include(cmake/external_projects/deepgemm.cmake)
    include(cmake/external_projects/flashmla.cmake)
+    include(cmake/external_projects/cutlass_fa3.cmake)
    include(cmake/external_projects/qutlass.cmake)

    # vllm-flash-attn should be last as it overwrites some CMake functions
@@ -0,0 +1,163 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# CUTLASS FA3 MLA Sparse Attention — requires CUDA >= 12.4, SM90a
+#
+# Vendors the sgl-attn CUTLASS FlashAttention3 kernel from SGLang into vLLM
+# as a self-contained extension (_cutlass_fa3_C). This provides a high-
+# performance sparse MLA attention kernel for SM90 (Hopper) GPUs.
+#
+# Source: https://github.com/sgl-project/sgl-attn (commit bcf72ccc)
+# CUTLASS: https://github.com/NVIDIA/cutlass (commit 57e3cfb4)
+
+# Guard: CUDA >= 12.4 required for SM90a features used by FA3
+if(NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "12.4")
+  message(STATUS "Skipping CUTLASS FA3: requires CUDA >= 12.4")
+  # Create empty target so setup.py doesn't fail on unsupported systems
+  add_custom_target(_cutlass_fa3_C)
+  return()
+endif()
+
+# Guard: SM90 architecture required
+set(CUTLASS_FA3_SUPPORT_ARCHS)
+list(APPEND CUTLASS_FA3_SUPPORT_ARCHS "9.0a")
+cuda_archs_loose_intersection(
+  CUTLASS_FA3_ARCHS "${CUTLASS_FA3_SUPPORT_ARCHS}" "${CUDA_ARCHS}")
+if(NOT CUTLASS_FA3_ARCHS)
+  message(STATUS "Skipping CUTLASS FA3: requires SM90 (CUDA_ARCHS=${CUDA_ARCHS})")
+  add_custom_target(_cutlass_fa3_C)
+  return()
+endif()
+
+include(FetchContent)
+
+# Fetch sgl-attn (Flash Attention 3 kernels from SGLang)
+# We only need the source files, not the build system, so we use
+# FetchContent_Populate to download without building.
+if (DEFINED ENV{SGL_ATTN_SRC_DIR})
+  set(SGL_ATTN_SRC_DIR $ENV{SGL_ATTN_SRC_DIR})
+endif()
+if(SGL_ATTN_SRC_DIR)
+  FetchContent_Declare(cutlass_fa3
+    SOURCE_DIR ${SGL_ATTN_SRC_DIR})
+else()
+  FetchContent_Declare(cutlass_fa3
+    GIT_REPOSITORY https://github.com/sgl-project/sgl-attn.git
+    GIT_TAG bcf72ccc6816b36a5fae2c5a3c027604629785e0
+    GIT_PROGRESS TRUE
+    GIT_SHALLOW FALSE)
+endif()
+FetchContent_GetProperties(cutlass_fa3)
+if(NOT cutlass_fa3_POPULATED)
+  FetchContent_Populate(cutlass_fa3)
+endif()
+message(STATUS "CUTLASS FA3 sgl-attn source: ${cutlass_fa3_SOURCE_DIR}")
+
+# Fetch CUTLASS for FA3 (headers only, separate from vLLM's main CUTLASS
+# to avoid version conflicts). Use FetchContent_Populate to avoid running
+# CUTLASS's own CMakeLists.txt which would create conflicting targets.
+if (DEFINED ENV{CUTLASS_FA3_CUTLASS_SRC_DIR})
+  set(CUTLASS_FA3_CUTLASS_SRC_DIR $ENV{CUTLASS_FA3_CUTLASS_SRC_DIR})
+endif()
+if(CUTLASS_FA3_CUTLASS_SRC_DIR)
+  FetchContent_Declare(cutlass_for_fa3
+    SOURCE_DIR ${CUTLASS_FA3_CUTLASS_SRC_DIR})
+else()
+  FetchContent_Declare(cutlass_for_fa3
+    GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git
+    GIT_TAG 57e3cfb47a2d9e0d46eb6335c3dc411498efa198
+    GIT_PROGRESS TRUE
+    GIT_SHALLOW FALSE)
+endif()
+FetchContent_GetProperties(cutlass_for_fa3)
+if(NOT cutlass_for_fa3_POPULATED)
+  FetchContent_Populate(cutlass_for_fa3)
+endif()
+message(STATUS "CUTLASS FA3 cutlass source: ${cutlass_for_fa3_SOURCE_DIR}")
+
+set(FA3_SRC "${cutlass_fa3_SOURCE_DIR}/hopper")
+
+# flash_api.cpp dispatches to all head dimensions + dtypes (BF16, FP16, FP8)
+# at compile time. With FLASHATTENTION_DISABLE_SM8x, only SM90 instantiations
+# are needed. We exclude hdimall_* (fails on CUDA 13+) and backward files.
+file(GLOB FA3_INSTANTIATION_SOURCES
+  # BF16 instantiations
+  "${FA3_SRC}/instantiations/flash_fwd_hdim64_bf16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim96_bf16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim128_bf16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim192_bf16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim256_bf16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdimdiff_bf16*_sm90.cu"
+  # FP16 instantiations
+  "${FA3_SRC}/instantiations/flash_fwd_hdim64_fp16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim96_fp16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim128_fp16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim192_fp16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim256_fp16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdimdiff_fp16*_sm90.cu"
+  # FP8 (e4m3) instantiations
+  "${FA3_SRC}/instantiations/flash_fwd_hdim64_e4m3*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim96_e4m3*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim128_e4m3*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim192_e4m3*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim256_e4m3*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdimdiff_e4m3*_sm90.cu")
+
+set(FA3_CORE_SOURCES
+  "${FA3_SRC}/flash_api.cpp"
+  "${FA3_SRC}/flash_prepare_scheduler.cu"
+  "${FA3_SRC}/flash_fwd_combine.cu")
+
+set(FA3_ALL_SOURCES
+  "${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_fa3_extension.cc"
+  ${FA3_CORE_SOURCES}
+  ${FA3_INSTANTIATION_SOURCES})
+
+set(FA3_INCLUDE_DIRS
+  ${FA3_SRC}
+  ${cutlass_fa3_SOURCE_DIR}/include
+  ${cutlass_for_fa3_SOURCE_DIR}/include
+  ${cutlass_for_fa3_SOURCE_DIR}/tools/util/include
+  ${CMAKE_CURRENT_SOURCE_DIR}/csrc)
+
+# Set SM90a gencode flags for all FA3 CUDA sources
+set_gencode_flags_for_srcs(
+  SRCS "${FA3_ALL_SOURCES}"
+  CUDA_ARCHS "${CUTLASS_FA3_ARCHS}")
+
+define_extension_target(_cutlass_fa3_C
+  DESTINATION vllm
+  LANGUAGE ${VLLM_GPU_LANG}
+  SOURCES ${FA3_ALL_SOURCES}
+  COMPILE_FLAGS ${VLLM_GPU_FLAGS}
+  ARCHITECTURES ${VLLM_GPU_ARCHES}
+  INCLUDE_DIRECTORIES ${FA3_INCLUDE_DIRS}
+  USE_SABI 3
+  WITH_SOABI)
+
+# FA3-specific compile options for CUDA and C++ source files:
+# - C++17 required by CUTLASS
+# - Fast math for performance
+# - Relaxed constexpr for CUTLASS template metaprogramming
+# - Disable backward pass, dropout, uneven K (not needed for inference)
+# - Enable varlen-only mode (all our use cases are variable-length)
+target_compile_options(_cutlass_fa3_C PRIVATE
+  $<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API>
+  $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>
+  $<$<COMPILE_LANGUAGE:CUDA>:-std=c++17>
+  $<$<COMPILE_LANGUAGE:CXX>:-std=c++17>
+  $<$<COMPILE_LANGUAGE:CUDA>:--use_fast_math>
+  $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>)
+
+target_compile_definitions(_cutlass_fa3_C PRIVATE
+  CUTE_USE_PACKED_TUPLE=1
+  CUTLASS_ENABLE_GDC_FOR_SM90
+  CUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED
+  CUTLASS_ENABLE_TENSOR_CORE_MMA=1
+  FLASHATTENTION_DISABLE_BACKWARD
+  FLASHATTENTION_DISABLE_DROPOUT
+  FLASHATTENTION_DISABLE_UNEVEN_K
+  FLASHATTENTION_DISABLE_SM8x
+  FLASHATTENTION_VARLEN_ONLY)
+
+message(STATUS "CUTLASS FA3 MLA Sparse: enabled for SM90 (${CUTLASS_FA3_ARCHS})")
@@ -0,0 +1,72 @@
+/* SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+ *
+ * Vendored CUTLASS FA3 MLA attention kernel binding for vLLM.
+ * Based on sgl-kernel/csrc/flash_extension.cc from SGLang.
+ *
+ * This registers the FA3 forward pass as a PyTorch C++ extension under
+ * the _cutlass_fa3_C namespace, enabling torch.ops._cutlass_fa3_C.fwd().
+ *
+ * Original source:
+ *   https://github.com/sgl-project/sgl-attn (commit bcf72ccc)
+ *   sgl-kernel/csrc/flash_extension.cc
+ */
+#include <Python.h>
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <torch/all.h>
+#include <torch/library.h>
+
+#include "sgl_flash_kernel_ops.h"
+
+TORCH_LIBRARY_FRAGMENT(_cutlass_fa3_C, m) {
+  /*
+   * CUTLASS FA3 MLA forward pass.
+   * Signature matches sgl-attn's mha_fwd() exactly.
+   */
+  m.def(
+      "fwd(Tensor   q,"
+      "    Tensor   k,"
+      "    Tensor   v,"
+      "    Tensor?  k_new,"
+      "    Tensor?  v_new,"
+      "    Tensor?  q_v,"
+      "    Tensor?  out,"
+      "    Tensor?  cu_seqlens_q,"
+      "    Tensor?  cu_seqlens_k,"
+      "    Tensor?  cu_seqlens_k_new,"
+      "    Tensor?  seqused_q,"
+      "    Tensor?  seqused_k,"
+      "    int?     max_seqlen_q,"
+      "    int?     max_seqlen_k,"
+      "    Tensor?  page_table,"
+      "    Tensor?  kv_batch_idx,"
+      "    Tensor?  leftpad_k,"
+      "    Tensor?  rotary_cos,"
+      "    Tensor?  rotary_sin,"
+      "    Tensor?  seqlens_rotary,"
+      "    Tensor?  q_descale,"
+      "    Tensor?  k_descale,"
+      "    Tensor?  v_descale,"
+      "    float?   softmax_scale,"
+      "    bool     is_causal,"
+      "    int      window_size_left,"
+      "    int      window_size_right,"
+      "    int      attention_chunk,"
+      "    float    softcap,"
+      "    bool     is_rotary_interleaved,"
+      "    Tensor?  scheduler_metadata,"
+      "    int      num_splits,"
+      "    bool?    pack_gqa,"
+      "    int      sm_margin,"
+      "    Tensor?  sinks"
+      ") -> (Tensor, Tensor, Tensor, Tensor)");
+
+  m.impl("fwd", torch::kCUDA, make_pytorch_shim(&mha_fwd));
+}
+
+// Python module initialization for _cutlass_fa3_C
+PyMODINIT_FUNC PyInit__cutlass_fa3_C() {
+  static struct PyModuleDef module = {PyModuleDef_HEAD_INIT, "_cutlass_fa3_C",
+                                      nullptr, 0, nullptr};
+  return PyModule_Create(&module);
+}
@@ -134,13 +134,4 @@ void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out,
                              torch::stable::Tensor& input,
                              torch::stable::Tensor& input_global_scale);

-void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
-                            const torch::stable::Tensor& a,
-                            const torch::stable::Tensor& b,
-                            const torch::stable::Tensor& a_blockscale,
-                            const torch::stable::Tensor& b_blockscales,
-                            const torch::stable::Tensor& problem_sizes,
-                            const torch::stable::Tensor& expert_offsets,
-                            const torch::stable::Tensor& sf_offsets);
-
 #endif
@@ -1,468 +0,0 @@
-/*
- * SPDX-License-Identifier: Apache-2.0
- * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
- *
- * MXFP4 x MXFP4 block-scaled grouped GEMM kernel for MoE on SM100.
- * Uses Cutlass mx_float4_t operands, E8M0 block scales, and 32-element groups.
- */
-
-#include <torch/csrc/stable/library.h>
-#include <torch/csrc/stable/tensor.h>
-#include "libtorch_stable/torch_utils.h"
-
-#include <cutlass/arch/arch.h>
-
-#include "cutlass_extensions/common.hpp"
-
-#include "cute/tensor.hpp"
-#include "cutlass/tensor_ref.h"
-#include "cutlass/epilogue/collective/default_epilogue.hpp"
-#include "cutlass/epilogue/thread/linear_combination.h"
-#include "cutlass/gemm/dispatch_policy.hpp"
-#include "cutlass/gemm/group_array_problem_shape.hpp"
-#include "cutlass/gemm/collective/collective_builder.hpp"
-#include "cutlass/epilogue/collective/collective_builder.hpp"
-#include "cutlass/gemm/device/gemm_universal_adapter.h"
-#include "cutlass/gemm/kernel/gemm_universal.hpp"
-
-#include "cutlass/util/packed_stride.hpp"
-#include <cassert>
-
-using namespace cute;
-
-// Offset-computation kernel for MXFP4 grouped GEMM (group size 32).
-template <typename ElementAB, typename ElementC, typename ElementSF,
-          typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
-__global__ void __mxfp4_get_group_gemm_starts(
-    ElementAB** a_offsets, ElementAB** b_offsets, ElementC** out_offsets,
-    ElementSF** a_scales_offsets, ElementSF** b_scales_offsets,
-    LayoutSFA* layout_sfa_base_as_int, LayoutSFB* layout_sfb_base_as_int,
-    ElementAB* a_base_as_int, ElementAB* b_base_as_int,
-    ElementC* out_base_as_int, ElementSF* a_scales_base_as_int,
-    ElementSF* b_scales_base_as_int, const int32_t* expert_offsets,
-    const int32_t* sf_offsets, const int32_t* problem_sizes_as_shapes,
-    int64_t* a_strides, int64_t* b_strides, int64_t* c_strides,
-    const int64_t a_stride_val, const int64_t b_stride_val,
-    const int64_t c_stride_val, const int K, const int N) {
-  int64_t expert_id = threadIdx.x;
-  if (expert_id >= gridDim.x * blockDim.x) {
-    return;
-  }
-  int64_t expert_offset = static_cast<int64_t>(expert_offsets[expert_id]);
-  int64_t sf_offset = static_cast<int64_t>(sf_offsets[expert_id]);
-  int64_t group_size = 32;
-  int64_t m = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3]);
-  int64_t n = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 1]);
-  int64_t k = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 2]);
-  assert((m >= 0 && n == N && k == K && k % 2 == 0) &&
-         "unexpected problem sizes");
-
-  int64_t half_k = static_cast<int64_t>(k / 2);
-  int64_t group_k = static_cast<int64_t>(k / group_size);
-  // Shape of A as uint8/byte = [M, K // 2]
-  a_offsets[expert_id] = a_base_as_int + expert_offset * half_k;
-  // Shape of B as uint8/byte = [E, N, K // 2]
-  b_offsets[expert_id] = b_base_as_int + expert_id * n * half_k;
-  // Shape of C = [M, N]
-  out_offsets[expert_id] = out_base_as_int + expert_offset * n;
-  // Shape of a_scale = [sum(sf_sizes), K // group_size]
-  a_scales_offsets[expert_id] = a_scales_base_as_int + sf_offset * group_k;
-
-  assert((reinterpret_cast<uintptr_t>(a_scales_offsets[expert_id]) % 128) ==
-             0 &&
-         "TMA requires 128-byte alignment");
-
-  // Shape of B scale = [E, N, K // group_size]
-  b_scales_offsets[expert_id] = b_scales_base_as_int + expert_id * n * group_k;
-  assert((reinterpret_cast<uintptr_t>(b_scales_offsets[expert_id]) % 128) ==
-             0 &&
-         "TMA requires 128-byte alignment");
-
-  // Initialize strides
-  a_strides[expert_id] = a_stride_val;
-  b_strides[expert_id] = b_stride_val;
-  c_strides[expert_id] = c_stride_val;
-
-  LayoutSFA* layout_sfa_ptr = layout_sfa_base_as_int + expert_id;
-  LayoutSFB* layout_sfb_ptr = layout_sfb_base_as_int + expert_id;
-
-  *layout_sfa_ptr = ScaleConfig::tile_atom_to_shape_SFA(cute::make_shape(
-      static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
-  *layout_sfb_ptr = ScaleConfig::tile_atom_to_shape_SFB(cute::make_shape(
-      static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
-}
-
-#define __CALL_MXFP4_GET_STARTS_KERNEL(ELEMENT_AB_TYPE, SF_TYPE,               \
-                                       TENSOR_C_TYPE, C_TYPE, LayoutSFA,       \
-                                       LayoutSFB, ScaleConfig)                 \
-  else if (out_tensors.scalar_type() == TENSOR_C_TYPE) {                       \
-    __mxfp4_get_group_gemm_starts<ELEMENT_AB_TYPE, C_TYPE, SF_TYPE, LayoutSFA, \
-                                  LayoutSFB, ScaleConfig>                      \
-        <<<1, num_experts, 0, stream>>>(                                       \
-            static_cast<ELEMENT_AB_TYPE**>(a_starts.data_ptr()),               \
-            static_cast<ELEMENT_AB_TYPE**>(b_starts.data_ptr()),               \
-            static_cast<C_TYPE**>(out_starts.data_ptr()),                      \
-            static_cast<SF_TYPE**>(a_scales_starts.data_ptr()),                \
-            static_cast<SF_TYPE**>(b_scales_starts.data_ptr()),                \
-            reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),               \
-            reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr()),               \
-            static_cast<ELEMENT_AB_TYPE*>(a_tensors.data_ptr()),               \
-            static_cast<ELEMENT_AB_TYPE*>(b_tensors.data_ptr()),               \
-            static_cast<C_TYPE*>(out_tensors.data_ptr()),                      \
-            static_cast<SF_TYPE*>(a_scales.data_ptr()),                        \
-            static_cast<SF_TYPE*>(b_scales.data_ptr()),                        \
-            static_cast<int32_t*>(expert_offsets.data_ptr()),                  \
-            static_cast<int32_t*>(sf_offsets.data_ptr()),                      \
-            static_cast<int32_t*>(problem_sizes.data_ptr()),                   \
-            static_cast<int64_t*>(a_strides.data_ptr()),                       \
-            static_cast<int64_t*>(b_strides.data_ptr()),                       \
-            static_cast<int64_t*>(c_strides.data_ptr()), a_stride_val,         \
-            b_stride_val, c_stride_val, K, N);                                 \
-  }
-
-template <typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
-void mxfp4_run_get_group_gemm_starts(
-    const torch::stable::Tensor& a_starts,
-    const torch::stable::Tensor& b_starts,
-    const torch::stable::Tensor& out_starts,
-    const torch::stable::Tensor& a_scales_starts,
-    const torch::stable::Tensor& b_scales_starts,
-    const torch::stable::Tensor& layout_sfa,
-    const torch::stable::Tensor& layout_sfb,
-    const torch::stable::Tensor& a_strides,
-    const torch::stable::Tensor& b_strides,
-    const torch::stable::Tensor& c_strides, int64_t a_stride_val,
-    int64_t b_stride_val, int64_t c_stride_val,
-    torch::stable::Tensor const& a_tensors,
-    torch::stable::Tensor const& b_tensors,
-    torch::stable::Tensor const& out_tensors,
-    torch::stable::Tensor const& a_scales,
-    torch::stable::Tensor const& b_scales,
-    torch::stable::Tensor const& expert_offsets,
-    torch::stable::Tensor const& sf_offsets,
-    torch::stable::Tensor const& problem_sizes, int M, int N, int K) {
-  int num_experts = (int)expert_offsets.size(0);
-  auto stream = get_current_cuda_stream(a_tensors.get_device_index());
-
-  STD_TORCH_CHECK(out_tensors.size(1) == N,
-                  "Output tensor shape doesn't match expected shape");
-  STD_TORCH_CHECK(K / 2 == b_tensors.size(2),
-                  "b_tensors(dim = 2) and a_tensors(dim = 1) trailing"
-                  " dimension must match");
-  if (false) {
-  }
-  // MXFP4 uses E8M0 (float_ue8m0_t) scale factors
-  __CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t,
-                                 torch::headeronly::ScalarType::BFloat16,
-                                 cutlass::bfloat16_t, LayoutSFA, LayoutSFB,
-                                 ScaleConfig)
-  __CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t,
-                                 torch::headeronly::ScalarType::Half, half,
-                                 LayoutSFA, LayoutSFB, ScaleConfig)
-  else {
-    STD_TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
-  }
-}
-
-template <typename OutType>
-void run_mxfp4_blockwise_scaled_group_mm_sm100(
-    torch::stable::Tensor& output, const torch::stable::Tensor& a,
-    const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
-    const torch::stable::Tensor& b_blockscales,
-    const torch::stable::Tensor& problem_sizes,
-    const torch::stable::Tensor& expert_offsets,
-    const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
-  using ProblemShape =
-      cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
-  using ElementType = cutlass::float_e2m1_t;
-  using ElementSFType = cutlass::float_ue8m0_t;
-  using ElementA = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
-  using ElementB = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
-
-  using ElementC = OutType;
-  using ElementD = ElementC;
-  using ElementAccumulator = float;
-  // Layout definitions
-  using LayoutA = cutlass::layout::RowMajor;
-  using LayoutB = cutlass::layout::ColumnMajor;
-  using LayoutC = cutlass::layout::RowMajor;
-  using LayoutD = LayoutC;
-
-  static constexpr int AlignmentA = 32;
-  static constexpr int AlignmentB = 32;
-  static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
-  static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
-
-  // Architecture definitions
-  using ArchTag = cutlass::arch::Sm100;
-  using EpilogueOperatorClass = cutlass::arch::OpClassTensorOp;
-  using MainloopOperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
-  using StageCountType = cutlass::gemm::collective::StageCountAuto;
-
-  using ClusterShape = Shape<_1, _1, _1>;
-  struct MMA1SMConfig {
-    using MmaTileShape = Shape<_128, _128, _128>;
-    using KernelSchedule =
-        cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmMxf4Sm100;
-    using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm;
-  };
-
-  using CollectiveEpilogue =
-      typename cutlass::epilogue::collective::CollectiveBuilder<
-          ArchTag, EpilogueOperatorClass, typename MMA1SMConfig::MmaTileShape,
-          ClusterShape, Shape<_128, _64>, ElementAccumulator,
-          ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementD,
-          LayoutC*, AlignmentD,
-          typename MMA1SMConfig::EpilogueSchedule>::CollectiveOp;
-
-  using CollectiveMainloop =
-      typename cutlass::gemm::collective::CollectiveBuilder<
-          ArchTag, MainloopOperatorClass, ElementA, LayoutA*, AlignmentA,
-          ElementB, LayoutB*, AlignmentB, ElementAccumulator,
-          typename MMA1SMConfig::MmaTileShape, ClusterShape,
-          cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
-              sizeof(typename CollectiveEpilogue::SharedStorage))>,
-          typename MMA1SMConfig::KernelSchedule>::CollectiveOp;
-
-  using GemmKernel =
-      cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop,
-                                           CollectiveEpilogue>;
-
-  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  using StrideA = typename Gemm::GemmKernel::InternalStrideA;
-  using StrideB = typename Gemm::GemmKernel::InternalStrideB;
-  using StrideC = typename Gemm::GemmKernel::InternalStrideC;
-  using StrideD = typename Gemm::GemmKernel::InternalStrideD;
-
-  using LayoutSFA =
-      typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA;
-  using LayoutSFB =
-      typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB;
-  using ScaleConfig =
-      typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
-
-  using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
-  int num_experts = static_cast<int>(expert_offsets.size(0));
-
-  torch::stable::Tensor a_ptrs =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-  torch::stable::Tensor b_ptrs =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-  torch::stable::Tensor out_ptrs =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-  torch::stable::Tensor a_scales_ptrs =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-  torch::stable::Tensor b_scales_ptrs =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-  torch::stable::Tensor layout_sfa = torch::stable::empty(
-      {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
-      a.device());
-  torch::stable::Tensor layout_sfb = torch::stable::empty(
-      {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
-      a.device());
-  torch::stable::Tensor a_strides1 =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-  torch::stable::Tensor b_strides1 =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-  torch::stable::Tensor c_strides1 =
-      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
-                           std::nullopt, a.device());
-
-  mxfp4_run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
-      a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, layout_sfa,
-      layout_sfb, a_strides1, b_strides1, c_strides1, a.stride(0) * 2,
-      b.stride(1) * 2, output.stride(0), a, b, output, a_blockscale,
-      b_blockscales, expert_offsets, sf_offsets, problem_sizes, M, N, K);
-
-  // Create an instance of the GEMM
-  Gemm gemm_op;
-
-  UnderlyingProblemShape* problem_sizes_as_shapes =
-      static_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
-
-  // Set the Scheduler info
-  cutlass::KernelHardwareInfo hw_info;
-  using RasterOrderOptions = typename cutlass::gemm::kernel::detail::
-      PersistentTileSchedulerSm100GroupParams<
-          typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions;
-  typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
-  scheduler.raster_order = RasterOrderOptions::AlongM;
-  hw_info.device_id = a.get_device_index();
-  static std::unordered_map<int, int> cached_sm_counts;
-  if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
-    cached_sm_counts[hw_info.device_id] =
-        cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
-            hw_info.device_id);
-  }
-  hw_info.sm_count = min(cached_sm_counts[hw_info.device_id], INT_MAX);
-
-  // Mainloop Arguments
-  typename GemmKernel::MainloopArguments mainloop_args{
-      static_cast<const ElementType**>(a_ptrs.data_ptr()),
-      static_cast<StrideA*>(a_strides1.data_ptr()),
-      static_cast<const ElementType**>(b_ptrs.data_ptr()),
-      static_cast<StrideB*>(b_strides1.data_ptr()),
-      static_cast<const ElementSFType**>(a_scales_ptrs.data_ptr()),
-      reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
-      static_cast<const ElementSFType**>(b_scales_ptrs.data_ptr()),
-      reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())};
-
-  // Epilogue Arguments
-  typename GemmKernel::EpilogueArguments epilogue_args{
-      {},  // epilogue.thread
-      nullptr,
-      static_cast<StrideC*>(c_strides1.data_ptr()),
-      static_cast<ElementD**>(out_ptrs.data_ptr()),
-      static_cast<StrideC*>(c_strides1.data_ptr())};
-  auto& fusion_args = epilogue_args.thread;
-  // Scalar epilogue (CUTLASS grouped GEMM): D = 1 * accum + 0 * C
-  fusion_args.alpha_ptr = nullptr;
-  fusion_args.beta_ptr = nullptr;
-  fusion_args.alpha = 1.0f;
-  fusion_args.alpha_ptr_array = nullptr;
-  fusion_args.dAlpha = {_0{}, _0{}, 0};
-  fusion_args.beta = 0.0f;
-  fusion_args.beta_ptr_array = nullptr;
-  fusion_args.dBeta = {_0{}, _0{}, 0};
-
-  // Gemm Arguments
-  typename GemmKernel::Arguments args{
-      cutlass::gemm::GemmUniversalMode::kGrouped,
-      {num_experts, problem_sizes_as_shapes, nullptr},
-      mainloop_args,
-      epilogue_args,
-      hw_info,
-      scheduler};
-
-  size_t workspace_size = Gemm::get_workspace_size(args);
-  auto workspace =
-      torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
-                           std::nullopt, a.device());
-  const cudaStream_t stream = get_current_cuda_stream(a.get_device_index());
-
-  auto can_implement_status = gemm_op.can_implement(args);
-  STD_TORCH_CHECK(
-      can_implement_status == cutlass::Status::kSuccess,
-      "Failed to implement MXFP4 GEMM: status=", (int)can_implement_status);
-
-  // Run the GEMM
-  auto status = gemm_op.initialize(args, workspace.data_ptr());
-  STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
-                  "Failed to initialize MXFP4 GEMM: status=", (int)status,
-                  " workspace_size=", workspace_size,
-                  " num_experts=", num_experts, " M=", M, " N=", N, " K=", K);
-
-  status = gemm_op.run(args, workspace.data_ptr(), stream);
-  STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
-                  "Failed to run MXFP4 GEMM");
-}
-
-template <typename OutType>
-void run_mxfp4_blockwise_scaled_group_mm(
-    torch::stable::Tensor& output, const torch::stable::Tensor& a,
-    const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
-    const torch::stable::Tensor& b_blockscales,
-    const torch::stable::Tensor& problem_sizes,
-    const torch::stable::Tensor& expert_offsets,
-    const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
-  int32_t version_num = get_sm_version_num();
-#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
-  if (version_num >= 100 && version_num < 120) {
-    run_mxfp4_blockwise_scaled_group_mm_sm100<OutType>(
-        output, a, b, a_blockscale, b_blockscales, problem_sizes,
-        expert_offsets, sf_offsets, M, N, K);
-    return;
-  }
-#endif
-  STD_TORCH_CHECK_NOT_IMPLEMENTED(
-      false,
-      "No compiled cutlass_mxfp4_group_mm kernel for CUDA device capability: ",
-      version_num, ". Required capability: 100");
-}
-
-#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
-constexpr auto MXFP4_FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte;
-// E8M0 scale factors stored as uint8
-constexpr auto MXFP4_SF_DTYPE = torch::headeronly::ScalarType::Byte;
-#endif
-
-#define CHECK_TYPE(x, st, m)             \
-  STD_TORCH_CHECK(x.scalar_type() == st, \
-                  ": Inconsistency of torch::stable::Tensor type:", m)
-#define CHECK_TH_CUDA(x, m) \
-  STD_TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor.")
-#define CHECK_CONTIGUOUS(x, m) \
-  STD_TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous.")
-#define CHECK_INPUT(x, st, m) \
-  CHECK_TH_CUDA(x, m);        \
-  CHECK_CONTIGUOUS(x, m);     \
-  CHECK_TYPE(x, st, m)
-
-void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
-                            const torch::stable::Tensor& a,
-                            const torch::stable::Tensor& b,
-                            const torch::stable::Tensor& a_blockscale,
-                            const torch::stable::Tensor& b_blockscales,
-                            const torch::stable::Tensor& problem_sizes,
-                            const torch::stable::Tensor& expert_offsets,
-                            const torch::stable::Tensor& sf_offsets) {
-#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
-  // Input validation
-  CHECK_INPUT(a, MXFP4_FLOAT4_E2M1X2, "a");
-  CHECK_INPUT(b, MXFP4_FLOAT4_E2M1X2, "b");
-  // MXFP4 uses E8M0 scale factors (stored as uint8)
-  CHECK_INPUT(a_blockscale, MXFP4_SF_DTYPE, "a_blockscale");
-  CHECK_INPUT(b_blockscales, MXFP4_SF_DTYPE, "b_blockscales");
-
-  STD_TORCH_CHECK(
-      a_blockscale.dim() == 2,
-      "expected a_blockscale to be of shape [num_experts, rounded_m,"
-      " k // group_size], observed rank: ",
-      a_blockscale.dim())
-  STD_TORCH_CHECK(b_blockscales.dim() == 3,
-                  "expected b_blockscale to be of shape: "
-                  " [num_experts, n, k // group_size], observed rank: ",
-                  b_blockscales.dim())
-  STD_TORCH_CHECK(problem_sizes.dim() == 2,
-                  "problem_sizes must be a 2D tensor");
-  STD_TORCH_CHECK(problem_sizes.size(1) == 3,
-                  "problem_sizes must have the shape (num_experts, 3)");
-  STD_TORCH_CHECK(
-      problem_sizes.size(0) == expert_offsets.size(0),
-      "Number of experts in problem_sizes must match expert_offsets");
-  STD_TORCH_CHECK(
-      problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
-      "problem_sizes must be int32.");
-
-  int M = static_cast<int>(a.size(0));
-  int N = static_cast<int>(b.size(1));
-  int E = static_cast<int>(b.size(0));
-  int K = static_cast<int>(2 * b.size(2));
-
-  if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
-    run_mxfp4_blockwise_scaled_group_mm<cutlass::bfloat16_t>(
-        output, a, b, a_blockscale, b_blockscales, problem_sizes,
-        expert_offsets, sf_offsets, M, N, K);
-  } else {
-    run_mxfp4_blockwise_scaled_group_mm<cutlass::half_t>(
-        output, a, b, a_blockscale, b_blockscales, problem_sizes,
-        expert_offsets, sf_offsets, M, N, K);
-  }
-#else
-  STD_TORCH_CHECK_NOT_IMPLEMENTED(
-      false,
-      "No compiled cutlass_mxfp4_group_mm kernel; build vLLM with "
-      "SM100 block-scaled FP4 MoE (ENABLE_NVFP4_SM100) and CUDA 12.8+.");
-#endif
-}
-
-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
-  m.impl("cutlass_mxfp4_group_mm", TORCH_BOX(&cutlass_mxfp4_group_mm));
-}
@@ -1,432 +0,0 @@
-/*
- * SPDX-License-Identifier: Apache-2.0
- * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
- *
- * MXFP4 activation quantization kernel for MoE experts.
- * Quantizes BF16/FP16 activations to MXFP4: E2M1 values with E8M0 block scales
- * over 32-element groups.
- *
- * Uses PACK16 E2M1 conversion helpers (nvfp4_utils.cuh) configured for:
- *   - Block size 32 (2 threads per SF in PACK16 mode)
- *   - E8M0 (power-of-two) scale factors
- *   - SF layout: [numMTiles, numKTiles, 32, 4, 4] where numKTiles=ceil(K/128)
- */
-
-// MXFP4 requires PACK16 mode (16 elements per thread) so that
-// 2 threads cover 32-element blocks. This requires CUDA >= 12.9.
-// Must be defined before any header that (transitively) includes
-// nvfp4_utils.cuh.
-#define NVFP4_ENABLE_ELTS16 1
-
-#include <cuda.h>
-#include <cuda_runtime_api.h>
-#include <cuda_runtime.h>
-#include <cuda_fp8.h>
-
-#include <torch/csrc/stable/library.h>
-#include <torch/csrc/stable/tensor.h>
-#include "libtorch_stable/torch_utils.h"
-#include "libtorch_stable/dispatch_utils.h"
-#include "cuda_vec_utils.cuh"
-#include "cuda_utils.h"
-
-#include "nvfp4_utils.cuh"
-static_assert(CVT_FP4_ELTS_PER_THREAD == 16,
-              "MXFP4 experts quant requires PACK16 mode (CUDA >= 12.9)");
-
-#include "launch_bounds_utils.h"
-
-namespace vllm {
-
-// MXFP4 block size constants
-static constexpr int MXFP4_SF_VEC_SIZE = 32;
-
-// For PACK16 mode (CVT_FP4_ELTS_PER_THREAD=16): 2 threads per SF
-// For PACK8 mode (CVT_FP4_ELTS_PER_THREAD=8): 4 threads per SF
-static constexpr int MXFP4_NUM_THREADS_PER_SF =
-    MXFP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
-
-// MXFP4 quantization kernel for experts.
-// Uses 32-element blocks with E8M0 (UE8M0) scale factors.
-// When FUSE_SILU_MUL=true, expects input with gate||up layout and fuses
-// SiLU(gate)*up before quantization.
-template <class Type, bool FUSE_SILU_MUL = false,
-          bool SMALL_NUM_EXPERTS = false>
-__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
-    mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
-                          fp4_packed_t* out, uint32_t* SFout,
-                          uint32_t* input_offset_by_experts,
-                          uint32_t* output_scale_offset_by_experts,
-                          int n_experts, bool low_latency) {
-  using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
-  static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
-                "Vec size is not matched.");
-
-  // MXFP4: numKTiles = ceil(numCols / 128) since block_size=32, 4 SFs/tile
-  int32_t const numKTiles = (numCols + 127) / 128;
-
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
-  int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
-
-  for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
-       globalIdx += gridDim.x * blockDim.x) {
-    int rowIdx = globalIdx / colsPerRow;
-    int colIdx = globalIdx % colsPerRow;
-
-    int rowIdx_in_expert = 0;
-    int expert_idx = 0;
-
-    if constexpr (SMALL_NUM_EXPERTS) {
-      for (int i = 0; i < n_experts; i++) {
-        uint32_t current_offset = __ldca(&input_offset_by_experts[i]);
-        uint32_t next_offset = __ldca(&input_offset_by_experts[i + 1]);
-        if (rowIdx >= current_offset && rowIdx < next_offset) {
-          rowIdx_in_expert = rowIdx - current_offset;
-          expert_idx = i;
-          break;
-        }
-      }
-    } else {
-      uint32_t local_offsets[17];
-      for (int chunk_start = 0; chunk_start < n_experts; chunk_start += 16) {
-        *reinterpret_cast<int4*>(local_offsets) =
-            __ldca(reinterpret_cast<const int4*>(
-                &input_offset_by_experts[chunk_start]));
-        *reinterpret_cast<int4*>(local_offsets + 4) =
-            __ldca(reinterpret_cast<const int4*>(
-                &input_offset_by_experts[chunk_start + 4]));
-        *reinterpret_cast<int4*>(local_offsets + 8) =
-            __ldca(reinterpret_cast<const int4*>(
-                &input_offset_by_experts[chunk_start + 8]));
-        *reinterpret_cast<int4*>(local_offsets + 12) =
-            __ldca(reinterpret_cast<const int4*>(
-                &input_offset_by_experts[chunk_start + 12]));
-        local_offsets[16] = __ldca(&input_offset_by_experts[chunk_start + 16]);
-
-#pragma unroll
-        for (int i = 0; i < 16; i++) {
-          if (rowIdx >= local_offsets[i] && rowIdx < local_offsets[i + 1]) {
-            rowIdx_in_expert = rowIdx - local_offsets[i];
-            expert_idx = chunk_start + i;
-            break;
-          }
-        }
-      }
-    }
-
-    // Load input and optionally apply fused SiLU+Mul
-    int64_t inOffset = rowIdx * inColsPerRow + colIdx;
-    PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
-    PackedVec quant_input;
-    if constexpr (FUSE_SILU_MUL) {
-      PackedVec in_vec_up =
-          reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
-      quant_input = compute_silu_mul(in_vec, in_vec_up);
-    } else {
-      quant_input = in_vec;
-    }
-
-    // In PACK16 mode, each thread outputs 16 E2M1 values = u32x2
-    int64_t outOffset = rowIdx * colsPerRow + colIdx;
-    auto& out_pos = out[outOffset];
-
-    uint32_t* SFout_in_expert =
-        SFout + output_scale_offset_by_experts[expert_idx] * numKTiles;
-
-    // Use MXFP4_NUM_THREADS_PER_SF (2 for PACK16) for 32-element blocks
-    auto sf_out =
-        cvt_quant_to_fp4_get_sf_out_offset<uint32_t, MXFP4_NUM_THREADS_PER_SF>(
-            rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
-
-    // Block E8M0 scales only; no extra tensor-level scale in this path
-    constexpr float SFScaleVal = 1.0f;
-    // UE8M0_SF=true for MXFP4 E8M0 scale factors
-    out_pos =
-        cvt_warp_fp16_to_fp4<Type, MXFP4_NUM_THREADS_PER_SF, /*UE8M0_SF=*/true>(
-            quant_input, SFScaleVal, sf_out);
-  }
-}
-
-// Large M_topk variant using shared memory for expert offsets
-template <class Type, bool FUSE_SILU_MUL = false,
-          bool SMALL_NUM_EXPERTS = false>
-__global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
-    mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
-                          fp4_packed_t* out, uint32_t* SFout,
-                          uint32_t* input_offset_by_experts,
-                          uint32_t* output_scale_offset_by_experts,
-                          int n_experts) {
-  using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
-  static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
-                "Vec size is not matched.");
-
-  // MXFP4: numKTiles = ceil(numCols / 128)
-  int32_t const numKTiles = (numCols + 127) / 128;
-
-  extern __shared__ uint32_t shared_input_offsets[];
-
-  if constexpr (SMALL_NUM_EXPERTS) {
-    for (int i = threadIdx.x; i < n_experts + 1; i += blockDim.x) {
-      shared_input_offsets[i] = input_offset_by_experts[i];
-    }
-  } else {
-    for (int i = threadIdx.x * 4; i < n_experts; i += blockDim.x * 4) {
-      *reinterpret_cast<int4*>(&shared_input_offsets[i]) =
-          *reinterpret_cast<const int4*>(&input_offset_by_experts[i]);
-    }
-    if (threadIdx.x == 0) {
-      shared_input_offsets[n_experts] = input_offset_by_experts[n_experts];
-    }
-  }
-
-  __syncthreads();
-
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
-  int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
-
-  for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
-       globalIdx += gridDim.x * blockDim.x) {
-    int rowIdx = globalIdx / colsPerRow;
-    int colIdx = globalIdx % colsPerRow;
-
-    int rowIdx_in_expert = 0;
-    int expert_idx = 0;
-
-    // Binary search through experts using shared memory
-    int left = 0, right = n_experts - 1;
-    while (left <= right) {
-      int mid = (left + right) / 2;
-      uint32_t mid_offset = shared_input_offsets[mid];
-      uint32_t next_offset = shared_input_offsets[mid + 1];
-
-      if (rowIdx >= mid_offset && rowIdx < next_offset) {
-        rowIdx_in_expert = rowIdx - mid_offset;
-        expert_idx = mid;
-        break;
-      } else if (rowIdx < mid_offset) {
-        right = mid - 1;
-      } else {
-        left = mid + 1;
-      }
-    }
-
-    int64_t inOffset = rowIdx * inColsPerRow + colIdx;
-    PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
-    PackedVec quant_input;
-    if constexpr (FUSE_SILU_MUL) {
-      PackedVec in_vec_up =
-          reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
-      quant_input = compute_silu_mul(in_vec, in_vec_up);
-    } else {
-      quant_input = in_vec;
-    }
-
-    int64_t outOffset = rowIdx * colsPerRow + colIdx;
-    auto& out_pos = out[outOffset];
-
-    // MXFP4 has no global scale - only block-level E8M0 scale factors
-    constexpr float SFScaleVal = 1.0f;
-
-    uint32_t* SFout_in_expert =
-        SFout + output_scale_offset_by_experts[expert_idx] * numKTiles;
-
-    auto sf_out =
-        cvt_quant_to_fp4_get_sf_out_offset<uint32_t, MXFP4_NUM_THREADS_PER_SF>(
-            rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
-
-    out_pos =
-        cvt_warp_fp16_to_fp4<Type, MXFP4_NUM_THREADS_PER_SF, /*UE8M0_SF=*/true>(
-            quant_input, SFScaleVal, sf_out);
-  }
-}
-
-template <typename T, bool FUSE_SILU_MUL = false>
-void mxfp4_quant_impl(void* output, void* output_scale, void* input,
-                      void* input_offset_by_experts,
-                      void* output_scale_offset_by_experts, int m_topk, int k,
-                      int n_experts, cudaStream_t stream) {
-  int multiProcessorCount =
-      get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
-
-  int const workSizePerRow = k / ELTS_PER_THREAD;
-  int const totalWorkSize = m_topk * workSizePerRow;
-  dim3 block(std::min(workSizePerRow, 512));
-  int const numBlocksPerSM =
-      vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
-  dim3 grid(std::min(static_cast<int>((totalWorkSize + block.x - 1) / block.x),
-                     multiProcessorCount * numBlocksPerSM));
-  while (grid.x <= multiProcessorCount && block.x > 64) {
-    grid.x *= 2;
-    block.x = (block.x + 1) / 2;
-  }
-
-  int const blockRepeat =
-      (totalWorkSize + block.x * grid.x - 1) / (block.x * grid.x);
-  if (blockRepeat > 1) {
-    size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t);
-    if (n_experts >= 4) {
-      mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false>
-          <<<grid, block, shared_mem_size, stream>>>(
-              m_topk, k, reinterpret_cast<T*>(input),
-              reinterpret_cast<fp4_packed_t*>(output),
-              reinterpret_cast<uint32_t*>(output_scale),
-              reinterpret_cast<uint32_t*>(input_offset_by_experts),
-              reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
-              n_experts);
-    } else {
-      mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true>
-          <<<grid, block, shared_mem_size, stream>>>(
-              m_topk, k, reinterpret_cast<T*>(input),
-              reinterpret_cast<fp4_packed_t*>(output),
-              reinterpret_cast<uint32_t*>(output_scale),
-              reinterpret_cast<uint32_t*>(input_offset_by_experts),
-              reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
-              n_experts);
-    }
-  } else {
-    if (n_experts >= 16) {
-      mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false>
-          <<<grid, block, 0, stream>>>(
-              m_topk, k, reinterpret_cast<T*>(input),
-              reinterpret_cast<fp4_packed_t*>(output),
-              reinterpret_cast<uint32_t*>(output_scale),
-              reinterpret_cast<uint32_t*>(input_offset_by_experts),
-              reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
-              n_experts, /* bool low_latency */ true);
-    } else {
-      mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true><<<grid, block, 0, stream>>>(
-          m_topk, k, reinterpret_cast<T*>(input),
-          reinterpret_cast<fp4_packed_t*>(output),
-          reinterpret_cast<uint32_t*>(output_scale),
-          reinterpret_cast<uint32_t*>(input_offset_by_experts),
-          reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
-          n_experts, /* bool low_latency */ true);
-    }
-  }
-}
-
-}  // namespace vllm
-
-/*Quantization entry for mxfp4 experts quantization*/
-#define CHECK_TH_CUDA(x, m) \
-  STD_TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
-#define CHECK_CONTIGUOUS(x, m) \
-  STD_TORCH_CHECK(x.is_contiguous(), m, "must be contiguous")
-#define CHECK_INPUT(x, m) \
-  CHECK_TH_CUDA(x, m);    \
-  CHECK_CONTIGUOUS(x, m);
-
-constexpr auto HALF = torch::headeronly::ScalarType::Half;
-constexpr auto BF16 = torch::headeronly::ScalarType::BFloat16;
-constexpr auto INT = torch::headeronly::ScalarType::Int;
-constexpr auto UINT8 = torch::headeronly::ScalarType::Byte;
-
-static constexpr int MXFP4_BLOCK_SIZE = 32;
-
-static void validate_mxfp4_experts_quant_inputs(
-    torch::stable::Tensor const& output,
-    torch::stable::Tensor const& output_scale,
-    torch::stable::Tensor const& input,
-    torch::stable::Tensor const& input_offset_by_experts,
-    torch::stable::Tensor const& output_scale_offset_by_experts,
-    int64_t n_experts, int64_t m_topk, int64_t k) {
-  CHECK_INPUT(output, "output");
-  CHECK_INPUT(output_scale, "output_scale");
-  CHECK_INPUT(input, "input");
-  CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts");
-  CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts");
-
-  STD_TORCH_CHECK(output.dim() == 2);
-  STD_TORCH_CHECK(output_scale.dim() == 2);
-  STD_TORCH_CHECK(input.dim() == 2);
-  STD_TORCH_CHECK(input_offset_by_experts.dim() == 1);
-  STD_TORCH_CHECK(output_scale_offset_by_experts.dim() == 1);
-
-  STD_TORCH_CHECK(input.scalar_type() == HALF || input.scalar_type() == BF16);
-  STD_TORCH_CHECK(input_offset_by_experts.scalar_type() == INT);
-  STD_TORCH_CHECK(output_scale_offset_by_experts.scalar_type() == INT);
-  // output is uint8 (two mxfp4 values packed into one uint8)
-  // output_scale is int32 (four E8M0 values packed into one int32)
-  STD_TORCH_CHECK(output.scalar_type() == UINT8);
-  STD_TORCH_CHECK(output_scale.scalar_type() == INT);
-
-  STD_TORCH_CHECK(k % MXFP4_BLOCK_SIZE == 0, "k must be a multiple of 32");
-  STD_TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1);
-  STD_TORCH_CHECK(output_scale_offset_by_experts.size(0) == n_experts + 1);
-  STD_TORCH_CHECK(output.size(0) == m_topk);
-  STD_TORCH_CHECK(output.size(1) == k / 2);
-  int scales_k = k / MXFP4_BLOCK_SIZE;
-  // K-dimension scale columns padded to a multiple of 4 for swizzle layout
-  int padded_k = (scales_k + (4 - 1)) / 4 * 4;
-  // 4 = 4 E8M0 values packed into one int32
-  STD_TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
-}
-
-void mxfp4_experts_quant(
-    torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
-    torch::stable::Tensor const& input,
-    torch::stable::Tensor const& input_offset_by_experts,
-    torch::stable::Tensor const& output_scale_offset_by_experts,
-    int64_t n_experts) {
-  auto m_topk = input.size(0);
-  auto k = input.size(1);
-
-  validate_mxfp4_experts_quant_inputs(
-      output, output_scale, input, input_offset_by_experts,
-      output_scale_offset_by_experts, n_experts, m_topk, k);
-
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      input.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
-
-  VLLM_STABLE_DISPATCH_HALF_TYPES(
-      input.scalar_type(), "mxfp4_experts_quant_kernel", [&] {
-        using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
-        vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/false>(
-            output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
-            input_offset_by_experts.data_ptr(),
-            output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
-            stream);
-      });
-}
-
-void silu_and_mul_mxfp4_experts_quant(
-    torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
-    torch::stable::Tensor const& input,
-    torch::stable::Tensor const& input_offset_by_experts,
-    torch::stable::Tensor const& output_scale_offset_by_experts,
-    int64_t n_experts) {
-  auto m_topk = input.size(0);
-  auto k_times_2 = input.size(1);
-  STD_TORCH_CHECK(k_times_2 % 2 == 0, "input width must be even (gate || up)");
-  auto k = k_times_2 / 2;
-
-  validate_mxfp4_experts_quant_inputs(
-      output, output_scale, input, input_offset_by_experts,
-      output_scale_offset_by_experts, n_experts, m_topk, k);
-
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      input.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
-
-  VLLM_STABLE_DISPATCH_HALF_TYPES(
-      input.scalar_type(), "silu_mul_mxfp4_experts_quant_kernel", [&] {
-        using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
-        vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/true>(
-            output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
-            input_offset_by_experts.data_ptr(),
-            output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
-            stream);
-      });
-}
-
-// Registered here (not torch_bindings.cpp) because VLLM_GPU_FLAGS is applied
-// only under COMPILE_LANGUAGE:CUDA, so ENABLE_NVFP4_SM100 is invisible to
-// .cpp files and cannot gate the registration from there.
-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
-  m.impl("mxfp4_experts_quant", TORCH_BOX(&mxfp4_experts_quant));
-  m.impl("silu_and_mul_mxfp4_experts_quant",
-         TORCH_BOX(&silu_and_mul_mxfp4_experts_quant));
-}
@@ -116,12 +116,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
      " Tensor a_blockscale, Tensor b_blockscales, Tensor alphas,"
      " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");

-  // cutlass mxfp4 block scaled group GEMM (MXFP4 x MXFP4 MoE)
-  ops.def(
-      "cutlass_mxfp4_group_mm(Tensor! out, Tensor a, Tensor b,"
-      " Tensor a_blockscale, Tensor b_blockscales,"
-      " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
-
  // Compute NVFP4 block quantized tensor.
  ops.def(
      "scaled_fp4_quant(Tensor input,"
@@ -155,19 +149,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
      "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
      "Tensor output_scale_offset_by_experts) -> ()");

-  // Compute MXFP4 experts quantization (32-element blocks, E8M0 SFs).
-  ops.def(
-      "mxfp4_experts_quant(Tensor! output, Tensor! output_scale,"
-      "Tensor input, Tensor input_offset_by_experts,"
-      "Tensor output_scale_offset_by_experts, int n_experts) -> ()");
-
-  // Fused SiLU+Mul+MXFP4 experts quantization.
-  ops.def(
-      "silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! "
-      "output_scale,"
-      "Tensor input, Tensor input_offset_by_experts,"
-      "Tensor output_scale_offset_by_experts, int n_experts) -> ()");
-
  // Fused SiLU+Mul+NVFP4 quantization.
  ops.def(
      "silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, "
@@ -252,8 +233,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
  ops.impl("silu_and_mul_scaled_fp4_experts_quant",
           TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant));
  ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant));
-  // mxfp4_experts_quant: registered in mxfp4_experts_quant.cu (SM100 only).
-  // W4A8 ops: registered in w4a8_mm_entry.cu / w4a8_grouped_mm_entry.cu.
+
+  // W4A8 ops: impl registrations are in the source files
+  // (w4a8_mm_entry.cu and w4a8_grouped_mm_entry.cu)
 #endif
 }

@@ -126,9 +126,7 @@ __launch_bounds__(TPB) __global__
    {
        const int idx = thread_row_offset + ii;
        const float val = toFloat(input[idx]);
-        float softmax_val = expf(val - float_max) * normalizing_factor;
-        // Clamp NaN/Inf to 0 to prevent duplicate expert IDs downstream.
-        if (isnan(softmax_val) || isinf(softmax_val)) softmax_val = 0.f;
+        const float softmax_val = expf(val - float_max) * normalizing_factor;
        output[idx] = softmax_val;
    }
 }
@@ -149,9 +147,7 @@ __launch_bounds__(TPB) __global__
    {
        const int idx = thread_row_offset + ii;
        const float val = toFloat(input[idx]);
-        float sigmoid_val = 1.0f / (1.0f + __expf(-val));
-        // Clamp NaN/Inf to 0 to prevent duplicate expert IDs downstream.
-        if (isnan(sigmoid_val) || isinf(sigmoid_val)) sigmoid_val = 0.f;
+        const float sigmoid_val = 1.0f / (1.0f + __expf(-val));
        output[idx] = sigmoid_val;
    }
 }
@@ -446,19 +442,6 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
      }
    }

-    // Fix: clamp NaN/Inf values to 0 to prevent duplicate expert IDs.
-    // NaN gating (from degenerate hidden states in CUDA graph padding) causes
-    // softmax to produce all-NaN, which makes the argmax loop always pick
-    // expert 0 for every top-k slot, producing duplicate expert IDs that
-    // crash FlashInfer's three-step MoE sort.
-    // With 0s, the argmax uses index tie-breaking to pick [0,1,2,...,k-1].
-#pragma unroll
-    for (int ii = 0; ii < VPT; ++ii) {
-      if (isnan(row_chunk[ii]) || isinf(row_chunk[ii])) {
-        row_chunk[ii] = 0.f;
-      }
-    }
-
    static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;

    // If bias is not null, use biased value for selection
@@ -0,0 +1,45 @@
+/* SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright 2025 SGLang Team. All Rights Reserved.
+ *
+ * Vendored from sgl-kernel/include/sgl_flash_kernel_ops.h (commit bcf72ccc).
+ * Declares the mha_fwd() C++ function signature for CUTLASS FA3 kernels.
+ * NO MODIFICATIONS from the original (except removing unused macros).
+ */
+
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/Tensor.h>
+#include <torch/library.h>
+#include <torch/torch.h>
+
+#include <vector>
+
+#include "sgl_kernel_torch_shim.h"
+
+/*
+ * From flash-attention (sgl-attn fork)
+ */
+std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> mha_fwd(
+    at::Tensor q,  // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
+    at::Tensor k,  // (b_k, s_k, h_k, d) or (total_k, h_k, d) or paged
+    at::Tensor v,  // (b_k, s_k, h_k, dv) or (total_k, h_k, dv) or paged
+    std::optional<at::Tensor> k_new_, std::optional<at::Tensor> v_new_,
+    std::optional<at::Tensor> q_v_,  // MLA value projection query
+    std::optional<at::Tensor> out_, std::optional<at::Tensor> cu_seqlens_q_,
+    std::optional<at::Tensor> cu_seqlens_k_,
+    std::optional<at::Tensor> cu_seqlens_k_new_,
+    std::optional<at::Tensor> seqused_q_, std::optional<at::Tensor> seqused_k_,
+    std::optional<int64_t> max_seqlen_q_, std::optional<int64_t> max_seqlen_k_,
+    std::optional<at::Tensor> page_table_,
+    std::optional<at::Tensor> kv_batch_idx_,
+    std::optional<at::Tensor> leftpad_k_, std::optional<at::Tensor> rotary_cos_,
+    std::optional<at::Tensor> rotary_sin_,
+    std::optional<at::Tensor> seqlens_rotary_,
+    std::optional<at::Tensor> q_descale_, std::optional<at::Tensor> k_descale_,
+    std::optional<at::Tensor> v_descale_, std::optional<double> softmax_scale_,
+    bool is_causal, int64_t window_size_left, int64_t window_size_right,
+    int64_t attention_chunk, double softcap, bool is_rotary_interleaved,
+    std::optional<at::Tensor> scheduler_metadata_, int64_t num_splits,
+    std::optional<bool> pack_gqa_, int64_t sm_margin,
+    std::optional<const at::Tensor>& sinks_);
@@ -0,0 +1,121 @@
+/* Adapted from:
+ * https://github.com/neuralmagic/vllm-flash-attention/blob/90eacc1af2a7c3de62ea249e929ed5faccf38954/csrc/common/pytorch_shim.h
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright 2025 SGLang Team. All Rights Reserved.
+ *
+ * Vendored from sgl-kernel/include/sgl_kernel_torch_shim.h (commit bcf72ccc).
+ * Provides make_pytorch_shim() template for PyTorch op registration type
+ * conversion. NO MODIFICATIONS from the original.
+ */
+
+#pragma once
+
+#include <torch/library.h>
+
+/**
+ * Unfortunately, the type signatures of the flash_attn ops are not compatible
+ * with the PyTorch library bindings. To get around that we use
+ * `make_pytorch_shim` which creates a lambda that exposes the API using
+ * PyTorch compatible types to the types, then converts them to the types
+ * expected by the flash_attn ops. This shims allows us to make minimal changes
+ * to `flash_api.cpp` making it easier to synchronize with upstream changes.
+ *
+ * The `pytorch_library_compatible_type` struct is used to map from the
+ * flash_attn ops types to a PyTorch library compatible one. The main issues is
+ * that the following types are not support by PyTorch library bindings:
+ *  - `int`
+ *  - `float`
+ *  - `std::optional<T> &`
+ *  - `std::optional<const at::Tensor> &`
+ * So we convert them to (respectively):
+ *  - `int64_t`
+ *  - `double`
+ *  - `const std::optional<T>&`
+ *  - `const std::optional<at::Tensor>&`
+ */
+
+template <typename T>
+struct pytorch_library_compatible_type {
+  using type = T;
+  static T convert_from_type(T arg) { return arg; }
+};
+
+template <typename T>
+using pytorch_library_compatible_type_t =
+    typename pytorch_library_compatible_type<T>::type;
+
+template <typename T>
+T convert_from_pytorch_compatible_type(
+    pytorch_library_compatible_type_t<T> arg) {
+  return pytorch_library_compatible_type<T>::convert_from_type(arg);
+}
+
+// Map `c10::optional<T> &` -> `const c10::optional<T>&`
+//  (NOTE: this is bit unsafe but non of the ops in flash_attn mutate
+//   the optional container)
+template <typename T>
+struct pytorch_library_compatible_type<c10::optional<T>&> {
+  using type = const c10::optional<T>&;
+  static c10::optional<T>& convert_from_type(const c10::optional<T>& arg) {
+    return const_cast<c10::optional<T>&>(arg);
+  }
+};
+
+// Map `c10::optional<T>` ->
+//          `c10::optional<pytorch_library_compatible_type_t<T>>`
+//  (NOTE: tested for `c10::optional<int>` -> `c10::optional<int64_t>`)
+template <typename T>
+struct pytorch_library_compatible_type<c10::optional<T>> {
+  using type = c10::optional<pytorch_library_compatible_type_t<T>>;
+  static c10::optional<pytorch_library_compatible_type_t<T>> convert_from_type(
+      c10::optional<T> arg) {
+    return arg;
+  }
+};
+
+// Map `c10::optional<const at::Tensor>&` -> `const c10::optional<at::Tensor>&`
+template <>
+struct pytorch_library_compatible_type<c10::optional<const at::Tensor>&> {
+  using type = const c10::optional<at::Tensor>&;
+  static c10::optional<const at::Tensor>& convert_from_type(
+      const c10::optional<at::Tensor>& arg) {
+    return const_cast<c10::optional<const at::Tensor>&>(
+        reinterpret_cast<const c10::optional<const at::Tensor>&>(arg));
+  }
+};
+
+// Map `int` -> `int64_t`
+template <>
+struct pytorch_library_compatible_type<int> {
+  using type = int64_t;
+  static int convert_from_type(int64_t arg) {
+    TORCH_CHECK(arg <= std::numeric_limits<int>::max(),
+                "int64_t value is too large to be converted to int");
+    TORCH_CHECK(arg >= std::numeric_limits<int>::min(),
+                "int64_t value is too small to be converted to int");
+    return arg;
+  }
+};
+
+// Map `float` -> `double`
+template <>
+struct pytorch_library_compatible_type<float> {
+  using type = double;
+  static float convert_from_type(double arg) {
+    TORCH_CHECK(std::abs(arg) <= std::numeric_limits<float>::max(),
+                "double value is too large to be converted to float");
+    return arg;
+  }
+};
+
+//
+//  Shim Utils
+//
+
+template <typename Ret, typename... Args>
+auto make_pytorch_shim(Ret (*fun)(Args... args)) {
+  return [fun](pytorch_library_compatible_type_t<Args>... args) {
+    return fun(convert_from_pytorch_compatible_type<Args>(args)...);
+  };
+}
@@ -192,10 +192,9 @@ RUN cd /opt/rixl && mkdir -p /app/install && \
 FROM base AS build_deep
 ARG ROCSHMEM_BRANCH="ba0bf0f3"
 ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
-ARG DEEPEP_BRANCH="5d90af8b"
+ARG DEEPEP_BRANCH="e84464ec"
 ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
 ARG DEEPEP_NIC="cx7"
-ARG DEEPEP_ROCM_ARCH="gfx942;gfx950"
 ENV ROCSHMEM_DIR=/opt/rocshmem

 RUN git clone ${ROCSHMEM_REPO} \
@@ -203,11 +202,13 @@ RUN git clone ${ROCSHMEM_REPO} \
 && git checkout ${ROCSHMEM_BRANCH} \
 && mkdir -p projects/rocshmem/build \
 && cd projects/rocshmem/build \
- && bash ../scripts/build_configs/all_backends \
-      -DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
-      -DROCM_PATH=/opt/rocm \
-      -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \
-      -DUSE_EXTERNAL_MPI=OFF
+ && cmake .. \
+    -DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
+    -DROCM_PATH=/opt/rocm \
+    -DCMAKE_POSITION_INDEPENDENT_CODE=ON \
+    -DUSE_EXTERNAL_MPI=OFF \
+ && make -j \
+ && make install

 # Build DeepEP wheel.
 # DeepEP looks for rocshmem at ROCSHMEM_DIR.
@@ -193,7 +193,7 @@ Provide a fast duration→token estimate to improve streaming usage statistics:

 The API server takes care of basic audio I/O and optional chunking before building prompts:

- Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate` using `AudioResampler`.
+- Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate` using `librosa`.
 - Chunking: If `SpeechToTextConfig.allow_audio_chunking` is True and the duration exceeds `max_audio_clip_s`, the server splits the audio into overlapping chunks and generates a prompt per chunk. Overlap is controlled by `overlap_chunk_second`.
 - Energy-aware splitting: When `min_energy_split_window_size` is set, the server finds low-energy regions to minimize cutting within words.

@@ -206,8 +206,8 @@ Relevant server logic:
    async def _preprocess_speech_to_text(...):
        language = self.model_cls.validate_language(request.language)
        ...
-        y, sr = load_audio(bytes_, sr=self.asr_config.sample_rate)
-        duration = get_audio_duration(y=y, sr=sr)
+        y, sr = librosa.load(bytes_, sr=self.asr_config.sample_rate)
+        duration = librosa.get_duration(y=y, sr=sr)
        do_split_audio = (self.asr_config.allow_audio_chunking
                        and duration > self.asr_config.max_audio_clip_s)
        chunks = [y] if not do_split_audio else self._split_audio(y, int(sr))
@@ -206,8 +206,8 @@ Both the `vllm.utils.profiling.cprofile` and `vllm.utils.profiling.cprofile_cont
 used to profile a section of code.

 !!! note
-    The `vllm.utils.profiling` helpers are deprecated and will be removed in
-    `v0.21`. Please use Python's `cProfile` module directly instead.
+    The legacy import paths `vllm.utils.cprofile` and `vllm.utils.cprofile_context` are deprecated.
+    Please use `vllm.utils.profiling.cprofile` and `vllm.utils.profiling.cprofile_context` instead.

 ### Example usage - decorator

@@ -132,16 +132,6 @@ Priority is **1 = highest** (tried first).
 | 6 | `FLASHINFER_MLA_SPARSE`**\*** |
 | 7 | `FLASHMLA_SPARSE` |

-**Ampere/Hopper (SM 8.x-9.x):**
-
-| Priority | Backend |
-| -------- | ------- |
-| 1 | `FLASH_ATTN_MLA` |
-| 2 | `FLASHMLA` |
-| 3 | `FLASHINFER_MLA` |
-| 4 | `TRITON_MLA` |
-| 5 | `FLASHMLA_SPARSE` |
-
 > **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise.
 >
 > **Note:** ROCm and CPU platforms have their own selection logic. See the platform-specific documentation for details.
@@ -209,6 +199,7 @@ configuration.

 | Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------ | --------- | ----------- | ---------- | ---- | ------ | --------- | --- | --------------- | ------------ |
+| `CUTLASS_FA3_MLA_SPARSE` | bf16 | `auto` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x |
 | `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
 | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
 | `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
@@ -21,7 +21,6 @@ or just on the low or high end.
 | Fusion                                                                         | `PassConfig` flag            | Fused operations                               | Default at                     | E2E Speedup        | Fullgraph | `num_tokens` |
 | ------------------------------------------------------------------------------ | ---------------------------- | ---------------------------------------------- | ------------------------------ | ------------------ | --------- | ------------ |
 | [AllReduce + RMSNorm](#allreduce--rmsnorm-fuse_allreduce_rms)                  | `fuse_allreduce_rms`         | All-reduce → RMSNorm (+residual_add) (→ quant) | O2 (Hopper/Blackwell + TP > 1) | 5-20%              | No        | Low          |
-| [MiniMax QK Norm](#minimax-qk-norm-fuse_minimax_qk_norm)                       | `fuse_minimax_qk_norm`       | Q/K variance all-reduce → Q/K RMSNorm          | Off by default                 | 2-3%               | No        | Low          |
 | [Attention + Quant](#attention--quantization-fuse_attn_quant)                  | `fuse_attn_quant`            | Attention output → FP8/NVFP4 quant             | Off by default                 | 3-7%               | Yes       | Always       |
 | [MLA Attention + Quant](#attention--quantization-fuse_attn_quant)              | `fuse_attn_quant`            | MLA Attention output → FP8/NVFP4 quant         | Off by default                 | TBD                | Yes       | Always       |
 | [RoPE + KV-Cache Update](#rope--kv-cache-update-fuse_rope_kvcache)             | `fuse_rope_kvcache`          | Rotary embedding → KV cache write              | O2 (ROCm/AITER only)           | 2-4%               | No        | Low          |
@@ -41,7 +40,6 @@ The table below lists the quantization schemes supported by each fusion on each
 | Fusion                       | SM100 (Blackwell)                        | SM90 (Hopper)                            | SM89 (Ada)                               | SM80 (Ampere) | ROCm                                     |
 | ---------------------------- | ---------------------------------------- | ---------------------------------------- | ---------------------------------------- | ------------- | ---------------------------------------- |
 | `fuse_allreduce_rms`         | FP16/BF16, FP8 static, NVFP4             | FP16/BF16, FP8 static                    | —                                        | —             | —                                        |
-| `fuse_minimax_qk_norm`\*     | FP16/BF16                                | FP16/BF16                                | FP16/BF16                                | FP16/BF16     | —                                        |
 | `fuse_attn_quant`\*          | FP8 static\*, NVFP4\*                    | FP8 static\*                             | FP8 static\*                             | —             | FP8 static\*                             |
 | `fuse_attn_quant` (MLA)\*    | FP8 static\*, NVFP4\*                    | FP8 static\*                             | FP8 static\*                             | —             | FP8 static(untested)\*                   |
 | `fuse_rope_kvcache`          | —                                        | —                                        | —                                        | —             | FP16/BF16                                |
@@ -56,9 +54,6 @@ The table below lists the quantization schemes supported by each fusion on each
 fused quantization output. See the [`fuse_attn_quant` section](#attention--quantization-fuse_attn_quant)
 for per-backend details.

-\* `fuse_minimax_qk_norm` is a model-specific pass for `MiniMaxM2ForCausalLM`. It also requires
-tensor parallelism (`tp_size > 1`) and the CUDA custom op `minimax_allreduce_rms_qk`.
-
 † `enable_sp` and `fuse_gemm_comms` are only autoconfigured for SM90 today;
 other architectures support requires setting `PassConfig.sp_min_token_num` explicitly.
 SM100 support also requires setting `VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel`.
@@ -189,35 +184,6 @@ If these conditions are set, the fusion is enabled automatically for optimizatio

 - Pass: [`vllm/compilation/passes/fusion/rope_kvcache_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rope_kvcache_fusion.py)

-### MiniMax QK Norm (`fuse_minimax_qk_norm`)
-
-!!! info
-    This is a MiniMax-specific compile pass. It is currently only enabled when all of the following hold:
-    the model architecture is `MiniMaxM2ForCausalLM`, tensor parallelism is enabled (`tp_size > 1`),
-    and the CUDA custom op `minimax_allreduce_rms_qk` is available. It is not enabled by default at any
-    optimization level.
-
-**What it fuses.** Fuses the MiniMax M2 Q/K normalization path that performs an all-reduce over the
-per-token Q/K variances before applying RMS normalization to Q and K.
-
-This pass is distinct from [`enable_qk_norm_rope_fusion`](#qk-norm--rope-enable_qk_norm_rope_fusion):
-`fuse_minimax_qk_norm` targets MiniMax M2's tensor-parallel all-reduce + RMSNorm sequence, while
-`enable_qk_norm_rope_fusion` targets the later Q/K RMSNorm + RoPE sequence used by several other models.
-
-Example:
-
-```bash
-vllm serve MiniMaxAI/MiniMax-M2.5 \
-  --tensor-parallel-size 4 \
-  --compilation-config '{"mode": 3, "pass_config": {"fuse_minimax_qk_norm": true}}'
-```
-
-**Code locations.**
-
- Pass: [`vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py)
- CUDA op: [`csrc/minimax_reduce_rms_kernel.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/minimax_reduce_rms_kernel.cu) (`minimax_allreduce_rms_qk`)
- Workspace helper: [`vllm/model_executor/layers/mamba/lamport_workspace.py`](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/mamba/lamport_workspace.py)
-
 ### Sequence Parallelism (`enable_sp`)

 **What it fuses.** Replaces all-reduce collectives with reduce-scatter + local RMSNorm + all-gather,
@@ -300,12 +300,12 @@ Full example: [examples/offline_inference/audio_language.py](../../examples/offl
 Speech-to-text models like Whisper have a maximum audio length they can process (typically 30 seconds). For longer audio files, vLLM provides a utility to intelligently split audio into chunks at quiet points to minimize cutting through speech.

 ```python
+import librosa
 from vllm import LLM, SamplingParams
 from vllm.multimodal.audio import split_audio
-from vllm.multimodal.media.audio import load_audio

 # Load long audio file
-audio, sr = load_audio("long_audio.wav", sr=16000)
+audio, sr = librosa.load("long_audio.wav", sr=16000)

 # Split into chunks at low-energy (quiet) regions
 chunks = split_audio(
@@ -832,7 +832,7 @@ Then, you can use the OpenAI client as follows:
        base_url=openai_api_base,
    )

-    # Any format supported by soundfile/PyAV is supported
+    # Any format supported by librosa is supported
    audio_url = AudioAsset("winning_call").url
    audio_base64 = encode_base64_content_from_url(audio_url)

@@ -9,14 +9,14 @@
 - Online APIs:
    - Pooling API (`/pooling`)

-The difference between the (sequence) embedding task and the token embedding task is that (sequence) embedding outputs one embedding for each sequence, while token embedding outputs an embedding for each token.
+The difference between the (sequence) embedding task and the token embedding task is that (sequence) embedding outputs one embedding for each sequence, while token embedding outputs a embedding for each token.

 Many embedding models support both (sequence) embedding and token embedding. For further details on (sequence) embedding, please refer to [this page](embed.md).

 !!! note

    Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task (embed) is not 
-    what you want, you need to manually specify it via `PoolerConfig(task="token_embed")` offline or
+    what you want, you need to manually specify it via via `PoolerConfig(task="token_embed")` offline or
    `--pooler-config.task token_embed` online.

 ## Typical Use Cases
@@ -682,24 +682,6 @@ Speech2Text models trained specifically for Automatic Speech Recognition.
 !!! note
    `VoxtralForConditionalGeneration` requires `mistral-common[audio]` to be installed.

-#### Realtime Transcription
-
-Speech models that support streaming transcription via the
-[`/v1/realtime`](../serving/openai_compatible_server.md#realtime-api)
-WebSocket endpoint.
-
-| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
-| ------------ | ------ | ----------------- | -------------------- | ------------------------- |
-| `VoxtralRealtimeGeneration` | Voxtral Realtime | `mistralai/Voxtral-Mini-4B-Realtime-2602` | | |
-| `Qwen3ASRRealtimeGeneration` | Qwen3-ASR Realtime | `Qwen/Qwen3-ASR-0.6B` | | |
-
-!!! note
-    `VoxtralRealtimeGeneration` requires `mistral-common[audio]` to be installed, and must be served with `--tokenizer-mode mistral`.
-
-    `Qwen3ASRRealtimeGeneration` is not auto-detected from `config.json`.
-    You must pass `--hf-overrides '{"architectures":["Qwen3ASRRealtimeGeneration"]}'`
-    when serving.
-
 ## Pooling Models

 See [this page](pooling_models/README.md) for more information on how to use pooling models.
@@ -60,7 +60,7 @@ We currently support the following OpenAI APIs:
 - [Translation API](#translations-api) (`/v1/audio/translations`)
    - Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#transcription).
 - [Realtime API](#realtime-api) (`/v1/realtime`)
-    - Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#realtime-transcription).
+    - Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#transcription).

 In addition, we have the following custom APIs:

@@ -1,117 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Disaggregated multimodal serving: render → generate round-trip.
-
-Demonstrates the two-phase disaggregated flow:
-  1. /v1/chat/completions/render  – preprocesses a multimodal chat request
-     into token IDs and serialized tensor features.
-  2. /inference/v1/generate       – runs inference on the preprocessed tokens.
-
-The render response is passed *directly* to generate with only
-``sampling_params`` added, showing that the two endpoints compose with
-zero client-side transformation.
-
-Launch the server first:
-
-    vllm serve Qwen/Qwen3-VL-2B-Instruct \
-        --dtype bfloat16 --max-model-len 4096 --enforce-eager
-
-Then run this script:
-
-    python example_mm_serve.py
-"""
-
-import io
-
-import pybase64 as base64
-import requests
-from PIL import Image
-from transformers import AutoTokenizer
-
-BASE_URL = "http://localhost:8000"
-MODEL_NAME = "Qwen/Qwen3-VL-2B-Instruct"
-
-
-def make_data_url(image: Image.Image) -> str:
-    """Encode a PIL image as a base64 data URL."""
-    buf = io.BytesIO()
-    image.save(buf, format="PNG")
-    b64 = base64.b64encode(buf.getvalue()).decode()
-    return f"data:image/png;base64,{b64}"
-
-
-def main():
-    # -- Step 1: Create a test image (solid red) -------------------------
-    image = Image.new("RGB", (224, 224), color=(255, 0, 0))
-    data_url = make_data_url(image)
-    print("Created 224x224 red test image")
-
-    # -- Step 2: Render (preprocess) -------------------------------------
-    render_payload = {
-        "model": MODEL_NAME,
-        "messages": [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image_url", "image_url": {"url": data_url}},
-                    {
-                        "type": "text",
-                        "text": "What color is this image? Answer in one word.",
-                    },
-                ],
-            }
-        ],
-    }
-
-    print("\n--- Render ---")
-    render_resp = requests.post(
-        f"{BASE_URL}/v1/chat/completions/render", json=render_payload
-    )
-    render_resp.raise_for_status()
-    render_data = render_resp.json()
-
-    print(f"Response keys: {list(render_data.keys())}")
-    print(f"Number of token_ids: {len(render_data['token_ids'])}")
-
-    features = render_data.get("features")
-    if features and features.get("kwargs_data"):
-        print(f"kwargs_data modalities: {list(features['kwargs_data'].keys())}")
-        for modality, items in features["kwargs_data"].items():
-            print(
-                f"  {modality}: {len(items)} item(s), "
-                f"first item type: {type(items[0])} length: {len(items[0])}"
-                if items
-                else "First item: (empty)"
-            )
-    else:
-        print("WARNING: no kwargs_data in render response")
-
-    # -- Step 3: Generate (inference) ------------------------------------
-    # Pass the render output directly — only add sampling_params.
-    generate_payload = render_data
-    generate_payload["sampling_params"] = {
-        "max_tokens": 20,
-        "temperature": 0.0,
-    }
-
-    print("\n--- Generate ---")
-    gen_resp = requests.post(f"{BASE_URL}/inference/v1/generate", json=generate_payload)
-    gen_resp.raise_for_status()
-    gen_data = gen_resp.json()
-
-    # -- Step 4: Decode & print ------------------------------------------
-    output_ids = gen_data["choices"][0]["token_ids"]
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    text = tokenizer.decode(output_ids, skip_special_tokens=True)
-
-    print(f"Output token count: {len(output_ids)}")
-    print(f"Generated text: {text!r}")
-
-    if "red" in text.lower():
-        print("\nModel correctly identified the red image.")
-    else:
-        print(f"\nWARNING: Expected 'red' in output, got: {text!r}")
-
-
-if __name__ == "__main__":
-    main()
@@ -267,7 +267,7 @@ def run_audio(model: str, max_completion_tokens: int) -> None:
                    {
                        "type": "input_audio",
                        "input_audio": {
-                            # Any format supported by soundfile/PyAV is supported
+                            # Any format supported by librosa is supported
                            "data": audio_base64,
                            "format": "wav",
                        },
@@ -292,7 +292,7 @@ def run_audio(model: str, max_completion_tokens: int) -> None:
                    {
                        "type": "audio_url",
                        "audio_url": {
-                            # Any format supported by soundfile/PyAV is supported
+                            # Any format supported by librosa is supported
                            "url": audio_url
                        },
                    },
@@ -316,7 +316,7 @@ def run_audio(model: str, max_completion_tokens: int) -> None:
                    {
                        "type": "audio_url",
                        "audio_url": {
-                            # Any format supported by soundfile/PyAV is supported
+                            # Any format supported by librosa is supported
                            "url": f"data:audio/ogg;base64,{audio_base64}"
                        },
                    },
@@ -12,6 +12,7 @@ model, for example:
 Requirements:
 - vllm with audio support
 - websockets
+- librosa
 - numpy

 The script:
@@ -25,12 +26,12 @@ import argparse
 import asyncio
 import json

+import librosa
 import numpy as np
 import pybase64 as base64
 import websockets

 from vllm.assets.audio import AudioAsset
-from vllm.multimodal.media.audio import load_audio


 def audio_to_pcm16_base64(audio_path: str) -> str:
@@ -38,7 +39,7 @@ def audio_to_pcm16_base64(audio_path: str) -> str:
    Load an audio file and convert it to base64-encoded PCM16 @ 16kHz.
    """
    # Load audio and resample to 16kHz mono
-    audio, _ = load_audio(audio_path, sr=16000, mono=True)
+    audio, _ = librosa.load(audio_path, sr=16000, mono=True)
    # Convert to PCM16
    pcm16 = (audio * 32767).astype(np.int16)
    # Encode as base64
@@ -170,7 +170,6 @@ eles = "eles"
 datas = "datas"
 ser = "ser"
 ure = "ure"
-VALU = "VALU"
 # Walsh-Hadamard Transform
 wht = "wht"
 WHT = "WHT"
@@ -32,7 +32,9 @@ pyzmq >= 25.0.0
 msgspec
 gguf >= 0.17.0
 mistral_common[image] >= 1.11.0
+av  # required for audio in video IO
 opencv-python-headless >= 4.13.0    # required for video IO
+soundfile  # required for audio IO
 pyyaml
 six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12
 setuptools>=77.0.3,<81.0.0; python_version > '3.11' # Setuptools is used by triton, we need to ensure a modern version is installed for 3.12+ so that it does not try to import distutils, which was removed in 3.12
@@ -20,4 +20,6 @@ conch-triton-kernels==1.2.1
 timm>=1.0.17
 # amd-quark: required for Quark quantization on ROCm 
 # To be consistent with test_quark.py
-amd-quark>=0.8.99
+amd-quark>=0.8.99
+# Required for faster safetensors model loading
+fastsafetensors >= 0.2.2
@@ -55,7 +55,7 @@ arctic-inference==0.1.1 # Required for suffix decoding test
 numba==0.61.2 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs,azure]==0.15.7
-fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@0.2.2 # PyPI only ships CUDA wheels
+fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
 instanttensor>=0.1.5
 pydantic>=2.12 # 2.11 leads to error on python 3.13
 decord==0.6.0
@@ -76,7 +76,9 @@ attrs==26.1.0
 audioread==3.0.1
    # via librosa
 av==16.1.0
-    # via -r requirements/test/rocm.in
+    # via
+    #   -r requirements/test/../common.txt
+    #   -r requirements/test/rocm.in
 azure-core==1.39.0
    # via
    #   azure-identity
@@ -275,8 +277,10 @@ fastar==0.10.0
    # via fastapi-cloud-cli
 fastparquet==2026.3.0
    # via genai-perf
-fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@65d80088fca7a8f567fba30415fbcc80f7d2259c
-    # via -r requirements/test/rocm.in
+fastsafetensors==0.2.2
+    # via
+    #   -c requirements/rocm.txt
+    #   -r requirements/test/rocm.in
 filelock==3.25.2
    # via
    #   -c requirements/common.txt
@@ -1329,6 +1333,7 @@ sortedcontainers==2.4.0
    # via hypothesis
 soundfile==0.13.1
    # via
+    #   -r requirements/test/../common.txt
    #   -r requirements/test/rocm.in
    #   genai-perf
    #   librosa
@@ -1085,18 +1085,14 @@ setup(
    install_requires=get_requirements(),
    extras_require={
        # AMD Zen CPU optimizations via zentorch
-        "zen": [
-            "zentorch-weekly==5.2.1.dev20260408"
-        ],  # Zentorch has weekly releases. This pulls the known-good version.
+        "zen": ["zentorch"],
        "bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
        "tensorizer": ["tensorizer==2.10.1"],
        "fastsafetensors": ["fastsafetensors >= 0.2.2"],
        "instanttensor": ["instanttensor >= 0.1.5"],
        "runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"],
        "audio": [
-            "av",
            "scipy",
-            "soundfile",
            "mistral_common[audio]",
        ],  # Required for audio processing
        "video": [],  # Kept for backwards compatibility
@@ -38,8 +38,6 @@ llm = LLM(
    distributed_executor_backend="external_launcher",
    gpu_memory_utilization=random.uniform(0.7, 0.9),
    seed=0,
-    max_model_len=1024,
-    max_num_seqs=16,
 )

 outputs = llm.generate(prompts, sampling_params)
@@ -13,6 +13,7 @@ import io
 import time
 from statistics import mean, median

+import librosa
 import pytest
 import soundfile
 import torch
@@ -20,7 +21,6 @@ from datasets import load_dataset
 from evaluate import load
 from transformers.models.whisper.english_normalizer import EnglishTextNormalizer

-from vllm.multimodal.audio import get_audio_duration
 from vllm.tokenizers import get_tokenizer

 from ....models.registry import HF_EXAMPLE_MODELS
@@ -84,7 +84,7 @@ async def process_dataset(model, client, data, concurrent_request):
        trust_remote_code=model_info.trust_remote_code,
    )

-    # Warmup call as the first `load_audio` server-side is quite slow.
+    # Warmup call as the first `librosa.load` server-side is quite slow.
    audio, sr = data[0]["audio"]["array"], data[0]["audio"]["sampling_rate"]
    _ = await bound_transcribe(sem, client, tokenizer, (audio, sr), "")

@@ -118,7 +118,7 @@ def print_performance_metrics(results, total_time):

 def add_duration(sample):
    y, sr = sample["audio"]["array"], sample["audio"]["sampling_rate"]
-    sample["duration_ms"] = get_audio_duration(y=y, sr=sr) * 1000
+    sample["duration_ms"] = librosa.get_duration(y=y, sr=sr) * 1000
    return sample


@@ -5,6 +5,7 @@ import asyncio
 import json
 import warnings

+import librosa
 import numpy as np
 import pybase64 as base64
 import pytest
@@ -13,7 +14,6 @@ import websockets
 from tests.entrypoints.openai.conftest import add_attention_backend
 from tests.utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
 from vllm.assets.audio import AudioAsset
-from vllm.multimodal.media.audio import load_audio

 # Increase engine iteration timeout for ROCm where first-use JIT compilation
 # can exceed the default 60s, causing a silent deadlock in feed_tokens.
@@ -56,7 +56,7 @@ async def send_event(ws, event: dict) -> None:
 def mary_had_lamb_audio_chunks() -> list[str]:
    """Audio split into ~1 second chunks for streaming."""
    path = AudioAsset("mary_had_lamb").get_local_path()
-    audio, _ = load_audio(str(path), sr=16000, mono=True)
+    audio, _ = librosa.load(str(path), sr=16000, mono=True)

    # Split into ~0.1 second chunks (1600 samples at 16kHz)
    chunk_size = 1600
@@ -6,6 +6,7 @@ import asyncio
 import io
 import json

+import librosa
 import numpy as np
 import openai
 import pytest
@@ -13,7 +14,6 @@ import pytest_asyncio
 import soundfile as sf

 from tests.utils import RemoteOpenAIServer
-from vllm.multimodal.media.audio import load_audio
 from vllm.platforms import current_platform

 MODEL_NAME = "openai/whisper-large-v3-turbo"
@@ -134,7 +134,7 @@ async def test_bad_requests(mary_had_lamb, whisper_client):
@pytest.mark.asyncio
 async def test_long_audio_request(mary_had_lamb, whisper_client):
    mary_had_lamb.seek(0)
-    audio, sr = load_audio(mary_had_lamb)
+    audio, sr = librosa.load(mary_had_lamb)
    # Add small silence after each audio for repeatability in the split process
    audio = np.pad(audio, (0, 1600))
    repeated_audio = np.tile(audio, 10)
@@ -7,6 +7,7 @@ import io
 import json

 import httpx
+import librosa
 import numpy as np
 import openai
 import pytest
@@ -16,7 +17,6 @@ import soundfile as sf
 from tests.entrypoints.openai.conftest import add_attention_backend
 from tests.utils import RemoteOpenAIServer
 from vllm.logger import init_logger
-from vllm.multimodal.media.audio import load_audio

 logger = init_logger(__name__)

@@ -264,7 +264,7 @@ async def test_long_audio_request(foscolo, client_and_model):
    if model_name == "google/gemma-3n-E2B-it":
        pytest.skip("Gemma3n does not support long audio requests")
    foscolo.seek(0)
-    audio, sr = load_audio(foscolo)
+    audio, sr = librosa.load(foscolo)
    repeated_audio = np.tile(audio, 2)
    # Repeated audio to buffer
    buffer = io.BytesIO()
@@ -1,111 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Roundtrip tests for multimodal serde used by the disagg generate endpoint."""
-
-import torch
-
-from vllm.entrypoints.serve.disagg.mm_serde import (
-    decode_mm_kwargs_item,
-    encode_mm_kwargs_item,
-)
-from vllm.entrypoints.serve.disagg.protocol import (
-    MultiModalFeatures,
-    PlaceholderRangeInfo,
-)
-from vllm.multimodal.inputs import (
-    MultiModalBatchedField,
-    MultiModalFieldElem,
-    MultiModalFlatField,
-    MultiModalKwargsItem,
-    MultiModalSharedField,
-)
-
-
-def test_mm_kwargs_item_roundtrip():
-    """Full roundtrip test with all three field types and multiple dtypes."""
-    e1 = MultiModalFieldElem(
-        data=torch.zeros(1000, dtype=torch.bfloat16),
-        field=MultiModalBatchedField(),
-    )
-    e2 = MultiModalFieldElem(
-        data=torch.ones(100, dtype=torch.int32),
-        field=MultiModalSharedField(batch_size=4),
-    )
-    e3 = MultiModalFieldElem(
-        data=torch.randn(20, dtype=torch.float32),
-        field=MultiModalFlatField(slices=[slice(0, 10), slice(10, 20)], dim=0),
-    )
-
-    item = MultiModalKwargsItem({"pixel_values": e1, "grid_thw": e2, "embeds": e3})
-    encoded = encode_mm_kwargs_item(item)
-
-    # Encoded result is a base64 string
-    assert isinstance(encoded, str)
-
-    decoded = decode_mm_kwargs_item(encoded)
-
-    assert set(decoded.keys()) == {"pixel_values", "grid_thw", "embeds"}
-    assert torch.equal(item["pixel_values"].data, decoded["pixel_values"].data)
-    assert torch.equal(item["grid_thw"].data, decoded["grid_thw"].data)
-    assert torch.equal(item["embeds"].data, decoded["embeds"].data)
-    assert isinstance(decoded["pixel_values"].field, MultiModalBatchedField)
-    assert isinstance(decoded["grid_thw"].field, MultiModalSharedField)
-    assert isinstance(decoded["embeds"].field, MultiModalFlatField)
-
-
-def test_mm_kwargs_item_none_data():
-    """Roundtrip with None data field."""
-    elem = MultiModalFieldElem(
-        data=None,
-        field=MultiModalSharedField(batch_size=2),
-    )
-    item = MultiModalKwargsItem({"empty": elem})
-    encoded = encode_mm_kwargs_item(item)
-    decoded = decode_mm_kwargs_item(encoded)
-
-    assert decoded["empty"].data is None
-    assert isinstance(decoded["empty"].field, MultiModalSharedField)
-
-
-def test_mm_kwargs_item_nested_tensors():
-    """Roundtrip with nested tensor data."""
-    nested = [torch.randn(3, 4), torch.randn(5, 4)]
-    elem = MultiModalFieldElem(
-        data=nested,
-        field=MultiModalBatchedField(),
-    )
-    item = MultiModalKwargsItem({"nested": elem})
-    encoded = encode_mm_kwargs_item(item)
-    decoded = decode_mm_kwargs_item(encoded)
-
-    decoded_data = decoded["nested"].data
-    assert len(decoded_data) == 2
-    assert torch.equal(nested[0], decoded_data[0])
-    assert torch.equal(nested[1], decoded_data[1])
-
-
-def test_mm_features_with_kwargs_data():
-    """Test that MultiModalFeatures can carry serialized tensor data."""
-    elem = MultiModalFieldElem(
-        data=torch.randn(5, 3, dtype=torch.float32),
-        field=MultiModalBatchedField(),
-    )
-    item = MultiModalKwargsItem({"pixel_values": elem})
-    encoded = encode_mm_kwargs_item(item)
-
-    features = MultiModalFeatures(
-        mm_hashes={"image": ["abc123"]},
-        mm_placeholders={"image": [PlaceholderRangeInfo(offset=0, length=10)]},
-        kwargs_data={"image": [encoded]},
-    )
-
-    # JSON roundtrip
-    json_str = features.model_dump_json()
-    features2 = MultiModalFeatures.model_validate_json(json_str)
-
-    assert features2.mm_hashes == {"image": ["abc123"]}
-    assert features2.kwargs_data is not None
-    assert len(features2.kwargs_data["image"]) == 1
-
-    decoded = decode_mm_kwargs_item(features2.kwargs_data["image"][0])
-    assert torch.equal(elem.data, decoded["pixel_values"].data)
@@ -1,158 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for multimodal features through the /inference/v1/generate endpoint.
-
-Mirrors test_serving_tokens.py but exercises the multimodal piping
-using Qwen/Qwen3-VL-2B-Instruct end-to-end via the server's /render ->
-/generate -> /detokenize path. Intentionally avoids running the HF
-processor in the pytest parent process to keep os.fork() in sibling
-tests (e.g. test_weight_transfer_llm.py) deadlock-free.
-"""
-
-import os
-
-import httpx
-import pytest
-import pytest_asyncio
-from PIL import Image
-
-from tests.utils import RemoteOpenAIServer
-from vllm.multimodal.utils import encode_image_url
-
-MODEL_NAME = "Qwen/Qwen3-VL-2B-Instruct"
-GEN_ENDPOINT = "/inference/v1/generate"
-RENDER_ENDPOINT = "/v1/chat/completions/render"
-DETOKENIZE_ENDPOINT = "/detokenize"
-
-
-@pytest.fixture(scope="module")
-def test_image():
-    return Image.new("RGB", (224, 224), color=(255, 0, 0))
-
-
-@pytest.fixture(scope="module")
-def server():
-    args = [
-        "--dtype",
-        "bfloat16",
-        "--max-model-len",
-        "4096",
-        "--enforce-eager",
-        "--no-enable-prefix-caching",
-    ]
-
-    envs = os.environ.copy()
-    envs["VLLM_ROCM_USE_SKINNY_GEMM"] = "0"
-
-    with RemoteOpenAIServer(MODEL_NAME, args, env_dict=envs) as remote_server:
-        yield remote_server
-
-
-@pytest_asyncio.fixture
-async def client(server: RemoteOpenAIServer):
-    transport = httpx.AsyncHTTPTransport(uds=server.uds) if server.uds else None
-    headers = {"Authorization": f"Bearer {server.DUMMY_API_KEY}"}
-    async with httpx.AsyncClient(
-        transport=transport,
-        base_url=server.url_root,
-        timeout=600,
-        headers=headers,
-    ) as c:
-        yield c
-
-
-@pytest.mark.asyncio
-async def test_render_to_generate_roundtrip(client, test_image):
-    """End-to-end: render a multimodal chat -> feed into generate -> decode.
-
-    All preprocessing and detokenization happens in the server subprocess;
-    the pytest parent never imports transformers or touches torch tensors.
-    """
-    data_url = encode_image_url(test_image, format="PNG")
-
-    render_payload = {
-        "model": MODEL_NAME,
-        "messages": [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image_url", "image_url": {"url": data_url}},
-                    {
-                        "type": "text",
-                        "text": "What color is this image? Answer in one word.",
-                    },
-                ],
-            }
-        ],
-    }
-
-    render_resp = await client.post(RENDER_ENDPOINT, json=render_payload)
-    render_resp.raise_for_status()
-    render_data = render_resp.json()
-
-    # Validate render output structure: keys exist and values are non-empty
-    # and well-typed.
-    assert "token_ids" in render_data
-    assert isinstance(render_data["token_ids"], list)
-    assert len(render_data["token_ids"]) > 0
-    assert all(isinstance(t, int) for t in render_data["token_ids"])
-
-    assert "features" in render_data
-    features = render_data["features"]
-    assert features is not None
-    assert isinstance(features, dict)
-
-    assert "mm_hashes" in features
-    assert "image" in features["mm_hashes"]
-    image_hashes = features["mm_hashes"]["image"]
-    assert isinstance(image_hashes, list)
-    assert len(image_hashes) > 0
-    assert all(isinstance(h, str) and h for h in image_hashes)
-
-    assert "mm_placeholders" in features
-    assert "image" in features["mm_placeholders"]
-    image_placeholders = features["mm_placeholders"]["image"]
-    assert isinstance(image_placeholders, list)
-    assert len(image_placeholders) > 0
-    for p in image_placeholders:
-        assert isinstance(p.get("offset"), int)
-        assert isinstance(p.get("length"), int)
-        assert p["length"] > 0
-
-    assert "kwargs_data" in features
-    assert "image" in features["kwargs_data"]
-    assert len(features["kwargs_data"]["image"]) > 0
-
-    # Build generate request from render output
-    generate_payload = render_data
-    generate_payload["sampling_params"] = {
-        "max_tokens": 10,
-        "temperature": 0.0,
-    }
-
-    gen_resp = await client.post(GEN_ENDPOINT, json=generate_payload)
-    gen_resp.raise_for_status()
-    gen_data = gen_resp.json()
-
-    assert "choices" in gen_data
-    assert isinstance(gen_data["choices"], list)
-    assert len(gen_data["choices"]) >= 1
-    choice = gen_data["choices"][0]
-    assert "token_ids" in choice
-    assert isinstance(choice["token_ids"], list)
-    assert len(choice["token_ids"]) > 0
-    assert all(isinstance(t, int) for t in choice["token_ids"])
-
-    detok_resp = await client.post(
-        DETOKENIZE_ENDPOINT,
-        json={"model": MODEL_NAME, "tokens": choice["token_ids"]},
-    )
-    detok_resp.raise_for_status()
-    detok_data = detok_resp.json()
-    assert "prompt" in detok_data
-    text = detok_data["prompt"]
-    assert isinstance(text, str)
-    assert len(text) > 0
-    assert "red" in text.lower(), (
-        f"Expected model to identify the red image, got: {text!r}"
-    )
@@ -135,70 +135,3 @@ def test_fused_topk_bias(
        topk_weights_ref.to(torch.float32), topk_weights, atol=1e-2, rtol=1e-2
    )
    torch.testing.assert_close(topk_ids_ref.to(torch.int32), topk_ids, atol=0, rtol=0)
-
-
-@pytest.mark.skipif(
-    not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
-)
-@pytest.mark.parametrize("num_experts", [6, 8, 16])
-@pytest.mark.parametrize("topk", [3, 4])
-@pytest.mark.parametrize("scoring_func", ["softmax", "sigmoid"])
-@pytest.mark.parametrize("bad_value", [float("nan"), float("inf")])
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32])
-def test_fused_topk_nan_inf_clamp(
-    num_experts: int,
-    topk: int,
-    scoring_func: str,
-    bad_value: float,
-    dtype: torch.dtype,
-):
-    """Regression test for the NaN/Inf clamp in topk_softmax_kernels.cu.
-
-    Degenerate hidden states (e.g., from CUDA graph padding) can produce
-    NaN/Inf gating logits. Without the clamp, softmax/sigmoid outputs are
-    NaN and the argmax loop picks expert 0 for every top-k slot (since
-    "NaN > NaN" is false per IEEE 754), yielding duplicate expert IDs that
-    crash downstream MoE sort kernels. The fix clamps NaN/Inf to 0 before
-    argmax so index tie-breaking selects unique experts [0, 1, ..., k-1].
-    """
-    torch.manual_seed(0)
-    num_tokens = 4
-    hidden_size = 1024
-    hidden_states = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda")
-
-    # Row 0: all normal. Rows 1-3: fully poisoned with NaN or Inf.
-    gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
-    gating_output[1:, :] = bad_value
-
-    topk_weights, topk_ids, _ = fused_topk(
-        hidden_states=hidden_states,
-        gating_output=gating_output,
-        topk=topk,
-        renormalize=False,
-        scoring_func=scoring_func,
-    )
-
-    # Normal row must still match the torch reference.
-    ref_weights, ref_ids = torch_topk(
-        gating_output=gating_output[:1],
-        topk=topk,
-        renormalize=False,
-        scoring_func=scoring_func,
-    )
-    torch.testing.assert_close(
-        ref_weights.to(torch.float32), topk_weights[:1], atol=1e-2, rtol=1e-2
-    )
-    torch.testing.assert_close(ref_ids.to(torch.int32), topk_ids[:1], atol=0, rtol=0)
-
-    # Poisoned rows: IDs must be unique (no duplicates) and weights must be
-    # finite (no NaN/Inf propagation into downstream MoE kernels).
-    for row in range(1, num_tokens):
-        row_ids = topk_ids[row]
-        assert row_ids.unique().numel() == topk, (
-            f"Row {row} has duplicate expert IDs {row_ids.tolist()} "
-            f"(bad_value={bad_value}, scoring_func={scoring_func})"
-        )
-        assert torch.isfinite(topk_weights[row]).all(), (
-            f"Row {row} has non-finite weights {topk_weights[row].tolist()} "
-            f"(bad_value={bad_value}, scoring_func={scoring_func})"
-        )
@@ -1,57 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import pytest
-import torch
-
-from vllm.model_executor.models.gemma4 import (
-    gemma4_fused_routing_kernel_triton,
-    gemma4_routing_function_torch,
-)
-
-
-def sort_by_id(w, ids):
-    order = ids.argsort(dim=-1)
-    return w.gather(1, order), ids.gather(1, order)
-
-
-# Gemma4 MoE Model has context length of 250K
-# the minus 1 is to ensure that edge cases are tested
-@pytest.mark.parametrize("num_tokens", [1, 2, 2048, 250000])
-@pytest.mark.parametrize("num_experts", [128])  # gemma4 moe experts
-@pytest.mark.parametrize("topk", [8])  # gemma4 topk
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32])
-def test_gemma4_routing_kernel_triton(
-    num_tokens: int,
-    num_experts: int,
-    topk: int,
-    dtype: torch.dtype,
-):
-    torch.manual_seed(0)
-
-    gating = torch.randn(num_tokens, num_experts, dtype=dtype, device="cuda")
-    scales = torch.rand(num_experts, dtype=torch.float32, device="cuda")
-
-    ref_w, ref_ids = gemma4_routing_function_torch(gating, topk, scales)
-    tri_w, tri_ids = gemma4_fused_routing_kernel_triton(gating, topk, scales)
-
-    # Sort by expert id — to remove tie-breaking differences
-    ref_ws, ref_is = sort_by_id(ref_w, ref_ids)
-    tri_ws, tri_is = sort_by_id(tri_w, tri_ids)
-
-    ids_match = (ref_is == tri_is).all().item()
-    weights_match = torch.allclose(ref_ws, tri_ws, atol=1e-2, rtol=1e-2)
-    all_match = ids_match and weights_match
-    max_err = (ref_ws - tri_ws).abs().max().item()
-    print(
-        f"T={num_tokens:5d} E={num_experts:4d} K={topk} "
-        f"{str(dtype).split('.')[-1]:7s} ids={ids_match} max_Δweight={max_err:.2e}"
-    )
-    if not all_match:
-        bad = (ref_is != tri_is).any(dim=-1).nonzero(as_tuple=True)[0]
-        if len(bad):
-            r = bad[0].item()
-            print(
-                f"  first bad row {r}: ref_ids={ref_ids[r].tolist()} "
-                f"tri_ids={tri_ids[r].tolist()}"
-            )
-        assert all_match
@@ -14,6 +14,8 @@ import pytest
 import torch
 from torch.nn import Parameter
 from torch.nn import functional as F
+from transformers import MixtralConfig
+from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock

 import vllm.model_executor.layers.fused_moe  # noqa
 from tests.kernels.moe.utils import (
@@ -22,7 +24,10 @@ from tests.kernels.moe.utils import (
    modular_triton_fused_moe,
 )
 from tests.kernels.utils import opcheck, stack_and_dev, torch_experts, torch_moe
+from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.distributed.parallel_state import init_distributed_environment
+from vllm.forward_context import get_forward_context, set_forward_context
 from vllm.model_executor.layers.fused_moe import (
    MoEActivation,
    fused_topk,
@@ -51,10 +56,12 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
    marlin_quantize,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_weights
+from vllm.model_executor.models.mixtral import MixtralMoE
 from vllm.platforms import current_platform
 from vllm.scalar_type import ScalarType, scalar_types
 from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import set_random_seed
+from vllm.v1.worker.workspace import init_workspace_manager


 def iterative_moe(
@@ -143,14 +150,12 @@ MOE_MARLIN_QUANT_TEST_CONFIGS = [
    {
        "a_type": [scalar_types.bfloat16],
        "b_type": scalar_types.float4_e2m1f,
-        "c_type": [scalar_types.bfloat16],
        "group_blocks": [2],
    },
    # MXFP8
    {
        "a_type": [scalar_types.bfloat16],
        "b_type": scalar_types.float8_e4m3fn,
-        "c_type": [scalar_types.bfloat16],
        "group_blocks": [2],
    },
    # AWQ-INT4 with INT8 activation
@@ -676,35 +681,154 @@ def test_fused_moe_wn16(
    torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0)


-MARLIN_MOE_SCENARIOS = [
-    # (m, n, k, e, topk, ep_size, act_order, is_k_full)
-    # No act_order: is_k_full=True matches usual case (marlin_is_k_full).
-    # N>=256 required for Marlin kernel thread config for MXFP8.
-    # Single token, small matrices
-    (1, 128, 256, 5, 2, 1, False, True),
-    # Single token, large matrices
-    (1, 1024, 2048, 5, 2, 1, False, True),
-    # Unaligned m, small matrices
-    (133, 256, 256, 5, 2, 1, False, True),
-    # Unaligned m, large matrices
-    (133, 1024, 2048, 12, 3, 1, False, True),
-    # Aligned batch, small matrices
-    (128, 256, 256, 5, 2, 1, False, True),
-    # Aligned batch, large matrices
-    (128, 1024, 2048, 12, 3, 1, False, True),
-    # Expert parallelism
-    (64, 1024, 2048, 12, 3, 4, False, True),
-    # Act order with is_k_full=True (no tensor parallelism)
-    (1, 1024, 2048, 5, 2, 1, True, True),
-    # Act order with is_k_full=False (tensor parallelism)
-    (133, 256, 256, 5, 2, 1, True, False),
-]
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("padding", [True, False])
+@pytest.mark.parametrize(
+    "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
+)
+@torch.inference_mode()
+def test_mixtral_moe(
+    default_vllm_config,
+    dist_init,
+    dtype: torch.dtype,
+    padding: bool,
+    use_rocm_aiter: bool,
+    monkeypatch,
+):
+    """Make sure our Mixtral MoE implementation agrees with the one from
+    huggingface."""
+
+    # Explicitly set AITER env var based on test parameter to ensure
+    # consistent behavior regardless of external environment
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_rocm_aiter else "0")
+    rocm_aiter_ops.refresh_env_variables()
+
+    if use_rocm_aiter and dtype == torch.float32:
+        pytest.skip("AITER ROCm test skip for float32")
+
+    monkeypatch.setenv("RANK", "0")
+    monkeypatch.setenv("LOCAL_RANK", "0")
+    monkeypatch.setenv("WORLD_SIZE", "1")
+    monkeypatch.setenv("MASTER_ADDR", "localhost")
+    monkeypatch.setenv("MASTER_PORT", "12345")
+    init_distributed_environment()
+    init_workspace_manager(torch.accelerator.current_device_index())
+
+    # Instantiate our and huggingface's MoE blocks
+    vllm_config.compilation_config.static_forward_context = dict()
+    with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config):
+        config = MixtralConfig()
+        hf_moe = MixtralSparseMoeBlock(config).to(dtype).to("cuda")
+        vllm_moe = MixtralMoE(
+            num_experts=config.num_local_experts,
+            top_k=config.num_experts_per_tok,
+            hidden_size=config.hidden_size,
+            intermediate_size=config.intermediate_size,
+            params_dtype=dtype,
+            tp_size=1,
+            dp_size=1,
+        ).cuda()
+
+        # Load the weights
+        vllm_moe.gate.weight.data[:] = hf_moe.gate.weight.data
+        if isinstance(hf_moe.experts, torch.nn.ModuleList):
+            # Transformers v4
+            for i in range(config.num_local_experts):
+                weights = (
+                    hf_moe.experts[i].w1.weight.data,
+                    hf_moe.experts[i].w3.weight.data,
+                )
+                vllm_moe.experts.w13_weight[i][:] = torch.cat(weights, dim=0)
+                vllm_moe.experts.w2_weight[i][:] = hf_moe.experts[i].w2.weight.data
+        else:
+            # Transformers v5
+            vllm_moe.experts.w13_weight.data[:] = hf_moe.experts.gate_up_proj.data
+            vllm_moe.experts.w2_weight.data[:] = hf_moe.experts.down_proj.data
+            # TODO: remove this line after https://github.com/huggingface/transformers/pull/43622
+            hf_moe.experts.config._experts_implementation = "eager"
+
+        # Generate input batch of dimensions [batch_size, seq_len, hidden_dim]
+        hf_inputs = torch.randn((1, 64, config.hidden_size)).to(dtype).to("cuda")
+        # vLLM uses 1D query [num_tokens, hidden_dim]
+        vllm_inputs = hf_inputs.flatten(0, 1)
+
+        # Pad the weight if moe padding is enabled
+        if padding:
+            vllm_moe.experts.w13_weight = Parameter(
+                F.pad(vllm_moe.experts.w13_weight, (0, 128), "constant", 0)[
+                    ..., 0:-128
+                ],
+                requires_grad=False,
+            )
+            vllm_moe.experts.w2_weight = Parameter(
+                F.pad(vllm_moe.experts.w2_weight, (0, 128), "constant", 0)[..., 0:-128],
+                requires_grad=False,
+            )
+            torch.accelerator.synchronize()
+            torch.accelerator.empty_cache()
+
+        # FIXME (zyongye) fix this after we move self.kernel
+        # assignment in FusedMoE.__init__
+
+        vllm_moe.experts.quant_method.process_weights_after_loading(vllm_moe.experts)
+
+        # need to override the forward context for unittests, otherwise it assumes
+        # we're running the model forward pass (the model specified in vllm_config)
+        get_forward_context().all_moe_layers = None
+
+        # Run forward passes for both MoE blocks
+        hf_states = hf_moe.forward(hf_inputs)
+        if isinstance(hf_states, tuple):
+            # Transformers v4
+            hf_states = hf_states[0]
+        vllm_states = vllm_moe.forward(vllm_inputs)
+
+    mixtral_moe_tol = {
+        torch.float32: 1e-3,
+        torch.float16: 1e-3,
+        torch.bfloat16: 1e-2,
+    }
+
+    if use_rocm_aiter:
+        # The values of rtol and atol are set based on the tests in ROCM AITER package.
+        # https://github.com/ROCm/aiter/blob/dfed377f4be7da96ca2d75ac0761f569676f7240/op_tests/test_moe.py#L174
+        torch.testing.assert_close(
+            hf_states.flatten(0, 1), vllm_states, rtol=0.01, atol=100
+        )
+    else:
+        torch.testing.assert_close(
+            hf_states.flatten(0, 1),
+            vllm_states,
+            rtol=mixtral_moe_tol[dtype],
+            atol=mixtral_moe_tol[dtype],
+        )


 def marlin_moe_generate_valid_test_cases():
    import itertools

-    def is_valid(
+    m_list = [1, 123, 666]
+    n_list = [128, 1024]
+    k_list = [256, 2048]
+    e_list = [5, 12]
+    topk_list = [2, 3]
+    ep_size_list = [1, 4]
+    act_order_list = [True, False]
+    is_k_full_list = [True, False]
+
+    all_combinations = itertools.product(
+        MOE_MARLIN_QUANT_TEST_CONFIGS,
+        m_list,
+        n_list,
+        k_list,
+        e_list,
+        topk_list,
+        ep_size_list,
+        act_order_list,
+        is_k_full_list,
+    )
+
+    def is_invalid(
        a_type,
        b_type,
        c_type,
@@ -721,27 +845,29 @@ def marlin_moe_generate_valid_test_cases():
        group_size = group_blocks if group_blocks <= 0 else group_blocks * 16
        if group_size > 0 and k % group_size != 0:
            return False
+
        if act_order and group_size in [-1, k, n]:
            return False
        if group_size in [k, n]:
            return False
-        if b_type == scalar_types.float8_e4m3fn and group_size == 32 and is_k_full:
+        if not act_order and is_k_full:
            return False
+
        return a_type.size_bits < 16 or a_type is c_type

    cases = []
-    for quant_test_config in MOE_MARLIN_QUANT_TEST_CONFIGS:
-        f16_types = [scalar_types.float16]
-        inner_combinations = list(
-            itertools.product(
-                quant_test_config.get("a_type", f16_types),
-                [quant_test_config["b_type"]],
-                quant_test_config.get("c_type", f16_types),
-                quant_test_config["group_blocks"],
-            )
-        )
+    for case in all_combinations:
+        quant_test_config, m, n, k, _, _, _, act_order, *_ = case
+        if act_order and not quant_test_config.get("support_act_order", False):
+            continue

-        supports_act_order = quant_test_config.get("support_act_order", False)
+        f16_types = [scalar_types.float16]
+        inner_combinations = itertools.product(
+            quant_test_config.get("a_type", f16_types),
+            [quant_test_config["b_type"]],
+            quant_test_config.get("c_type", f16_types),
+            quant_test_config["group_blocks"],
+        )

        for sub_case in inner_combinations:
            if (
@@ -749,14 +875,9 @@ def marlin_moe_generate_valid_test_cases():
                and current_platform.get_device_capability() not in [89, 120]
            ):
                continue
-
-            for scenario in MARLIN_MOE_SCENARIOS:
-                m, n, k, e, topk, ep_size, act_order, is_k_full = scenario
-                if act_order and not supports_act_order:
-                    continue
-                args = sub_case + (m, n, k, e, topk, ep_size, act_order, is_k_full)
-                if is_valid(*args):
-                    cases.append(args)
+            args = sub_case + (m, n, k) + case[4:]
+            if is_invalid(*args):
+                cases.append(args)
    return cases


@@ -1,248 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Tests for SM100 CUTLASS MXFP4 x MXFP4 grouped MoE kernels."""
-
-import random
-
-import pytest
-import torch
-
-from tests.kernels.utils import torch_moe_single
-from vllm import _custom_ops as ops
-from vllm.platforms import current_platform
-from vllm.utils.torch_utils import set_random_seed
-
-random.seed(42)
-set_random_seed(42)
-
-MXFP4_BLOCK_SIZE = 32
-
-
-def align(val: int, alignment: int = 128) -> int:
-    return int((val + alignment - 1) // alignment * alignment)
-
-
-def calc_diff(x, y):
-    x, y = x.double(), y.double()
-    denominator = (x * x + y * y).sum()
-    sim = 2 * (x * y).sum() / denominator
-    return 1 - sim
-
-
-def is_sm100_supported() -> bool:
-    return current_platform.is_cuda() and current_platform.is_device_capability_family(
-        100
-    )
-
-
-def compute_ref_output(
-    input_tensor: torch.Tensor,
-    weight_list: list[torch.Tensor],
-    expert_offsets: list[int],
-    expert_offset: int,
-    num_experts: int,
-) -> torch.Tensor:
-    """Reference output using torch_moe_single with top-1 routing."""
-    score = torch.full(
-        (expert_offset, num_experts),
-        -1e9,
-        device=input_tensor.device,
-        dtype=torch.float32,
-    )
-    for g in range(num_experts):
-        start = expert_offsets[g]
-        end = expert_offsets[g + 1] if g + 1 < num_experts else expert_offset
-        score[start:end, g] = 0.0
-
-    return torch_moe_single(
-        input_tensor, torch.stack(weight_list, dim=0), score, topk=1
-    )
-
-
-@pytest.mark.skipif(
-    not is_sm100_supported(),
-    reason="cutlass_mxfp4_group_mm requires CUDA SM100",
-)
-@pytest.mark.parametrize("num_experts", [8, 16, 32])
-@pytest.mark.parametrize("out_dtype", [torch.bfloat16])
-def test_cutlass_mxfp4_grouped_mm(num_experts, out_dtype):
-    """
-    Test the MXFP4 grouped GEMM kernel by:
-    1. Creating random per-expert inputs and weights
-    2. Quantizing both to MXFP4 using the CUDA kernel
-    3. Running the CUTLASS grouped GEMM
-    4. Comparing against BF16 reference
-    """
-    device = "cuda"
-    alignment = 128
-    # N and K must be multiples of 128 for clean swizzle layout
-    n_g = random.randint(1, 16) * alignment
-    k_g = random.randint(1, 16) * alignment
-
-    expert_offset = 0
-    expert_offsets_input = []
-    problem_sizes = []
-    input_list = []
-    weight_list = []
-
-    for g in range(num_experts):
-        m_g = random.randint(1, 256)
-        expert_offsets_input.append(expert_offset)
-        expert_offset += m_g
-        problem_sizes.append([m_g, n_g, k_g])
-
-        input_list.append(
-            torch.normal(0.0, std=0.5, size=(m_g, k_g), device=device, dtype=out_dtype)
-        )
-        weight_list.append(
-            torch.normal(0.0, std=0.5, size=(n_g, k_g), device=device, dtype=out_dtype)
-        )
-
-    input_tensor = torch.concat(input_list, dim=0)  # [M_total, K]
-
-    # --- Quantize INPUTS via mxfp4_experts_quant ---
-    input_bs_offsets = []
-    tot = 0
-    for g in range(num_experts):
-        input_bs_offsets.append(tot)
-        tot += align(problem_sizes[g][0], 128)
-    input_bs_offsets.append(tot)
-
-    _inp_expert_offsets = torch.tensor(
-        expert_offsets_input + [expert_offset], device=device, dtype=torch.int32
-    )
-    _inp_bs_offsets = torch.tensor(input_bs_offsets, device=device, dtype=torch.int32)
-
-    input_quant, input_sf = ops.mxfp4_experts_quant(
-        input_tensor,
-        _inp_expert_offsets,
-        _inp_bs_offsets,
-        num_experts,
-        topk=1,
-    )
-
-    # --- Quantize WEIGHTS via mxfp4_experts_quant ---
-    # Treat each expert's N weight rows as an "expert" with N tokens
-    weight_tensor = torch.concat(weight_list, dim=0)  # [E*N, K]
-    weight_expert_offsets = [g * n_g for g in range(num_experts)] + [num_experts * n_g]
-    # N is always multiple of 128, so blockscale offsets are clean
-    weight_bs_offsets = [g * n_g for g in range(num_experts)] + [num_experts * n_g]
-
-    _wt_expert_offsets = torch.tensor(
-        weight_expert_offsets, device=device, dtype=torch.int32
-    )
-    _wt_bs_offsets = torch.tensor(weight_bs_offsets, device=device, dtype=torch.int32)
-
-    weight_quant, weight_sf = ops.mxfp4_experts_quant(
-        weight_tensor,
-        _wt_expert_offsets,
-        _wt_bs_offsets,
-        num_experts,
-        topk=1,
-    )
-
-    # Reshape weight quantized data to [E, N, K//2]
-    weight_quant = weight_quant[: num_experts * n_g].view(num_experts, n_g, k_g // 2)
-
-    # Reshape weight scale factors to [E, N, K//32]
-    # The quant kernel produces uint8 SF buffer. Each row has K//32 SFs.
-    scales_per_row = k_g // MXFP4_BLOCK_SIZE
-    weight_sf_flat = weight_sf.view(-1)[: num_experts * n_g * scales_per_row]
-    weight_sf_3d = weight_sf_flat.view(num_experts, n_g, scales_per_row)
-
-    # Output
-    output = torch.empty((expert_offset, n_g), device=device, dtype=out_dtype)
-
-    _problem_sizes = torch.tensor(problem_sizes, device=device, dtype=torch.int32)
-    _expert_offsets = torch.tensor(
-        expert_offsets_input, device=device, dtype=torch.int32
-    )
-    _input_bs = torch.tensor(input_bs_offsets[:-1], device=device, dtype=torch.int32)
-
-    # Run the MXFP4 grouped GEMM
-    ops.cutlass_mxfp4_moe_mm(
-        output,
-        input_quant,
-        weight_quant,
-        input_sf,
-        weight_sf_3d,
-        _problem_sizes,
-        _expert_offsets,
-        _input_bs,
-    )
-
-    # Reference: BF16 matmul
-    ref_output = compute_ref_output(
-        input_tensor=input_tensor,
-        weight_list=weight_list,
-        expert_offsets=expert_offsets_input,
-        expert_offset=expert_offset,
-        num_experts=num_experts,
-    )
-
-    # Compare per-expert
-    for g in range(num_experts):
-        start = expert_offsets_input[g]
-        end = expert_offsets_input[g + 1] if g + 1 < num_experts else expert_offset
-        if start == end:
-            continue
-        baseline = ref_output[start:end]
-        actual = output[start:end]
-        diff = calc_diff(actual, baseline)
-        print(
-            f"m_g={end - start} n_g={n_g} k_g={k_g} "
-            f"num_experts={num_experts}, "
-            f"out_dtype={out_dtype}, diff={diff:.5f}"
-        )
-        # FP4 quantization is very lossy (~4 bits precision)
-        # Comparing quantized vs full-precision gives cosine diff of 0.05-0.15
-        assert diff < 0.15, f"Expert {g}: diff={diff:.5f} exceeds threshold"
-
-
-@pytest.mark.skipif(
-    not is_sm100_supported(),
-    reason="mxfp4_experts_quant requires CUDA SM100",
-)
-def test_mxfp4_experts_quant_basic():
-    """
-    Basic smoke test for the MXFP4 experts quantization kernel.
-    """
-    device = "cuda"
-    num_experts = 4
-    k = 256
-    tokens_per_expert = 16
-
-    total_tokens = tokens_per_expert * num_experts
-    input_tensor = torch.randn(total_tokens, k, device=device, dtype=torch.bfloat16) / 5
-
-    expert_offsets = [i * tokens_per_expert for i in range(num_experts + 1)]
-    blockscale_offsets = [
-        align(i * tokens_per_expert, 128) for i in range(num_experts + 1)
-    ]
-
-    _expert_offsets = torch.tensor(expert_offsets, device=device, dtype=torch.int32)
-    _blockscale_offsets = torch.tensor(
-        blockscale_offsets, device=device, dtype=torch.int32
-    )
-
-    output, output_sf = ops.mxfp4_experts_quant(
-        input_tensor,
-        _expert_offsets,
-        _blockscale_offsets,
-        num_experts,
-        topk=1,
-    )
-
-    assert output.shape == (total_tokens, k // 2)
-    assert output.dtype == torch.uint8
-    assert output_sf.dtype == torch.uint8
-    assert output.any(), "Quantized output is all zeros"
-    print(
-        f"MXFP4 experts quant: output shape={output.shape}, sf shape={output_sf.shape}"
-    )
-    print("PASSED")
-
-
-if __name__ == "__main__":
-    pytest.main([__file__, "-v", "-s"])
@@ -11,11 +11,6 @@ from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
 )
 from vllm.platforms import current_platform

-skipif_not_cuda_rocm = pytest.mark.skipif(
-    not (current_platform.is_cuda() or current_platform.is_rocm()),
-    reason="Only supported on CUDA/ROCm platforms.",
-)
-

@pytest.mark.parametrize(
    "platform_method,expected_backend",
@@ -195,83 +190,3 @@ def test_select_cuda_flashinfer_cutlass_backend(

        assert selected_backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS
        assert experts_cls is not None
-
-
-@skipif_not_cuda_rocm
-def test_select_lora_backend_prefers_triton():
-    """LoRA-enabled unquantized MoE should select Triton backend."""
-    moe_config = make_dummy_moe_config()
-    moe_config.is_lora_enabled = True
-    selected_backend, experts_cls = select_unquantized_moe_backend(
-        moe_config=moe_config
-    )
-
-    assert selected_backend == UnquantizedMoeBackend.TRITON
-    assert experts_cls is not None
-
-
-@skipif_not_cuda_rocm
-def test_select_lora_explicit_non_triton_backend():
-    """LoRA should override explicit non-Triton backend to Triton."""
-    moe_config = make_dummy_moe_config()
-    moe_config.is_lora_enabled = True
-
-    # Use string from mapping in function map_unquantized_backend()
-    moe_config.moe_backend = "flashinfer_cutlass"
-
-    selected_backend, experts_cls = select_unquantized_moe_backend(
-        moe_config=moe_config
-    )
-
-    assert selected_backend == UnquantizedMoeBackend.TRITON
-    assert experts_cls is not None
-
-
-@skipif_not_cuda_rocm
-@pytest.mark.parametrize("is_lora_enabled", [False, True])
-def test_select_explicit_triton_backend(is_lora_enabled):
-    """Explicit triton backend selection should return Triton."""
-    moe_config = make_dummy_moe_config()
-    moe_config.is_lora_enabled = is_lora_enabled
-    moe_config.moe_backend = "triton"
-
-    selected_backend, experts_cls = select_unquantized_moe_backend(
-        moe_config=moe_config
-    )
-
-    assert selected_backend == UnquantizedMoeBackend.TRITON
-    assert experts_cls is not None
-
-
-@skipif_not_cuda_rocm
-def test_select_explicit_triton_ignores_flashinfer_env(monkeypatch):
-    """Explicit triton backend should override FlashInfer env selection."""
-    monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1")
-    monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput")
-
-    moe_config = make_dummy_moe_config()
-    moe_config.is_lora_enabled = False
-    moe_config.moe_backend = "triton"
-
-    selected_backend, experts_cls = select_unquantized_moe_backend(
-        moe_config=moe_config
-    )
-
-    assert selected_backend == UnquantizedMoeBackend.TRITON
-    assert experts_cls is not None
-
-
-@skipif_not_cuda_rocm
-def test_select_lora_ignores_flashinfer_env(monkeypatch):
-    """LoRA path should still choose Triton even if FlashInfer env is on."""
-    monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1")
-    monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput")
-
-    moe_config = make_dummy_moe_config()
-    moe_config.is_lora_enabled = True
-    selected_backend, experts_cls = select_unquantized_moe_backend(
-        moe_config=moe_config
-    )
-
-    assert selected_backend == UnquantizedMoeBackend.TRITON
-    assert experts_cls is not None
@@ -4,6 +4,7 @@
 import os
 from collections.abc import Sequence

+import librosa
 import pytest
 import regex as re
 from huggingface_hub import snapshot_download
@@ -13,7 +14,6 @@ from vllm.assets.image import ImageAsset
 from vllm.logprobs import SampleLogprobs
 from vllm.lora.request import LoRARequest
 from vllm.multimodal.image import convert_image_mode, rescale_image_size
-from vllm.multimodal.media.audio import load_audio

 from ....conftest import (
    IMAGE_ASSETS,
@@ -290,7 +290,7 @@ def test_vision_speech_models(
    num_logprobs: int,
 ) -> None:
    # use the example speech question so that the model outputs are reasonable
-    audio = load_audio(speech_question, sr=None)
+    audio = librosa.load(speech_question, sr=None)
    image = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")

    inputs_vision_speech = [
@@ -25,7 +25,6 @@ if TYPE_CHECKING:

 PIXTRAL_ID = "mistralai/Pixtral-12B-2409"
 MISTRAL_SMALL_3_1_ID = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
-MINISTRAL_3B_ID = "mistralai/Ministral-3-3B-Instruct-2512"

 MODELS = [PIXTRAL_ID, MISTRAL_SMALL_3_1_ID]

@@ -117,7 +116,6 @@ assert FIXTURES_PATH.exists()
 FIXTURE_LOGPROBS_CHAT = {
    PIXTRAL_ID: FIXTURES_PATH / "pixtral_chat.json",
    MISTRAL_SMALL_3_1_ID: FIXTURES_PATH / "mistral_small_3_chat.json",
-    MINISTRAL_3B_ID: FIXTURES_PATH / "ministral_3b_chat.json",
 }

 OutputsLogprobs = list[tuple[list[int], str, SampleLogprobs | None]]
@@ -211,41 +209,3 @@ def test_chat(
        name_0="h100_ref",
        name_1="output",
    )
-
-
-@large_gpu_test(min_gb=16)
-@pytest.mark.parametrize("dtype", ["bfloat16"])
-def test_chat_consolidated(vllm_runner, dtype: str, local_asset_server) -> None:
-    EXPECTED_CHAT_LOGPROBS = load_outputs_w_logprobs(
-        FIXTURE_LOGPROBS_CHAT[MINISTRAL_3B_ID]
-    )
-    with vllm_runner(
-        MINISTRAL_3B_ID,
-        dtype=dtype,
-        tokenizer_mode="mistral",
-        load_format="mistral",
-        config_format="mistral",
-        max_model_len=8192,
-        limit_mm_per_prompt=LIMIT_MM_PER_PROMPT,
-    ) as vllm_model:
-        outputs = []
-        urls_all = [local_asset_server.url_for(u) for u in IMG_URLS]
-        msgs = [
-            _create_msg_format(urls_all[:1]),
-            _create_msg_format(urls_all[:2]),
-            _create_msg_format(urls_all),
-        ]
-        for msg in msgs:
-            output = vllm_model.llm.chat(msg, sampling_params=SAMPLING_PARAMS)
-            outputs.extend(output)
-
-    logprobs = vllm_runner._final_steps_generate_w_logprobs(outputs)
-    for i in range(len(logprobs)):
-        assert logprobs[i][-1] is None
-        logprobs[i] = logprobs[i][:-1]
-    check_logprobs_close(
-        outputs_0_lst=EXPECTED_CHAT_LOGPROBS,
-        outputs_1_lst=logprobs,
-        name_0="h100_ref",
-        name_1="output",
-    )
@@ -4,11 +4,11 @@
 from collections.abc import Sequence
 from typing import Any

+import librosa
 import pytest
 from transformers import AutoModelForSpeechSeq2Seq

 from vllm.assets.audio import AudioAsset
-from vllm.multimodal.audio import AudioResampler
 from vllm.platforms import current_platform

 from ....conftest import HfRunner, PromptAudioInput, VllmRunner
@@ -93,12 +93,13 @@ def run_test(
 def resampled_assets() -> list[tuple[Any, int]]:
    audio_assets = [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")]
    sampled_assets = []
-    resampler = AudioResampler(target_sr=WHISPER_SAMPLE_RATE)
    for asset in audio_assets:
        audio, orig_sr = asset.audio_and_sample_rate
        # Resample to Whisper's expected sample rate (16kHz)
        if orig_sr != WHISPER_SAMPLE_RATE:
-            audio = resampler.resample(audio, orig_sr=orig_sr)
+            audio = librosa.resample(
+                audio, orig_sr=orig_sr, target_sr=WHISPER_SAMPLE_RATE
+            )
        sampled_assets.append(
            (audio, WHISPER_SAMPLE_RATE),
        )
@@ -3,12 +3,12 @@
 from pathlib import Path
 from unittest.mock import patch

+import librosa
 import numpy as np
 import pybase64 as base64
 import pytest

 from vllm.multimodal.media import AudioMediaIO
-from vllm.multimodal.media.audio import load_audio

 from ...conftest import AudioTestAssets

@@ -73,6 +73,6 @@ def test_audio_media_io_from_video(video_assets):
    video_path = video_assets[0].video_path
    with open(video_path, "rb") as f:
        audio, sr = audio_io.load_bytes(f.read())
-    audio_ref, sr_ref = load_audio(video_path, sr=None)
+    audio_ref, sr_ref = librosa.load(video_path, sr=None)
    assert sr == sr_ref
    np.testing.assert_allclose(audio_ref, audio, atol=1e-4)
@@ -26,8 +26,11 @@ def test_placeholder_range_get_num_embeds(is_embed, expected):
    "is_embed,expected",
    [
        (None, None),
-        (torch.tensor([False, True, False, True, True]), [0, 1, 1, 2, 3]),
-        (torch.tensor([True, True, True]), [1, 2, 3]),
+        (
+            torch.tensor([False, True, False, True, True]),
+            torch.tensor([0, 1, 1, 2, 3]),
+        ),
+        (torch.tensor([True, True, True]), torch.tensor([1, 2, 3])),
    ],
 )
 def test_placeholder_range_embeds_cumsum(is_embed, expected):
@@ -38,6 +41,6 @@ def test_placeholder_range_embeds_cumsum(is_embed, expected):
        assert pr.embeds_cumsum is None
        return

-    assert pr.embeds_cumsum == expected
+    assert torch.equal(pr.embeds_cumsum, expected)
    # cached_property should return the same object on repeated access
    assert pr.embeds_cumsum is pr.embeds_cumsum
@@ -18,7 +18,9 @@ from vllm.model_executor.layers.quantization.turboquant.config import (
    TQ_PRESETS,
    TurboQuantConfig,
 )
-from vllm.platforms import current_platform
+from vllm.model_executor.layers.quantization.turboquant.quantizer import (
+    generate_wht_signs,
+)
 from vllm.utils.math_utils import next_power_of_2

 # ============================================================================
@@ -343,8 +345,7 @@ class TestLloydMax:
 # Rotation matrix tests (GPU required)
 # ============================================================================

-GPGPU_AVAILABLE = torch.cuda.is_available() or torch.xpu.is_available()
-DEVICE_TYPE = current_platform.device_type
+CUDA_AVAILABLE = torch.cuda.is_available()


 def generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.Tensor:
@@ -359,16 +360,16 @@ def generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.Te
    return Q.to(device)


-@pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available")
+@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
 class TestRotationMatrix:
    """Tests for the QR-based rotation (standalone benchmarks only)."""

    @pytest.mark.parametrize("dim", [64, 96, 128, 256])
    def test_rotation_matrix_shape_and_orthogonal(self, dim):
-        Pi = generate_rotation_matrix(dim, seed=42, device=DEVICE_TYPE)
+        Pi = generate_rotation_matrix(dim, seed=42, device="cuda")
        assert Pi.shape == (dim, dim)
        eye = Pi @ Pi.T
-        assert torch.allclose(eye, torch.eye(dim, device=DEVICE_TYPE), atol=1e-5), (
+        assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), (
            f"Pi not orthogonal for dim={dim}"
        )

@@ -384,13 +385,13 @@ class TestRotationMatrix:

    def test_rotation_matrix_det_is_pm1(self):
        """Orthogonal matrix determinant must be +1 or -1."""
-        Pi = generate_rotation_matrix(128, seed=42, device=DEVICE_TYPE)
+        Pi = generate_rotation_matrix(128, seed=42, device="cuda")
        det = torch.linalg.det(Pi)
        assert abs(abs(det.item()) - 1.0) < 1e-4


 # ============================================================================
-# Hadamard rotation tests (serving path: _build_hadamard)
+# WHT rotation tests (serving path: generate_wht_signs + _build_hadamard)
 # ============================================================================


@@ -402,34 +403,58 @@ def _build_hadamard(d: int, device: str = "cpu") -> torch.Tensor:
    return (H / math.sqrt(d)).to(torch.device(device))


-@pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available")
-class TestHadamardRotation:
-    """Tests for the Hadamard rotation used in serving."""
+@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
+class TestWHTRotation:
+    """Tests for the WHT rotation actually used in serving."""

    @pytest.mark.parametrize("dim", [64, 128, 256])
-    def test_hadamard_orthonormal(self, dim):
-        """H must be orthonormal: H @ H^T = I."""
-        H = _build_hadamard(dim, DEVICE_TYPE)
-        eye = H @ H.T
-        assert torch.allclose(eye, torch.eye(dim, device=DEVICE_TYPE), atol=1e-5), (
-            f"Hadamard not orthonormal for dim={dim}"
+    def test_wht_orthonormal(self, dim):
+        """signs * H must be orthonormal: (signs*H) @ (signs*H)^T = I."""
+        signs = generate_wht_signs(dim, seed=42, device="cuda")
+        H = _build_hadamard(dim, "cuda")
+        PiT = (signs.unsqueeze(1) * H).contiguous()
+        eye = PiT @ PiT.T
+        assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), (
+            f"WHT rotation not orthonormal for dim={dim}"
        )

    @pytest.mark.parametrize("dim", [64, 128, 256])
-    def test_hadamard_symmetric(self, dim):
-        """Sylvester Hadamard must be symmetric: H = H^T."""
-        H = _build_hadamard(dim, DEVICE_TYPE)
-        assert torch.allclose(H, H.T, atol=1e-6), (
-            f"Hadamard not symmetric for dim={dim}"
+    def test_wht_self_inverse(self, dim):
+        """PiT should be self-inverse: PiT @ PiT = I (up to sign flip)."""
+        signs = generate_wht_signs(dim, seed=42, device="cuda")
+        H = _build_hadamard(dim, "cuda")
+        PiT = (signs.unsqueeze(1) * H).contiguous()
+        Pi = PiT.T.contiguous()
+        # Pi @ PiT should be identity (rotation then inverse)
+        result = Pi @ PiT
+        assert torch.allclose(result, torch.eye(dim, device="cuda"), atol=1e-5), (
+            f"WHT rotation not self-inverse for dim={dim}"
        )

+    def test_wht_signs_deterministic(self):
+        """Same seed must produce identical signs."""
+        s1 = generate_wht_signs(128, seed=42)
+        s2 = generate_wht_signs(128, seed=42)
+        assert torch.equal(s1, s2)
+
+    def test_wht_signs_different_seeds(self):
+        """Different seeds must produce different signs."""
+        s1 = generate_wht_signs(128, seed=42)
+        s2 = generate_wht_signs(128, seed=99)
+        assert not torch.equal(s1, s2)
+
+    def test_wht_signs_are_pm1(self):
+        """All sign values must be exactly +1 or -1."""
+        signs = generate_wht_signs(128, seed=42)
+        assert torch.all(signs.abs() == 1.0)
+

 # ============================================================================
 # Store → Decode round-trip test (GPU + Triton required)
 # ============================================================================


-@pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available")
+@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
 class TestStoreDecodeRoundTrip:
    """End-to-end: store KV into TQ cache, decode, compare vs fp16 ref."""

@@ -462,12 +487,13 @@ class TestStoreDecodeRoundTrip:
        block_size = 16
        num_blocks = 1

-        device = torch.device(DEVICE_TYPE)
+        device = torch.device("cuda")

-        # Pure Hadamard rotation (symmetric: H = H^T, so Pi = PiT = H)
-        H = _build_hadamard(D, DEVICE_TYPE)
-        PiT = H
-        Pi = H
+        # Generate rotation
+        signs = generate_wht_signs(D, seed=42, device=device)
+        H = _build_hadamard(D, "cuda")
+        PiT = (signs.unsqueeze(1) * H).contiguous().float()
+        Pi = PiT.T.contiguous()

        # Generate centroids
        centroids, _ = solve_lloyd_max(D, cfg.centroid_bits)
@@ -17,6 +17,22 @@ from vllm import LLM, SamplingParams
 from vllm.platforms import current_platform


+@pytest.mark.skip(reason="In V1, we reject tokens > max_seq_len")
+def test_duplicated_ignored_sequence_group():
+    """https://github.com/vllm-project/vllm/issues/1655"""
+
+    sampling_params = SamplingParams(temperature=0.01, top_p=0.1, max_tokens=256)
+    llm = LLM(
+        model="distilbert/distilgpt2",
+        max_num_batched_tokens=4096,
+        tensor_parallel_size=1,
+    )
+    prompts = ["This is a short prompt", "This is a very long prompt " * 1000]
+    outputs = llm.generate(prompts, sampling_params=sampling_params)
+
+    assert len(prompts) == len(outputs)
+
+
@pytest.mark.parametrize(
    "model",
    [
@@ -1,116 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Regression tests for Responses API tool-calling request adjustment.
-
-Covers two bugs on the ``/v1/responses`` path that broke streaming tool
-calling for parsers relying on special-token delimiters (Gemma4):
-
-1. :class:`Gemma4ToolParser.adjust_request` used an
-   ``isinstance(request, ChatCompletionRequest)`` guard, so a
-   :class:`ResponsesRequest` with tools never had
-   ``skip_special_tokens`` flipped to ``False``. The default (``True``)
-   stripped ``<|tool_call>`` / ``<tool_call|>`` delimiters, causing
-   :meth:`Gemma4ToolParser.extract_tool_calls_streaming` to fall through
-   to the content branch and leak the raw ``call:fn{...}`` body via
-   ``response.output_text.delta``.
-
-2. :meth:`ToolParser.adjust_request` built
-   :class:`ResponseTextConfig` in two steps (bare constructor then
-   ``.format = ...``). Under Pydantic v2 the later assignment is not
-   tracked in ``__fields_set__``, which can drop the nested config from
-   ``model_dump``. It also passed a ``description`` kwarg carrying the
-   wrong-purpose string ``"Response format for tool calling"``.
-"""
-
-from __future__ import annotations
-
-from typing import Any
-
-from openai.types.responses.tool_param import FunctionToolParam
-
-from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
-from vllm.tool_parsers.abstract_tool_parser import ToolParser
-from vllm.tool_parsers.gemma4_tool_parser import Gemma4ToolParser
-
-
-def _get_weather_tool() -> FunctionToolParam:
-    return FunctionToolParam(
-        type="function",
-        name="get_weather",
-        description="Get current weather for a city",
-        parameters={
-            "type": "object",
-            "properties": {"city": {"type": "string"}},
-            "required": ["city"],
-        },
-        strict=True,
-    )
-
-
-def _build_responses_request(*, tool_choice: str) -> ResponsesRequest:
-    return ResponsesRequest(
-        model="gemma4-test",
-        input=[{"role": "user", "content": "What is the weather in Hanoi?"}],
-        tools=[_get_weather_tool()],
-        tool_choice=tool_choice,
-        stream=True,
-        max_output_tokens=200,
-    )
-
-
-class _StubTokenizer:
-    """Minimal tokenizer stub to satisfy ``Gemma4ToolParser.__init__``."""
-
-    def get_vocab(self) -> dict[str, int]:
-        return {"<|tool_call>": 256_000, "<tool_call|>": 256_001, '<|"|>': 52}
-
-
-def test_gemma4_adjust_request_sets_skip_special_tokens_on_responses() -> None:
-    """``Gemma4ToolParser.adjust_request`` must flip
-    ``skip_special_tokens=False`` for both ``ChatCompletionRequest`` and
-    ``ResponsesRequest`` so that ``<|tool_call>`` delimiters reach the
-    streaming extractor. The previous
-    ``isinstance(ChatCompletionRequest)`` guard omitted the Responses
-    path, causing raw ``call:fn{...}`` text to leak via
-    ``response.output_text.delta``.
-    """
-    parser = Gemma4ToolParser.__new__(Gemma4ToolParser)
-    parser.model_tokenizer = _StubTokenizer()
-
-    request = _build_responses_request(tool_choice="auto")
-    assert request.skip_special_tokens is True, (
-        "Precondition: ResponsesRequest.skip_special_tokens default is True"
-    )
-
-    Gemma4ToolParser.adjust_request(parser, request)
-
-    assert request.skip_special_tokens is False
-
-
-def test_tool_parser_adjust_request_builds_valid_response_text_config() -> None:
-    """``ToolParser.adjust_request`` must produce a ``ResponseTextConfig``
-    whose dumped form contains the JSON schema under the ``schema`` alias
-    and does not leak the unrelated ``"Response format for tool calling"``
-    description string that the previous two-step construction injected.
-    """
-    parser = ToolParser.__new__(ToolParser)
-    parser.model_tokenizer = None
-
-    request = _build_responses_request(tool_choice="required")
-    ToolParser.adjust_request(parser, request)
-
-    assert request.text is not None
-    assert request.text.format is not None
-    assert request.text.format.type == "json_schema"
-
-    dump: dict[str, Any] = request.text.model_dump(mode="json", by_alias=True)
-    fmt = dump.get("format") or {}
-    assert fmt.get("type") == "json_schema"
-    assert fmt.get("name") == "tool_calling_response"
-    assert fmt.get("strict") is True
-    # Nested config must be present under the alias. Two-step Pydantic v2
-    # construction could drop it from __fields_set__.
-    assert "schema" in fmt and isinstance(fmt["schema"], dict)
-    # The old code passed a wrong-purpose string; valid field should now
-    # either be absent or None (the openai-python default).
-    assert fmt.get("description") in (None, "")
@@ -7,39 +7,17 @@ from tests.models.utils import check_embeddings_close
 from vllm.utils.serial_utils import (
    EMBED_DTYPES,
    ENDIANNESS,
-    MM_METADATA_DTYPES,
    EmbedDType,
    Endianness,
-    MmMetadataDType,
    binary2tensor,
    tensor2binary,
 )

-FLOAT_EMBED_DTYPES = tuple(EMBED_DTYPES.keys())
-INTEGER_EMBED_DTYPES = tuple(MM_METADATA_DTYPES.keys())
-
-
-def _build_integer_tensor(
-    embed_dtype: MmMetadataDType, shape: tuple[int, ...]
-) -> torch.Tensor:
-    torch_dtype = MM_METADATA_DTYPES[embed_dtype].torch_dtype
-
-    if torch_dtype is torch.bool:
-        return torch.randint(0, 2, shape, dtype=torch.int32).to(torch.bool)
-    if torch_dtype is torch.uint8:
-        return torch.randint(0, 256, shape, dtype=torch.uint8)
-    if torch_dtype is torch.int32:
-        return torch.randint(-(2**20), 2**20, shape, dtype=torch.int32)
-    if torch_dtype is torch.int64:
-        return torch.randint(-(2**62), 2**62, shape, dtype=torch.int64)
-
-    raise AssertionError(f"Unsupported non-floating embed dtype: {embed_dtype}")
-

@pytest.mark.parametrize("endianness", ENDIANNESS)
-@pytest.mark.parametrize("embed_dtype", FLOAT_EMBED_DTYPES)
+@pytest.mark.parametrize("embed_dtype", EMBED_DTYPES.keys())
@torch.inference_mode()
-def test_encode_and_decode_floats(embed_dtype: EmbedDType, endianness: Endianness):
+def test_encode_and_decode(embed_dtype: EmbedDType, endianness: Endianness):
    for i in range(10):
        tensor = torch.rand(2, 3, 5, 7, 11, 13, device="cpu", dtype=torch.float32)
        shape = tensor.shape
@@ -62,20 +40,3 @@ def test_encode_and_decode_floats(embed_dtype: EmbedDType, endianness: Endiannes
            name_1="new",
            tol=1e-2,
        )
-
-
-@pytest.mark.parametrize("endianness", ENDIANNESS)
-@pytest.mark.parametrize("embed_dtype", INTEGER_EMBED_DTYPES)
-@torch.inference_mode()
-def test_encode_and_decode_integers(
-    embed_dtype: MmMetadataDType, endianness: Endianness
-):
-    shape = (2, 3, 5, 7, 11, 13)
-
-    for i in range(10):
-        tensor = _build_integer_tensor(embed_dtype, shape)
-        binary = tensor2binary(tensor, embed_dtype, endianness)
-        new_tensor = binary2tensor(binary, shape, embed_dtype, endianness)
-
-        assert new_tensor.dtype == MM_METADATA_DTYPES[embed_dtype].torch_dtype
-        torch.testing.assert_close(tensor, new_tensor, atol=0, rtol=0)
@@ -0,0 +1,968 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Backend integration tests for CUTLASS FA3 sparse MLA attention.
+
+Tests verify:
+  - Backend class properties
+  - Metadata builder (decode, prefill, mixed, topk clipping)
+  - KV cache write/read consistency
+  - Backend registration and selection
+"""
+
+import pytest
+import torch
+
+from vllm.v1.attention.ops.cutlass_fa3 import is_cutlass_fa3_available
+
+pytestmark = pytest.mark.skipif(
+    not is_cutlass_fa3_available(),
+    reason="CUTLASS FA3 not available (requires CUDA >= 12.4, SM90)",
+)
+
+
+# ─── TEST 2.1: Backend Class Properties ──────────────────────────────
+
+
+def test_backend_class_properties():
+    """Verify CutlassFA3MLASparseBackend class attributes."""
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        CutlassFA3MLASparseBackend,
+    )
+
+    assert CutlassFA3MLASparseBackend.get_name() == "CUTLASS_FA3_MLA_SPARSE"
+    assert CutlassFA3MLASparseBackend.is_mla() is True
+    assert CutlassFA3MLASparseBackend.is_sparse() is True
+    assert CutlassFA3MLASparseBackend.get_supported_head_sizes() == [576]
+    assert CutlassFA3MLASparseBackend.supported_kv_cache_dtypes == ["auto"]
+    assert CutlassFA3MLASparseBackend.get_supported_kernel_block_sizes() == [64]
+
+
+def test_backend_compute_capability():
+    """Verify SM90-only support."""
+    from vllm.platforms.interface import DeviceCapability
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        CutlassFA3MLASparseBackend,
+    )
+
+    assert CutlassFA3MLASparseBackend.supports_compute_capability(
+        DeviceCapability(major=9, minor=0)
+    )
+    assert not CutlassFA3MLASparseBackend.supports_compute_capability(
+        DeviceCapability(major=8, minor=0)
+    )
+    assert not CutlassFA3MLASparseBackend.supports_compute_capability(
+        DeviceCapability(major=10, minor=0)
+    )
+
+
+def test_backend_kv_cache_shape():
+    """Verify KV cache shape for BF16 format."""
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        CutlassFA3MLASparseBackend,
+    )
+
+    shape = CutlassFA3MLASparseBackend.get_kv_cache_shape(
+        num_blocks=100,
+        block_size=64,
+        num_kv_heads=1,
+        head_size=576,
+        cache_dtype_str="auto",
+    )
+    assert shape == (100, 64, 576)
+
+
+# ─── TEST 2.2: Backend Registration ──────────────────────────────────
+
+
+def test_backend_enum_registered():
+    """Verify CUTLASS_FA3_MLA_SPARSE is in the backend enum."""
+    from vllm.v1.attention.backends.registry import AttentionBackendEnum
+
+    assert hasattr(AttentionBackendEnum, "CUTLASS_FA3_MLA_SPARSE")
+    backend_enum = AttentionBackendEnum.CUTLASS_FA3_MLA_SPARSE
+    assert "cutlass_fa3_sparse" in backend_enum.get_path()
+
+
+def test_backend_class_loadable():
+    """Verify the backend class can be loaded from the enum."""
+    from vllm.v1.attention.backends.registry import AttentionBackendEnum
+
+    backend_cls = AttentionBackendEnum.CUTLASS_FA3_MLA_SPARSE.get_class()
+    assert backend_cls.get_name() == "CUTLASS_FA3_MLA_SPARSE"
+
+
+# ─── TEST 2.3: KV Cache Write/Read ───────────────────────────────────
+
+
+def test_kv_cache_write_read_consistency():
+    """Verify do_kv_cache_update writes match what forward_mqa would read."""
+    device = "cuda"
+    num_blocks = 4
+    block_size = 64
+    head_size = 576
+    kv_lora_rank = 512
+    qk_rope_head_dim = 64
+
+    # Create BF16 cache
+    cache = torch.zeros(
+        num_blocks, block_size, head_size, dtype=torch.bfloat16, device=device
+    )
+
+    # Write known values
+    T = 3
+    kv_c_normed = torch.randn(T, kv_lora_rank, dtype=torch.bfloat16, device=device)
+    k_pe = torch.randn(T, 1, qk_rope_head_dim, dtype=torch.bfloat16, device=device)
+    slot_mapping = torch.tensor([0, 1, 2], dtype=torch.int64, device=device)
+    k_scale = torch.ones(1, dtype=torch.float32, device=device)
+
+    from vllm import _custom_ops as ops
+
+    ops.concat_and_cache_mla(
+        kv_c_normed,
+        k_pe.squeeze(1),
+        cache,
+        slot_mapping,
+        kv_cache_dtype="auto",
+        scale=k_scale,
+    )
+
+    # Read back via flatten + split (same as forward_mqa does)
+    S = num_blocks * block_size
+    kv_flat = cache.reshape(S, head_size)
+    c_kv_read = kv_flat[:T, :kv_lora_rank]
+    k_rope_read = kv_flat[:T, kv_lora_rank:]
+
+    # Verify consistency
+    torch.testing.assert_close(c_kv_read, kv_c_normed, rtol=1e-3, atol=1e-3)
+    torch.testing.assert_close(k_rope_read, k_pe.squeeze(1), rtol=1e-3, atol=1e-3)
+
+
+def test_kv_cache_dtype_auto():
+    """Verify kv_cache_dtype='auto' uses BF16 direct copy."""
+    device = "cuda"
+    cache = torch.zeros(1, 64, 576, dtype=torch.bfloat16, device=device)
+
+    kv_c = torch.randn(1, 512, dtype=torch.bfloat16, device=device)
+    k_pe = torch.randn(1, 1, 64, dtype=torch.bfloat16, device=device)
+    slot_mapping = torch.tensor([0], dtype=torch.int64, device=device)
+    k_scale = torch.ones(1, dtype=torch.float32, device=device)
+
+    from vllm import _custom_ops as ops
+
+    ops.concat_and_cache_mla(
+        kv_c, k_pe.squeeze(1), cache, slot_mapping, kv_cache_dtype="auto", scale=k_scale
+    )
+
+    assert cache.dtype == torch.bfloat16
+
+
+# ─── TEST 2.4: Edge Cases ────────────────────────────────────────────
+
+
+def test_empty_kv_cache():
+    """Verify do_kv_cache_update handles empty cache gracefully."""
+
+    kv_cache = torch.empty(0, device="cuda")
+    # Should return without error (numel() == 0 check)
+    # We call the static method from parent class directly
+    from vllm.v1.attention.backend import SparseMLAAttentionImpl
+
+    SparseMLAAttentionImpl.do_kv_cache_update(
+        None,
+        kv_c_normed=torch.empty(0),
+        k_pe=torch.empty(0),
+        kv_cache=kv_cache,
+        slot_mapping=torch.empty(0),
+        kv_cache_dtype="auto",
+        k_scale=torch.ones(1),
+    )
+
+
+# ─── TEST 2.5: Valid Counts from Index Conversion ───────────────────
+
+
+def test_triton_convert_valid_counts():
+    """Verify triton_convert_req_index_to_global_index with return_valid_counts.
+
+    This tests the core fix mechanism: the Triton kernel atomically counts
+    valid (non -1) entries per row while converting indices.
+    """
+    from vllm.v1.attention.backends.mla.sparse_utils import (
+        triton_convert_req_index_to_global_index,
+    )
+
+    device = "cuda"
+    T = 4
+    topk = 128
+    num_blocks = 16
+    block_size = 64
+
+    req_id = torch.zeros(T, dtype=torch.int32, device=device)
+    block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze(
+        0
+    )  # [1, num_blocks]
+
+    # Create topk_indices with varying valid entries per token
+    topk_indices = torch.full((T, topk), -1, dtype=torch.int32, device=device)
+    expected_valid = [1, 10, 50, 100]
+    for i in range(T):
+        nv = expected_valid[i]
+        # Use indices within the valid range
+        topk_indices[i, :nv] = torch.randint(
+            0,
+            num_blocks * block_size,
+            (nv,),
+            dtype=torch.int32,
+            device=device,
+        )
+
+    global_idx, valid_counts = triton_convert_req_index_to_global_index(
+        req_id,
+        block_table,
+        topk_indices,
+        BLOCK_SIZE=block_size,
+        NUM_TOPK_TOKENS=topk,
+        return_valid_counts=True,
+    )
+
+    # Verify valid counts match expected
+    for i in range(T):
+        assert valid_counts[i].item() == expected_valid[i], (
+            f"Token {i}: expected {expected_valid[i]} valid, "
+            f"got {valid_counts[i].item()}"
+        )
+
+    # Verify -1 propagation
+    for i in range(T):
+        nv = expected_valid[i]
+        # Entries beyond valid should be -1
+        assert (global_idx[i, nv:] == -1).all(), (
+            f"Token {i}: entries beyond valid count should be -1"
+        )
+
+
+# ─── TEST 2.6: Prefill Metadata Correctness ─────────────────────────
+
+
+def test_prefill_cache_seqlens_vs_valid_counts():
+    """Verify metadata cache_seqlens = min(seq_len, topk) and that the
+    forward_mqa fix overrides with valid_counts.
+
+    The metadata builder computes cache_seqlens as min(seq_len, topk).
+    For prefill tokens, this can exceed the actual valid topk entries.
+    The fix in forward_mqa uses valid_counts instead.
+    """
+    import numpy as np
+
+    device = "cuda"
+
+    # Simulate a prefill batch: 1 request, 4 tokens, seq_len=4
+    num_reqs = 1
+    T = 4
+    topk = 2048
+    seq_len = 4
+
+    # The metadata builder's logic (simplified):
+    starts = np.array([0, T], dtype=np.int32)
+    seg_lens = np.diff(starts)  # [4]
+    seq_lens_np = np.array([seq_len], dtype=np.int32)
+
+    per_tok_seqlens = np.minimum(np.repeat(seq_lens_np, seg_lens), topk)  # [4, 4, 4, 4]
+
+    # This is what the metadata builder produces:
+    assert all(per_tok_seqlens == 4), (
+        "Metadata cache_seqlens should be min(seq_len, topk) = 4"
+    )
+
+    # But the actual valid entries per token (with causal masking):
+    # Token 0: 1 valid entry, Token 1: 2, Token 2: 3, Token 3: 4
+    expected_valid = [1, 2, 3, 4]
+
+    # The fix in forward_mqa computes valid_counts from the page_table
+    # and uses those as cache_seqlens. Verify the fix produces correct
+    # valid counts:
+    from vllm.v1.attention.backends.mla.sparse_utils import (
+        triton_convert_req_index_to_global_index,
+    )
+
+    req_id = torch.zeros(T, dtype=torch.int32, device=device)
+    block_table = torch.arange(32, dtype=torch.int32, device=device).unsqueeze(0)
+
+    topk_indices = torch.full((T, topk), -1, dtype=torch.int32, device=device)
+    for i in range(T):
+        nv = expected_valid[i]
+        topk_indices[i, :nv] = torch.arange(nv, dtype=torch.int32, device=device)
+
+    _, valid_counts = triton_convert_req_index_to_global_index(
+        req_id,
+        block_table,
+        topk_indices,
+        BLOCK_SIZE=64,
+        NUM_TOPK_TOKENS=topk,
+        return_valid_counts=True,
+    )
+
+    for i in range(T):
+        assert valid_counts[i].item() == expected_valid[i], (
+            f"Token {i}: valid_counts should be {expected_valid[i]}, "
+            f"got {valid_counts[i].item()}"
+        )
+
+
+# ─── TEST 2.7: Clamp -1 to 0 Safety ─────────────────────────────────
+
+
+def test_global_idx_clamp_safety():
+    """Verify clamping -1 page indices to 0 prevents OOB access."""
+    device = "cuda"
+
+    # Create a page_table with -1 entries
+    page_table = torch.tensor(
+        [[5, 10, -1, -1], [3, -1, -1, -1]],
+        dtype=torch.int32,
+        device=device,
+    )
+
+    # Clamp -1 to 0
+    clamped = page_table.clamp(min=0)
+
+    # Verify
+    expected = torch.tensor(
+        [[5, 10, 0, 0], [3, 0, 0, 0]],
+        dtype=torch.int32,
+        device=device,
+    )
+    assert torch.equal(clamped, expected), (
+        f"Clamped page_table doesn't match expected: {clamped} vs {expected}"
+    )
+
+
+# ─── TEST 2.8: In-place clamp correctness ───────────────────────────
+
+
+def test_inplace_clamp_no_negative_indices():
+    """Verify in-place clamp_(min=0) on global_idx leaves no -1 entries.
+
+    The review-fixed code uses clamp_() (in-place) instead of clamp()
+    to avoid unnecessary tensor allocations during CUDA graph capture.
+    """
+    device = "cuda"
+
+    # Create a global_idx tensor with -1 entries
+    global_idx = torch.tensor(
+        [[100, 200, -1, -1, -1], [50, -1, -1, -1, -1]],
+        dtype=torch.int32,
+        device=device,
+    )
+
+    # In-place clamp
+    global_idx.clamp_(min=0)
+
+    # Verify no -1 entries remain
+    assert (global_idx >= 0).all(), (
+        f"In-place clamp should remove all -1 entries: {global_idx}"
+    )
+    # Verify valid entries are preserved
+    assert global_idx[0, 0].item() == 100
+    assert global_idx[0, 1].item() == 200
+    assert global_idx[1, 0].item() == 50
+
+
+# ─── TEST 2.9: Full fix flow with index conversion ──────────────────
+
+
+def test_full_fix_flow_valid_counts_and_clamp():
+    """End-to-end test of the complete fix flow:
+    1. triton_convert_req_index_to_global_index with return_valid_counts=True
+    2. In-place clamp global_idx to replace -1 with 0
+    3. In-place clamp valid_counts to min=1
+    4. Use valid_counts as cache_seqlens
+
+    This simulates what forward_mqa does after the fix.
+    """
+    from vllm.v1.attention.backends.mla.sparse_utils import (
+        triton_convert_req_index_to_global_index,
+    )
+
+    device = "cuda"
+    T = 4
+    topk = 128
+    num_blocks = 16
+    block_size = 64
+
+    req_id = torch.zeros(T, dtype=torch.int32, device=device)
+    block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze(
+        0
+    )
+
+    # Simulate causal prefill: token i has (i+1) valid entries
+    topk_indices = torch.full((T, topk), -1, dtype=torch.int32, device=device)
+    expected_valid = [1, 2, 3, 4]
+    for i in range(T):
+        nv = expected_valid[i]
+        topk_indices[i, :nv] = torch.arange(nv, dtype=torch.int32, device=device)
+
+    # Step 1: Convert with valid counts
+    global_idx, valid_counts = triton_convert_req_index_to_global_index(
+        req_id,
+        block_table,
+        topk_indices,
+        BLOCK_SIZE=block_size,
+        NUM_TOPK_TOKENS=topk,
+        return_valid_counts=True,
+    )
+
+    # Step 2: In-place clamp global_idx (no -1 entries after)
+    global_idx.clamp_(min=0)
+    assert (global_idx >= 0).all(), "No -1 entries should remain after clamp_"
+
+    # Step 3: In-place clamp valid_counts to min=1
+    valid_counts.clamp_(min=1)
+    cache_seqlens = valid_counts
+
+    # Step 4: Verify valid counts match expected
+    for i in range(T):
+        assert cache_seqlens[i].item() == expected_valid[i], (
+            f"Token {i}: expected cache_seqlens={expected_valid[i]}, "
+            f"got {cache_seqlens[i].item()}"
+        )
+
+    # Step 5: Verify that for each token, entries 0..cache_seqlens-1 in
+    # global_idx are valid (non-zero, since we clamped -1 to 0 for the
+    # entries beyond valid_counts, the valid entries at positions 0..nv-1
+    # should be the actual converted indices)
+    for i in range(T):
+        nv = expected_valid[i]
+        valid_region = global_idx[i, :nv]
+        # Valid region should have specific converted values from block_table
+        # (not just zeros from clamping)
+        # For indices [0, 1, ..., nv-1] with block_size=64:
+        #   block_id = index // 64, inblock_off = index % 64
+        #   out = block_table[0, block_id] * 64 + inblock_off
+        for j in range(nv):
+            block_id = j // block_size
+            inblock_off = j % block_size
+            expected_val = block_table[0, block_id].item() * block_size + inblock_off
+            assert valid_region[j].item() == expected_val, (
+                f"Token {i}, position {j}: expected {expected_val}, "
+                f"got {valid_region[j].item()}"
+            )
+
+
+# ─── TEST 2.10: CUDA Graph Padding Fix ─────────────────────────────
+# These tests verify the fix for Issue 2: RuntimeError when
+# num_actual_tokens (padded) != sum(seg_lens) (real tokens).
+# This is the core bug that caused the crash during lm_eval with
+# 32 concurrent requests on DeepSeek-V3.2.
+
+
+def _make_mock_vllm_config(max_tokens=512):
+    """Create a mock VllmConfig for metadata builder tests."""
+    from unittest.mock import MagicMock
+
+    vllm_config = MagicMock()
+    vllm_config.scheduler_config.max_num_batched_tokens = max_tokens
+    vllm_config.speculative_config = None
+    vllm_config.parallel_config.decode_context_parallel_size = 1
+    return vllm_config
+
+
+def test_metadata_builder_cuda_graph_padding():
+    """Verify build() handles CUDA graph padding (T > actual_tokens).
+
+    Reproduces the exact crash from Issue 2:
+      RuntimeError: The size of tensor a (32) must match the size
+      of tensor b (31) at non-singleton dimension 0
+
+    This happens when num_actual_tokens=32 (padded for CUDA graph)
+    but only 31 real tokens exist (one request completed mid-batch).
+    """
+    from unittest.mock import MagicMock
+
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        CutlassFA3MLASparseMetadataBuilder,
+    )
+
+    device = "cuda"
+    max_tokens = 512
+    block_size = 64
+    topk = 2048
+
+    # Mock kv_cache_spec
+    kv_cache_spec = MagicMock()
+    kv_cache_spec.block_size = block_size
+
+    # Mock vllm_config
+    vllm_config = _make_mock_vllm_config(max_tokens)
+
+    builder = CutlassFA3MLASparseMetadataBuilder(
+        kv_cache_spec=kv_cache_spec,
+        layer_names=["layers.0.self_attn"],
+        vllm_config=vllm_config,
+        device=torch.device(device),
+    )
+    builder.topk_tokens = topk
+
+    # Simulate the crash scenario: 31 real tokens padded to 32
+    padded_T = 32
+    real_tokens = 31
+    num_reqs_padded = 32  # padded request count
+
+    # Accurately mock gpu_model_runner.py's padding behavior:
+    # query_start_loc.cpu[:num_reqs_padded+1] = [:33], 33 entries
+    # Real entries: [0,1,...,31], Padding: [31] (repeats last value)
+    query_start_loc_cpu = list(range(real_tokens + 1)) + [real_tokens]
+    # seq_lens_cpu[:num_reqs_padded] = [:32], 32 entries
+    # Real entries: [100]*31, Padding: [0] (stale/zero for padding slot)
+    seq_lens_cpu = [100] * real_tokens + [0]
+
+    # Build the mock CommonAttentionMetadata
+    cm = MagicMock()
+    cm.num_actual_tokens = padded_T  # PADDED to 32
+    cm.query_start_loc_cpu = query_start_loc_cpu
+    cm.seq_lens_cpu = seq_lens_cpu
+    cm.num_reqs = num_reqs_padded  # gpu_model_runner passes padded count
+    cm.max_query_len = 1
+    cm.max_seq_len = 100
+    cm.query_start_loc = torch.tensor(
+        query_start_loc_cpu, dtype=torch.int32, device=device
+    )
+    cm.slot_mapping = torch.zeros(padded_T, dtype=torch.int64, device=device)
+    cm.block_table_tensor = torch.zeros(
+        num_reqs_padded, 4, dtype=torch.int32, device=device
+    )
+
+    # This should NOT raise RuntimeError
+    metadata = builder.build(
+        common_prefix_len=0,
+        common_attn_metadata=cm,
+    )
+
+    # Verify metadata shapes match padded T
+    assert metadata.req_id_per_token.shape[0] == padded_T, (
+        f"req_id_per_token should have padded size {padded_T}, "
+        f"got {metadata.req_id_per_token.shape[0]}"
+    )
+    assert metadata.cache_seqlens.shape[0] == padded_T, (
+        f"cache_seqlens should have padded size {padded_T}, "
+        f"got {metadata.cache_seqlens.shape[0]}"
+    )
+    assert metadata.cu_seqlens_q.shape[0] == padded_T + 1
+    assert metadata.cu_seqlens_k.shape[0] == padded_T + 1
+
+    # Verify real data portion is correct
+    for i in range(real_tokens):
+        assert metadata.req_id_per_token[i].item() == i, (
+            f"Token {i}: req_id should be {i}, "
+            f"got {metadata.req_id_per_token[i].item()}"
+        )
+        assert metadata.cache_seqlens[i].item() == 100, (
+            f"Token {i}: cache_seqlens should be 100, "
+            f"got {metadata.cache_seqlens[i].item()}"
+        )
+
+    # Verify padding tokens have safe defaults
+    assert metadata.req_id_per_token[real_tokens].item() == 0, (
+        "Padding token req_id should be 0"
+    )
+    assert metadata.cache_seqlens[real_tokens].item() >= 1, (
+        "Padding token cache_seqlens should be >= 1 (safe minimum)"
+    )
+
+    # Verify cu_seqlens_q is [0, 1, 2, ..., padded_T] (always correct)
+    for i in range(padded_T + 1):
+        assert metadata.cu_seqlens_q[i].item() == i, (
+            f"cu_seqlens_q[{i}] should be {i}, got {metadata.cu_seqlens_q[i].item()}"
+        )
+
+    # Verify cu_seqlens_k is monotonically non-decreasing
+    for i in range(padded_T):
+        assert metadata.cu_seqlens_k[i + 1].item() >= metadata.cu_seqlens_k[i].item(), (
+            f"cu_seqlens_k must be non-decreasing at index {i}: "
+            f"{metadata.cu_seqlens_k[i].item()} -> {metadata.cu_seqlens_k[i + 1].item()}"
+        )
+    # Verify cu_seqlens_k at the real/padding boundary
+    assert metadata.cu_seqlens_k[real_tokens].item() == real_tokens * 100, (
+        f"cu_seqlens_k[{real_tokens}] should be {real_tokens * 100}, "
+        f"got {metadata.cu_seqlens_k[real_tokens].item()}"
+    )
+
+
+@pytest.mark.parametrize(
+    "real_tokens,padded_T",
+    [
+        (1, 2),  # minimal padding
+        (3, 32),  # large padding gap
+        (7, 8),  # small batch
+        (15, 16),  # medium batch
+        (31, 32),  # the exact crash scenario
+        (100, 104),  # larger padding gap
+    ],
+)
+def test_metadata_builder_cuda_graph_padding_various(real_tokens, padded_T):
+    """Verify build() handles various CUDA graph padding scenarios.
+
+    Uses accurate mock that matches gpu_model_runner.py's padding behavior:
+    - query_start_loc_cpu has num_reqs_padded+1 entries (with padded suffix)
+    - seq_lens_cpu has num_reqs_padded entries (with stale padding entries)
+    """
+    from unittest.mock import MagicMock
+
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        CutlassFA3MLASparseMetadataBuilder,
+    )
+
+    device = "cuda"
+    max_tokens = max(512, padded_T + 1)  # ensure buffer large enough
+    block_size = 64
+    topk = 2048
+    num_reqs_padded = padded_T  # For decode-only, padded_T == num_reqs_padded
+
+    kv_cache_spec = MagicMock()
+    kv_cache_spec.block_size = block_size
+
+    vllm_config = _make_mock_vllm_config(max_tokens)
+
+    builder = CutlassFA3MLASparseMetadataBuilder(
+        kv_cache_spec=kv_cache_spec,
+        layer_names=["layers.0.self_attn"],
+        vllm_config=vllm_config,
+        device=torch.device(device),
+    )
+    builder.topk_tokens = topk
+
+    # Accurate mock: query_start_loc_cpu[:num_reqs_padded+1]
+    # Real entries [0,1,...,real_tokens], then (num_reqs_padded - real_tokens)
+    # padding entries all equal to real_tokens (flat, non-decreasing)
+    query_start_loc_cpu = list(range(real_tokens + 1))
+    num_padding_reqs = num_reqs_padded - real_tokens
+    query_start_loc_cpu += [real_tokens] * num_padding_reqs
+
+    # seq_lens_cpu[:num_reqs_padded] — padding entries are stale (zero)
+    seq_lens_cpu = [200] * real_tokens + [0] * num_padding_reqs
+
+    cm = MagicMock()
+    cm.num_actual_tokens = padded_T
+    cm.query_start_loc_cpu = query_start_loc_cpu
+    cm.seq_lens_cpu = seq_lens_cpu
+    cm.num_reqs = num_reqs_padded
+    cm.max_query_len = 1
+    cm.max_seq_len = 200
+    cm.query_start_loc = torch.tensor(
+        query_start_loc_cpu, dtype=torch.int32, device=device
+    )
+    cm.slot_mapping = torch.zeros(padded_T, dtype=torch.int64, device=device)
+    cm.block_table_tensor = torch.zeros(
+        max(num_reqs_padded, 1), 4, dtype=torch.int32, device=device
+    )
+
+    # Should NOT raise any errors
+    metadata = builder.build(
+        common_prefix_len=0,
+        common_attn_metadata=cm,
+    )
+
+    # Verify shapes match padded T
+    assert metadata.req_id_per_token.shape[0] == padded_T
+    assert metadata.cache_seqlens.shape[0] == padded_T
+    assert metadata.cu_seqlens_q.shape[0] == padded_T + 1
+    assert metadata.cu_seqlens_k.shape[0] == padded_T + 1
+    assert metadata.num_actual_tokens == padded_T
+
+    # Verify real portion
+    for i in range(real_tokens):
+        assert metadata.req_id_per_token[i].item() == i
+        assert metadata.cache_seqlens[i].item() == 200
+
+    # Verify padding
+    for i in range(real_tokens, padded_T):
+        assert metadata.req_id_per_token[i].item() == 0
+        assert metadata.cache_seqlens[i].item() >= 1
+
+    # Verify cu_seqlens_q is [0, 1, ..., padded_T]
+    for i in range(padded_T + 1):
+        assert metadata.cu_seqlens_q[i].item() == i
+
+    # Verify cu_seqlens_k monotonicity
+    for i in range(padded_T):
+        assert metadata.cu_seqlens_k[i + 1].item() >= metadata.cu_seqlens_k[i].item()
+    # Verify cu_seqlens_k at boundary
+    assert metadata.cu_seqlens_k[real_tokens].item() == real_tokens * 200
+
+
+def test_metadata_builder_no_padding():
+    """Verify build() still works correctly when T == actual_tokens (no padding)."""
+    from unittest.mock import MagicMock
+
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        CutlassFA3MLASparseMetadataBuilder,
+    )
+
+    device = "cuda"
+    max_tokens = 512
+    block_size = 64
+
+    kv_cache_spec = MagicMock()
+    kv_cache_spec.block_size = block_size
+
+    vllm_config = _make_mock_vllm_config(max_tokens)
+
+    builder = CutlassFA3MLASparseMetadataBuilder(
+        kv_cache_spec=kv_cache_spec,
+        layer_names=["layers.0.self_attn"],
+        vllm_config=vllm_config,
+        device=torch.device(device),
+    )
+    builder.topk_tokens = 2048
+
+    # No padding: T == real tokens
+    T = 4
+    query_start_loc_cpu = [0, 1, 2, 3, 4]  # 4 decode tokens
+    seq_lens_cpu = [50, 100, 150, 200]
+
+    cm = MagicMock()
+    cm.num_actual_tokens = T
+    cm.query_start_loc_cpu = query_start_loc_cpu
+    cm.seq_lens_cpu = seq_lens_cpu
+    cm.num_reqs = 4
+    cm.max_query_len = 1
+    cm.max_seq_len = 200
+    cm.query_start_loc = torch.tensor(
+        query_start_loc_cpu, dtype=torch.int32, device=device
+    )
+    cm.slot_mapping = torch.zeros(T, dtype=torch.int64, device=device)
+    cm.block_table_tensor = torch.zeros(4, 4, dtype=torch.int32, device=device)
+
+    metadata = builder.build(
+        common_prefix_len=0,
+        common_attn_metadata=cm,
+    )
+
+    assert metadata.req_id_per_token.shape[0] == T
+    assert metadata.cache_seqlens.shape[0] == T
+    assert metadata.num_actual_tokens == T
+
+    # Verify exact values
+    assert metadata.req_id_per_token[0].item() == 0
+    assert metadata.req_id_per_token[1].item() == 1
+    assert metadata.req_id_per_token[2].item() == 2
+    assert metadata.req_id_per_token[3].item() == 3
+    assert metadata.cache_seqlens[0].item() == 50
+    assert metadata.cache_seqlens[1].item() == 100
+    assert metadata.cache_seqlens[2].item() == 150
+    assert metadata.cache_seqlens[3].item() == 200
+
+
+def test_metadata_builder_mixed_prefill_decode_with_padding():
+    """Verify build() handles mixed prefill+decode with CUDA graph padding.
+
+    This tests a more complex scenario: 2 decode tokens + 3 prefill tokens
+    from 3 requests, padded from 5 to 8 tokens.
+    """
+    from unittest.mock import MagicMock
+
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        CutlassFA3MLASparseMetadataBuilder,
+    )
+
+    device = "cuda"
+    max_tokens = 512
+    block_size = 64
+
+    kv_cache_spec = MagicMock()
+    kv_cache_spec.block_size = block_size
+
+    vllm_config = _make_mock_vllm_config(max_tokens)
+
+    builder = CutlassFA3MLASparseMetadataBuilder(
+        kv_cache_spec=kv_cache_spec,
+        layer_names=["layers.0.self_attn"],
+        vllm_config=vllm_config,
+        device=torch.device(device),
+    )
+    builder.topk_tokens = 2048
+
+    # 3 real requests: req0 (1 decode token), req1 (1 decode token),
+    #                  req2 (3 prefill tokens)
+    # Total: 5 real tokens, padded to 8 tokens, 8 padded request slots
+    real_tokens = 5
+    num_real_reqs = 3
+    padded_T = 8
+    num_reqs_padded = 8  # padded request count
+
+    # Accurate: query_start_loc_cpu[:num_reqs_padded+1] = 9 entries
+    # Real: [0, 1, 2, 5], Padding: [5, 5, 5, 5, 5]
+    query_start_loc_cpu = [0, 1, 2, 5] + [5] * (num_reqs_padded - num_real_reqs)
+    # seq_lens_cpu[:num_reqs_padded] = 8 entries
+    seq_lens_cpu = [100, 200, 3] + [0] * (num_reqs_padded - num_real_reqs)
+
+    cm = MagicMock()
+    cm.num_actual_tokens = padded_T
+    cm.query_start_loc_cpu = query_start_loc_cpu
+    cm.seq_lens_cpu = seq_lens_cpu
+    cm.num_reqs = num_reqs_padded
+    cm.max_query_len = 3
+    cm.max_seq_len = 200
+    cm.query_start_loc = torch.tensor(
+        query_start_loc_cpu, dtype=torch.int32, device=device
+    )
+    cm.slot_mapping = torch.zeros(padded_T, dtype=torch.int64, device=device)
+    cm.block_table_tensor = torch.zeros(
+        num_reqs_padded, 4, dtype=torch.int32, device=device
+    )
+
+    metadata = builder.build(
+        common_prefix_len=0,
+        common_attn_metadata=cm,
+    )
+
+    # Verify shapes
+    assert metadata.req_id_per_token.shape[0] == padded_T
+    assert metadata.cache_seqlens.shape[0] == padded_T
+
+    # Verify req_id mapping
+    assert metadata.req_id_per_token[0].item() == 0  # req0, decode
+    assert metadata.req_id_per_token[1].item() == 1  # req1, decode
+    assert metadata.req_id_per_token[2].item() == 2  # req2, prefill tok0
+    assert metadata.req_id_per_token[3].item() == 2  # req2, prefill tok1
+    assert metadata.req_id_per_token[4].item() == 2  # req2, prefill tok2
+    # Padding tokens
+    assert metadata.req_id_per_token[5].item() == 0
+    assert metadata.req_id_per_token[6].item() == 0
+    assert metadata.req_id_per_token[7].item() == 0
+
+    # Verify cache_seqlens
+    assert metadata.cache_seqlens[0].item() == 100  # req0 seq_len
+    assert metadata.cache_seqlens[1].item() == 200  # req1 seq_len
+    assert metadata.cache_seqlens[2].item() == 3  # req2 seq_len
+    assert metadata.cache_seqlens[3].item() == 3  # req2 seq_len
+    assert metadata.cache_seqlens[4].item() == 3  # req2 seq_len
+    # Padding (default = 1)
+    assert metadata.cache_seqlens[5].item() >= 1
+    assert metadata.cache_seqlens[6].item() >= 1
+    assert metadata.cache_seqlens[7].item() >= 1
+
+
+# ─── TEST 2.11: Zero Real Tokens Edge Case (Review Issue #3) ─────
+# Tests the edge case where ALL tokens are padding (actual_tokens=0).
+# This can happen during CUDA graph warmup/capture with dummy batches.
+
+
+def test_metadata_builder_zero_real_tokens():
+    """Verify build() handles the case where all tokens are padding.
+
+    This edge case can occur during CUDA graph warmup or capture where
+    dummy batches may have zero real tokens but T > 0 (padded size).
+    """
+    from unittest.mock import MagicMock
+
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        CutlassFA3MLASparseMetadataBuilder,
+    )
+
+    device = "cuda"
+    max_tokens = 512
+    block_size = 64
+
+    kv_cache_spec = MagicMock()
+    kv_cache_spec.block_size = block_size
+
+    vllm_config = _make_mock_vllm_config(max_tokens)
+
+    builder = CutlassFA3MLASparseMetadataBuilder(
+        kv_cache_spec=kv_cache_spec,
+        layer_names=["layers.0.self_attn"],
+        vllm_config=vllm_config,
+        device=torch.device(device),
+    )
+    builder.topk_tokens = 2048
+
+    # Zero real tokens, padded to 4
+    # This happens when query_start_loc = [0] only (1 entry, no requests)
+    # and num_actual_tokens is still the padded count.
+    padded_T = 4
+    real_tokens = 0
+    # query_start_loc_cpu with a single entry means 0 requests
+    query_start_loc_cpu = [0]
+    seq_lens_cpu = []
+
+    cm = MagicMock()
+    cm.num_actual_tokens = padded_T
+    cm.query_start_loc_cpu = query_start_loc_cpu
+    cm.seq_lens_cpu = seq_lens_cpu
+    cm.num_reqs = 0
+    cm.max_query_len = 0
+    cm.max_seq_len = 0
+    cm.query_start_loc = torch.tensor(
+        query_start_loc_cpu, dtype=torch.int32, device=device
+    )
+    cm.slot_mapping = torch.zeros(padded_T, dtype=torch.int64, device=device)
+    cm.block_table_tensor = torch.zeros(1, 4, dtype=torch.int32, device=device)
+
+    # Should NOT raise any errors
+    metadata = builder.build(
+        common_prefix_len=0,
+        common_attn_metadata=cm,
+    )
+
+    # Verify shapes match padded T
+    assert metadata.req_id_per_token.shape[0] == padded_T
+    assert metadata.cache_seqlens.shape[0] == padded_T
+    assert metadata.cu_seqlens_q.shape[0] == padded_T + 1
+    assert metadata.cu_seqlens_k.shape[0] == padded_T + 1
+
+    # All tokens are padding — verify safe defaults
+    for i in range(padded_T):
+        assert metadata.req_id_per_token[i].item() == 0
+        assert metadata.cache_seqlens[i].item() >= 1
+
+    # cu_seqlens_k should be monotonically non-decreasing
+    for i in range(padded_T):
+        assert metadata.cu_seqlens_k[i + 1].item() >= metadata.cu_seqlens_k[i].item()
+
+
+# ─── TEST 2.12: Batch Size Gating Constant ───────────────────────────
+
+
+def test_batch_size_gating_threshold():
+    """Verify MAX_BATCH_SIZE_FOR_FA3 is 16 and controls routing."""
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        MAX_BATCH_SIZE_FOR_FA3,
+        _flashmla_sparse_available,
+    )
+
+    assert MAX_BATCH_SIZE_FOR_FA3 == 16
+    # On SM90 builds, FlashMLA fallback should be available
+    # (unless FlashMLA was explicitly excluded from the build)
+    assert isinstance(_flashmla_sparse_available, bool)
+
+
+# ─── TEST 2.13: FlashMLA Fallback Head Padding ──────────────────────
+
+
+def test_flashmla_fallback_head_padding():
+    """Verify FlashMLA fallback head padding constant is 64 for SM90."""
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        _FLASHMLA_SM90_HEAD_PADDING,
+    )
+
+    assert _FLASHMLA_SM90_HEAD_PADDING == 64, (
+        f"SM90 head padding should be 64, got {_FLASHMLA_SM90_HEAD_PADDING}"
+    )
+
+
+# ─── TEST 2.14: Forward MQA Dispatch Verification ────────────────────
+
+
+def test_forward_mqa_has_fa3_and_fallback_methods():
+    """Verify CutlassFA3MLASparseImpl has both kernel dispatch methods."""
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        CutlassFA3MLASparseImpl,
+    )
+
+    assert hasattr(CutlassFA3MLASparseImpl, "_forward_fa3"), (
+        "CutlassFA3MLASparseImpl should have _forward_fa3 method"
+    )
+    assert hasattr(CutlassFA3MLASparseImpl, "_forward_flashmla_bf16_fallback"), (
+        "CutlassFA3MLASparseImpl should have _forward_flashmla_bf16_fallback method"
+    )
@@ -1,105 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Test batch-invariant matmul against torch.matmul for various shape combinations.
-
-Tests correctness (matches torch.matmul) and batch invariance (result for one
-item doesn't change based on other items in the batch).
-"""
-
-import pytest
-import torch
-from utils import skip_unsupported
-
-from vllm.model_executor.layers.batch_invariant import matmul_batch_invariant
-from vllm.platforms import current_platform
-
-DEVICE_TYPE = current_platform.device_type
-
-
-@skip_unsupported
-@pytest.mark.parametrize(
-    "a_shape,b_shape",
-    [
-        # 2D x 2D
-        ((32, 64), (64, 16)),
-        # 2D x 3D
-        ((64, 16), (4, 16, 32)),
-        # 3D x 2D
-        ((4, 32, 64), (64, 16)),
-        # 4D x 2D
-        ((1, 4, 32, 64), (64, 16)),
-        # 3D x 3D
-        ((4, 32, 64), (4, 64, 16)),
-        # 3D x 4D
-        ((2, 32, 64), (1, 2, 64, 16)),
-        # 4D x 3D (Gemma4 pattern)
-        ((1, 2, 32, 64), (2, 64, 16)),
-        # 4D x 4D
-        ((1, 2, 32, 64), (4, 2, 64, 16)),
-        # 2D x 4D
-        ((32, 64), (1, 2, 64, 16)),
-        # 2D x 5D
-        ((32, 64), (1, 2, 2, 64, 16)),
-        # 5D x 2D
-        ((1, 2, 2, 32, 64), (64, 16)),
-        # 5D x 5D
-        ((1, 2, 4, 32, 64), (1, 2, 4, 64, 16)),
-    ],
-)
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-def test_matmul_correctness(a_shape, b_shape, dtype):
-    """
-    Compare matmul_batch_invariant against torch.matmul for various shapes.
-    """
-    device = torch.device(DEVICE_TYPE)
-
-    torch.manual_seed(42)
-    a = torch.rand(a_shape, dtype=dtype, device=device)
-    b = torch.rand(b_shape, dtype=dtype, device=device)
-
-    # Standard implementation (CUDA ops)
-    standard_output = torch.matmul(a, b)
-
-    # Batch-invariant implementation (Triton)
-    triton_output = matmul_batch_invariant(a, b)
-
-    # Compare outputs
-    # Use looser tolerance for bfloat16 due to its lower precision
-    if dtype == torch.bfloat16:
-        rtol, atol = 1e-1, 1e-1  # 10% relative tolerance for bfloat16
-    else:
-        rtol, atol = 1e-2, 1e-2  # 1% for float16/float32
-
-    torch.testing.assert_close(
-        triton_output,
-        standard_output,
-        rtol=rtol,
-        atol=atol,
-        msg=f"matmul mismatch for a ndim={a.ndim}, b ndim={b.ndim},",
-    )
-
-
-@skip_unsupported
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-def test_matmul_batch_invariance(dtype):
-    """
-    Verify that the result for one item is bitwise identical regardless
-    of what other items are in the batch.
-    """
-
-    device = torch.device(DEVICE_TYPE)
-
-    torch.manual_seed(42)
-    a_single = torch.rand((1, 64, 32), dtype=dtype, device=device)
-    b = torch.rand((32, 128), dtype=dtype, device=device)
-
-    standard_output = matmul_batch_invariant(a_single, b)
-
-    a_batch = torch.rand((8, 64, 32), dtype=dtype, device=device)
-    a_batch[3] = a_single[0]
-
-    batch_output = matmul_batch_invariant(a_batch, b)
-    batch_output_a = batch_output[3]
-
-    assert torch.equal(standard_output[0], batch_output_a)
@@ -32,8 +32,8 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
    # 3 blocks, store just the middle block (skip first and last)
    # blocks = [0, 1, 2], [3, 4, 5], [6, 7, 8]
    runner.new_request(token_ids=[0] * offloaded_block_size * 3)
-    runner.manager.prepare_store.side_effect = (
-        lambda keys, req_context: generate_store_output(list(keys)[1:2])
+    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(
+        list(keys)[1:2]
    )
    runner.run(decoded_tokens=[0])

@@ -45,22 +45,18 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
    runner.manager.prepare_store.assert_not_called()

    # +1 token -> single block, fail prepare_store
-    runner.manager.prepare_store.side_effect = lambda keys, req_context: None
+    runner.manager.prepare_store.side_effect = lambda keys: None
    runner.run(decoded_tokens=[0])
    runner.manager.prepare_store.assert_called()

    # 1 more block (+ token for async scheduling)
    # now set block_hashes_to_store = []
-    runner.manager.prepare_store.side_effect = (
-        lambda keys, req_context: generate_store_output([])
-    )
+    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
    runner.run(decoded_tokens=[0] * (offloaded_block_size + 1))

    # 1 more block (+ token for kicking off offloading)
    # now check touch was called with all 6 blocks
-    runner.manager.prepare_store.side_effect = (
-        lambda keys, req_context: generate_store_output(keys)
-    )
+    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
    runner.run(
        decoded_tokens=[0] * (offloaded_block_size + 1),
        expected_stored_gpu_block_indexes=(15, 16, 17),
@@ -93,17 +89,13 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
    runner.new_request(
        token_ids=[0] * gpu_block_size + [1] * (offloaded_block_size - gpu_block_size)
    )
-    runner.manager.prepare_store.side_effect = (
-        lambda keys, req_context: generate_store_output([])
-    )
+    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
    runner.run(decoded_tokens=[EOS_TOKEN_ID])
    runner.manager.lookup.assert_not_called()

    # single block lookup with no hits
    runner.new_request(token_ids=[1] * offloaded_block_size)
-    runner.manager.prepare_store.side_effect = (
-        lambda keys, req_context: generate_store_output([])
-    )
+    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
    runner.run(decoded_tokens=[EOS_TOKEN_ID])
    runner.manager.lookup.assert_called()
    assert len(list(runner.manager.lookup.call_args.args[0])) == 1
@@ -111,9 +103,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
    # single block lookup with a hit
    runner.scheduler.reset_prefix_cache()
    runner.new_request(token_ids=[0] * offloaded_block_size)
-    runner.manager.prepare_store.side_effect = (
-        lambda keys, req_context: generate_store_output([])
-    )
+    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
    runner.manager.lookup.return_value = 1
    runner.run(
        decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(0, 1, 2)
@@ -123,9 +113,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
    runner.new_request(
        token_ids=[0] * offloaded_block_size * 2 + [1] * offloaded_block_size
    )
-    runner.manager.prepare_store.side_effect = (
-        lambda keys, req_context: generate_store_output([])
-    )
+    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
    runner.manager.lookup.return_value = 1
    runner.run(
        decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(3, 4, 5)
@@ -176,18 +164,14 @@ def test_request_preemption(request_runner, async_scheduling: bool):
    # 2 blocks, store all, without flushing
    # blocks = [0, 1, 2], [3, 4, 5]
    runner.new_request(token_ids=[0] * offloaded_block_size * 2)
-    runner.manager.prepare_store.side_effect = (
-        lambda keys, req_context: generate_store_output(keys)
-    )
+    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
    runner.run(
        decoded_tokens=[0],
        complete_transfers=False,
    )

    # decode 2 more blocks - 1 gpu block, storing [6, 7, 8] (no flush)
-    runner.manager.prepare_store.side_effect = (
-        lambda keys, req_context: generate_store_output(keys)
-    )
+    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
    runner.run(
        decoded_tokens=[0] * (2 * offloaded_block_size - gpu_block_size),
        complete_transfers=False,
@@ -211,9 +195,7 @@ def test_request_preemption(request_runner, async_scheduling: bool):
    # request should now return from preemption
    # re-load [0, ..., 8] from the CPU and store [9, 10, 11]
    runner.manager.lookup.return_value = 3
-    runner.manager.prepare_store.side_effect = (
-        lambda keys, req_context: generate_store_output(keys)
-    )
+    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
    runner.run(
        decoded_tokens=[0] * gpu_block_size,
        expected_loaded_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
@@ -240,9 +222,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:

    # store 1 blocks
    runner.new_request(token_ids=[0] * offloaded_block_size)
-    runner.manager.prepare_store.side_effect = (
-        lambda keys, req_context: generate_store_output(keys)
-    )
+    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
    runner.run(
        decoded_tokens=[EOS_TOKEN_ID],
        expected_stored_gpu_block_indexes=(0, 1, 2),
@@ -273,9 +253,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
    assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)

    # complete transfers
-    runner.manager.prepare_store.side_effect = (
-        lambda keys, req_context: generate_store_output([])
-    )
+    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
    runner.run(
        decoded_tokens=[EOS_TOKEN_ID],
        expected_loaded_gpu_block_indexes=(0, 1, 2),
@@ -300,9 +278,7 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):

    # store 1 blocks
    runner.new_request(token_ids=[0] * offloaded_block_size)
-    runner.manager.prepare_store.side_effect = (
-        lambda keys, req_context: generate_store_output(keys)
-    )
+    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
    runner.run(
        decoded_tokens=[EOS_TOKEN_ID],
        expected_stored_gpu_block_indexes=(0, 1, 2),
@@ -115,7 +115,7 @@ class MockOffloadingSpec(OffloadingSpec):

        self.manager = MagicMock(spec=OffloadingManager)
        self.manager.lookup.return_value = 0
-        self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys)
+        self.manager.prepare_load = lambda keys: MockLoadStoreSpec(keys)
        self.handler = MockOffloadingHandler()

    def get_manager(self) -> OffloadingManager:
@@ -11,7 +11,6 @@ from vllm.v1.kv_offload.abstract import (
    OffloadingEvent,
    OffloadKey,
    PrepareStoreOutput,
-    ReqContext,
    make_offload_key,
 )
 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
@@ -20,14 +19,6 @@ from vllm.v1.kv_offload.mediums import CPULoadStoreSpec
 from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager


-def make_req_context(kv_transfer_params: dict | None = None) -> ReqContext:
-    """Create a ReqContext as production code would, from a request's params."""
-    return ReqContext(kv_transfer_params=kv_transfer_params)
-
-
-_EMPTY_REQ_CTX = make_req_context()
-
-
@dataclass
 class ExpectedPrepareStoreOutput:
    keys_to_store: list[int]
@@ -112,7 +103,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
    )

    # store [1, 2] and complete
-    manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
+    manager.prepare_store(to_keys([1, 2]))
    manager.complete_store(to_keys([1, 2]))

    # touch [1] to make block 2 the LRU candidate
@@ -122,7 +113,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
    #   - block 2 is already stored -> filtered out of keys_to_store
    #   - block 2 must NOT be evicted even though it is the LRU candidate
    #   - block 1 (ID 0) is evicted instead; new blocks [3,4,5] get IDs 2,3,0
-    prepare_store_output = manager.prepare_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX)
+    prepare_store_output = manager.prepare_store(to_keys([2, 3, 4, 5]))
    verify_store_output(
        prepare_store_output,
        ExpectedPrepareStoreOutput(
@@ -136,7 +127,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
    manager.complete_store(to_keys([2, 3, 4, 5]))

    # block 2 must still be present in the cache
-    assert manager.lookup(to_keys([2]), _EMPTY_REQ_CTX) == 1
+    assert manager.lookup(to_keys([2])) == 1


 def test_cpu_manager():
@@ -149,7 +140,7 @@ def test_cpu_manager():
    )

    # prepare store [1, 2]
-    prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
+    prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2]))
    verify_store_output(
        prepare_store_output,
        ExpectedPrepareStoreOutput(
@@ -160,7 +151,7 @@ def test_cpu_manager():
    )

    # lookup [1, 2] -> not ready
-    assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
+    assert cpu_manager.lookup(to_keys([1, 2])) == 0

    # no events so far
    assert list(cpu_manager.take_events()) == []
@@ -170,14 +161,12 @@ def test_cpu_manager():
    verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))

    # lookup [1, 2]
-    assert cpu_manager.lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1
-    assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2
-    assert cpu_manager.lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2
+    assert cpu_manager.lookup(to_keys([1])) == 1
+    assert cpu_manager.lookup(to_keys([1, 2])) == 2
+    assert cpu_manager.lookup(to_keys([1, 2, 3])) == 2

    # prepare store [2, 3, 4, 5] -> evicts [1]
-    prepare_store_output = cpu_manager.prepare_store(
-        to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX
-    )
+    prepare_store_output = cpu_manager.prepare_store(to_keys([2, 3, 4, 5]))
    verify_store_output(
        prepare_store_output,
        ExpectedPrepareStoreOutput(
@@ -191,23 +180,23 @@ def test_cpu_manager():
    verify_events(cpu_manager.take_events(), expected_evictions=({1},))

    # prepare store with no space
-    assert cpu_manager.prepare_store(to_keys([1, 6]), _EMPTY_REQ_CTX) is None
+    assert cpu_manager.prepare_store(to_keys([1, 6])) is None

    # complete store [2, 3, 4, 5]
    cpu_manager.complete_store(to_keys([2, 3, 4, 5]))

    # prepare load [2, 3]
-    prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX)
+    prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]))
    verify_load_output(prepare_load_output, [1, 2])

    # prepare store with no space ([2, 3] is being loaded)
-    assert cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) is None
+    assert cpu_manager.prepare_store(to_keys([6, 7, 8])) is None

    # complete load [2, 3]
    cpu_manager.complete_load(to_keys([2, 3]))

    # prepare store [6, 7, 8] -> evicts [2, 3, 4] (oldest)
-    prepare_store_output = cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX)
+    prepare_store_output = cpu_manager.prepare_store(to_keys([6, 7, 8]))
    verify_store_output(
        prepare_store_output,
        ExpectedPrepareStoreOutput(
@@ -224,7 +213,7 @@ def test_cpu_manager():
    cpu_manager.touch(to_keys([5, 6, 7]))

    # prepare store [7, 9] -> evicts [8] (oldest following previous touch)
-    prepare_store_output = cpu_manager.prepare_store(to_keys([9]), _EMPTY_REQ_CTX)
+    prepare_store_output = cpu_manager.prepare_store(to_keys([9]))
    verify_store_output(
        prepare_store_output,
        ExpectedPrepareStoreOutput(
@@ -238,8 +227,8 @@ def test_cpu_manager():
    cpu_manager.complete_store(to_keys([7, 9]), success=False)

    # assert [7] is still stored, but [9] is not
-    assert cpu_manager.lookup(to_keys([7]), _EMPTY_REQ_CTX) == 1
-    assert cpu_manager.lookup(to_keys([9]), _EMPTY_REQ_CTX) == 0
+    assert cpu_manager.lookup(to_keys([7])) == 1
+    assert cpu_manager.lookup(to_keys([9])) == 0

    verify_events(
        cpu_manager.take_events(),
@@ -271,9 +260,7 @@ class TestARCPolicy:
        cpu_manager, arc_policy = self._make_manager()

        # prepare store [1, 2]
-        prepare_store_output = cpu_manager.prepare_store(
-            to_keys([1, 2]), _EMPTY_REQ_CTX
-        )
+        prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2]))
        verify_store_output(
            prepare_store_output,
            ExpectedPrepareStoreOutput(
@@ -284,7 +271,7 @@ class TestARCPolicy:
        )

        # lookup [1, 2] -> not ready
-        assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
+        assert cpu_manager.lookup(to_keys([1, 2])) == 0

        # no events so far
        assert list(cpu_manager.take_events()) == []
@@ -294,9 +281,9 @@ class TestARCPolicy:
        verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))

        # lookup [1, 2]
-        assert cpu_manager.lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1
-        assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2
-        assert cpu_manager.lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2
+        assert cpu_manager.lookup(to_keys([1])) == 1
+        assert cpu_manager.lookup(to_keys([1, 2])) == 2
+        assert cpu_manager.lookup(to_keys([1, 2, 3])) == 2

        # blocks should be in T1 (recent)
        assert len(arc_policy.t1) == 2
@@ -310,7 +297,7 @@ class TestARCPolicy:
        cpu_manager, arc_policy = self._make_manager(enable_events=False)

        # store and complete block 1
-        cpu_manager.prepare_store(to_keys([1]), _EMPTY_REQ_CTX)
+        cpu_manager.prepare_store(to_keys([1]))
        cpu_manager.complete_store(to_keys([1]))

        # block 1 starts in T1 (recent)
@@ -332,9 +319,7 @@ class TestARCPolicy:
        cpu_manager, _ = self._make_manager()

        # prepare and complete store [1, 2, 3, 4]
-        prepare_store_output = cpu_manager.prepare_store(
-            to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX
-        )
+        prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2, 3, 4]))
        verify_store_output(
            prepare_store_output,
            ExpectedPrepareStoreOutput(
@@ -346,21 +331,19 @@ class TestARCPolicy:
        cpu_manager.complete_store(to_keys([1, 2, 3, 4]))

        # prepare load [2, 3] (increases ref_cnt)
-        prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX)
+        prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]))
        verify_load_output(prepare_load_output, [1, 2])

        # prepare store [5, 6, 7] with [2, 3] being loaded
        # should fail because [2, 3] have ref_cnt > 0
-        assert cpu_manager.prepare_store(to_keys([5, 6, 7]), _EMPTY_REQ_CTX) is None
+        assert cpu_manager.prepare_store(to_keys([5, 6, 7])) is None

        # complete load [2, 3]
        cpu_manager.complete_load(to_keys([2, 3]))

        # now prepare store [5, 6, 7] should succeed
        # ARC will evict blocks one at a time from T1 as needed
-        prepare_store_output = cpu_manager.prepare_store(
-            to_keys([5, 6, 7]), _EMPTY_REQ_CTX
-        )
+        prepare_store_output = cpu_manager.prepare_store(to_keys([5, 6, 7]))
        assert prepare_store_output is not None
        # Should successfully evict enough blocks to make room (at least 1)
        assert len(prepare_store_output.evicted_keys) >= 1
@@ -374,13 +357,13 @@ class TestARCPolicy:
        cpu_manager, arc_policy = self._make_manager(num_blocks=2, enable_events=False)

        # store blocks 1, 2 (fills cache)
-        cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
+        cpu_manager.prepare_store(to_keys([1, 2]))
        cpu_manager.complete_store(to_keys([1, 2]))

        initial_target = arc_policy.target_t1_size

        # store block 3, evicting block 1 (moves to B1 ghost list)
-        cpu_manager.prepare_store(to_keys([3]), _EMPTY_REQ_CTX)
+        cpu_manager.prepare_store(to_keys([3]))
        cpu_manager.complete_store(to_keys([3]))

        # block 1 should be in B1 (ghost list)
@@ -401,7 +384,7 @@ class TestARCPolicy:
        cpu_manager, arc_policy = self._make_manager(enable_events=False)

        # store blocks 1, 2, 3, 4
-        cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
+        cpu_manager.prepare_store(to_keys([1, 2, 3, 4]))
        cpu_manager.complete_store(to_keys([1, 2, 3, 4]))

        # promote blocks 3, 4 to T2 by touching them
@@ -416,7 +399,7 @@ class TestARCPolicy:
        arc_policy.target_t1_size = 1

        # store block 5, should evict from T1 (block 1, LRU in T1)
-        output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX)
+        output = cpu_manager.prepare_store(to_keys([5]))
        assert output is not None
        assert to_keys([1]) == output.evicted_keys

@@ -435,12 +418,12 @@ class TestARCPolicy:
        cpu_manager, arc_policy = self._make_manager(num_blocks=2, enable_events=False)

        # fill cache with blocks 1, 2
-        cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
+        cpu_manager.prepare_store(to_keys([1, 2]))
        cpu_manager.complete_store(to_keys([1, 2]))

        # store many blocks to fill ghost lists
        for i in range(3, 20):
-            cpu_manager.prepare_store(to_keys([i]), _EMPTY_REQ_CTX)
+            cpu_manager.prepare_store(to_keys([i]))
            cpu_manager.complete_store(to_keys([i]))

        # ghost lists should not exceed cache_capacity
@@ -455,7 +438,7 @@ class TestARCPolicy:
        cpu_manager, arc_policy = self._make_manager()

        # store blocks 1, 2, 3, 4
-        cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
+        cpu_manager.prepare_store(to_keys([1, 2, 3, 4]))
        cpu_manager.complete_store(to_keys([1, 2, 3, 4]))

        # promote 3, 4 to T2
@@ -470,7 +453,7 @@ class TestARCPolicy:
        assert len(arc_policy.t2) == 3

        # store block 5, should evict from T1 (block 2, only one in T1)
-        prepare_store_output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX)
+        prepare_store_output = cpu_manager.prepare_store(to_keys([5]))
        verify_store_output(
            prepare_store_output,
            ExpectedPrepareStoreOutput(
@@ -488,11 +471,11 @@ class TestARCPolicy:
        cpu_manager, arc_policy = self._make_manager()

        # store blocks 1, 2, 3, 4
-        cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
+        cpu_manager.prepare_store(to_keys([1, 2, 3, 4]))
        cpu_manager.complete_store(to_keys([1, 2, 3, 4]))

        # prepare store block 5 (will evict block 1)
-        prepare_store_output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX)
+        prepare_store_output = cpu_manager.prepare_store(to_keys([5]))
        assert prepare_store_output is not None
        assert len(prepare_store_output.evicted_keys) == 1

@@ -500,7 +483,7 @@ class TestARCPolicy:
        cpu_manager.complete_store(to_keys([5]), success=False)

        # block 5 should not be in cache
-        assert cpu_manager.lookup(to_keys([5]), _EMPTY_REQ_CTX) == 0
+        assert cpu_manager.lookup(to_keys([5])) == 0
        # block 5 should not be in T1 or T2
        assert to_keys([5])[0] not in arc_policy.t1
        assert to_keys([5])[0] not in arc_policy.t2
@@ -517,13 +500,11 @@ class TestARCPolicy:
        cpu_manager, arc_policy = self._make_manager()

        # store [1, 2]
-        cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
+        cpu_manager.prepare_store(to_keys([1, 2]))
        cpu_manager.complete_store(to_keys([1, 2]))

        # store [3, 4, 5] -> evicts [1]
-        prepare_store_output = cpu_manager.prepare_store(
-            to_keys([3, 4, 5]), _EMPTY_REQ_CTX
-        )
+        prepare_store_output = cpu_manager.prepare_store(to_keys([3, 4, 5]))
        assert prepare_store_output is not None
        assert len(prepare_store_output.evicted_keys) == 1
        cpu_manager.complete_store(to_keys([3, 4, 5]))
@@ -536,13 +517,13 @@ class TestARCPolicy:
        assert len(arc_policy.t2) == 2

        # store [6] -> should evict from T1 (4 is oldest in T1)
-        prepare_store_output = cpu_manager.prepare_store(to_keys([6]), _EMPTY_REQ_CTX)
+        prepare_store_output = cpu_manager.prepare_store(to_keys([6]))
        assert prepare_store_output is not None
        cpu_manager.complete_store(to_keys([6]))

        # verify blocks 2, 3 (in T2) are still present
-        assert cpu_manager.lookup(to_keys([2]), _EMPTY_REQ_CTX) == 1
-        assert cpu_manager.lookup(to_keys([3]), _EMPTY_REQ_CTX) == 1
+        assert cpu_manager.lookup(to_keys([2])) == 1
+        assert cpu_manager.lookup(to_keys([3])) == 1

        # verify events
        events = list(cpu_manager.take_events())
@@ -562,34 +543,34 @@ def test_filter_reused_manager():
    )

    # Lookup [1, 2] -> 1st time, added to tracker but not eligible for store yet
-    assert manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
+    assert manager.lookup(to_keys([1, 2])) == 0

    # prepare store [1, 2] -> should be filtered
-    prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
+    prepare_store_output = manager.prepare_store(to_keys([1, 2]))
    assert prepare_store_output is not None
    assert prepare_store_output.keys_to_store == []

    # Lookup [1] -> 2nd time, eligible now
-    assert manager.lookup(to_keys([1]), _EMPTY_REQ_CTX) == 0
+    assert manager.lookup(to_keys([1])) == 0

    # prepare store [1, 2] -> [1] should be eligible, [2] should be filtered
-    prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
+    prepare_store_output = manager.prepare_store(to_keys([1, 2]))
    assert prepare_store_output is not None
    assert prepare_store_output.keys_to_store == to_keys([1])

    # Lookup [3, 4] -> 1st time
    # (evicts [2] from tracker since max_size is 3 and tracker has [1])
-    assert manager.lookup(to_keys([3, 4]), _EMPTY_REQ_CTX) == 0
+    assert manager.lookup(to_keys([3, 4])) == 0
    # Verify [2] was evicted from the tracker (tracker now has: [1], [3], [4])
    assert to_keys([2])[0] not in manager.counts

    # Lookup [2] again -> (this adds [2] back to the tracker as 1st time)
-    assert manager.lookup(to_keys([2]), _EMPTY_REQ_CTX) == 0
+    assert manager.lookup(to_keys([2])) == 0
    # Verify [2] was re-added with count=1 (not eligible yet)
    assert manager.counts.get(to_keys([2])[0]) == 1

    # prepare store [2] -> should still be filtered out since count was reset
-    prepare_store_output = manager.prepare_store(to_keys([2]), _EMPTY_REQ_CTX)
+    prepare_store_output = manager.prepare_store(to_keys([2]))
    assert prepare_store_output is not None
    assert prepare_store_output.keys_to_store == []

@@ -1150,38 +1150,6 @@ def cutlass_fp4_moe_mm(
    )


-def cutlass_mxfp4_moe_mm(
-    out_tensors: torch.Tensor,
-    a_tensors: torch.Tensor,
-    b_tensors: torch.Tensor,
-    a_scales: torch.Tensor,
-    b_scales: torch.Tensor,
-    problem_sizes: torch.Tensor,
-    expert_offsets: torch.Tensor,
-    sf_offsets: torch.Tensor,
-):
-    """
-    An MXFP4 Blockscaled Group Gemm for MoE (MXFP4 x MXFP4).
-
-    Uses mx_float4_t types with E8M0 scale factors and 32-element blocks.
-    - a/b_tensors: MXFP4 packed activations/weights (uint8, 2 E2M1 per byte)
-    - a_/b_scales: E8M0 blockscales (uint8, stored in swizzled layout)
-    - Epilogue uses scalar alpha=1, beta=0 inside the CUDA op (no global scales).
-    - expert_offsets/sf_offsets: expert boundary indices
-    - problem_sizes: (num_experts, 3) with (M, N, K) per expert
-    """
-    return torch.ops._C.cutlass_mxfp4_group_mm(
-        out_tensors,
-        a_tensors,
-        b_tensors,
-        a_scales,
-        b_scales,
-        problem_sizes,
-        expert_offsets,
-        sf_offsets,
-    )
-
-
 def mxfp8_experts_quant(
    input_tensor: torch.Tensor,
    problem_sizes: torch.Tensor,
@@ -1880,109 +1848,6 @@ def silu_and_mul_scaled_fp4_experts_quant(
    return output, output_scales


-def mxfp4_experts_quant(
-    input_tensor: torch.Tensor,
-    expert_offsets: torch.Tensor,
-    blockscale_offsets: torch.Tensor,
-    n_experts: int,
-    topk: int,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Quantize input tensor to MXFP4 for packed MoE inputs.
-    Uses 32-element blocks with E8M0 (power-of-two) scale factors.
-    MXFP4 has no global scale - only block-level E8M0 scale factors.
-
-    Args:
-        input_tensor: [m_topk, k] BF16/FP16 activations
-        expert_offsets: [n_experts+1] token boundaries per expert
-        blockscale_offsets: [n_experts+1] SF row boundaries per expert
-        n_experts: number of experts
-        topk: number of top-k experts
-    Returns:
-        output: [m_topk, k//2] packed E2M1 values (uint8)
-        output_scales: E8M0 blockscales in swizzled layout (uint8 view)
-    """
-    assert not current_platform.is_rocm()
-    assert input_tensor.ndim == 2
-
-    MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
-    m_numtopk, k = input_tensor.shape
-
-    assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, (
-        f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT("
-        f"{MAX_TOKENS_PER_EXPERT})"
-        f" for cutlass_moe_mxfp4, observed m_numtopk = {m_numtopk}. Use"
-        f" VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value."
-    )
-    scales_k = k // 32
-    padded_k = (scales_k + (4 - 1)) // 4
-
-    output = torch.empty(
-        m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
-    )
-    output_scales = torch.empty(
-        MAX_TOKENS_PER_EXPERT * topk,
-        padded_k,
-        dtype=torch.int32,
-        device=input_tensor.device,
-    )
-    torch.ops._C.mxfp4_experts_quant(
-        output,
-        output_scales,
-        input_tensor,
-        expert_offsets,
-        blockscale_offsets,
-        n_experts,
-    )
-    # E8M0 SFs are stored as uint8
-    output_scales = output_scales.view(torch.uint8)
-    return output, output_scales
-
-
-def silu_and_mul_mxfp4_experts_quant(
-    input_tensor: torch.Tensor,
-    expert_offsets: torch.Tensor,
-    blockscale_offsets: torch.Tensor,
-    n_experts: int,
-    topk: int,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Fused SiLU+Mul+MXFP4 quantization for MoE intermediate activations.
-    MXFP4 has no global scale - only block-level E8M0 scale factors.
-    """
-    assert not current_platform.is_rocm()
-    assert input_tensor.ndim == 2
-
-    MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
-    m_numtopk, k_times_2 = input_tensor.shape
-    assert k_times_2 % 2 == 0, "input width must be even (gate || up layout)"
-    k = k_times_2 // 2
-
-    assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk
-    scales_k = k // 32
-    padded_k = (scales_k + (4 - 1)) // 4
-
-    output = torch.empty(
-        m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
-    )
-    output_scales = torch.empty(
-        MAX_TOKENS_PER_EXPERT * topk,
-        padded_k,
-        dtype=torch.int32,
-        device=input_tensor.device,
-    )
-    torch.ops._C.silu_and_mul_mxfp4_experts_quant(
-        output,
-        output_scales,
-        input_tensor,
-        expert_offsets,
-        blockscale_offsets,
-        n_experts,
-    )
-    output_scales = output_scales.view(torch.uint8)
-    return output, output_scales
-
-
 # fp8
 def scaled_fp8_quant(
    input: torch.Tensor,
@@ -22,23 +22,6 @@ else:
    except ImportError:
        from torch.library import impl_abstract as register_fake

-if hasattr(torch.ops._xpu_C, "fp8_gemm"):
-
-    @register_fake("_xpu_C::fp8_gemm")
-    def _fp8_gemm_fake(
-        q_input: torch.Tensor,
-        q_weight: torch.Tensor,
-        out_dtype: torch.dtype,
-        input_scales: torch.Tensor,
-        weight_scale: torch.Tensor,
-        bias: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        input_2d = q_input.view(-1, q_input.shape[-1])
-        M = input_2d.size(0)
-        N = q_weight.size(1)
-        return torch.empty((M, N), dtype=out_dtype, device=q_input.device)
-
-
 if hasattr(torch.ops._xpu_C, "fp8_gemm_w8a16"):

    @register_fake("_xpu_C::fp8_gemm_w8a16")
@@ -47,10 +47,9 @@ class XpuCommunicator(DeviceCommunicatorBase):
                self.all2all_manager = AgRsAll2AllManager(self.cpu_group)
                logger.info("Using AgRs manager on XPU device.")

-    def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
-        output = input_.clone() if torch.compiler.is_compiling() else input_
-        dist.all_reduce(output, group=self.device_group)
-        return output
+    def all_reduce(self, input_) -> torch.Tensor:
+        dist.all_reduce(input_, group=self.device_group)
+        return input_

    def reduce_scatter(self, input_: torch.Tensor, dim: int = -1):
        world_size = self.world_size
@@ -548,13 +548,7 @@ class MultiConnector(KVConnectorBase_V1):
            if stats_by_connector is None:
                # Lazy init to allow optional return value.
                stats_by_connector = MultiKVConnectorStats()
-            connector_id = c.__class__.__name__
-            if connector_id in stats_by_connector.data:
-                stats_by_connector[connector_id] = stats_by_connector[
-                    connector_id
-                ].aggregate(stats)
-            else:
-                stats_by_connector[connector_id] = stats
+            stats_by_connector[c.__class__.__name__] = stats
        return stats_by_connector

    @classmethod
@@ -566,13 +560,9 @@ class MultiConnector(KVConnectorBase_V1):
        per_engine_labelvalues: dict[int, list[object]],
    ) -> KVConnectorPromMetrics:
        prom_metrics: dict[str, KVConnectorPromMetrics] = {}
-        seen_classes: set[type] = set()
        for connector_cls, temp_config in cls._get_connector_classes_and_configs(
            vllm_config
        ):
-            if connector_cls in seen_classes:
-                continue
-            seen_classes.add(connector_cls)
            connector_prom = connector_cls.build_prom_metrics(
                temp_config, metric_types, labelnames, per_engine_labelvalues
            )
@@ -19,7 +19,6 @@ from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.kv_offload.abstract import (
    OffloadingManager,
    OffloadKey,
-    ReqContext,
    get_offload_block_hash,
    make_offload_key,
 )
@@ -75,7 +74,6 @@ class RequestOffloadState:
    config: SchedulerOffloadConfig
    req: Request
    group_states: tuple[RequestGroupState, ...] = field(init=False)
-    req_context: ReqContext = field(init=False)
    # number of hits in the GPU cache
    num_locally_computed_tokens: int = 0

@@ -83,7 +81,6 @@ class RequestOffloadState:
        self.group_states = tuple(
            RequestGroupState() for _ in self.config.kv_group_configs
        )
-        self.req_context = ReqContext(kv_transfer_params=self.req.kv_transfer_params)

    def update_offload_keys(self) -> None:
        for group_config, group_state in zip(
@@ -184,10 +181,7 @@ class OffloadingConnectorScheduler:
            return 0, False

        start_block_idx = num_computed_tokens // group_config.offloaded_block_size
-        hits = self.manager.lookup(
-            offload_keys[start_block_idx:],
-            req_status.req_context,
-        )
+        hits = self.manager.lookup(offload_keys[start_block_idx:])
        if hits is None:
            # indicates a lookup that should be tried later
            return None, False
@@ -255,7 +249,7 @@ class OffloadingConnectorScheduler:
        assert len(request.block_hashes) // self.config.block_size_factor >= num_blocks
        offload_keys = group_state.offload_keys[start_block_idx:num_blocks]

-        src_spec = self.manager.prepare_load(offload_keys, req_status.req_context)
+        src_spec = self.manager.prepare_load(offload_keys)
        dst_spec = GPULoadStoreSpec(
            block_ids[num_computed_gpu_blocks:],
            group_sizes=(num_pending_gpu_blocks,),
@@ -310,9 +304,7 @@ class OffloadingConnectorScheduler:
            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
-            )
+            store_output = self.manager.prepare_store(new_offload_keys)
            if store_output is None:
                logger.warning(
                    "Request %s: cannot store %s blocks", req_id, num_new_blocks
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # Adapted from
-# https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/openai/chat_completion/serving.py
+# https://github.com/vllm/vllm/entrypoints/openai/serving_chat.py

 """Anthropic Messages API serving handler"""

@@ -1638,17 +1638,6 @@ class LLM:
        seq_params = self._params_to_seq(params, len(seq_convs))
        seq_lora_requests = self._lora_request_to_seq(lora_request, len(seq_convs))

-        # When thinking is enabled or tools are provided, and the model
-        # uses special tokens for structured output (e.g. Gemma4's
-        # <|channel>, <|tool_call>, <|"|>), automatically set
-        # skip_special_tokens=False so these tokens are preserved in
-        # output.text for downstream parsing.
-        needs_parsing = (
-            chat_template_kwargs and chat_template_kwargs.get("enable_thinking")
-        ) or tools
-        if needs_parsing:
-            self._adjust_params_for_parsing(seq_params)
-
        return self._render_and_run_requests(
            prompts=(
                self._preprocess_chat_one(
@@ -1674,53 +1663,6 @@ class LLM:
            use_tqdm=use_tqdm,
        )

-    def _adjust_params_for_parsing(
-        self, params: Sequence[SamplingParams | PoolingParams]
-    ) -> None:
-        """Set ``skip_special_tokens=False`` when the model encodes
-        structured output syntax as special tokens.
-
-        Models like Gemma4 register thinking delimiters
-        (``<|channel>``/``<channel|>``) and tool call tokens
-        (``<|tool_call>``/``<tool_call|>``/``<|"|>``) as special tokens.
-        The default ``skip_special_tokens=True`` strips them from
-        ``output.text``, breaking parsing of both reasoning blocks and
-        tool calls.
-
-        This is a no-op for models whose structured tokens are regular
-        text tokens (e.g. DeepSeek's ``<think>``/``</think>``).
-        """
-        # The offline API currently lacks a unified rendering pipeline.
-        # Until the planned Renderer refactor is complete, we hardcode
-        # this token preservation logic specifically for Gemma4 models
-        # to avoid regressions on other models.
-        hf_config = getattr(self.model_config, "hf_config", None)
-        architectures = getattr(hf_config, "architectures", [])
-
-        if any("Gemma4" in arch for arch in architectures):
-            tokenizer = self.renderer.get_tokenizer()
-            vocab = tokenizer.get_vocab()
-            special_ids = set(getattr(tokenizer, "all_special_ids", []))
-
-            # Tokens used for thinking delimiters and tool call syntax
-            # that some models (Gemma4) register as special tokens.
-            structured_tokens = (
-                "<|channel>",
-                "<channel|>",  # thinking delimiters
-                "<|tool_call>",
-                "<tool_call|>",  # tool call delimiters
-                '<|"|>',  # string quoting in tool args
-            )
-            needs_special = any(
-                vocab.get(tok) in special_ids
-                for tok in structured_tokens
-                if tok in vocab
-            )
-            if needs_special:
-                for sp in params:
-                    if isinstance(sp, SamplingParams) and sp.skip_special_tokens:
-                        sp.skip_special_tokens = False
-
    def _render_and_run_requests(
        self,
        prompts: Iterable[EngineInput],
@@ -557,20 +557,6 @@ class OpenAIServingChat(OpenAIServing):
            and self._should_stream_with_auto_tool_parsing(request)
        )

-        # Determine whether required/named tool_choice should fall back to
-        # the auto tool_parser path instead of the standard JSON-based parsing.
-        # This happens when the parser declares supports_required_and_named=False
-        # (e.g. GLM models that output XML instead of JSON).
-        tool_choice_uses_parser = (
-            self.tool_parser is not None
-            and not self.tool_parser.supports_required_and_named
-            and request.tools
-            and (
-                request.tool_choice == "required"
-                or isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam)
-            )
-        )
-
        all_previous_token_ids: list[list[int]] | None
        function_name_returned = [False] * num_choices
        if self.tool_call_id_type == "kimi_k2":
@@ -583,12 +569,7 @@ class OpenAIServingChat(OpenAIServing):

        # Only one of these will be used, thus previous_texts and
        # all_previous_token_ids will not be used twice in the same iteration.
-        if (
-            is_mistral_grammar_path
-            or tool_choice_auto
-            or tool_choice_uses_parser
-            or reasoning_parser
-        ):
+        if is_mistral_grammar_path or tool_choice_auto or reasoning_parser:
            # These are only required in "auto" tool choice case
            all_previous_token_ids = [[] for _ in range(num_choices)]
            reasoning_end_arr = [False] * num_choices
@@ -783,12 +764,7 @@ class OpenAIServingChat(OpenAIServing):
                    delta_message: DeltaMessage | None

                    # just update previous_texts and previous_token_ids
-                    if (
-                        is_mistral_grammar_path
-                        or tool_choice_auto
-                        or tool_choice_uses_parser
-                        or reasoning_parser
-                    ):
+                    if is_mistral_grammar_path or tool_choice_auto or reasoning_parser:
                        assert previous_texts is not None
                        assert all_previous_token_ids is not None
                        previous_text = previous_texts[i]
@@ -837,9 +813,7 @@ class OpenAIServingChat(OpenAIServing):
                        if result.tools_called:
                            tools_streamed[i] = True
                    # handle streaming deltas for tools with named tool_choice
-                    # Skip when tool_choice_uses_parser so it falls through
-                    # to the auto tool_parser branches below.
-                    elif tool_choice_function_name and not tool_choice_uses_parser:
+                    elif tool_choice_function_name:
                        # When encountering think end id in prompt_token_ids
                        # i.e {"enable_thinking": False},
                        # check BEFORE calling the parser to avoid a spurious
@@ -877,6 +851,7 @@ class OpenAIServingChat(OpenAIServing):
                            ):
                                reasoning_end_arr[i] = True
                                if delta_message and delta_message.content:
+                                    # This need to be added to next `delta_text`
                                    current_text = delta_message.content
                                    delta_message.content = None
                                else:
@@ -921,12 +896,7 @@ class OpenAIServingChat(OpenAIServing):
                            )
                            tools_streamed[i] = True

-                    # Skip when tool_choice_uses_parser so it falls through
-                    # to the auto tool_parser branches below.
-                    elif (
-                        request.tool_choice == "required"
-                        and not tool_choice_uses_parser
-                    ):
+                    elif request.tool_choice == "required":
                        assert previous_texts is not None
                        previous_text = previous_texts[i]
                        current_text = previous_text + delta_text
@@ -996,10 +966,7 @@ class OpenAIServingChat(OpenAIServing):

                    # update the previous values for the next iteration
                    if (
-                        is_mistral_grammar_path
-                        or tool_choice_auto
-                        or tool_choice_uses_parser
-                        or reasoning_parser
+                        is_mistral_grammar_path or tool_choice_auto or reasoning_parser
                    ) and not self.use_harmony:
                        assert previous_texts is not None
                        assert all_previous_token_ids is not None
@@ -627,7 +627,7 @@ class OpenAIServing:
            and isinstance(request.tool_choice, ToolChoiceFunction)
        ):
            assert content is not None
-            # Forced Function Call (Responses API)
+            # Forced Function Call
            function_calls.append(
                FunctionCall(name=request.tool_choice.name, arguments=content)
            )
@@ -636,20 +636,14 @@ class OpenAIServing:
            not use_mistral_tool_parser
            and request.tool_choice
            and isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam)
-            and (tool_parser_cls is None or tool_parser_cls.supports_required_and_named)
        ):
-            # Named function with standard JSON-based parsing
            assert content is not None
+            # Forced Function Call
            function_calls.append(
                FunctionCall(name=request.tool_choice.function.name, arguments=content)
            )
            content = None  # Clear content since tool is called.
-        elif (
-            not use_mistral_tool_parser
-            and request.tool_choice == "required"
-            and (tool_parser_cls is None or tool_parser_cls.supports_required_and_named)
-        ):
-            # "required" with standard JSON-based parsing
+        elif not use_mistral_tool_parser and request.tool_choice == "required":
            tool_calls = []
            with contextlib.suppress(ValidationError):
                content = content or ""
@@ -668,30 +662,15 @@ class OpenAIServing:
            use_mistral_tool_parser
            or (
                enable_auto_tools
-                and (
-                    request.tool_choice == "auto"
-                    or request.tool_choice is None
-                    or (
-                        not tool_parser_cls.supports_required_and_named
-                        and request.tools
-                        and (
-                            request.tool_choice == "required"
-                            or isinstance(
-                                request.tool_choice,
-                                ChatCompletionNamedToolChoiceParam,
-                            )
-                        )
-                    )
-                )
+                and (request.tool_choice == "auto" or request.tool_choice is None)
            )
        ):
-            # Automatic Tool Call Parsing (also used as fallback for
-            # required/named when supports_required_and_named=False)
            if tokenizer is None:
                raise ValueError(
                    "Tokenizer not available when `skip_tokenizer_init=True`"
                )

+            # Automatic Tool Call Parsing
            try:
                tool_parser = tool_parser_cls(tokenizer, request.tools)
            except RuntimeError as e:
@@ -264,7 +264,7 @@ def convert_tool_responses_to_completions_format(tool: dict) -> dict:
 def construct_tool_dicts(
    tools: list[Tool], tool_choice: ToolChoice
 ) -> list[dict[str, Any]] | None:
-    if not tools or (tool_choice == "none"):
+    if tools is None or (tool_choice == "none"):
        tool_dicts = None
    else:
        tool_dicts = [
@@ -1,27 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Encode/decode utilities for multimodal tensors and field metadata
-over JSON/HTTP, used by the disaggregated generate endpoint."""
-
-from __future__ import annotations
-
-import pybase64
-
-from vllm.multimodal.inputs import MultiModalKwargsItem
-from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
-
-_encoder = MsgpackEncoder(size_threshold=2**62)  # force all tensors inline
-_decoder = MsgpackDecoder(t=MultiModalKwargsItem)
-
-
-def encode_mm_kwargs_item(item: MultiModalKwargsItem) -> str:
-    """Serialize a MultiModalKwargsItem to a base64 string."""
-    bufs = _encoder.encode(item)
-    assert len(bufs) == 1, "All tensors should be inline"
-    return pybase64.b64encode(bufs[0]).decode("ascii")
-
-
-def decode_mm_kwargs_item(data: str) -> MultiModalKwargsItem:
-    """Deserialize a base64 string back to a MultiModalKwargsItem."""
-    raw = pybase64.b64decode(data)
-    return _decoder.decode(raw)
@@ -35,6 +35,14 @@ class MultiModalFeatures(BaseModel):
    Carries hashes (for cache lookup / identification) and placeholder
    positions so the downstream `/generate` service knows *where* in
    the token sequence each multimodal item lives.
+
+    Note:
+        Phase 1 — metadata only.
+        Phase 2 should add `mm_kwargs` (processed tensor data) using a
+        binary transport so the ``/generate` side can skip re-processing.
+        The `/generate` endpoint must also be updated to inject these
+        features into `EngineInput` before passing to
+        `InputProcessor.process_inputs`.
    """

    mm_hashes: dict[str, list[str]]
@@ -43,15 +51,6 @@ class MultiModalFeatures(BaseModel):
    mm_placeholders: dict[str, list[PlaceholderRangeInfo]]
    """Per-modality placeholder ranges in the token sequence."""

-    kwargs_data: dict[str, list[str | None]] | None = None
-    """Per-modality serialized tensor data.
-
-    Each value is a list parallel to ``mm_hashes[modality]``.  A ``str``
-    entry is a base64-encoded ``MultiModalKwargsItem``; ``None`` means
-    the item should be resolved from cache.  The entire field is
-    ``None`` for metadata-only (cache-hit) responses.
-    """
-

 class GenerateRequest(BaseModel):
    request_id: str = Field(
@@ -25,7 +25,6 @@ from vllm.entrypoints.openai.engine.protocol import (
 )
 from vllm.entrypoints.openai.engine.serving import OpenAIServing, clamp_prompt_logprobs
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
-from vllm.entrypoints.serve.disagg.mm_serde import decode_mm_kwargs_item
 from vllm.entrypoints.serve.disagg.protocol import (
    GenerateRequest,
    GenerateResponse,
@@ -35,14 +34,8 @@ from vllm.entrypoints.serve.disagg.protocol import (
 )
 from vllm.entrypoints.serve.render.serving import OpenAIServingRender
 from vllm.entrypoints.utils import should_include_usage
-from vllm.inputs import EngineInput, mm_input
 from vllm.logger import init_logger
 from vllm.logprobs import Logprob
-from vllm.multimodal.inputs import (
-    MultiModalKwargsItem,
-    MultiModalKwargsItems,
-    PlaceholderRange,
-)
 from vllm.outputs import RequestOutput
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 from vllm.utils.collection_utils import as_list
@@ -110,42 +103,11 @@ class ServingTokens(OpenAIServing):
        if raw_request:
            raw_request.state.request_metadata = request_metadata

-        engine_input: EngineInput
-        if features := request.features:
-            # Convert PlaceholderRangeInfo → PlaceholderRange per modality.
-            mm_placeholders: dict[str, list[PlaceholderRange]] = {
-                modality: [
-                    PlaceholderRange(offset=p.offset, length=p.length) for p in ranges
-                ]
-                for modality, ranges in features.mm_placeholders.items()
-            }
-
-            # Deserialize tensor data when present; None → cache hit.
-            mm_kwargs: dict[str, list[MultiModalKwargsItem | None]] = {}
-            if features.kwargs_data is not None:
-                for modality, items in features.kwargs_data.items():
-                    mm_kwargs[modality] = [
-                        decode_mm_kwargs_item(item) if item is not None else None
-                        for item in items
-                    ]
-            else:
-                for modality, hashes in features.mm_hashes.items():
-                    mm_kwargs[modality] = [None] * len(hashes)
-
-            engine_input = mm_input(
-                prompt_token_ids=request.token_ids,
-                mm_kwargs=MultiModalKwargsItems(mm_kwargs),
-                mm_hashes=features.mm_hashes,
-                mm_placeholders=mm_placeholders,
-                cache_salt=request.cache_salt,
-            )
-        else:
-            (engine_input,) = await self.openai_serving_render.preprocess_completion(
-                request,
-                prompt_input=request.token_ids,
-                prompt_embeds=None,
-                skip_mm_cache=True,
-            )
+        (engine_input,) = await self.openai_serving_render.preprocess_completion(
+            request,
+            prompt_input=request.token_ids,
+            prompt_embeds=None,
+        )

        # Schedule the request and get the result generator.
        result_generator: AsyncGenerator[RequestOutput, None] | None = None
@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Sequence
 from http import HTTPStatus
-from typing import Any, cast
+from typing import Any

 from openai_harmony import Message as OpenAIMessage

@@ -25,7 +25,6 @@ from vllm.entrypoints.openai.parser.harmony_utils import (
    render_for_completion,
 )
 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
-from vllm.entrypoints.serve.disagg.mm_serde import encode_mm_kwargs_item
 from vllm.entrypoints.serve.disagg.protocol import (
    GenerateRequest,
    MultiModalFeatures,
@@ -38,7 +37,6 @@ from vllm.entrypoints.utils import (
 from vllm.inputs import (
    EngineInput,
    MultiModalHashes,
-    MultiModalInput,
    MultiModalPlaceholders,
    PromptType,
    SingletonPrompt,
@@ -253,7 +251,6 @@ class OpenAIServingRender:
                default_template_kwargs=self.default_chat_template_kwargs,
                tool_dicts=tool_dicts,
                tool_parser=tool_parser,
-                skip_mm_cache=True,
                reasoning_parser=self.reasoning_parser,
            )
        else:
@@ -345,7 +342,6 @@ class OpenAIServingRender:
            request,
            prompt_input=request.prompt,
            prompt_embeds=request.prompt_embeds,
-            skip_mm_cache=True,
        )

        return engine_inputs
@@ -361,10 +357,9 @@ class OpenAIServingRender:
        if engine_input.get("type") != "multimodal":
            return None

-        # At this point engine_input is a MultiModalInput TypedDict.
-        mm_engine_input = cast(MultiModalInput, engine_input)
-        mm_hashes: MultiModalHashes = mm_engine_input["mm_hashes"]
-        raw_placeholders: MultiModalPlaceholders = mm_engine_input["mm_placeholders"]
+        # At this point engine_input is a MultiModalInputs TypedDict.
+        mm_hashes: MultiModalHashes = engine_input["mm_hashes"]  # type: ignore[typeddict-item]
+        raw_placeholders: MultiModalPlaceholders = engine_input["mm_placeholders"]  # type: ignore[typeddict-item]

        mm_placeholders = {
            modality: [
@@ -373,20 +368,9 @@ class OpenAIServingRender:
            for modality, ranges in raw_placeholders.items()
        }

-        # Serialize tensor data per modality.
-        kwargs_data: dict[str, list[str | None]] | None = None
-        if raw_mm_kwargs := mm_engine_input.get("mm_kwargs"):
-            kwargs_data = {}
-            for modality, items in raw_mm_kwargs.items():
-                kwargs_data[modality] = [
-                    encode_mm_kwargs_item(item) if item is not None else None
-                    for item in items
-                ]
-
        return MultiModalFeatures(
            mm_hashes=mm_hashes,
            mm_placeholders=mm_placeholders,
-            kwargs_data=kwargs_data,
        )

    def _make_request_with_harmony(
@@ -53,15 +53,14 @@ if not has_helion():
    )

 import helion
+from helion._compat import requires_torch_version
 from helion.autotuner.base_search import BaseAutotuner
 from helion.runtime.config import Config
 from helion.runtime.settings import default_autotuner_fn

 # TODO(gmagogsfm): Remove CustomOp fallback path (_get_or_register_custom_op,
 # vllm_helion_lib, direct_register_custom_op) once vLLM requires PyTorch >= 2.11.
-# FIXME(gmagogsfm): Re-enable HOP path once performance regression is fixed.
-# _HOP_AVAILABLE = requires_torch_version("2.11")
-_HOP_AVAILABLE = False
+_HOP_AVAILABLE = requires_torch_version("2.11")

 if _HOP_AVAILABLE:
    from helion._compat import supports_torch_compile_fusion
@@ -27,42 +27,9 @@ def bgmv_expand(
    lora_indices_tensor: torch.Tensor,
    add_inputs: bool = True,
 ) -> None:
-    weight_out_dim = lora_b_weights.size(-2)
-    output_dim = output_tensor.size(1)
-
-    if weight_out_dim == output_dim:
-        torch.ops._xpu_C.bgmv_expand(
-            output_tensor,
-            inputs,
-            lora_b_weights,
-            lora_indices_tensor,
-            add_inputs,
-        )
-    elif weight_out_dim < output_dim:
-        # LoRA weight output dim can be smaller than the output tensor
-        # (e.g. vocab_size vs padded logits). Use expand_slice to write
-        # only the matching portion, mirroring torch_ops common_len logic.
-        torch.ops._xpu_C.bgmv_expand_slice(
-            output_tensor,
-            inputs,
-            lora_b_weights,
-            lora_indices_tensor,
-            0,
-            weight_out_dim,
-            add_inputs,
-        )
-    else:
-        # Weight output dim larger than output tensor: truncate weights.
-        lora_b_weights = lora_b_weights[..., :output_dim, :].contiguous()
-        torch.ops._xpu_C.bgmv_expand_slice(
-            output_tensor,
-            inputs,
-            lora_b_weights,
-            lora_indices_tensor,
-            0,
-            output_dim,
-            add_inputs,
-        )
+    torch.ops._xpu_C.bgmv_expand(
+        output_tensor, inputs, lora_b_weights, lora_indices_tensor, add_inputs
+    )


 def bgmv_expand_slice(
@@ -406,16 +406,33 @@ class Attention(nn.Module, AttentionLayerBase):
    def _init_turboquant_buffers(
        self, cache_dtype: str, head_size: int, prefix: str
    ) -> None:
-        """Initialize TurboQuant centroids for Lloyd-Max quantization."""
+        """Initialize TurboQuant rotation/projection matrices and centroids."""
        from vllm.model_executor.layers.quantization.turboquant.centroids import (
            get_centroids,
        )
        from vllm.model_executor.layers.quantization.turboquant.config import (
            TurboQuantConfig,
        )
+        from vllm.model_executor.layers.quantization.turboquant.quantizer import (
+            generate_wht_signs,
+        )

        tq_config = TurboQuantConfig.from_cache_dtype(cache_dtype, head_size)

+        # Each layer needs a unique rotation matrix so quantization errors
+        # don't correlate across layers. Stride must exceed max head_dim to
+        # ensure non-overlapping RNG streams between adjacent layers.
+        _TQ_LAYER_SEED_STRIDE = 1337
+
+        from vllm.model_executor.models.utils import extract_layer_index
+
+        layer_idx = extract_layer_index(prefix)
+        seed = tq_config.seed + layer_idx * _TQ_LAYER_SEED_STRIDE
+
+        self.register_buffer(
+            "_tq_signs",
+            generate_wht_signs(head_size, seed=seed),
+        )
        self.register_buffer(
            "_tq_centroids",
            get_centroids(head_size, tq_config.centroid_bits),
@@ -367,6 +367,16 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                "KV cache format, please set `--attention-backend FLASHMLA_SPARSE`"
            )

+        # CUTLASS FA3 MLA Sparse requires BF16 KV cache — force "auto" dtype
+        if self.attn_backend.get_name() == "CUTLASS_FA3_MLA_SPARSE":
+            if cache_config is not None:
+                cache_config.cache_dtype = "auto"
+            kv_cache_dtype = "auto"
+            logger.info_once(
+                "CUTLASS FA3 MLA Sparse backend requires BF16 KV cache. "
+                "Setting kv_cache_dtype to 'auto' (BF16)."
+            )
+
        # Initialize KV cache quantization attributes
        self.kv_cache_dtype = kv_cache_dtype
        self.calculate_kv_scales = calculate_kv_scales
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import math
 import os
 from collections.abc import Callable
 from typing import Any
@@ -612,43 +611,51 @@ def matmul_batch_invariant(a, b, *, out=None):
            out.copy_(result)
            return out
        return result
-    elif b.ndim == 2:
-        # Handle ND x 2D: Common for linear layers
-        # (..., batch, seq, hidden) @ (hidden, out) -> (..., batch, seq, out)
-        batch_dims = a.shape[:-1]
-        hidden = a.shape[-1]
-        out_dim = b.shape[-1]
+    elif a.ndim == 3 and b.ndim == 3:
+        # Handle batched case like bmm
+        return bmm_batch_invariant(a, b, out=out)
+    elif a.ndim == 3 and b.ndim == 2:
+        # Handle 3D x 2D: common for linear layers
+        # (batch, seq, hidden) @ (hidden, out) -> (batch, seq, out)
+        # Reshape to 2D, do mm, reshape back
+        batch, seq, hidden = a.shape
        a_2d = a.reshape(-1, hidden)
        result_2d = matmul_persistent(a_2d, b)
-        result = result_2d.reshape(batch_dims + (out_dim,))
+        result = result_2d.reshape(batch, seq, -1)
        if out is not None:
            out.copy_(result)
            return out
        return result
-    elif a.ndim >= 2 and b.ndim >= 3:
-        # Generic handler for 2D x ND and ND x ND (except 1D)
-        # Broadcast dims to ensure both matrices have the same shape
-        # If 2D x ND, then unsqueeze to add a dim to a
-        if a.ndim == 2:
-            a = a.unsqueeze(0)
-        broadcast_shape = torch.broadcast_shapes(a.shape[:-2], b.shape[:-2])
-        a = a.expand(broadcast_shape + a.shape[-2:])
-        b = b.expand(broadcast_shape + b.shape[-2:])
-        batch_dim = math.prod(broadcast_shape)
-        # Reuse broadcast shape to get all dims except mm dims
-        a_3d = a.reshape(batch_dim, a.shape[-2], a.shape[-1])
-        b_3d = b.reshape(batch_dim, b.shape[-2], b.shape[-1])
+    elif a.ndim == 2 and b.ndim == 3:
+        # Handle 2D x 3D: (M, K) @ (B, K, N) -> (B, M, N)
+        # By broadcasting `a` to 3D, we can reuse the batched matrix
+        # multiplication logic.
+        a_expanded = a.unsqueeze(0).expand(b.shape[0], -1, -1)
+        return bmm_batch_invariant(a_expanded, b, out=out)
+    elif a.ndim == 4 and b.ndim == 4:
+        # Handle 4D attention tensors: [batch, heads, seq, dim]
+        # Reshape to 3D, process, reshape back
+        batch, heads, seq_a, dim_a = a.shape
+        _, _, dim_b, seq_b = b.shape
+
+        # Reshape to [batch*heads, seq_a, dim_a]
+        a_3d = a.reshape(batch * heads, seq_a, dim_a)
+        b_3d = b.reshape(batch * heads, dim_b, seq_b)
+
        # Do batched matmul
        result_3d = bmm_batch_invariant(a_3d, b_3d)
-        # Reshape back to [broadcast_shape, seq_a, seq_b]
-        result = result_3d.reshape(broadcast_shape + (a.shape[-2], b.shape[-1]))
+
+        # Reshape back to [batch, heads, seq_a, seq_b]
+        result = result_3d.reshape(batch, heads, seq_a, seq_b)
+
        if out is not None:
            out.copy_(result)
            return out
        return result
    else:
        raise ValueError(
-            f"matmul_batch_invariant requires both inputs be at least 2D "
+            f"matmul_batch_invariant currently only supports 2D x 2D, 3D x 3D, "
+            f"3D x 2D, 2D x 3D, and 4D x 4D, "
            f"got shapes {a.shape} and {b.shape}"
        )

@@ -762,25 +762,6 @@ def nvfp4_moe_quant_config(
    )


-def mxfp4_moe_quant_config(
-    w1_scale: torch.Tensor,
-    w2_scale: torch.Tensor,
-) -> FusedMoEQuantConfig:
-    """
-    Construct a quant config for MXFP4 x MXFP4 MoE.
-    MXFP4 uses block scaling only (E8M0 scales, 32-element groups), with no
-    separate alphas / global activation scales in this config.
-    """
-    return FusedMoEQuantConfig.make(
-        "mxfp4",
-        w1_scale=w1_scale,
-        w2_scale=w2_scale,
-        per_act_token_quant=False,
-        per_out_ch_quant=False,
-        block_shape=None,
-    )
-
-
 def nvfp4_w4a16_moe_quant_config(
    g1_alphas: torch.Tensor,
    g2_alphas: torch.Tensor,
@@ -36,8 +36,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
    kFp8DynamicTokenSym,
    kFp8StaticChannelSym,
    kFp8StaticTensorSym,
-    kMxfp4Dynamic,
-    kMxfp4Static,
    kNvfp4Dynamic,
    kNvfp4Static,
 )
@@ -797,299 +795,6 @@ class CutlassExpertsFp4(mk.FusedMoEExpertsModular):
        )


-def run_cutlass_moe_mxfp4(
-    output: torch.Tensor,
-    a: torch.Tensor,
-    w1_fp4: torch.Tensor,
-    w1_blockscale: torch.Tensor,
-    w2_fp4: torch.Tensor,
-    w2_blockscale: torch.Tensor,
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    activation: MoEActivation,
-    workspace13: torch.Tensor,
-    workspace2: torch.Tensor,
-    m: int,
-    n: int,
-    k: int,
-    e: int,
-    device: torch.device,
-    apply_router_weight_on_input: bool = False,
-) -> None:
-    """MXFP4 x MXFP4 MoE implementation using CUTLASS grouped GEMM."""
-    is_gated = activation.is_gated
-    w1_n = n * 2 if is_gated else n
-
-    assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
-    assert w1_fp4.dtype == torch.uint8, "weight 1 must be uint8"
-    assert w2_fp4.dtype == torch.uint8, "weight 2 must be uint8"
-    assert (
-        w1_fp4.ndim == 3
-        and w2_fp4.ndim == 3
-        and w1_blockscale.ndim == 3
-        and w2_blockscale.ndim == 3
-    ), "All Weights must be of rank 3 for cutlass_moe_mxfp4"
-    m_a, k_a = a.shape
-    e_w1, w1_n_actual, half_k_w1 = w1_fp4.shape
-    e_w2, k_w2, half_n_w2 = w2_fp4.shape
-
-    assert e_w1 == e_w2 and e_w1 == e
-    assert k_a == half_k_w1 * 2 and k == k_w2
-    assert w1_n_actual == w1_n and half_n_w2 * 2 == n
-    assert m == m_a
-    assert 2 * half_k_w1 == k_w2
-    assert a.dtype in [torch.half, torch.bfloat16], "Invalid input dtype"
-    assert topk_weights.size(0) == m and topk_ids.size(0) == m
-
-    topk = topk_ids.size(1)
-    out_dtype = a.dtype
-    num_topk = topk_ids.size(1)
-
-    expert_offsets = torch.empty((e + 1), dtype=torch.int32, device=device)
-    blockscale_offsets = torch.empty((e + 1), dtype=torch.int32, device=device)
-    problem_sizes1 = torch.empty((e, 3), dtype=torch.int32, device=device)
-    problem_sizes2 = torch.empty((e, 3), dtype=torch.int32, device=device)
-
-    a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
-    c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
-
-    if apply_router_weight_on_input:
-        assert num_topk == 1, (
-            "apply_router_weight_on_input is only implemented for topk=1"
-        )
-        a.mul_(topk_weights.to(out_dtype))
-
-    ops.get_cutlass_moe_mm_data(
-        topk_ids,
-        expert_offsets,
-        problem_sizes1,
-        problem_sizes2,
-        a_map,
-        c_map,
-        e,
-        n,
-        k,
-        blockscale_offsets,
-        is_gated=is_gated,
-    )
-
-    a = ops.shuffle_rows(a, a_map)
-    rep_a_fp4, rep_a_blockscale = ops.mxfp4_experts_quant(
-        a,
-        expert_offsets,
-        blockscale_offsets,
-        e,
-        num_topk,
-    )
-    c1 = _resize_cache(workspace13, (m * topk, w1_n))
-    c2 = _resize_cache(workspace2, (m * topk, n))
-    c3 = _resize_cache(workspace13, (m * topk, k))
-
-    ops.cutlass_mxfp4_moe_mm(
-        c1,
-        rep_a_fp4,
-        w1_fp4,
-        rep_a_blockscale,
-        w1_blockscale,
-        problem_sizes1,
-        expert_offsets[:-1],
-        blockscale_offsets[:-1],
-    )
-    del rep_a_fp4, rep_a_blockscale
-    if activation == MoEActivation.SILU:
-        int_fp4, int_blockscale = ops.silu_and_mul_mxfp4_experts_quant(
-            c1, expert_offsets, blockscale_offsets, e, num_topk
-        )
-    else:
-        apply_moe_activation(activation, c2, c1)
-        int_fp4, int_blockscale = ops.mxfp4_experts_quant(
-            c2, expert_offsets, blockscale_offsets, e, num_topk
-        )
-
-    ops.cutlass_mxfp4_moe_mm(
-        c3,
-        int_fp4,
-        w2_fp4,
-        int_blockscale,
-        w2_blockscale,
-        problem_sizes2,
-        expert_offsets[:-1],
-        blockscale_offsets[:-1],
-    )
-    del int_fp4, int_blockscale
-
-    c3 = ops.shuffle_rows(c3, c_map)
-
-    assert output.dtype == out_dtype
-    if not apply_router_weight_on_input:
-        output.copy_(
-            (
-                c3.view(m, num_topk, k)
-                * topk_weights.view(m, num_topk, 1).to(out_dtype)
-            ).sum(dim=1),
-            non_blocking=True,
-        )
-    else:
-        output.copy_(c3.view(m, num_topk, k).sum(dim=1), non_blocking=True)
-    return
-
-
-def swizzle_mxfp4_scales(
-    scales: torch.Tensor,
-    N: int,
-    K: int,
-) -> torch.Tensor:
-    """Swizzle flat [N, K//32] E8M0 scales to CUTLASS tiled layout.
-
-    CUTLASS expects MX scale factors in a tiled layout:
-        [numMTiles, numKTiles, 32, 4, 4]
-    where numMTiles = ceil(N/128), numKTiles = ceil(K/128),
-    and the inner dimensions correspond to the swizzle pattern:
-        mTileIdx = mIdx / 128
-        outerMIdx = mIdx % 32
-        innerMIdx = (mIdx / 32) % 4
-        kTileIdx = kIdx / 4
-        innerKIdx = kIdx % 4
-    with kIdx = col_in_scale_space (i.e., index into K//32).
-    """
-    assert scales.dtype == torch.uint8
-    num_scale_cols = K // 32  # number of E8M0 scale values per row
-
-    num_m_tiles = (N + 127) // 128
-    num_k_tiles = (num_scale_cols + 3) // 4
-
-    # Pad N to multiple of 128 and scale_cols to multiple of 4
-    padded_N = num_m_tiles * 128
-    padded_scale_cols = num_k_tiles * 4
-
-    # Start with flat scales, pad if needed
-    padded = torch.zeros(
-        padded_N, padded_scale_cols, dtype=torch.uint8, device=scales.device
-    )
-    padded[:N, :num_scale_cols] = scales
-
-    # Reshape to tile structure:
-    # [numMTiles, 4, 32, numKTiles, 4]
-    #  mTileIdx, innerMIdx, outerMIdx, kTileIdx, innerKIdx
-    tiled = padded.reshape(num_m_tiles, 4, 32, num_k_tiles, 4)
-    # Permute to [numMTiles, numKTiles, 32, 4, 4]
-    #            (outerMIdx, innerMIdx, innerKIdx)
-    tiled = tiled.permute(0, 3, 2, 1, 4).contiguous()
-    return tiled.reshape(-1)
-
-
-class CutlassExpertsMxfp4(mk.FusedMoEExpertsModular):
-    """CUTLASS MXFP4 x MXFP4 fused MoE expert implementation."""
-
-    @property
-    def expects_unquantized_inputs(self) -> bool:
-        return True
-
-    @staticmethod
-    def _supports_current_device() -> bool:
-        p = current_platform
-        return p.is_cuda() and p.is_device_capability_family(100)
-
-    @staticmethod
-    def _supports_no_act_and_mul() -> bool:
-        return True
-
-    @staticmethod
-    def _supports_quant_scheme(
-        weight_key: QuantKey | None,
-        activation_key: QuantKey | None,
-    ) -> bool:
-        return (weight_key, activation_key) == (kMxfp4Static, kMxfp4Dynamic)
-
-    @staticmethod
-    def _supports_activation(activation: MoEActivation) -> bool:
-        return activation in [
-            MoEActivation.SILU,
-            MoEActivation.GELU,
-            MoEActivation.SWIGLUOAI,
-            MoEActivation.SWIGLUSTEP,
-            MoEActivation.SILU_NO_MUL,
-            MoEActivation.GELU_NO_MUL,
-            MoEActivation.RELU2_NO_MUL,
-        ]
-
-    @staticmethod
-    def _supports_parallel_config(
-        moe_parallel_config: FusedMoEParallelConfig,
-    ) -> bool:
-        return moe_parallel_config.ep_size == 1
-
-    @staticmethod
-    def activation_format() -> mk.FusedMoEActivationFormat:
-        return mk.FusedMoEActivationFormat.Standard
-
-    def supports_expert_map(self) -> bool:
-        return False
-
-    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
-        return TopKWeightAndReduceNoOP()
-
-    def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:
-        return act_dtype
-
-    def workspace_shapes(
-        self,
-        M: int,
-        N: int,
-        K: int,
-        topk: int,
-        global_num_experts: int,
-        local_num_experts: int,
-        expert_tokens_meta: mk.ExpertTokensMetadata | None,
-        activation: MoEActivation,
-    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
-        workspace1 = (M * topk, max(2 * N, K))
-        workspace2 = (M * topk, N)
-        output = (M, K)
-        return (workspace1, workspace2, output)
-
-    def apply(
-        self,
-        output: torch.Tensor,
-        hidden_states: torch.Tensor,
-        w1: torch.Tensor,
-        w2: torch.Tensor,
-        topk_weights: torch.Tensor,
-        topk_ids: torch.Tensor,
-        activation: MoEActivation,
-        global_num_experts: int,
-        expert_map: torch.Tensor | None,
-        a1q_scale: torch.Tensor | None,
-        a2_scale: torch.Tensor | None,
-        workspace13: torch.Tensor | None,
-        workspace2: torch.Tensor | None,
-        expert_tokens_meta: mk.ExpertTokensMetadata | None,
-        apply_router_weight_on_input: bool,
-    ):
-        e, m, n, k, _ = self.moe_problem_size(hidden_states, w1, w2, topk_ids)
-        n = w2.shape[2] * 2
-
-        run_cutlass_moe_mxfp4(
-            output=output,
-            a=hidden_states,
-            w1_fp4=w1,
-            w1_blockscale=self.w1_scale,
-            w2_fp4=w2,
-            w2_blockscale=self.w2_scale,
-            topk_weights=topk_weights,
-            topk_ids=topk_ids,
-            activation=activation,
-            workspace13=workspace13,
-            workspace2=workspace2,
-            m=m,
-            n=n,
-            k=k,
-            e=e,
-            device=hidden_states.device,
-            apply_router_weight_on_input=apply_router_weight_on_input,
-        )
-
-
 # W4A8
 def run_cutlass_moe_w4a8_fp8(
    output: torch.Tensor,
@@ -163,11 +163,6 @@ def select_unquantized_moe_backend(
    if current_platform.is_out_of_tree():
        return UnquantizedMoeBackend.OOT, None

-    if moe_config.is_lora_enabled:
-        return UnquantizedMoeBackend.TRITON, backend_to_kernel_cls(
-            UnquantizedMoeBackend.TRITON
-        )
-
    # NOTE: the kernels are selected in the following order.
    AVAILABLE_BACKENDS = _get_priority_backends(moe_config)

@@ -478,12 +478,9 @@ class RMSNormGated(CustomOp):
        weight = self.weight.float()
        z = z.float() if z is not None else None

-        assert self.activation in ["silu", "sigmoid", "swish"]
-        act_fn = F.sigmoid if self.activation == "sigmoid" else F.silu
-
        # Apply gating before normalization if needed
        if z is not None and not self.norm_before_gate:
-            x = x * act_fn(z)
+            x = x * F.silu(z)

        # RMS Normalization
        if self.group_size is None:
@@ -502,7 +499,7 @@ class RMSNormGated(CustomOp):

        # Apply gating after normalization if needed
        if z is not None and self.norm_before_gate:
-            out = out * act_fn(z)
+            out = out * F.silu(z)

        return out.to(orig_dtype)

@@ -916,15 +916,9 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
                            loaded_weight=loaded_weight, shard_id=idx
                        )
                else:
-                    # When weights are already fused on disk (e.g. Phi-3's
-                    # gate_up_proj), there is only a single scale for the
-                    # entire fused matrix. Fill all slots with this scale
-                    # to ensure that any subsequent reduction (like .max())
-                    # works correctly while preserving the parameter shape.
-                    for idx in range(param.data.shape[0]):
-                        param.load_merged_column_weight(
-                            loaded_weight=loaded_weight, shard_id=idx
-                        )
+                    param.load_merged_column_weight(
+                        loaded_weight=loaded_weight, shard_id=0
+                    )
                return
            elif type(param) in (RowvLLMParameter, BasevLLMParameter):
                param.load_merged_column_weight(loaded_weight=loaded_weight)
@@ -1136,15 +1130,9 @@ class QKVParallelLinear(ColumnParallelLinear):
        self.validate_shard_id(loaded_shard_id)
        if loaded_shard_id is None:  # special case for certain models
            if isinstance(param, PerTensorScaleParameter):
-                # When weights are already fused on disk (e.g. Phi-3's
-                # qkv_proj), there is only a single scale for the entire
-                # fused matrix. Fill all slots (q, k, v) with this scale
-                # to ensure that any subsequent reduction (like .max())
-                # works correctly while preserving the parameter shape.
-                for idx in range(param.data.shape[0]):
-                    param.load_qkv_weight(
-                        loaded_weight=loaded_weight, shard_id=idx, tp_rank=self.tp_rank
-                    )
+                param.load_qkv_weight(
+                    loaded_weight=loaded_weight, shard_id=0, tp_rank=self.tp_rank
+                )
                return
            elif type(param) in (RowvLLMParameter, BasevLLMParameter):
                param.load_qkv_weight(loaded_weight=loaded_weight, tp_rank=self.tp_rank)
@@ -357,19 +357,11 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
        set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(0)})
        set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})

-        output_gate_type = getattr(config, "output_gate_type", "silu")
-        if output_gate_type == "swish":
-            output_gate_type = "silu"
-        assert output_gate_type in ["silu", "swish", "sigmoid"], (
-            f"unsupported {output_gate_type=}"
-        )
-
        self.norm = RMSNormGated(
            self.head_v_dim,
            eps=self.layer_norm_epsilon,
            group_size=None,
            norm_before_gate=True,
-            activation=output_gate_type,
            device=current_platform.current_device(),
        )

@@ -4,7 +4,6 @@

 import torch

-import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe import (
    FusedMoE,
@@ -12,10 +11,6 @@ from vllm.model_executor.layers.fused_moe import (
 )
 from vllm.model_executor.layers.fused_moe.config import (
    FusedMoEQuantConfig,
-    mxfp4_moe_quant_config,
-)
-from vllm.model_executor.layers.fused_moe.cutlass_moe import (
-    CutlassExpertsMxfp4,
 )
 from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
    MarlinExperts,
@@ -41,14 +36,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
        super().__init__(moe)
        self.group_size = 32
        self.mxfp4_backend = Mxfp4MoeBackend.MARLIN
-        self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device()
-        self.experts_cls: type[mk.FusedMoEExperts]
-        if self.use_cutlass_mxfp4:
-            logger.info_once("Using CutlassExpertsMxfp4 for MXFP4 MoE", scope="local")
-            self.experts_cls = CutlassExpertsMxfp4
-        else:
-            logger.info_once("Using MarlinExperts for MXFP4 MoE", scope="local")
-            self.experts_cls = MarlinExperts
+        self.experts_cls = MarlinExperts

    def create_weights(
        self,
@@ -121,19 +109,11 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
    def get_fused_moe_quant_config(
        self, layer: torch.nn.Module
    ) -> FusedMoEQuantConfig | None:
-        if self.use_cutlass_mxfp4:
-            # W4A4: both weights and activations quantized to MXFP4
-            return mxfp4_moe_quant_config(
-                w1_scale=layer.w13_weight_scale,
-                w2_scale=layer.w2_weight_scale,
-            )
-        else:
-            # W4A16: weight-only via Marlin
-            return make_mxfp4_moe_quant_config(
-                mxfp4_backend=self.mxfp4_backend,
-                w1_scale=layer.w13_weight_scale,
-                w2_scale=layer.w2_weight_scale,
-            )
+        return make_mxfp4_moe_quant_config(
+            mxfp4_backend=self.mxfp4_backend,
+            w1_scale=layer.w13_weight_scale,
+            w2_scale=layer.w2_weight_scale,
+        )

    def process_weights_after_loading(self, layer: FusedMoE) -> None:
        layer.w13_weight = torch.nn.Parameter(
@@ -146,45 +126,13 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
        )
        delattr(layer, "w2_weight_packed")

-        if self.use_cutlass_mxfp4:
-            # Swizzle weight scales from flat checkpoint layout [E, N, K//32]
-            # to CUTLASS tiled layout [E, numMTiles*numKTiles*512].
-            from vllm.model_executor.layers.fused_moe.cutlass_moe import (
-                swizzle_mxfp4_scales,
-            )
-
-            E = layer.w13_weight_scale.shape[0]
-            w13_N = layer.w13_weight_scale.shape[1]
-            w13_scale_K = layer.w13_weight_scale.shape[2]
-            w13_K = w13_scale_K * 32
-
-            w2_M = layer.w2_weight_scale.shape[1]
-            w2_scale_N = layer.w2_weight_scale.shape[2]
-            w2_N = w2_scale_N * 32
-
-            swizzled_w13 = []
-            swizzled_w2 = []
-            for e_idx in range(E):
-                s13 = layer.w13_weight_scale[e_idx]
-                sw13 = swizzle_mxfp4_scales(s13, w13_N, w13_K)
-                swizzled_w13.append(sw13.reshape(w13_N, w13_scale_K))
-                s2 = layer.w2_weight_scale[e_idx]
-                sw2 = swizzle_mxfp4_scales(s2, w2_M, w2_N)
-                swizzled_w2.append(sw2.reshape(w2_M, w2_scale_N))
-            layer.w13_weight_scale = torch.nn.Parameter(
-                torch.stack(swizzled_w13), requires_grad=False
-            )
-            layer.w2_weight_scale = torch.nn.Parameter(
-                torch.stack(swizzled_w2), requires_grad=False
-            )
-        else:
-            logger.warning_once(
-                "Your GPU does not have native support for FP4 computation "
-                "but FP4 quantization is being used. Weight-only FP4 "
-                "compression will be used leveraging the Marlin kernel. "
-                "This may degrade performance for compute-heavy workloads."
-            )
-            prepare_moe_fp4_layer_for_marlin(layer)
+        logger.warning_once(
+            "Your GPU does not have native support for FP4 computation but "
+            "FP4 quantization is being used. Weight-only FP4 compression "
+            "will be used leveraging the Marlin kernel. This may degrade "
+            "performance for compute-heavy workloads."
+        )
+        prepare_moe_fp4_layer_for_marlin(layer)

        self.moe_quant_config = self.get_fused_moe_quant_config(layer)
        if self.moe_quant_config is not None:
@@ -1,19 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""TurboQuant: KV-cache quantization for vLLM.
+"""TurboQuant: Near-optimal KV-cache quantization for vLLM.

-Hadamard rotation + per-coordinate Lloyd-Max scalar quantization for
-keys, uniform quantization for values.
+PolarQuant compression: random rotation + per-coordinate Lloyd-Max
+scalar quantization for keys, uniform quantization for values.

-The technique implemented here consists of the scalar case of the HIGGS
-quantization method (Malinovskii et al., "Pushing the Limits of Large
-Language Model Quantization via the Linearity Theorem", NAACL 2025;
-preprint arXiv:2411.17525): rotation + optimized grid + optional
-re-normalization, applied to KV cache compression. A first application
-of this approach to KV-cache compression is in "Cache Me If You Must:
-Adaptive Key-Value Quantization for Large Language Models" (Shutova
-et al., ICML 2025; preprint arXiv:2501.19392). Both these references
-pre-date the TurboQuant paper (Zandieh et al., ICLR 2026).
+Reference: "TurboQuant: Online Vector Quantization with Near-optimal
+Distortion Rate" (ICLR 2026), Zandieh et al.
 """

 from vllm.model_executor.layers.quantization.turboquant.config import TurboQuantConfig
@@ -36,22 +36,10 @@ TQ_PRESETS: dict[str, dict] = {
 class TurboQuantConfig:
    """Configuration for TurboQuant KV-cache quantization.

-    Applies Hadamard rotation followed by per-coordinate Lloyd-Max scalar
-    quantization for keys, and uniform quantization for values.
-
-    Historical note: this is the scalar case of the HIGGS quantization
-    method (Malinovskii et al., "Pushing the Limits of Large Language Model
-    Quantization via the Linearity Theorem", NAACL 2025; preprint
-    arXiv:2411.17525): rotation + optimized grid + optional re-normalization,
-    applied to KV cache compression. A first application of this approach to
-    KV-cache compression is in "Cache Me If You Must: Adaptive Key-Value
-    Quantization for Large Language Models" (Shutova et al., ICML 2025;
-    preprint arXiv:2501.19392). Both these references pre-date the
-    TurboQuant paper.
-
-    QJL is intentionally omitted — community consensus (5+ independent
-    groups) found it hurts attention quality by amplifying variance through
-    softmax.
+    Uses PolarQuant (WHT rotation + Lloyd-Max scalar quantization) for keys
+    and uniform quantization for values. QJL is intentionally omitted —
+    community consensus (5+ independent groups) found it hurts attention
+    quality by amplifying variance through softmax.

    Named presets (use via --kv-cache-dtype):
        turboquant_k8v4:   FP8 keys + 4-bit values, 2.6x, +1.17% PPL
@@ -65,6 +53,8 @@ class TurboQuantConfig:
            rotation/MSE). 3-4 = Lloyd-Max MSE quantized keys.
        value_quant_bits: Bits per value dimension for uniform quantization.
            3 = 8 levels, 4 = 16 levels (default).
+        seed: Base seed for deterministic random matrix generation.
+            Actual seed per layer = seed + layer_idx * 1337.
        norm_correction: Re-normalize centroid vectors to unit norm before
            inverse rotation during dequant. Fixes quantization-induced norm
            distortion, improving PPL by ~0.8% at 4-bit.
@@ -73,7 +63,7 @@ class TurboQuantConfig:
    head_dim: int = 128
    key_quant_bits: int = 3  # 3-4 = MSE keys, 8 = FP8 keys
    value_quant_bits: int = 4  # 3-4 = uniform quantized values
-    seed: int = 42  # kept for backward compatibility; no longer used internally
+    seed: int = 42
    norm_correction: bool = False

    @property
@@ -2,5 +2,23 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """TurboQuant quantizer utilities.

+Serving path uses generate_wht_signs() for WHT rotation sign buffers.
 Triton kernels handle all quantization, packing, and dequantization on GPU.
 """
+
+import torch
+
+_CPU = torch.device("cpu")
+
+
+def generate_wht_signs(d: int, seed: int, device: torch.device = _CPU) -> torch.Tensor:
+    """Generate deterministic random ±1 signs for WHT rotation.
+
+    Used with Walsh-Hadamard Transform for per-layer rotation randomization.
+    Same seed derivation as QR (per-layer via seed + layer_idx * stride).
+    """
+    gen = torch.Generator(device="cpu")
+    gen.manual_seed(seed)
+    bits = torch.randint(0, 2, (d,), generator=gen, device="cpu")
+    signs = bits.float() * 2 - 1
+    return signs.to(device)
@@ -57,9 +57,7 @@ from vllm.model_executor.model_loader.weight_utils import (
    default_weight_loader,
    maybe_remap_kv_scale_name,
 )
-from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
-from vllm.triton_utils import tl, triton
 from vllm.v1.attention.backends.utils import KVSharingFastPrefillMetadata

 from .interfaces import (
@@ -81,120 +79,6 @@ from .utils import (
 logger = init_logger(__name__)


-@triton.jit
-def _gemma4_routing_kernel(
-    gating_ptr,
-    per_expert_scale_ptr,
-    topk_weights_ptr,
-    topk_ids_ptr,
-    E: tl.constexpr,
-    K: tl.constexpr,
-    BLOCK_E: tl.constexpr,
-):
-    pid = tl.program_id(0)
-    offs_e = tl.arange(0, BLOCK_E)
-    valid = offs_e < E
-
-    logits = tl.load(
-        gating_ptr + pid * E + offs_e,
-        mask=valid,
-        other=-float("inf"),
-    ).to(tl.float32)
-
-    max_l = tl.max(logits, axis=0)
-
-    # Float32 → ascending-sortable bijection
-    MIN32 = -2147483648
-    logit_bits = logits.to(tl.int32, bitcast=True)
-    sign_b = logit_bits >> 31
-    key = tl.where(sign_b == 0, logit_bits ^ -1, logit_bits ^ MIN32)
-    key = tl.where(valid, key, 0x7FFFFFFF)
-    sk64 = key.to(tl.int64) & 0x00000000FFFFFFFF
-    packed = (sk64 << 32) | offs_e.to(tl.int64)
-    sorted_p = tl.sort(packed, descending=False)
-
-    # Vectorized extraction of ALL sorted elements — no K-loop, no cross-lane reductions
-    all_keys = ((sorted_p >> 32) & 0x00000000FFFFFFFF).to(tl.int32)
-    all_ids = (sorted_p & 0x00000000FFFFFFFF).to(tl.int32)
-
-    # Inverse bijection: recover original logit bits
-    sign_k = all_keys >> 31
-    all_bits = tl.where(sign_k < 0, all_keys ^ -1, all_keys ^ MIN32)
-    all_logits = all_bits.to(tl.float32, bitcast=True)
-
-    # Compute raw_exp for ALL BLOCK_E elements — vectorized, ~2 VALU clocks
-    all_raw_exp = tl.math.exp2((all_logits - max_l) * 1.4426950408889634)
-
-    # Sum only top-K for renorm — ONE masked reduction
-    top_mask = offs_e < K
-    renorm_raw = tl.sum(tl.where(top_mask, all_raw_exp, 0.0), axis=0)
-    renorm_raw = tl.where(renorm_raw > 0.0, renorm_raw, 1.0)
-    inv_renorm = 1.0 / renorm_raw
-
-    # Load scales for top-K only (masked gather; scale array is tiny → L1 cached)
-    all_scales = tl.load(
-        per_expert_scale_ptr + all_ids.to(tl.int64),
-        mask=top_mask,
-        other=1.0,
-    ).to(tl.float32)
-
-    # Final weights: vectorized multiply (only top-K will be stored)
-    all_weights = (all_raw_exp * inv_renorm * all_scales).to(tl.float32)
-
-    # Write results with TWO masked stores — replaces K × 2 serial scalar stores
-    base_off = pid * K + offs_e
-    tl.store(topk_ids_ptr + base_off, all_ids, mask=top_mask)
-    tl.store(topk_weights_ptr + base_off, all_weights, mask=top_mask)
-
-
-def gemma4_fused_routing_kernel_triton(
-    gating_output: torch.Tensor,
-    topk: int,
-    per_expert_scale: torch.Tensor,
-    num_warps: int = 1,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    gating_output = gating_output.contiguous()
-    per_expert_scale = per_expert_scale.contiguous()
-    T, E = gating_output.shape
-    weights = torch.empty(T, topk, dtype=torch.float32, device=gating_output.device)
-    ids = torch.empty(T, topk, dtype=torch.int32, device=gating_output.device)
-    BLOCK_E = triton.next_power_of_2(E)
-    _gemma4_routing_kernel[(T,)](
-        gating_output,
-        per_expert_scale,
-        weights,
-        ids,
-        E,
-        topk,
-        BLOCK_E,
-        num_warps=num_warps,
-    )
-    return weights, ids
-
-
-def gemma4_routing_function_torch(
-    gating_output: torch.Tensor,
-    topk: int,
-    per_expert_scale: torch.Tensor,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    _, topk_ids = torch.topk(gating_output, k=topk, dim=-1)
-    router_probabilities = torch.nn.functional.softmax(gating_output, dim=-1)
-    indicator = torch.nn.functional.one_hot(
-        topk_ids, num_classes=gating_output.size(-1)
-    ).sum(dim=-2)
-    gate_weights = indicator * router_probabilities
-    renorm_factor = torch.sum(gate_weights, dim=-1, keepdim=True)
-    renorm_factor = torch.where(renorm_factor > 0.0, renorm_factor, 1.0)
-    dispatch_weights = gate_weights / renorm_factor
-
-    topk_weights = dispatch_weights.gather(1, topk_ids)
-
-    # Fold per_expert_scale into routing weights
-    expert_scales = per_expert_scale[topk_ids].to(topk_weights.dtype)
-    topk_weights = topk_weights * expert_scales
-    return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
-
-
 def _get_text_config(config):
    """Dereference text_config if config is a nested Gemma4Config.

@@ -332,12 +216,22 @@ class Gemma4MoE(nn.Module):
            topk: int,
            renormalize: bool,
        ) -> tuple[torch.Tensor, torch.Tensor]:
-            if current_platform.is_cuda_alike() or current_platform.is_xpu():
-                return gemma4_fused_routing_kernel_triton(
-                    gating_output, topk, per_expert_scale
-                )
+            _, topk_ids = torch.topk(gating_output, k=topk, dim=-1)
+            router_probabilities = torch.nn.functional.softmax(gating_output, dim=-1)
+            indicator = torch.nn.functional.one_hot(
+                topk_ids, num_classes=gating_output.size(-1)
+            ).sum(dim=-2)
+            gate_weights = indicator * router_probabilities
+            renorm_factor = torch.sum(gate_weights, dim=-1, keepdim=True)
+            renorm_factor = torch.where(renorm_factor > 0.0, renorm_factor, 1.0)
+            dispatch_weights = gate_weights / renorm_factor

-            return gemma4_routing_function_torch(gating_output, topk, per_expert_scale)
+            topk_weights = dispatch_weights.gather(1, topk_ids)
+
+            # Fold per_expert_scale into routing weights
+            expert_scales = per_expert_scale[topk_ids].to(topk_weights.dtype)
+            topk_weights = topk_weights * expert_scales
+            return topk_weights.to(torch.float32), topk_ids.to(torch.int32)

        # FusedMoE experts with custom Gemma4 routing
        self.experts = FusedMoE(
@@ -67,7 +67,6 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape
 from .interfaces import (
    MultiModalEmbeddings,
    SupportsEagle3,
-    SupportsLoRA,
    SupportsMultiModal,
    SupportsPP,
 )
@@ -881,7 +880,6 @@ class Gemma4ForConditionalGeneration(
    nn.Module,
    SupportsMultiModal,
    SupportsPP,
-    SupportsLoRA,
    SupportsEagle3,
 ):
    packed_modules_mapping = {
@@ -1360,16 +1358,10 @@ class Gemma4ForConditionalGeneration(

    def get_mm_mapping(self) -> MultiModelKeys:
        """Get the module prefix mapping for multimodal models."""
-        connectors = ["embed_vision"]
-        tower_models = ["vision_tower"]
-        if self.audio_tower is not None:
-            connectors.append("embed_audio")
-            tower_models.append("audio_tower")
-
        return MultiModelKeys.from_string_field(
            language_model="language_model",
-            connector=connectors,
-            tower_model=tower_models,
+            connector=["embed_vision", "embed_audio"],
+            tower_model=["vision_tower", "audio_tower"],
        )

    @classmethod
@@ -66,7 +66,7 @@ from .interfaces import (
    SupportsTranscription,
 )
 from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
-from .whisper import ISO639_1_SUPPORTED_LANGS, _create_fake_bias_for_k_proj
+from .whisper import ISO639_1_SUPPORTED_LANGS


 class GlmAsrEncoderRotaryEmbedding(nn.Module):
@@ -499,8 +499,6 @@ class GlmAsrEncoder(nn.Module):
        """Custom weight loading to handle q_proj/k_proj/v_proj -> qkv_proj mapping."""
        from vllm.model_executor.model_loader.weight_utils import default_weight_loader

-        weights = _create_fake_bias_for_k_proj(weights, ".k_proj.weight")
-
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("qkv_proj", "q_proj", "q"),
@@ -458,27 +458,13 @@ class PixtralForConditionalGeneration(
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
        _vision_encoder_stacked_params = [
            # (param_name, shard_name, shard_id)
-            # HF format
            (".qkv_proj", ".q_proj", "q"),
            (".qkv_proj", ".k_proj", "k"),
            (".qkv_proj", ".v_proj", "v"),
            (".gate_up_proj", ".gate_proj", 0),
            (".gate_up_proj", ".up_proj", 1),
-            # Mistral native (consolidated) format
-            (".qkv_proj", ".wq", "q"),
-            (".qkv_proj", ".wk", "k"),
-            (".qkv_proj", ".wv", "v"),
-            (".gate_up_proj", ".w1", 0),
-            (".gate_up_proj", ".w3", 1),
        ]

-        # Remap Mistral native names to HF-style names
-        # used by the vLLM vision encoder modules.
-        _vision_encoder_name_remap = {
-            ".wo.": ".o_proj.",
-            ".w2.": ".down_proj.",
-        }
-
        def is_vision_encoder_weights(weight: tuple[str, torch.Tensor]):
            return weight[0].startswith(("vision_encoder", "vision_tower"))

@@ -532,11 +518,6 @@ class PixtralForConditionalGeneration(
                            weight_loader(param, w, shard_id)
                            break
                    else:
-                        for old, new in _vision_encoder_name_remap.items():
-                            if old in trimmed_name:
-                                trimmed_name = trimmed_name.replace(old, new)
-                                break
-
                        param = vision_encoder_dict.get(trimmed_name)
                        if param is not None:
                            weight_loader = getattr(
@@ -145,15 +145,14 @@ class PlaceholderRange:
    """

    @cached_property
-    def embeds_cumsum(self) -> list[int] | None:
-        # python list so python indexing avoids torch C++ overhead/conversions/deallocs
-        return None if self.is_embed is None else self.is_embed.cumsum(dim=0).tolist()
+    def embeds_cumsum(self) -> torch.Tensor | None:
+        return None if self.is_embed is None else self.is_embed.cumsum(dim=0)

    def get_num_embeds(self) -> int:
        if self.embeds_cumsum is None:
            return self.length

-        return self.embeds_cumsum[-1] if self.embeds_cumsum else 0
+        return int(self.embeds_cumsum[-1])

    def get_embeds_indices_in_range(
        self, start_idx: int, end_idx: int
@@ -171,8 +170,10 @@ class PlaceholderRange:
        if self.embeds_cumsum is None:
            return start_idx, end_idx

-        embeds_start_idx = self.embeds_cumsum[start_idx - 1] if start_idx > 0 else 0
-        embeds_end_idx = self.embeds_cumsum[end_idx - 1] if end_idx > 0 else 0
+        embeds_start_idx = (
+            int(self.embeds_cumsum[start_idx - 1]) if start_idx > 0 else 0
+        )
+        embeds_end_idx = int(self.embeds_cumsum[end_idx - 1])

        return embeds_start_idx, embeds_end_idx

@@ -29,9 +29,9 @@ except ImportError:
    soundfile = PlaceholderModule("soundfile")  # type: ignore[assignment]


-# Public libsndfile error codes exposed via `soundfile.LibsndfileError.code`,
-# soundfile being the main audio loading backend. Used to validate if an audio
-# loading error is due to a server error vs a client error (invalid audio file).
+# Public libsndfile error codes exposed via `soundfile.LibsndfileError.code`, soundfile
+# being librosa's main backend. Used to validate if an audio loading error is due to a
+# server error vs a client error (invalid audio file).
 # 0 = sf_error(NULL) race condition: when multiple threads fail sf_open_virtual
 #     concurrently, one thread may clear the global error before another reads it,
 #     producing code=0 ("Garbled error message from libsndfile" in soundfile).
@@ -108,6 +108,22 @@ def _get_backend_priorities(
                        AttentionBackendEnum.FLASHINFER_MLA_SPARSE,
                    ]

+            return [
+                AttentionBackendEnum.FLASHINFER_MLA,
+                AttentionBackendEnum.CUTLASS_MLA,
+                AttentionBackendEnum.FLASH_ATTN_MLA,
+                AttentionBackendEnum.FLASHMLA,
+                AttentionBackendEnum.TRITON_MLA,
+                *sparse_backends,
+            ]
+        elif device_capability.major == 9:
+            # Hopper (SM90) — CUTLASS FA3 is highest priority for sparse MLA
+            # with BF16 KV cache. Falls back to FlashMLA Sparse for FP8.
+            sparse_backends = [
+                AttentionBackendEnum.CUTLASS_FA3_MLA_SPARSE,
+                AttentionBackendEnum.FLASHINFER_MLA_SPARSE,
+                AttentionBackendEnum.FLASHMLA_SPARSE,
+            ]
            return [
                AttentionBackendEnum.FLASHINFER_MLA,
                AttentionBackendEnum.CUTLASS_MLA,
@@ -382,7 +382,6 @@ def _get_backend_priorities(
    if is_aiter_found_and_supported():
        backends.append(AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN)
    backends.append(AttentionBackendEnum.TRITON_ATTN)
-    backends.append(AttentionBackendEnum.TURBOQUANT)

    return backends

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

-import torch
-
 from vllm.logger import init_logger
 from vllm.platforms.cpu import CpuPlatform

@@ -24,9 +22,3 @@ class ZenCpuPlatform(CpuPlatform):
    def is_zen_cpu(self) -> bool:
        # is_cpu() also returns True for this platform (inherited from CpuPlatform).
        return True
-
-    # Currently, AMD CPUs do not support float16 compute.
-    # Hence explicitly return bfloat16 and float32.
-    @property
-    def supported_dtypes(self) -> list[torch.dtype]:
-        return [torch.bfloat16, torch.float32]
--- a/Show More
+++ b/Show More