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

CMakeLists.txt

@@ -1240,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/sgl_flash_kernel_ops.h

@@ -0,0 +1,58 @@
+/* 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,116 @@
+/* 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)...);
+  };
+}

tests/v1/attention/test_cutlass_fa3_sparse_backend.py

@@ -0,0 +1,972 @@
+# 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."""
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        CutlassFA3MLASparseImpl,
+    )
+
+    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).
+    """
+    import numpy as np
+    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}, "
+            f"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)
+    """
+    import numpy as np
+    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)."""
+    import numpy as np
+    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.
+    """
+    import numpy as np
+    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).
+    """
+    import numpy as np
+    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/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/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/v1/attention/backends/mla/cutlass_fa3_sparse.py

@@ -0,0 +1,643 @@
+# 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
+
+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
+
+import logging
+
+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 N < padded_heads:
+            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,10 @@ 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/ops/cutlass_fa3.py

@@ -0,0 +1,163 @@
+# 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)
+"""
+from typing import Optional
+
+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: Optional[torch.Tensor] = None,
+    v: Optional[torch.Tensor] = None,
+    qv: Optional[torch.Tensor] = None,
+    rotary_cos: Optional[torch.Tensor] = None,
+    rotary_sin: Optional[torch.Tensor] = None,
+    cache_seqlens: Optional[torch.Tensor] = None,
+    cache_batch_idx: Optional[torch.Tensor] = None,
+    cache_leftpad: Optional[torch.Tensor] = None,
+    page_table: Optional[torch.Tensor] = None,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k_new: Optional[torch.Tensor] = None,
+    max_seqlen_q: Optional[int] = None,
+    rotary_seqlens: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    window_size: tuple[int, int] = (-1, -1),
+    attention_chunk: Optional[int] = None,
+    softcap: float = 0.0,
+    rotary_interleaved: bool = True,
+    scheduler_metadata: Optional[torch.Tensor] = None,
+    num_splits: int = 0,
+    pack_gqa: Optional[bool] = None,
+    sm_margin: int = 0,
+    return_softmax_lse: bool = False,
+    sinks: Optional[torch.Tensor] = 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