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,27 +134,4 @@ void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out,
                               torch::stable::Tensor& input,
                               torch::stable::Tensor& input_global_scale);
 
-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);
-
-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);
-
-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,422 +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/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);
-      });
-}

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,9 +233,6 @@ 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));
-  ops.impl("mxfp4_experts_quant", TORCH_BOX(&mxfp4_experts_quant));
-  ops.impl("silu_and_mul_mxfp4_experts_quant",
-           TORCH_BOX(&silu_and_mul_mxfp4_experts_quant));
 
   // W4A8 ops: impl registrations are in the source files
   // (w4a8_mm_entry.cu and w4a8_grouped_mm_entry.cu)

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

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

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

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

@@ -21,7 +21,6 @@ from vllm.model_executor.layers.quantization.turboquant.config import (
 from vllm.model_executor.layers.quantization.turboquant.quantizer import (
     generate_wht_signs,
 )
-from vllm.platforms import current_platform
 from vllm.utils.math_utils import next_power_of_2
 
 # ============================================================================
@@ -346,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:
@@ -362,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}"
         )
 
@@ -387,7 +385,7 @@ 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
 
@@ -405,31 +403,31 @@ 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")
+@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_wht_orthonormal(self, dim):
         """signs * H must be orthonormal: (signs*H) @ (signs*H)^T = I."""
-        signs = generate_wht_signs(dim, seed=42, device=DEVICE_TYPE)
-        H = _build_hadamard(dim, DEVICE_TYPE)
+        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=DEVICE_TYPE), atol=1e-5), (
+        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_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=DEVICE_TYPE)
-        H = _build_hadamard(dim, DEVICE_TYPE)
+        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=DEVICE_TYPE), atol=1e-5), (
+        assert torch.allclose(result, torch.eye(dim, device="cuda"), atol=1e-5), (
             f"WHT rotation not self-inverse for dim={dim}"
         )
 
@@ -456,7 +454,7 @@ class TestWHTRotation:
 # ============================================================================
 
 
-@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."""
 
@@ -489,11 +487,11 @@ class TestStoreDecodeRoundTrip:
         block_size = 16
         num_blocks = 1
 
-        device = torch.device(DEVICE_TYPE)
+        device = torch.device("cuda")
 
         # Generate rotation
         signs = generate_wht_signs(D, seed=42, device=device)
-        H = _build_hadamard(D, DEVICE_TYPE)
+        H = _build_hadamard(D, "cuda")
         PiT = (signs.unsqueeze(1) * H).contiguous().float()
         Pi = PiT.T.contiguous()
 

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

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

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

@@ -507,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,
@@ -517,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
@@ -554,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,
@@ -566,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()
@@ -728,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,
@@ -740,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`