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

Signed-off-by: Alexander Matveev <amatveev@redhat.com>

.buildkite/test-amd.yaml

@@ -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
 
 

.buildkite/test_areas/kernels.yaml

@@ -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

CMakeLists.txt

@@ -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

cmake/external_projects/cutlass_fa3.cmake

@@ -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})")

csrc/cutlass_fa3_extension.cc

@@ -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);
+}

csrc/libtorch_stable/ops.h

@@ -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

csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu

@@ -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));
-}

csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu

@@ -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));
-}

csrc/libtorch_stable/torch_bindings.cpp

@@ -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
 }
 

csrc/sgl_flash_kernel_ops.h

@@ -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_);

csrc/sgl_kernel_torch_shim.h

@@ -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)...);
+  };
+}

docker/Dockerfile

@@ -330,7 +330,7 @@ WORKDIR /workspace
 # Build DeepEP wheels
 COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh
 # Defaults moved here from tools/ep_kernels/install_python_libraries.sh for centralized version management
-ARG DEEPEP_COMMIT_HASH=9249c25
+ARG DEEPEP_COMMIT_HASH=73b6ea4
 ARG NVSHMEM_VER
 RUN --mount=type=cache,target=/root/.cache/uv \
     mkdir -p /tmp/ep_kernels_workspace/dist && \

docker/Dockerfile.rocm

@@ -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.
@@ -261,6 +262,30 @@ RUN --mount=type=bind,source=.git,target=vllm/.git \
     && echo "Detected vLLM version: ${VLLM_VERSION}" \
     && echo "${VLLM_VERSION}" > /tmp/vllm_version.txt
 
+# Fail if git-based package dependencies are found in requirements files
+# (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI)
+# Extra notes: pip install is able to handle git+ URLs, but uv doesn't.
+RUN echo "Checking for git-based packages in requirements files..." \
+    && echo "Checking common.txt for git-based packages:" \
+    && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \
+         echo "ERROR: Git-based packages found in common.txt:"; \
+         grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \
+         echo "Please publish these packages to PyPI instead of using git dependencies."; \
+         exit 1; \
+       else \
+         echo "  ✓ No git-based packages found in common.txt"; \
+       fi \
+    && echo "Checking rocm.txt for git-based packages:" \
+    && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \
+         echo "ERROR: Git-based packages found in rocm.txt:"; \
+         grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \
+         echo "Please publish these packages to PyPI instead of using git dependencies."; \
+         exit 1; \
+       else \
+         echo "  ✓ No git-based packages found in rocm.txt"; \
+       fi \
+    && echo "All requirements files are clean - no git-based packages found"
+
 # Pin vLLM dependencies to exact versions of custom ROCm wheels
 # This ensures 'pip install vllm' automatically installs correct torch/triton/torchvision/amdsmi
 COPY tools/vllm-rocm/pin_rocm_dependencies.py /tmp/pin_rocm_dependencies.py

docker/versions.json

@@ -53,7 +53,7 @@
       "default": "cuda"
     },
     "DEEPEP_COMMIT_HASH": {
-      "default": "9249c25"
+      "default": "73b6ea4"
     },
     "GIT_REPO_CHECK": {
       "default": "0"

docs/contributing/model/transcription.md

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

docs/design/attention_backends.md

@@ -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 |

docs/design/fusions.md

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

docs/features/multimodal_inputs.md

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

docs/models/supported_models.md

@@ -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.

docs/serving/openai_compatible_server.md

@@ -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:
 

examples/online_serving/disaggregated_serving/example_mm_serve.py

@@ -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()

examples/online_serving/openai_chat_completion_client_for_multimodal.py

@@ -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}"
                         },
                     },

examples/online_serving/openai_realtime_client.py

@@ -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

requirements/rocm.txt

@@ -22,4 +22,4 @@ timm>=1.0.17
 # To be consistent with test_quark.py
 amd-quark>=0.8.99
 # Required for faster safetensors model loading
-fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@0.2.2
+fastsafetensors >= 0.2.2

requirements/test/rocm.in

@@ -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

requirements/test/rocm.txt

@@ -277,7 +277,7 @@ 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
+fastsafetensors==0.2.2
     # via
     #   -c requirements/rocm.txt
     #   -r requirements/test/rocm.in

setup.py

@@ -1085,9 +1085,7 @@ 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"],

tests/entrypoints/openai/correctness/test_transcription_api_correctness.py

@@ -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
 
 

tests/entrypoints/openai/realtime/test_realtime_validation.py

@@ -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

tests/entrypoints/openai/speech_to_text/test_transcription_validation_whisper.py

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

tests/entrypoints/openai/speech_to_text/test_translation_validation.py

@@ -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()

tests/entrypoints/openai/test_mm_serde.py

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

tests/entrypoints/serve/disagg/test_serving_multimodal_tokens.py

@@ -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}"
-    )

tests/kernels/moe/test_moe.py

@@ -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
 
 

tests/kernels/moe/test_mxfp4_moe.py

@@ -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"])

tests/kernels/moe/test_nvfp4_backend_selection.py

@@ -1,202 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from unittest.mock import patch
-
-import torch
-
-import vllm.model_executor.layers.fused_moe.modular_kernel as mk
-from tests.kernels.moe.utils import make_dummy_moe_config
-from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
-    NvFp4MoeBackend,
-    select_nvfp4_moe_backend,
-)
-from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided import (  # noqa: E501
-    _group_rank_batched_inputs_by_local_expert,
-    _reduce_local_expert_outputs_to_rank_batched_payload,
-)
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
-    kNvfp4Dynamic,
-    kNvfp4Static,
-)
-
-
-class _StandardNvFp4Kernel:
-    @staticmethod
-    def is_supported_config(
-        _cls,
-        moe_config,
-        weight_key,
-        activation_key,
-        activation_format,
-    ):
-        if activation_format == mk.FusedMoEActivationFormat.Standard:
-            return True, None
-        return False, f"{activation_format.value} activation format"
-
-
-class _BatchedNvFp4Kernel:
-    @staticmethod
-    def is_supported_config(
-        _cls,
-        moe_config,
-        weight_key,
-        activation_key,
-        activation_format,
-    ):
-        if activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
-            return True, None
-        return False, f"{activation_format.value} activation format"
-
-
-class _UnsupportedNvFp4Kernel:
-    @staticmethod
-    def is_supported_config(
-        _cls,
-        moe_config,
-        weight_key,
-        activation_key,
-        activation_format,
-    ):
-        return False, "unsupported"
-
-
-def _make_nvfp4_config(all2all_backend: str):
-    moe_config = make_dummy_moe_config(num_experts=8, hidden_dim=16)
-    moe_config.moe_backend = "flashinfer_cutedsl"
-    moe_config.moe_parallel_config.dp_size = 2
-    moe_config.moe_parallel_config.use_ep = True
-    moe_config.moe_parallel_config.all2all_backend = all2all_backend
-    return moe_config
-
-
-def _fake_backend_to_kernel_cls(backend: NvFp4MoeBackend):
-    if backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL:
-        return [_StandardNvFp4Kernel]
-    if backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED:
-        return [_BatchedNvFp4Kernel]
-    return [_UnsupportedNvFp4Kernel]
-
-
-@patch(
-    "vllm.model_executor.layers.fused_moe.oracle.nvfp4.backend_to_kernel_cls",
-    side_effect=_fake_backend_to_kernel_cls,
-)
-def test_select_nvfp4_backend_uses_standard_cutedsl_for_standard_all2all(
-    mock_backend_to_kernel_cls,
-):
-    moe_config = _make_nvfp4_config("allgather_reducescatter")
-
-    backend, experts_cls = select_nvfp4_moe_backend(
-        moe_config,
-        kNvfp4Static,
-        kNvfp4Dynamic,
-    )
-
-    assert backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL
-    assert experts_cls is _StandardNvFp4Kernel
-
-
-@patch(
-    "vllm.model_executor.layers.fused_moe.oracle.nvfp4.backend_to_kernel_cls",
-    side_effect=_fake_backend_to_kernel_cls,
-)
-def test_select_nvfp4_backend_promotes_cutedsl_for_batched_all2all(
-    mock_backend_to_kernel_cls,
-):
-    moe_config = _make_nvfp4_config("flashinfer_nvlink_one_sided")
-
-    backend, experts_cls = select_nvfp4_moe_backend(
-        moe_config,
-        kNvfp4Static,
-        kNvfp4Dynamic,
-    )
-
-    assert moe_config.moe_parallel_config.use_batched_activation_format
-    assert backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED
-    assert experts_cls is _BatchedNvFp4Kernel
-
-
-def test_flashinfer_one_sided_rank_batched_regroup_and_reduce():
-    hidden_states = torch.tensor(
-        [
-            [[10.0, 11.0], [20.0, 21.0], [30.0, 31.0], [0.0, 0.0]],
-            [[40.0, 41.0], [50.0, 51.0], [0.0, 0.0], [0.0, 0.0]],
-        ]
-    )
-    hidden_scales = torch.tensor(
-        [
-            [[1.0], [2.0], [3.0], [0.0]],
-            [[4.0], [5.0], [0.0], [0.0]],
-        ]
-    )
-    topk_ids = torch.tensor(
-        [
-            [[2, 0], [3, 2], [1, 0], [0, 0]],
-            [[3, 1], [2, 3], [0, 0], [0, 0]],
-        ],
-        dtype=torch.int32,
-    )
-    topk_weights = torch.tensor(
-        [
-            [[0.70, 0.30], [0.40, 0.60], [0.50, 0.50], [0.0, 0.0]],
-            [[0.80, 0.20], [0.25, 0.75], [0.0, 0.0], [0.0, 0.0]],
-        ],
-        dtype=torch.float32,
-    )
-
-    (
-        batched_hidden_states,
-        batched_hidden_scales,
-        expert_tokens_meta,
-        dispatch_metadata,
-    ) = _group_rank_batched_inputs_by_local_expert(
-        hidden_states,
-        hidden_scales,
-        topk_ids,
-        topk_weights,
-        num_local_experts=2,
-        first_local_expert=2,
-        num_dispatchers=2,
-        runtime_max_tokens_per_rank=4,
-        source_num_tokens=[3, 2],
-    )
-
-    assert expert_tokens_meta.expert_num_tokens.tolist() == [3, 3]
-    torch.testing.assert_close(
-        batched_hidden_states[0, :3],
-        torch.tensor([[10.0, 11.0], [20.0, 21.0], [50.0, 51.0]]),
-    )
-    torch.testing.assert_close(
-        batched_hidden_states[1, :3],
-        torch.tensor([[20.0, 21.0], [40.0, 41.0], [50.0, 51.0]]),
-    )
-    assert batched_hidden_scales is not None
-    torch.testing.assert_close(
-        batched_hidden_scales[0, :3],
-        torch.tensor([[1.0], [2.0], [5.0]]),
-    )
-    torch.testing.assert_close(
-        batched_hidden_scales[1, :3],
-        torch.tensor([[2.0], [4.0], [5.0]]),
-    )
-
-    fused_expert_output = torch.tensor(
-        [
-            [[1.0, 1.0], [2.0, 2.0], [3.0, 3.0], [0.0, 0.0]],
-            [[10.0, 10.0], [20.0, 20.0], [30.0, 30.0], [0.0, 0.0]],
-        ]
-    )
-    combine_payload = _reduce_local_expert_outputs_to_rank_batched_payload(
-        fused_expert_output,
-        dispatch_metadata,
-        num_dispatchers=2,
-        runtime_max_tokens_per_rank=4,
-        apply_router_weight_on_input=False,
-    )
-
-    expected = torch.zeros((2, 4, 2))
-    expected[0, 0] = torch.tensor([0.70, 0.70])
-    expected[0, 1] = torch.tensor([5.20, 5.20])
-    expected[1, 0] = torch.tensor([16.0, 16.0])
-    expected[1, 1] = torch.tensor([23.25, 23.25])
-    torch.testing.assert_close(combine_payload, expected)

tests/models/multimodal/generation/test_phi4mm.py

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

tests/models/multimodal/generation/test_whisper.py

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

tests/multimodal/media/test_audio.py

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

tests/multimodal/test_inputs.py

@@ -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

tests/quantization/test_turboquant.py

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

tests/test_regression.py

@@ -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",
     [

tests/utils_/test_serial_utils.py

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

tests/v1/attention/test_cutlass_fa3_sparse_backend.py

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

tests/v1/determinism/test_matmul_batch_invariant.py

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

tools/ep_kernels/install_python_libraries.sh

@@ -8,7 +8,7 @@ set -ex
 #   --nvshmem-ver <ver>  NVSHMEM version 
 
 CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
-DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"9249c25"}
+DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"73b6ea4"}
 NVSHMEM_VER=${NVSHMEM_VER:-"3.3.24"}  # Default supports both CUDA 12 and 13
 WORKSPACE=${WORKSPACE:-$(pwd)/ep_kernels_workspace}
 MODE=${MODE:-install}

vllm/_custom_ops.py

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

vllm/_xpu_ops.py

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

vllm/distributed/device_communicators/all2all.py

@@ -414,11 +414,9 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):
             num_qps_per_rank=num_qps_per_rank,
         )
         if not current_platform.is_rocm():
-            use_mnnvl = envs.VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL
             kwargs.update(
                 allow_nvlink_for_low_latency_mode=True,
-                allow_mnnvl=use_mnnvl,
-                use_fabric=use_mnnvl,
+                allow_mnnvl=envs.VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL,
                 explicitly_destroy=True,
             )
         return kwargs

vllm/distributed/device_communicators/xpu_communicator.py

@@ -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

vllm/entrypoints/anthropic/serving.py

@@ -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"""
 

vllm/entrypoints/openai/chat_completion/serving.py

@@ -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

vllm/entrypoints/openai/engine/serving.py

@@ -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:

vllm/entrypoints/serve/disagg/mm_serde.py

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

vllm/entrypoints/serve/disagg/protocol.py

@@ -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(

vllm/entrypoints/serve/disagg/serving.py

@@ -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

vllm/entrypoints/serve/render/serving.py

@@ -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(

vllm/kernels/helion/register.py

@@ -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

vllm/model_executor/layers/attention/attention.py

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

vllm/model_executor/layers/attention/mla_attention.py

@@ -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

vllm/model_executor/layers/batch_invariant.py

@@ -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}"
         )
 

vllm/model_executor/layers/fused_moe/config.py

@@ -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,
@@ -990,11 +971,7 @@ class FusedMoEParallelConfig:
 
     @property
     def use_batched_activation_format(self):
-        return (
-            self.use_deepep_ll_kernels
-            or self.use_fi_nvl_one_sided_kernels
-            or self.use_nixl_ep_kernels
-        )
+        return self.use_deepep_ll_kernels
 
     @property
     def use_ag_rs_all2all_kernels(self):

vllm/model_executor/layers/fused_moe/cutlass_moe.py

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

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

@@ -4,6 +4,7 @@
 import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm import envs
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (
@@ -125,7 +126,7 @@ class FlashInferCuteDSLBatchedExperts(mk.FusedMoEExpertsModular):
 
         # We use global_num_experts due to how moe_align_block_size handles
         # expert_maps.
-        K_dim = self.moe_config.hidden_dim if self.moe_config.hidden_dim == K * 2 else K
+        K_dim = K * 2 if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else K
         output_shape = (local_num_experts, M, K_dim)
         workspace2 = (local_num_experts, M, N)
         workspace1 = output_shape
@@ -162,12 +163,13 @@ class FlashInferCuteDSLBatchedExperts(mk.FusedMoEExpertsModular):
         assert self.w1_scale.ndim == 3
         assert self.w2_scale.ndim == 3
 
-        use_prequantized_inputs = (
-            a1q_scale is not None and hidden_states.dtype == torch.uint8
+        input_global_scale = (
+            None if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else self.a1_gscale
         )
-        input_global_scale = None if use_prequantized_inputs else self.a1_gscale
         flashinfer_hidden_states = (
-            (hidden_states, a1q_scale) if use_prequantized_inputs else hidden_states
+            (hidden_states, a1q_scale)
+            if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH
+            else hidden_states
         )
         flashinfer_cutedsl_moe_masked(
             hidden_states=flashinfer_hidden_states,

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

@@ -162,9 +162,10 @@ def select_nvfp4_moe_backend(
     # NOTE(rob): this is kind of a hack. We need to peak into
     # the prepare-finalize selection to determine if we are using
     # the batched or standard expert format.
+    use_batched = config.moe_parallel_config.use_deepep_ll_kernels
     activation_format = (
         mk.FusedMoEActivationFormat.BatchedExperts
-        if config.moe_parallel_config.use_batched_activation_format
+        if use_batched
         else mk.FusedMoEActivationFormat.Standard
     )
 
@@ -204,22 +205,16 @@ def select_nvfp4_moe_backend(
 
         raise ValueError(_make_log_unsupported(backend, reason))
 
-    def _resolve_requested_backend(
-        backend: NvFp4MoeBackend,
-    ) -> NvFp4MoeBackend:
-        if (
-            activation_format == mk.FusedMoEActivationFormat.BatchedExperts
-            and backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL
-        ):
-            return NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED
-        return backend
-
     # Handle explicit moe_backend from user.
     runner_backend = config.moe_backend
     if runner_backend != "auto":
-        requested_backend = _resolve_requested_backend(
-            map_nvfp4_backend(runner_backend)
-        )
+        requested_backend = map_nvfp4_backend(runner_backend)
+        # For batched activation format, use batched variant if available.
+        if (
+            activation_format == mk.FusedMoEActivationFormat.BatchedExperts
+            and requested_backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL
+        ):
+            requested_backend = NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED
         return _return_or_raise(
             requested_backend, config, weight_key, activation_key, activation_format
         )
@@ -232,9 +227,7 @@ def select_nvfp4_moe_backend(
 
         elif envs.is_set("VLLM_FLASHINFER_MOE_BACKEND"):
             # If user is explicit about backend, validate it.
-            backend = _resolve_requested_backend(
-                fi_2_vllm_backend_map[get_flashinfer_moe_backend()]
-            )
+            backend = fi_2_vllm_backend_map[get_flashinfer_moe_backend()]
             return _return_or_raise(
                 backend, config, weight_key, activation_key, activation_format
             )

vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ll.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import inspect
 from collections.abc import Callable
 
 import deep_ep
@@ -121,20 +120,6 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
         # time. This setting is handled by post_init_setup.
         self.use_ue8m0_dispatch = False
 
-        # Check if DeepEP supports use_nvfp4 in low_latency_dispatch.
-        # This requires the hybrid-ep branch of DeepEP.
-        self.has_nvfp4_support = (
-            "use_nvfp4" in inspect.signature(buffer.low_latency_dispatch).parameters
-        )
-        if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH and not self.has_nvfp4_support:
-            logger.warning_once(
-                "VLLM_DEEPEPLL_NVFP4_DISPATCH=1 but DeepEP does not support "
-                "use_nvfp4 in low_latency_dispatch. Falling back to FP8/BF16 "
-                "dispatch. Install DeepEP from the hybrid-ep branch for "
-                "NvFP4 dispatch support: "
-                "https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep"
-            )
-
     def post_init_setup(self, fused_experts: mk.FusedMoEExperts):
         if not fused_experts.supports_packed_ue8m0_act_scales():
             # Early exit.
@@ -198,29 +183,27 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
         assert isinstance(x, (torch.Tensor, tuple))
         q_dtype = quant_config.quant_dtype
 
-        nvfp4_quant_path = (
-            q_dtype == "nvfp4"
-            and envs.VLLM_DEEPEPLL_NVFP4_DISPATCH
-            and self.has_nvfp4_support
-        )
-        if nvfp4_quant_path:
+        if q_dtype == "nvfp4" and envs.VLLM_DEEPEPLL_NVFP4_DISPATCH:
             logger.info_once(
-                "Quantization is fused with DeepEP nvfp4 dispatch (hybrid-ep branch)"
+                "Since VLLM_DEEPEPLL_NVFP4_DISPATCH==1, make sure "
+                "using the hybrid-ep branch of DeepEP"
+                "(https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep)"
             )
             assert isinstance(x, tuple)
             x_scales = x[1]
             x = x[0].permute(2, 0, 1)
             num_experts, max_tokens, hidden_dim_by_2 = x.shape
             hidden_dim = hidden_dim_by_2 * 2
+            logger.info_once(
+                "Quantization is fused with DeepEP nvfp4 dispatch for "
+                "FlashInfer CUTEDSL as VLLM_DEEPEPLL_NVFP4_DISPATCH==1"
+            )
         else:
             if q_dtype == "nvfp4":
                 q_dtype = None
                 logger.info_once(
-                    "Using DeepEP bfloat16 dispatch for FlashInfer CUTEDSL "
-                    "(nvfp4 dispatch %s)",
-                    "not supported by this DeepEP build"
-                    if not self.has_nvfp4_support
-                    else "disabled",
+                    "Using DeepEP bfloat16 dispatch for FlashInfer CUTEDSL as "
+                    "VLLM_DEEPEPLL_NVFP4_DISPATCH==0"
                 )
             assert isinstance(x, torch.Tensor)
             num_experts, max_tokens, hidden_dim = x.size()
@@ -277,9 +260,7 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
 
         use_nvfp4 = False
         nvfp4_dispatch = (
-            quant_config.quant_dtype == "nvfp4"
-            and envs.VLLM_DEEPEPLL_NVFP4_DISPATCH
-            and self.has_nvfp4_support
+            quant_config.quant_dtype == "nvfp4" and envs.VLLM_DEEPEPLL_NVFP4_DISPATCH
         )
         if nvfp4_dispatch:
             use_nvfp4 = True

vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py

@@ -1,8 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from collections.abc import Sequence
-from dataclasses import dataclass
-
 import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -17,156 +14,6 @@ def get_local_sizes():
     return get_forward_context().dp_metadata.get_chunk_sizes_across_dp_rank()
 
 
-@dataclass
-class _OneSidedDispatchMetadata:
-    combine_indices: list[torch.Tensor]
-    combine_weights: list[torch.Tensor]
-
-
-def _normalize_source_num_tokens(
-    source_num_tokens: Sequence[int] | None,
-    num_dispatchers: int,
-    runtime_max_tokens_per_rank: int,
-) -> list[int]:
-    if source_num_tokens is None:
-        return [runtime_max_tokens_per_rank] * num_dispatchers
-
-    normalized = [int(x) for x in source_num_tokens[:num_dispatchers]]
-    if len(normalized) < num_dispatchers:
-        normalized.extend(
-            [runtime_max_tokens_per_rank] * (num_dispatchers - len(normalized))
-        )
-    return [min(max(x, 0), runtime_max_tokens_per_rank) for x in normalized]
-
-
-def _group_rank_batched_inputs_by_local_expert(
-    hidden_states: torch.Tensor,
-    hidden_scales: torch.Tensor | None,
-    topk_ids: torch.Tensor,
-    topk_weights: torch.Tensor,
-    *,
-    num_local_experts: int,
-    first_local_expert: int,
-    num_dispatchers: int,
-    runtime_max_tokens_per_rank: int,
-    source_num_tokens: Sequence[int] | None,
-) -> tuple[
-    torch.Tensor,
-    torch.Tensor | None,
-    mk.ExpertTokensMetadata,
-    _OneSidedDispatchMetadata,
-]:
-    hidden_dim = hidden_states.shape[-1]
-    batched_hidden_states = hidden_states.new_empty(
-        (num_local_experts, runtime_max_tokens_per_rank, hidden_dim)
-    )
-    batched_hidden_scales = (
-        None
-        if hidden_scales is None
-        else hidden_scales.new_empty(
-            (num_local_experts, runtime_max_tokens_per_rank, hidden_scales.shape[-1])
-        )
-    )
-    tokens_per_expert = torch.zeros(
-        num_local_experts, dtype=torch.int32, device=hidden_states.device
-    )
-
-    combine_indices: list[torch.Tensor] = []
-    combine_weights: list[torch.Tensor] = []
-    valid_source_num_tokens = _normalize_source_num_tokens(
-        source_num_tokens, num_dispatchers, runtime_max_tokens_per_rank
-    )
-
-    for local_expert in range(num_local_experts):
-        global_expert = first_local_expert + local_expert
-        expert_indices: list[torch.Tensor] = []
-        expert_weights: list[torch.Tensor] = []
-        cursor = 0
-
-        for dispatcher, num_tokens in enumerate(valid_source_num_tokens):
-            if num_tokens == 0:
-                continue
-
-            token_idx, topk_slot_idx = torch.where(
-                topk_ids[dispatcher, :num_tokens] == global_expert
-            )
-            rows = token_idx.numel()
-            if rows == 0:
-                continue
-
-            batched_hidden_states[local_expert, cursor : cursor + rows] = hidden_states[
-                dispatcher, token_idx
-            ]
-            if batched_hidden_scales is not None:
-                assert hidden_scales is not None
-                batched_hidden_scales[local_expert, cursor : cursor + rows] = (
-                    hidden_scales[dispatcher, token_idx]
-                )
-
-            expert_indices.append(
-                dispatcher * runtime_max_tokens_per_rank + token_idx.to(torch.int64)
-            )
-            expert_weights.append(
-                topk_weights[dispatcher, token_idx, topk_slot_idx].contiguous()
-            )
-            cursor += rows
-
-        tokens_per_expert[local_expert] = cursor
-        combine_indices.append(
-            torch.cat(expert_indices)
-            if expert_indices
-            else torch.empty(0, dtype=torch.int64, device=hidden_states.device)
-        )
-        combine_weights.append(
-            torch.cat(expert_weights)
-            if expert_weights
-            else torch.empty(0, dtype=topk_weights.dtype, device=topk_weights.device)
-        )
-
-    expert_tokens_meta = mk.ExpertTokensMetadata(
-        expert_num_tokens=tokens_per_expert, expert_num_tokens_cpu=None
-    )
-    dispatch_metadata = _OneSidedDispatchMetadata(
-        combine_indices=combine_indices,
-        combine_weights=combine_weights,
-    )
-    return (
-        batched_hidden_states,
-        batched_hidden_scales,
-        expert_tokens_meta,
-        dispatch_metadata,
-    )
-
-
-def _reduce_local_expert_outputs_to_rank_batched_payload(
-    fused_expert_output: torch.Tensor,
-    dispatch_metadata: _OneSidedDispatchMetadata,
-    *,
-    num_dispatchers: int,
-    runtime_max_tokens_per_rank: int,
-    apply_router_weight_on_input: bool,
-) -> torch.Tensor:
-    hidden_dim = fused_expert_output.shape[-1]
-    combine_payload = fused_expert_output.new_zeros(
-        (num_dispatchers, runtime_max_tokens_per_rank, hidden_dim)
-    )
-    flat_payload = combine_payload.view(-1, hidden_dim)
-
-    for local_expert, linear_indices in enumerate(dispatch_metadata.combine_indices):
-        rows = linear_indices.numel()
-        if rows == 0:
-            continue
-
-        expert_output = fused_expert_output[local_expert, :rows]
-        if not apply_router_weight_on_input:
-            expert_output = expert_output * dispatch_metadata.combine_weights[
-                local_expert
-            ].to(expert_output.dtype).unsqueeze(-1)
-        flat_payload.index_add_(0, linear_indices, expert_output)
-
-    return combine_payload
-
-
 class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
     """FlashInfer implementation using the Moe AlltoAll kernel."""
 
@@ -184,14 +31,8 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
         self.num_experts = num_experts
         self.hidden_size = hidden_size
         self.num_dispatchers_ = num_dispatchers
-        self.all2all_manager = get_ep_group().device_communicator.all2all_manager
-        assert self.num_experts % self.num_dispatchers_ == 0, (
-            "flashinfer_nvlink_one_sided requires evenly sharded local experts."
-        )
-        self.num_local_experts = self.num_experts // self.num_dispatchers_
-        self.first_local_expert = self.all2all_manager.rank * self.num_local_experts
-        self.dispatch_metadata: _OneSidedDispatchMetadata | None = None
 
+        self.all2all_manager = get_ep_group().device_communicator.all2all_manager
         self.all2all_manager.initialize(
             max_num_tokens=self.max_num_tokens,
             top_k=self.top_k,
@@ -201,10 +42,10 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
 
     @property
     def activation_format(self) -> mk.FusedMoEActivationFormat:
-        return mk.FusedMoEActivationFormat.BatchedExperts
+        return mk.FusedMoEActivationFormat.Standard
 
     def max_num_tokens_per_rank(self) -> int | None:
-        return self.max_num_tokens
+        return None
 
     def num_dispatchers(self) -> int:
         return self.num_dispatchers_
@@ -226,10 +67,6 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
         quant_config: FusedMoEQuantConfig,
         defer_input_quant: bool = False,
     ) -> mk.PrepareResultType:
-        if defer_input_quant:
-            raise NotImplementedError(
-                f"{self.__class__.__name__} does not support defer_input_quant=True."
-            )
         if apply_router_weight_on_input:
             topk = topk_ids.size(1)
             assert topk == 1, (
@@ -267,7 +104,7 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
         )
         if a1q_scale is not None:
             a1q_recv, a1q_scale_recv, topk_ids_recv, topk_weights_recv = recv_payloads
-            # Swizzle after dispatch when the selected MoE kernel expects it.
+            # Apply scale interleaving only for CUTLASS (not TRT-LLM)
             if (
                 quant_config.quant_dtype == "nvfp4"
                 and quant_config.is_nvfp4_scale_swizzled
@@ -275,32 +112,15 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
                 a1q_scale_recv = a1q_scale_recv.view(-1, a1q_scale_recv.shape[-1])
                 a1q_scale_recv = a1q_scale_recv.view(torch.uint8)
                 a1q_scale_recv = nvfp4_block_scale_interleave(a1q_scale_recv)
-            a1q_scale_recv = a1q_scale_recv.view(
-                self.num_dispatchers_,
-                self.runtime_max_tokens_per_rank,
-                self.hidden_size // 16,
-            )
+            a1q_scale_recv = a1q_scale_recv.view(-1, self.hidden_size // 16)
         else:
             a1q_recv, topk_ids_recv, topk_weights_recv = recv_payloads
             a1q_scale_recv = None
-        (
-            a1q_recv,
-            a1q_scale_recv,
-            expert_tokens_meta,
-            self.dispatch_metadata,
-        ) = _group_rank_batched_inputs_by_local_expert(
-            a1q_recv,
-            a1q_scale_recv,
-            topk_ids_recv,
-            topk_weights_recv,
-            num_local_experts=self.num_local_experts,
-            first_local_expert=self.first_local_expert,
-            num_dispatchers=self.num_dispatchers_,
-            runtime_max_tokens_per_rank=self.runtime_max_tokens_per_rank,
-            source_num_tokens=global_num_tokens_cpu,
-        )
+        a1q_recv = a1q_recv.view(-1, a1q_recv.shape[-1])
+        topk_ids_recv = topk_ids_recv.view(-1, topk_ids_recv.shape[-1])
+        topk_weights_recv = topk_weights_recv.view(-1, topk_weights_recv.shape[-1])
 
-        return a1q_recv, a1q_scale_recv, expert_tokens_meta, None, None
+        return a1q_recv, a1q_scale_recv, None, topk_ids_recv, topk_weights_recv
 
     def finalize(
         self,
@@ -312,20 +132,15 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
         weight_and_reduce_impl: mk.TopKWeightAndReduce,
     ) -> None:
         assert self.all2all_manager.moe_alltoall is not None
-        assert self.dispatch_metadata is not None, (
-            "flashinfer_nvlink_one_sided finalize called before prepare"
-        )
-        combine_payload = _reduce_local_expert_outputs_to_rank_batched_payload(
-            fused_expert_output,
-            self.dispatch_metadata,
-            num_dispatchers=self.num_dispatchers_,
-            runtime_max_tokens_per_rank=self.runtime_max_tokens_per_rank,
-            apply_router_weight_on_input=apply_router_weight_on_input,
+
+        ep_size = self.all2all_manager.world_size
+        hidden_size = fused_expert_output.shape[-1]
+        fused_expert_output = fused_expert_output.view(
+            ep_size, self.runtime_max_tokens_per_rank, hidden_size
         )
 
         combined_output = self.all2all_manager.moe_alltoall.combine(
-            payload=combine_payload,
+            payload=fused_expert_output,
             runtime_max_tokens_per_rank=self.runtime_max_tokens_per_rank,
         )
-        self.dispatch_metadata = None
         output.copy_(combined_output)

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

@@ -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:

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

@@ -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

vllm/model_executor/layers/quantization/turboquant/config.py

@@ -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

vllm/model_executor/layers/quantization/turboquant/quantizer.py

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

vllm/model_executor/models/gemma4_mm.py

@@ -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

vllm/model_executor/models/glmasr.py

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

vllm/multimodal/inputs.py

@@ -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
 

vllm/multimodal/media/audio.py

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

vllm/platforms/cuda.py

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

vllm/platforms/rocm.py

@@ -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
 

vllm/platforms/zen_cpu.py

@@ -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]

vllm/reasoning/__init__.py

@@ -60,10 +60,6 @@ _REASONING_PARSERS_TO_REGISTER = {
         "kimi_k2_reasoning_parser",
         "KimiK2ReasoningParser",
     ),
-    "mimo": (
-        "qwen3_reasoning_parser",
-        "Qwen3ReasoningParser",
-    ),
     "minimax_m2": (
         "minimax_m2_reasoning_parser",
         "MiniMaxM2ReasoningParser",

vllm/tool_parsers/__init__.py

@@ -94,10 +94,6 @@ _TOOL_PARSERS_TO_REGISTER = {
         "longcat_tool_parser",
         "LongcatFlashToolParser",
     ),
-    "mimo": (
-        "qwen3xml_tool_parser",
-        "Qwen3XMLToolParser",
-    ),
     "minimax_m2": (
         "minimax_m2_tool_parser",
         "MinimaxM2ToolParser",

vllm/tool_parsers/abstract_tool_parser.py

@@ -44,17 +44,6 @@ class ToolParser:
     derived classes.
     """
 
-    # When True (default), the serving layer uses the standard JSON-based
-    # parsing for tool_choice="required" and named function tool_choice,
-    # which works for models where guided decoding produces well-formed
-    # JSON output (e.g. Hermes).
-    # Subclasses set False when the standard parsing does not work for
-    # their model's output format (e.g. GLM models that use XML).  When
-    # False, the serving layer falls back to the tool_parser's
-    # extract_tool_calls / extract_tool_calls_streaming methods for
-    # required/named tool_choice, treating them the same as "auto".
-    supports_required_and_named: bool = True
-
     def __init__(
         self,
         tokenizer: TokenizerLike,

vllm/tool_parsers/glm47_moe_tool_parser.py

@@ -23,8 +23,6 @@ logger = init_logger(__name__)
 
 
 class Glm47MoeModelToolParser(Glm4MoeModelToolParser):
-    supports_required_and_named = False
-
     def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
         super().__init__(tokenizer, tools)
         # GLM-4.7 format: <tool_call>func_name[<arg_key>...]*</tool_call>

vllm/tool_parsers/glm4_moe_tool_parser.py

@@ -20,7 +20,6 @@ import regex as re
 
 from vllm.entrypoints.chat_utils import make_tool_call_id
 from vllm.entrypoints.openai.chat_completion.protocol import (
-    ChatCompletionNamedToolChoiceParam,
     ChatCompletionRequest,
 )
 from vllm.entrypoints.openai.engine.protocol import (
@@ -51,8 +50,6 @@ class Glm4MoeModelToolParser(ToolParser):
     call, and diffs against what was previously sent to emit only new content.
     """
 
-    supports_required_and_named = False
-
     def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
         super().__init__(tokenizer, tools)
         # Stateful streaming fields
@@ -159,25 +156,7 @@ class Glm4MoeModelToolParser(ToolParser):
     def adjust_request(
         self, request: ChatCompletionRequest | ResponsesRequest
     ) -> ChatCompletionRequest | ResponsesRequest:
-        """Adjust request parameters for tool call token handling.
-
-        For required/named tool_choice, skip setting structured_outputs
-        because GLM models output tool calls in XML format (per chat
-        template).  Guided decoding would force JSON output, conflicting
-        with the XML format and causing parsing failures.
-        """
-        if request.tools:
-            tc = request.tool_choice
-            if tc == "required" or isinstance(tc, ChatCompletionNamedToolChoiceParam):
-                # Do NOT call super().adjust_request() for required/named,
-                # because it would set structured_outputs and force JSON
-                # output via guided decoding.  GLM models use XML tool-call
-                # syntax (defined in the chat template), so guided decoding
-                # must be skipped to let the model output XML freely.
-                # The tool_parser handles extraction from XML output.
-                if request.tool_choice != "none":
-                    request.skip_special_tokens = False
-                return request
+        """Adjust request parameters for tool call token handling."""
         request = super().adjust_request(request)
         if request.tools and request.tool_choice != "none":
             # Ensure tool call tokens (<tool_call>, </tool_call>) are not skipped

vllm/transformers_utils/processors/cohere_asr.py

@@ -4,11 +4,11 @@ import logging
 import math
 import random
 
+import librosa
 import numpy as np
 import torch
 import torch.nn.functional as F
 from torch import nn
-from torchaudio.functional import melscale_fbanks
 from transformers import AutoFeatureExtractor, AutoProcessor, BatchFeature
 from transformers.feature_extraction_sequence_utils import (
     SequenceFeatureExtractor,
@@ -129,15 +129,17 @@ class FilterbankFeatures(nn.Module):
         self.pad_min_duration = 0.0
         self.pad_direction = "both"
 
-        filterbanks = melscale_fbanks(
-            n_freqs=self.n_fft // 2 + 1,
-            f_min=lowfreq,
-            f_max=highfreq,
-            n_mels=nfilt,
-            sample_rate=sample_rate,
-            norm=mel_norm,
-            mel_scale="slaney",
-        ).T.unsqueeze(0)
+        filterbanks = torch.tensor(
+            librosa.filters.mel(
+                sr=sample_rate,
+                n_fft=self.n_fft,
+                n_mels=nfilt,
+                fmin=lowfreq,
+                fmax=highfreq,
+                norm=mel_norm,
+            ),
+            dtype=torch.float,
+        ).unsqueeze(0)
         self.register_buffer("fb", filterbanks)
 
         # Calculate maximum sequence length

vllm/utils/serial_utils.py

@@ -27,7 +27,6 @@ class DTypeInfo:
 
 
 EmbedDType = Literal["float32", "float16", "bfloat16", "fp8_e4m3", "fp8_e5m2"]
-MmMetadataDType = Literal["int32", "int64", "uint8", "bool"]
 Endianness = Literal["native", "big", "little"]
 EncodingFormat = Literal["float", "base64", "bytes", "bytes_only"]
 
@@ -43,15 +42,6 @@ EMBED_DTYPES: Mapping[EmbedDType, DTypeInfo] = {
     "fp8_e4m3": DTypeInfo(torch.float8_e4m3fn, torch.uint8, np.uint8),
     "fp8_e5m2": DTypeInfo(torch.float8_e5m2, torch.uint8, np.uint8),
 }
-MM_METADATA_DTYPES: Mapping[MmMetadataDType, DTypeInfo] = {
-    "int32": DTypeInfo(torch.int32, torch.int32, np.int32),
-    "int64": DTypeInfo(torch.int64, torch.int64, np.int64),
-    "uint8": DTypeInfo(torch.uint8, torch.uint8, np.uint8),
-    "bool": DTypeInfo(torch.bool, torch.uint8, np.uint8),
-}
-_ALL_SERIAL_DTYPES: Mapping[str, DTypeInfo] = {
-    k: v for d in (EMBED_DTYPES, MM_METADATA_DTYPES) for k, v in d.items()
-}
 ENDIANNESS: tuple[Endianness, ...] = get_args(Endianness)
 
 
@@ -66,14 +56,14 @@ def tensor2base64(x: torch.Tensor) -> str:
 
 def tensor2binary(
     tensor: torch.Tensor,
-    embed_dtype: "EmbedDType | MmMetadataDType",
+    embed_dtype: EmbedDType,
     endianness: Endianness,
 ) -> bytes:
     assert isinstance(tensor, torch.Tensor)
-    assert embed_dtype in _ALL_SERIAL_DTYPES
+    assert embed_dtype in EMBED_DTYPES
     assert endianness in ENDIANNESS
 
-    dtype_info = _ALL_SERIAL_DTYPES[embed_dtype]
+    dtype_info = EMBED_DTYPES[embed_dtype]
 
     np_array = (
         tensor.to(dtype_info.torch_dtype)
@@ -92,13 +82,13 @@ def tensor2binary(
 def binary2tensor(
     binary: bytes,
     shape: tuple[int, ...],
-    embed_dtype: "EmbedDType | MmMetadataDType",
+    embed_dtype: EmbedDType,
     endianness: Endianness,
 ) -> torch.Tensor:
-    assert embed_dtype in _ALL_SERIAL_DTYPES
+    assert embed_dtype in EMBED_DTYPES
     assert endianness in ENDIANNESS
 
-    dtype_info = _ALL_SERIAL_DTYPES[embed_dtype]
+    dtype_info = EMBED_DTYPES[embed_dtype]
 
     np_array = np.frombuffer(binary, dtype=dtype_info.numpy_view_dtype).reshape(shape)
 

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

@@ -0,0 +1,633 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CUTLASS FA3 Sparse MLA Attention Backend for vLLM.
+
+This backend uses the vendored CUTLASS FlashAttention3 Sm90 kernel from
+sgl-attn to implement sparse MLA attention for DeepSeek-V3.2 and similar
+models on SM90 (Hopper) GPUs.
+
+Key differences from FlashMLASparseBackend:
+  - Uses BF16 KV cache (576 bytes/token) instead of FP8 (656 bytes/token)
+  - No head padding needed (FA3 handles arbitrary head counts natively)
+  - Accepts Q_rope and Q_nope (qv) separately (no ConcatMLAQ kernel)
+  - 3 sub-kernels: scheduler + main attention + combine
+  - ~4x faster per transformer block (~16us vs ~64us)
+
+All execution modes (decode, prefill, mixed) are handled identically:
+each token is treated as an independent batch element with seqlen=1.
+This simplifies metadata building and CUDA graph support.
+
+Backend priority: Highest for SM90 with kv_cache_dtype="auto".
+Graceful fallback to FlashMLA Sparse when FP8 cache requested or non-SM90.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, ClassVar
+
+import numpy as np
+import torch
+
+from vllm.v1.attention.backend import (
+    AttentionBackend,
+    AttentionCGSupport,
+    AttentionLayer,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+    CommonAttentionMetadata,
+    SparseMLAAttentionImpl,
+)
+from vllm.v1.attention.ops.cutlass_fa3 import is_cutlass_fa3_available
+
+logger = logging.getLogger(__name__)
+
+# Maximum batch size (number of tokens) for which CUTLASS FA3 is used.
+# For larger batch sizes, fall back to FlashMLA BF16 sparse prefill kernel.
+# FA3 is ~4x faster than FlashMLA for small batches (bs<=16) but regresses
+# for larger batches due to higher per-token overhead from the 3-kernel
+# launch pattern (scheduler + main + combine) and page_size=1 layout.
+MAX_BATCH_SIZE_FOR_FA3 = 16
+
+# FlashMLA sparse prefill kernel requires num_heads padded to this multiple
+# on SM90 (Hopper). SM100 (Blackwell) requires 128.
+_FLASHMLA_SM90_HEAD_PADDING = 64
+
+# Check if FlashMLA BF16 sparse kernel is available for fallback
+_flashmla_sparse_available = False
+try:
+    from vllm.v1.attention.ops.flashmla import flash_mla_sparse_fwd
+
+    _flashmla_sparse_available = True
+except (ImportError, Exception):
+    pass
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.config.cache import CacheDType
+    from vllm.model_executor.layers.linear import ColumnParallelLinear
+    from vllm.platforms.interface import DeviceCapability
+    from vllm.v1.kv_cache_interface import AttentionSpec
+
+
+# ─── Backend Class ────────────────────────────────────────────────────
+
+
+class CutlassFA3MLASparseBackend(AttentionBackend):
+    """CUTLASS FA3 sparse MLA for SM90 (Hopper). BF16 KV cache only.
+
+    When FP8 cache is requested, vLLM's backend selection falls back to
+    FlashMLASparseBackend automatically since this backend only supports
+    kv_cache_dtype="auto" (which maps to BF16 for MLA).
+    """
+
+    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto"]
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int]:
+        return [64]
+
+    @staticmethod
+    def get_name() -> str:
+        return "CUTLASS_FA3_MLA_SPARSE"
+
+    @staticmethod
+    def get_builder_cls() -> type[CutlassFA3MLASparseMetadataBuilder]:
+        return CutlassFA3MLASparseMetadataBuilder
+
+    @staticmethod
+    def get_impl_cls() -> type[CutlassFA3MLASparseImpl]:
+        return CutlassFA3MLASparseImpl
+
+    @classmethod
+    def get_supported_head_sizes(cls) -> list[int]:
+        return [576]
+
+    @classmethod
+    def is_mla(cls) -> bool:
+        return True
+
+    @classmethod
+    def is_sparse(cls) -> bool:
+        return True
+
+    @classmethod
+    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
+        return capability.major == 9
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+        cache_dtype_str: str = "auto",
+    ) -> tuple[int, ...]:
+        # BF16 cache: 576 bf16 elements per token = 1152 bytes
+        # Layout per token: [kv_c_normed(512 bf16) | k_pe(64 bf16)]
+        return (num_blocks, block_size, head_size)
+
+    @classmethod
+    def validate_configuration(
+        cls,
+        head_size: int,
+        dtype: torch.dtype,
+        kv_cache_dtype: CacheDType | None,
+        block_size: int | None,
+        use_mla: bool,
+        has_sink: bool,
+        use_sparse: bool,
+        use_mm_prefix: bool,
+        use_per_head_quant_scales: bool,
+        device_capability: DeviceCapability,
+        attn_type: str,
+        use_non_causal: bool = False,
+    ) -> list[str]:
+        invalid = super().validate_configuration(
+            head_size=head_size,
+            dtype=dtype,
+            kv_cache_dtype=kv_cache_dtype,
+            block_size=block_size,
+            use_mla=use_mla,
+            has_sink=has_sink,
+            use_sparse=use_sparse,
+            use_mm_prefix=use_mm_prefix,
+            use_per_head_quant_scales=use_per_head_quant_scales,
+            device_capability=device_capability,
+            attn_type=attn_type,
+            use_non_causal=use_non_causal,
+        )
+        if not is_cutlass_fa3_available():
+            invalid.append("_cutlass_fa3_C not available (requires CUDA >= 12.4, SM90)")
+        return invalid
+
+
+# ─── Metadata ─────────────────────────────────────────────────────────
+
+
+@dataclass
+class CutlassFA3MLASparseMetadata(AttentionMetadata):
+    """Flat metadata for CUTLASS FA3 sparse MLA attention.
+
+    ALL tokens (decode/prefill/mixed) are treated as independent batch
+    elements with seqlen=1. There are no nested Decode/Prefill sub-objects.
+
+    This simplification is valid because:
+    - Sparse MLA always routes through forward_mqa (not forward_mha)
+    - Each token independently selects its top-K KV positions
+    - The FA3 kernel handles variable-length sequences via cu_seqlens
+    """
+
+    num_reqs: int
+    max_query_len: int
+    max_seq_len: int
+    num_actual_tokens: int
+    query_start_loc: torch.Tensor
+    slot_mapping: torch.Tensor
+    block_table: torch.Tensor  # [num_reqs, max_blocks_per_req] int32
+    req_id_per_token: torch.Tensor  # [T] int32
+
+    block_size: int = 64
+    topk_tokens: int = 2048
+
+    # FA3-specific metadata (pre-allocated for CUDA graph safety)
+    cache_seqlens: torch.Tensor | None = None  # [T] int32
+    cu_seqlens_q: torch.Tensor | None = None  # [T+1] int32
+    cu_seqlens_k: torch.Tensor | None = None  # [T+1] int32
+
+    # For MLAAttention.forward_impl() routing: sparse -> all MQA
+    # Setting num_decodes = num_reqs ensures all tokens go through
+    # the forward_mqa path (no MHA prefill path).
+    num_decodes: int | None = 0
+    num_decode_tokens: int | None = 0
+    num_prefills: int | None = 0
+    num_prefill_tokens: int | None = 0
+
+
+# ─── Metadata Builder ─────────────────────────────────────────────────
+
+
+class CutlassFA3MLASparseMetadataBuilder(
+    AttentionMetadataBuilder[CutlassFA3MLASparseMetadata]
+):
+    """Builds CutlassFA3MLASparseMetadata from CommonAttentionMetadata.
+
+    Key design choices:
+    - Pre-allocates GPU buffers in __init__ for CUDA graph compatibility
+    - All tokens (decode + prefill) treated as independent seqlen=1 elements
+    - Uses in-place .copy_() for buffer updates (safe for CUDA graph replay)
+    """
+
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
+
+    def __init__(
+        self,
+        kv_cache_spec: AttentionSpec,
+        layer_names: list[str],
+        vllm_config: VllmConfig,
+        device: torch.device,
+    ) -> None:
+        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+        self.topk_tokens = 2048
+        max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+        self.block_size = kv_cache_spec.block_size
+
+        # Enable speculative decoding support
+        self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
+
+        # Pre-allocate GPU buffers (persist across CUDA graph replays).
+        # These are updated in-place via .copy_() before each replay.
+        self.req_id_buf = torch.zeros(max_tokens, dtype=torch.int32, device=device)
+        self.cache_seqlens_buf = torch.ones(
+            max_tokens, dtype=torch.int32, device=device
+        )
+        self.cu_seqlens_q_buf = torch.arange(
+            0, max_tokens + 1, dtype=torch.int32, device=device
+        )
+        self.cu_seqlens_k_buf = torch.zeros(
+            max_tokens + 1, dtype=torch.int32, device=device
+        )
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+    ) -> CutlassFA3MLASparseMetadata:
+        """Build metadata from common attention metadata.
+
+        Converts the request-level metadata into per-token flat metadata:
+        - req_id_per_token: maps each token to its request index
+        - cache_seqlens: min(seq_len, topk) per token (topk-clipped)
+        - cu_seqlens_q: [0, 1, 2, ..., T] (each token = seqlen 1)
+        - cu_seqlens_k: cumsum of cache_seqlens
+        """
+        cm = common_attn_metadata
+        T = cm.num_actual_tokens
+        starts = np.asarray(cm.query_start_loc_cpu, dtype=np.int32)
+        seg_lens = np.diff(starts)
+
+        # req_id_per_token: map each token -> request index
+        req_ids = np.repeat(np.arange(len(seg_lens), dtype=np.int32), seg_lens)
+        # CUDA graph padding fix: T = cm.num_actual_tokens may include
+        # padding tokens (e.g., T=32 when only 31 real tokens exist).
+        # The computed req_ids array has sum(seg_lens) elements which
+        # equals the real (unpadded) token count. We must:
+        # 1) Zero-fill the entire buffer first (safe default for padding)
+        # 2) Copy only the actual data using req_ids.shape[0]
+        # 3) Slice to padded T for the metadata return
+        # This matches the pattern used by FlashMLASparseMetadataBuilder,
+        # FlashInferMLASparseMetadataBuilder, and all other sparse backends.
+        actual_tokens = req_ids.shape[0]
+        self.req_id_buf.fill_(0)
+        self.req_id_buf[:actual_tokens].copy_(
+            torch.from_numpy(req_ids), non_blocking=True
+        )
+
+        # cache_seqlens: UPPER BOUND = min(seq_len, topk) per token.
+        # NOTE: This is a per-REQUEST uniform value, NOT the correct
+        # per-token causal seqlen. For prefill, token i at position p
+        # can only attend to min(p+1, topk) entries, but this gives
+        # all tokens min(seq_len, topk). The actual per-token
+        # cache_seqlens is computed in forward_mqa() using valid_counts
+        # from the index conversion kernel, which correctly reflects
+        # the number of valid KV entries per token.
+        seq_lens_np = np.asarray(cm.seq_lens_cpu, dtype=np.int32)
+        per_tok_seqlens = np.minimum(np.repeat(seq_lens_np, seg_lens), self.topk_tokens)
+        # Same CUDA graph padding fix: zero-fill then copy actual data.
+        # Default to 1 (safe minimum seqlen for FA3 kernel).
+        self.cache_seqlens_buf.fill_(1)
+        self.cache_seqlens_buf[:actual_tokens].copy_(
+            torch.from_numpy(per_tok_seqlens), non_blocking=True
+        )
+
+        # cu_seqlens_q: [0, 1, 2, ..., T] — each token is seqlen=1
+        cu_q = self.cu_seqlens_q_buf[: T + 1]
+
+        # cu_seqlens_k: cumsum(cache_seqlens)
+        self.cu_seqlens_k_buf[0] = 0
+        self.cu_seqlens_k_buf[1 : T + 1].copy_(
+            torch.cumsum(self.cache_seqlens_buf[:T], dim=0)
+        )
+        cu_k = self.cu_seqlens_k_buf[: T + 1]
+
+        return CutlassFA3MLASparseMetadata(
+            num_reqs=cm.num_reqs,
+            max_query_len=cm.max_query_len,
+            max_seq_len=cm.max_seq_len,
+            num_actual_tokens=T,
+            query_start_loc=cm.query_start_loc,
+            slot_mapping=cm.slot_mapping,
+            block_table=cm.block_table_tensor,
+            req_id_per_token=self.req_id_buf[:T],
+            block_size=self.block_size,
+            topk_tokens=self.topk_tokens,
+            cache_seqlens=self.cache_seqlens_buf[:T],
+            cu_seqlens_q=cu_q,
+            cu_seqlens_k=cu_k,
+            # Route ALL tokens through MQA in forward_impl
+            num_decodes=cm.num_reqs,
+            num_decode_tokens=T,
+            num_prefills=0,
+            num_prefill_tokens=0,
+        )
+
+
+# ─── Implementation ───────────────────────────────────────────────────
+
+
+class CutlassFA3MLASparseImpl(SparseMLAAttentionImpl[CutlassFA3MLASparseMetadata]):
+    """CUTLASS FA3 sparse MLA attention implementation.
+
+    This implementation replaces the FlashMLA C sparse_attn_fwd_kernel
+    with the CUTLASS FA3 Sm90 kernel from sgl-attn, providing ~4x speedup
+    per transformer block on Hopper GPUs.
+
+    Key advantages over FlashMLASparseImpl:
+    - No head padding (FA3 handles arbitrary head counts natively)
+    - No Q concatenation kernel (FA3 accepts q_rope and qv separately)
+    - BF16 KV cache (smaller footprint, no dequantization overhead)
+    - SM90 warpgroup MMA + TMA for higher compute efficiency
+    """
+
+    supports_quant_query_input: bool = False
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int,
+        alibi_slopes: list[float] | None,
+        sliding_window: int | None,
+        kv_cache_dtype: str,
+        logits_soft_cap: float | None,
+        attn_type: str,
+        kv_sharing_target_layer_name: str | None,
+        # MLA Specific Arguments
+        q_lora_rank: int | None = None,
+        kv_lora_rank: int = 512,
+        qk_nope_head_dim: int = 128,
+        qk_rope_head_dim: int = 64,
+        qk_head_dim: int = 192,
+        v_head_dim: int = 128,
+        kv_b_proj: ColumnParallelLinear | None = None,
+        indexer: object | None = None,
+        q_pad_num_heads: int | None = None,
+        **kwargs,
+    ) -> None:
+        self.num_heads = num_heads  # 16 (per GPU for TP=8)
+        self.head_size = head_size  # 576 (kv_lora_rank + qk_rope_head_dim)
+        self.scale = float(scale)  # 192**-0.5
+        self.num_kv_heads = num_kv_heads  # 1 (MQA)
+        self.kv_cache_dtype = kv_cache_dtype  # "auto" (maps to BF16)
+        self.kv_lora_rank = kv_lora_rank
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.softmax_scale = scale
+        self.topk_tokens = 2048
+        self.num_splits = 0  # auto; CUDA-graph safe (deterministic per bs)
+        self.logits_soft_cap = float(logits_soft_cap) if logits_soft_cap else 0.0
+
+        # The indexer provides topk_indices_buffer shared across layers
+        assert indexer is not None, (
+            "CutlassFA3MLASparseImpl requires an indexer "
+            "for sparse top-K index selection"
+        )
+        self.topk_indices_buffer = indexer.topk_indices_buffer
+
+        # DCP (Decode Context Parallelism) requires softmax LSE from the
+        # attention kernel. FA3's return_softmax_lse=True is not yet wired
+        # through this backend. When DCP is needed, fall back to FlashMLA.
+        # TODO: Wire return_softmax_lse=True through forward_mqa for DCP.
+
+    def forward_mqa(
+        self,
+        q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        kv_c_and_k_pe_cache: torch.Tensor,
+        attn_metadata: CutlassFA3MLASparseMetadata,
+        layer: AttentionLayer,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """FA3 sparse MLA attention with batch size gating.
+
+        For batch sizes <= MAX_BATCH_SIZE_FOR_FA3 (16), uses the fast
+        CUTLASS FA3 kernel. For larger batch sizes, falls back to the
+        FlashMLA BF16 sparse prefill kernel which handles larger batches
+        more efficiently.
+
+        All execution modes (decode/prefill/mixed) are handled identically:
+        each token is an independent batch element with seqlen=1.
+
+        Input:  q = tuple(ql_nope[T, N, 512], q_pe[T, N, 64])
+        Output: (attn_out[T, N, 512], None)
+
+        The _v_up_proj in MLAAttention.forward_impl() handles the
+        subsequent .view(-1, N, kv_lora_rank) correctly for 3D output.
+        """
+        # FA3 does not yet return LSE; DCP requires it.
+        assert self.dcp_world_size <= 1, (
+            "CutlassFA3MLASparseImpl does not support DCP (dcp_world_size > 1). "
+            "Use FlashMLA Sparse instead."
+        )
+
+        # 1) Unpack Q components
+        if isinstance(q, tuple):
+            ql_nope, q_pe = q  # [T, N, 512], [T, N, 64]
+        else:
+            ql_nope = q[..., : self.kv_lora_rank]  # [T, N, 512]
+            q_pe = q[..., self.kv_lora_rank :]  # [T, N, 64]
+        T = ql_nope.shape[0]
+
+        # 2) Convert topk_indices -> global cache slot indices
+        #    Reuses vLLM's existing Triton kernel (no changes needed)
+        from vllm.v1.attention.backends.mla.sparse_utils import (
+            triton_convert_req_index_to_global_index,
+        )
+
+        global_idx, valid_counts = triton_convert_req_index_to_global_index(
+            attn_metadata.req_id_per_token,  # [T] int32
+            attn_metadata.block_table,  # [R, max_blocks] int32
+            self.topk_indices_buffer[:T],  # [T, 2048] int32
+            BLOCK_SIZE=attn_metadata.block_size,  # 64
+            NUM_TOPK_TOKENS=self.topk_tokens,  # 2048
+            return_valid_counts=True,
+        )
+        # global_idx: [T, 2048] int32 — flat cache slot IDs
+        # valid_counts: [T] int32 — number of valid (non -1) entries per token
+
+        # Replace -1 (invalid) page indices with 0 (a safe, valid page index)
+        # IN-PLACE for CUDA graph friendliness (no extra allocation).
+        global_idx.clamp_(min=0)
+
+        # Use valid_counts as cache_seqlens instead of metadata.cache_seqlens.
+        # CRITICAL FIX (Issue #1): metadata cache_seqlens = min(seq_len, topk)
+        # can exceed actual valid topk entries for prefill tokens.
+        valid_counts.clamp_(min=1)  # in-place; min=1 for seqlen_k safety
+        cache_seqlens = valid_counts
+
+        # 3) Route to FA3 or FlashMLA based on batch size
+        #    FA3 is faster for small batches (bs<=16) but regresses for
+        #    larger batches. FlashMLA BF16 sparse prefill handles larger
+        #    batches more efficiently.
+        use_fa3 = (T <= MAX_BATCH_SIZE_FOR_FA3) or not _flashmla_sparse_available
+        if use_fa3:
+            attn_out = self._forward_fa3(
+                ql_nope,
+                q_pe,
+                kv_c_and_k_pe_cache,
+                global_idx,
+                cache_seqlens,
+                attn_metadata,
+            )
+        else:
+            attn_out = self._forward_flashmla_bf16_fallback(
+                ql_nope,
+                q_pe,
+                kv_c_and_k_pe_cache,
+                global_idx,
+                cache_seqlens,
+            )
+
+        # Output: [T, N, 512] — already 3D
+        return attn_out, None
+
+    def _forward_fa3(
+        self,
+        ql_nope: torch.Tensor,  # [T, N, 512]
+        q_pe: torch.Tensor,  # [T, N, 64]
+        kv_c_and_k_pe_cache: torch.Tensor,
+        global_idx: torch.Tensor,  # [T, 2048]
+        cache_seqlens: torch.Tensor,  # [T]
+        attn_metadata: CutlassFA3MLASparseMetadata,
+    ) -> torch.Tensor:
+        """CUTLASS FA3 kernel path — fast for small batch sizes (bs<=16).
+
+        Accepts Q_rope and Q_nope (qv) separately, no head padding needed.
+        Uses page_size=1 paged KV format with split-KV parallelism.
+        """
+        T = ql_nope.shape[0]
+        S = kv_c_and_k_pe_cache.shape[0] * kv_c_and_k_pe_cache.shape[1]
+        kv_flat = kv_c_and_k_pe_cache.reshape(S, self.head_size)  # [S, 576]
+
+        # Split NoPE and RoPE, reshape for FA3 paged format (page_size=1)
+        c_kv = kv_flat[:, : self.kv_lora_rank].reshape(
+            S, 1, 1, self.kv_lora_rank
+        )  # [S, 1, 1, 512]
+        k_rope = kv_flat[:, self.kv_lora_rank :].reshape(
+            S, 1, 1, self.qk_rope_head_dim
+        )  # [S, 1, 1, 64]
+
+        from vllm.v1.attention.ops.cutlass_fa3 import flash_attn_with_kvcache
+
+        attn_out = flash_attn_with_kvcache(
+            q=q_pe,  # [T, N, 64]
+            k_cache=k_rope,  # [S, 1, 1, 64]
+            v_cache=c_kv,  # [S, 1, 1, 512]
+            qv=ql_nope,  # [T, N, 512]
+            page_table=global_idx,  # [T, 2048]
+            cache_seqlens=cache_seqlens,  # [T]
+            cu_seqlens_q=attn_metadata.cu_seqlens_q,  # [T+1]
+            cu_seqlens_k_new=None,
+            max_seqlen_q=1,
+            softmax_scale=self.softmax_scale,  # 192**-0.5
+            causal=True,
+            window_size=(-1, -1),
+            softcap=self.logits_soft_cap,
+            num_splits=self.num_splits,
+        )
+        return attn_out  # [T, N, 512]
+
+    def _forward_flashmla_bf16_fallback(
+        self,
+        ql_nope: torch.Tensor,  # [T, N, 512]
+        q_pe: torch.Tensor,  # [T, N, 64]
+        kv_c_and_k_pe_cache: torch.Tensor,
+        global_idx: torch.Tensor,  # [T, 2048]
+        cache_seqlens: torch.Tensor,  # [T]
+    ) -> torch.Tensor:
+        """FlashMLA BF16 sparse prefill fallback — for larger batch sizes.
+
+        Used when T > MAX_BATCH_SIZE_FOR_FA3 (16). The FlashMLA BF16 sparse
+        prefill kernel handles larger batches more efficiently than FA3's
+        3-kernel launch pattern (scheduler + main + combine).
+
+        This path:
+        1. Concatenates Q components: [ql_nope | q_pe] -> [T, N, 576]
+        2. Pads heads to 64 (SM90 FlashMLA requirement)
+        3. Reshapes KV cache to [S, 1, 576] (flattened, MQA format)
+        4. Reshapes indices to [T, 1, topk] (MQA format)
+        5. Calls flash_mla_sparse_fwd with topk_length for valid bounds
+        6. Unpads output heads back to N
+
+        The BF16 KV cache format [kv_c_normed(512) | k_pe(64)] is identical
+        between FA3 and FlashMLA, so no cache format conversion is needed.
+        """
+        T = ql_nope.shape[0]
+        N = self.num_heads
+
+        # 1) Concatenate Q: [ql_nope(512) | q_pe(64)] -> [T, N, 576]
+        q_concat = torch.cat([ql_nope, q_pe], dim=-1)  # [T, N, 576]
+
+        # 2) Pad heads to _FLASHMLA_SM90_HEAD_PADDING (64 for SM90)
+        padded_heads = _FLASHMLA_SM90_HEAD_PADDING
+        if padded_heads > N:
+            q_padded = q_concat.new_zeros((T, padded_heads, q_concat.shape[-1]))
+            q_padded[:, :N, :] = q_concat
+            q_concat = q_padded
+
+        # 3) Reshape KV cache: (num_blocks, block_size, 576) -> (S, 1, 576)
+        S = kv_c_and_k_pe_cache.shape[0] * kv_c_and_k_pe_cache.shape[1]
+        kv = kv_c_and_k_pe_cache.reshape(S, 1, self.head_size)  # [S, 1, 576]
+
+        # 4) Reshape indices for MQA: (T, 2048) -> (T, 1, 2048)
+        indices = global_idx.unsqueeze(1)  # [T, 1, 2048]
+
+        # 5) Call FlashMLA BF16 sparse prefill kernel
+        #    NOTE: Unlike FlashMLASparseImpl._bf16_flash_mla_kernel which does
+        #    not pass topk_length (it relies on all indices being valid), we
+        #    pass topk_length=valid_counts because our indices have been
+        #    clamped (global_idx.clamp_(min=0)), so entries beyond valid_counts
+        #    are 0 (valid but irrelevant data). topk_length prevents the kernel
+        #    from processing these clamped entries, saving compute and ensuring
+        #    correctness.
+        output = flash_mla_sparse_fwd(
+            q_concat,  # [T, padded_heads, 576]
+            kv,  # [S, 1, 576]
+            indices,  # [T, 1, 2048]
+            self.softmax_scale,  # 192**-0.5
+            d_v=self.kv_lora_rank,  # 512
+            topk_length=cache_seqlens,  # [T] valid entry counts
+        )[0]  # extract output tensor from (output, max_logits, lse) tuple
+
+        # 6) Unpad heads: (T, padded_heads, 512) -> (T, N, 512)
+        return output[:, :N, :]
+
+    def do_kv_cache_update(
+        self,
+        kv_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        kv_cache: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        kv_cache_dtype: str,
+        k_scale: torch.Tensor,
+    ) -> None:
+        """BF16 KV cache write using existing vLLM kernel.
+
+        kv_cache_dtype MUST be "auto" which maps to Fp8KVCacheDataType::kAuto
+        in the C++ dispatch, performing a direct BF16 copy (no quantization).
+        Passing "bfloat16" would crash because concat_and_cache_mla expects
+        the "auto" string for the non-quantized path.
+        """
+        if kv_cache.numel() == 0:
+            return
+        from vllm import _custom_ops as ops
+
+        ops.concat_and_cache_mla(
+            kv_c_normed,
+            k_pe.squeeze(1),
+            kv_cache,
+            slot_mapping.flatten(),
+            kv_cache_dtype="auto",
+            scale=k_scale,
+        )

vllm/v1/attention/backends/registry.py

@@ -73,6 +73,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
     FLASHMLA_SPARSE = (
         "vllm.v1.attention.backends.mla.flashmla_sparse.FlashMLASparseBackend"
     )
+    CUTLASS_FA3_MLA_SPARSE = (
+        "vllm.v1.attention.backends.mla.cutlass_fa3_sparse.CutlassFA3MLASparseBackend"
+    )
     FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend"
     NO_ATTENTION = "vllm.v1.attention.backends.no_attention.NoAttentionBackend"
     FLEX_ATTENTION = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend"

vllm/v1/attention/backends/turboquant_attn.py

@@ -279,25 +279,29 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
         )
 
     def _ensure_on_device(self, layer, device):
-        """One-time derivation of TQ buffers (rotation matrix, midpoints).
+        """One-time derivation of TQ buffers (rotation matrices, midpoints).
 
-        The Hadamard rotation is shared across all layers: random sign
-        flips do not improve Lloyd-Max quantization quality because the
-        quantizer is symmetric around zero (sign-flipping a coordinate
-        maps it to the mirror centroid with identical distortion).
+        Registered buffers (_tq_signs, _tq_centroids) are already on the
+        correct device via register_buffer + model.to(device).
         """
         if not hasattr(layer, "_tq_cached"):
-            D = self.head_size
+            D = layer._tq_signs.shape[0]
+            signs = layer._tq_signs.to(device=device, dtype=torch.float32)
 
-            # Pure Hadamard: orthonormal + symmetric (H = H^T), enabling
+            # WHT rotation: orthonormal + self-inverse, enabling future
             # in-kernel butterfly fusion and trivial inverse for continuation.
             H = _build_hadamard(D, str(device))
-            layer._tq_PiT = H
-            layer._tq_Pi = H
+            layer._tq_PiT = (signs.unsqueeze(1) * H).contiguous()
+            layer._tq_Pi = layer._tq_PiT.T.contiguous()
 
             c = layer._tq_centroids.to(device=device, dtype=torch.float32)
+            # Precompute midpoints for threshold-based quantization
             c_sorted, _ = c.sort()
             layer._tq_midpoints = (c_sorted[:-1] + c_sorted[1:]) / 2
+            # Decode buffers (_tq_mid_o_buf, _tq_output_buf, _tq_lse_buf)
+            # are pre-allocated via register_buffer in Attention.__init__
+            # and moved to GPU by model.to(device) — no allocation needed
+            # here.  The memory profiler sees them before KV cache sizing.
             layer._tq_cached = True
 
     def do_kv_cache_update(
@@ -503,7 +507,8 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
         # max_query_len == max_seq_len means no request has prior cached KV.
         # Both are Python ints — no GPU sync.
         if _HAS_FLASH_ATTN and attn_metadata.max_query_len == attn_metadata.max_seq_len:
-            return flash_attn_varlen_func(
+            output = torch.empty(N, Hq, D, device=query.device, dtype=query.dtype)
+            flash_attn_varlen_func(
                 q=query,
                 k=key,
                 v=value,
@@ -513,7 +518,9 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
                 max_seqlen_k=attn_metadata.max_query_len,
                 softmax_scale=self.scale,
                 causal=True,
+                out=output,
             )
+            return output
 
         # Continuation or no flash_attn: per-request attention.
         # For continuation chunks (seq_len > q_len), we must attend to
@@ -550,9 +557,10 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
             if q_len == seq_len:
                 # First-chunk prefill: all K/V are in the current batch.
                 if _HAS_FLASH_ATTN:
+                    out = torch.empty_like(q_seq)
                     _cu_2[1] = q_len
                     cu = _cu_2
-                    out = flash_attn_varlen_func(
+                    flash_attn_varlen_func(
                         q=q_seq,
                         k=k_seq,
                         v=v_seq,
@@ -562,6 +570,7 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
                         max_seqlen_k=q_len,
                         softmax_scale=self.scale,
                         causal=True,
+                        out=out,
                     )
                 else:
                     q_t = q_seq.transpose(0, 1).contiguous()
@@ -724,9 +733,10 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
 
         # Attention: q_len queries attending to seq_len K/V with causal mask
         if _HAS_FLASH_ATTN:
+            output = torch.empty(q_len, Hq, D, device=device, dtype=query.dtype)
             cu_seqlens_q = torch.tensor([0, q_len], device=device, dtype=torch.int32)
             cu_seqlens_k = torch.tensor([0, seq_len], device=device, dtype=torch.int32)
-            return flash_attn_varlen_func(
+            flash_attn_varlen_func(
                 q=query,
                 k=k_full,
                 v=v_full,
@@ -736,7 +746,9 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
                 max_seqlen_k=seq_len,
                 softmax_scale=self.scale,
                 causal=True,
+                out=output,
             )
+            return output
         else:
             # SDPA fallback: expand KV for GQA, build causal mask
             q_t = query.transpose(0, 1).unsqueeze(0)  # (1, Hq, q_len, D)

vllm/v1/attention/ops/cutlass_fa3.py

@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Vendored CUTLASS FA3 MLA attention kernel wrapper.
+
+This module wraps the CUTLASS FlashAttention3 Sm90 kernel from sgl-attn,
+providing a Python interface compatible with vLLM's sparse MLA attention
+backend. The kernel is vendored as a self-contained C++ extension
+(_cutlass_fa3_C) and does NOT depend on sglang, sgl_kernel, or any sgl*
+modules.
+
+The FA3 kernel supports MLA (Multi-head Latent Attention) with:
+  - Separate Q_rope and QV (Q_nope) components
+  - Paged KV cache with page_size=1
+  - Variable-length sequences via cu_seqlens
+  - Split-KV parallelism with automatic split count
+  - SM90 (Hopper) CUTLASS warpgroup MMA + TMA
+
+Source: https://github.com/sgl-project/sgl-attn (commit bcf72ccc)
+"""
+
+import torch
+
+from vllm.platforms import current_platform
+
+_cutlass_fa3_available = False
+if current_platform.is_cuda():
+    try:
+        import vllm._cutlass_fa3_C  # noqa: F401
+
+        _cutlass_fa3_available = True
+    except ImportError:
+        pass
+
+
+def is_cutlass_fa3_available() -> bool:
+    """Check if the CUTLASS FA3 extension is available.
+
+    Requires CUDA >= 12.4 and SM90 (Hopper) GPU.
+    """
+    return _cutlass_fa3_available
+
+
+def flash_attn_with_kvcache(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    k: torch.Tensor | None = None,
+    v: torch.Tensor | None = None,
+    qv: torch.Tensor | None = None,
+    rotary_cos: torch.Tensor | None = None,
+    rotary_sin: torch.Tensor | None = None,
+    cache_seqlens: torch.Tensor | None = None,
+    cache_batch_idx: torch.Tensor | None = None,
+    cache_leftpad: torch.Tensor | None = None,
+    page_table: torch.Tensor | None = None,
+    cu_seqlens_q: torch.Tensor | None = None,
+    cu_seqlens_k_new: torch.Tensor | None = None,
+    max_seqlen_q: int | None = None,
+    rotary_seqlens: torch.Tensor | None = None,
+    q_descale: torch.Tensor | None = None,
+    k_descale: torch.Tensor | None = None,
+    v_descale: torch.Tensor | None = None,
+    softmax_scale: float | None = None,
+    causal: bool = False,
+    window_size: tuple[int, int] = (-1, -1),
+    attention_chunk: int | None = None,
+    softcap: float = 0.0,
+    rotary_interleaved: bool = True,
+    scheduler_metadata: torch.Tensor | None = None,
+    num_splits: int = 0,
+    pack_gqa: bool | None = None,
+    sm_margin: int = 0,
+    return_softmax_lse: bool = False,
+    sinks: torch.Tensor | None = None,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    """CUTLASS FA3 attention with paged KV cache for MLA.
+
+    MLA mode shapes (DeepSeek-V3.2, varlen mode with cu_seqlens_q):
+      q:       [T, N, 64]      -- RoPE query (total_q=T, heads=N, dim=64)
+      qv:      [T, N, 512]     -- NoPE query (total_q=T, heads=N, dim_v=512)
+      k_cache: [S, 1, 1, 64]   -- Paged RoPE keys (pages, pg_sz=1, kv_h=1, d)
+      v_cache: [S, 1, 1, 512]  -- Paged NoPE latent (pages, pg_sz=1, kv_h=1, dv)
+      page_table: [T, topk]    -- Global cache slot indices per token
+
+    The FA3 kernel internally computes:
+      score = Q_rope @ K_rope^T + QV(Q_nope) @ V_cache(C_KV)^T
+      output = softmax(score * scale) @ V_cache(C_KV)
+
+    FA3 MLA constraints (from flash_api.cpp):
+      headdim_qk <= 64, headdim_v >= 256, SM90 only, BF16/FP16 only
+
+    FA3 produces 3 sub-kernels:
+      1. prepare_varlen_num_blocks_kernel (scheduler)
+      2. FlashAttnFwdSm90 (main attention, TMA+WGMMA)
+      3. FlashAttnFwdCombine (split-KV merge, when num_splits > 1)
+
+    Args:
+        q: Query tensor for RoPE component.
+        k_cache: Paged K cache (RoPE keys).
+        v_cache: Paged V cache (NoPE latent).
+        qv: Query tensor for NoPE/value component (MLA specific).
+        page_table: Page table mapping tokens to cache slots.
+        cache_seqlens: Number of valid KV entries per batch element.
+        cu_seqlens_q: Cumulative query sequence lengths.
+        cu_seqlens_k_new: Cumulative KV sequence lengths.
+        max_seqlen_q: Maximum query sequence length.
+        softmax_scale: Softmax scale factor (default: q.shape[-1]**-0.5).
+        causal: Whether to apply causal masking.
+        window_size: (left, right) attention window sizes.
+        softcap: Logits soft cap value (0.0 = disabled).
+        num_splits: Number of split-KV splits (0 = auto).
+        return_softmax_lse: Whether to return log-sum-exp values.
+
+    Returns:
+        Attention output tensor, or tuple of (output, softmax_lse).
+    """
+    assert _cutlass_fa3_available, (
+        "CUTLASS FA3 requires CUDA >= 12.4 and SM90 (Hopper) GPU. "
+        "The _cutlass_fa3_C extension was not compiled or could not be loaded."
+    )
+    if softmax_scale is None:
+        softmax_scale = q.shape[-1] ** (-0.5)
+    attention_chunk_val = 0 if attention_chunk is None else int(attention_chunk)
+
+    out, softmax_lse, *rest = torch.ops._cutlass_fa3_C.fwd.default(
+        q,  # 0: q
+        k_cache,  # 1: k (paged KV cache)
+        v_cache,  # 2: v (paged KV cache)
+        k,  # 3: k_new
+        v,  # 4: v_new
+        qv,  # 5: q_v (MLA NoPE query)
+        None,  # 6: out buffer
+        cu_seqlens_q,  # 7: cu_seqlens_q
+        None,  # 8: cu_seqlens_k
+        cu_seqlens_k_new,  # 9: cu_seqlens_k_new
+        None,  # 10: seqused_q
+        cache_seqlens,  # 11: seqused_k
+        max_seqlen_q,  # 12: max_seqlen_q
+        None,  # 13: max_seqlen_k
+        page_table,  # 14: page_table
+        cache_batch_idx,  # 15: kv_batch_idx
+        cache_leftpad,  # 16: leftpad_k
+        rotary_cos,  # 17: rotary_cos
+        rotary_sin,  # 18: rotary_sin
+        rotary_seqlens,  # 19: seqlens_rotary
+        q_descale,  # 20: q_descale
+        k_descale,  # 21: k_descale
+        v_descale,  # 22: v_descale
+        softmax_scale,  # 23: softmax_scale
+        causal,  # 24: is_causal
+        window_size[0],  # 25: window_size_left
+        window_size[1],  # 26: window_size_right
+        attention_chunk_val,  # 27: attention_chunk
+        softcap,  # 28: softcap
+        rotary_interleaved,  # 29: is_rotary_interleaved
+        scheduler_metadata,  # 30: scheduler_metadata
+        num_splits,  # 31: num_splits
+        pack_gqa,  # 32: pack_gqa
+        sm_margin,  # 33: sm_margin
+        sinks,  # 34: sinks
+    )
+    return (out, softmax_lse) if return_softmax_lse else out

vllm/v1/attention/ops/triton_turboquant_decode.py

@@ -13,7 +13,6 @@ from typing import Any
 
 import torch
 
-from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.v1.attention.ops.triton_decode_attention import (
     _fwd_kernel_stage2,
@@ -23,15 +22,10 @@ _FP8_E4B15: dict[int, int] = {}
 
 
 def _use_fp8_e4b15(device: int = 0) -> int:
-    """Return 1 if device needs fp8e4b15 (Ampere/Ada, SM < 8.9), else 0.
-    On non-CUDA platforms (e.g. XPU), always returns 0 (use e4nv format).
-    """
+    """Return 1 if device needs fp8e4b15 (Ampere/Ada, SM < 8.9), else 0."""
     if device not in _FP8_E4B15:
-        if current_platform.is_cuda_alike():
-            cap = torch.cuda.get_device_capability(device)
-            _FP8_E4B15[device] = 1 if cap < (8, 9) else 0
-        else:
-            _FP8_E4B15[device] = 0
+        cap = torch.cuda.get_device_capability(device)
+        _FP8_E4B15[device] = 1 if cap < (8, 9) else 0
     return _FP8_E4B15[device]
 
 
@@ -143,12 +137,12 @@ def _tq_decode_stage1(
             Block_table_ptr + bt_base + page_idx,
             mask=kv_mask,
             other=0,
-        ).to(tl.int64)
+        )
 
         slot_bases = (
             block_nums * stride_cache_block
-            + page_off.to(tl.int64) * stride_cache_pos
-            + tl.cast(kv_head, tl.int64) * stride_cache_head
+            + page_off * stride_cache_pos
+            + kv_head * stride_cache_head
         )
 
         # ============================================================
@@ -356,11 +350,11 @@ def _tq_full_dequant_kv(
 
     page_idx = pos // BLOCK_SIZE
     page_off = pos % BLOCK_SIZE
-    block_num = tl.load(Block_table_ptr + bid * stride_bt_b + page_idx).to(tl.int64)
+    block_num = tl.load(Block_table_ptr + bid * stride_bt_b + page_idx)
     slot_base = (
         block_num * stride_cache_block
-        + tl.cast(page_off, tl.int64) * stride_cache_pos
-        + tl.cast(hid, tl.int64) * stride_cache_head
+        + page_off * stride_cache_pos
+        + hid * stride_cache_head
     )
 
     d_offs = tl.arange(0, BLOCK_D)

vllm/v1/attention/ops/triton_turboquant_store.py

@@ -174,13 +174,10 @@ def _tq_fused_store_fp8(
     slot = tl.load(Slot_mapping_ptr + token_idx)
     if slot < 0:
         return
-    blk = (slot // BLOCK_SIZE).to(tl.int64)
-    off = (slot % BLOCK_SIZE).to(tl.int64)
-    head_idx_i64 = tl.cast(head_idx, tl.int64)
+    blk = slot // BLOCK_SIZE
+    off = slot % BLOCK_SIZE
     slot_base = (
-        blk * stride_cache_block
-        + off * stride_cache_pos
-        + head_idx_i64 * stride_cache_head
+        blk * stride_cache_block + off * stride_cache_pos + head_idx * stride_cache_head
     )
 
     base = pid * D
@@ -262,13 +259,10 @@ def _tq_fused_store_mse(
     slot = tl.load(Slot_mapping_ptr + token_idx)
     if slot < 0:
         return
-    blk = (slot // BLOCK_SIZE).to(tl.int64)
-    off = (slot % BLOCK_SIZE).to(tl.int64)
-    head_idx_i64 = tl.cast(head_idx, tl.int64)
+    blk = slot // BLOCK_SIZE
+    off = slot % BLOCK_SIZE
     slot_base = (
-        blk * stride_cache_block
-        + off * stride_cache_pos
-        + head_idx_i64 * stride_cache_head
+        blk * stride_cache_block + off * stride_cache_pos + head_idx * stride_cache_head
     )
 
     base = pid * D

vllm/v1/worker/gpu_model_runner.py

@@ -3758,15 +3758,6 @@ class GPUModelRunner(
 
         return slot_mappings_by_gid, slot_mappings_by_layer
 
-    def _is_all_reqs_chunked_prefill(self) -> bool:
-        """Check if all scheduled requests are marked to discard sampled tokens.
-
-        This is true when `discard_request_mask` is set for every scheduled
-        request (e.g., for chunked prefill requests that are not the last
-        prefill chunk)."""
-        num_reqs = self.input_batch.num_reqs
-        return bool(self.discard_request_mask.np[:num_reqs].all())
-
     @torch.inference_mode()
     def execute_model(
         self,
@@ -4370,12 +4361,9 @@ class GPUModelRunner(
         assert sampled_token_ids.dim() == 2 and sampled_token_ids.shape[-1] == 1, (
             "PP+async expects sampled_token_ids to have shape [num_reqs, 1]"
         )
-        # Skip for chunked prefill: sampled tokens are dummy
-        # and will be discarded, no need to broadcast.
-        if not self._is_all_reqs_chunked_prefill():
-            torch.distributed.broadcast(
-                sampled_token_ids, src=pp.rank, group=pp.device_group
-            )
+        torch.distributed.broadcast(
+            sampled_token_ids, src=pp.rank, group=pp.device_group
+        )
 
     def _pp_receive_prev_sampled_token_ids_to_input_batch(self) -> None:
         """Receive sampled token ids broadcast from last PP stage"""
@@ -4384,9 +4372,7 @@ class GPUModelRunner(
         num_reqs = self.input_batch.num_reqs
         # `prev_sampled_token_ids` is expected to have shape [num_reqs, 1].
         recv = torch.empty((num_reqs, 1), dtype=torch.int32, device=self.device)
-        # skip for chunked prefill.
-        if not self._is_all_reqs_chunked_prefill():
-            torch.distributed.broadcast(recv, src=pp.last_rank, group=pp.device_group)
+        torch.distributed.broadcast(recv, src=pp.last_rank, group=pp.device_group)
         self.input_batch.prev_sampled_token_ids = recv
 
         # construct `prev_req_id_to_index` here so `_prepare_input_ids`