From a60418e6fb4e2ababec06b5614c370f323bbe7be Mon Sep 17 00:00:00 2001 From: Alexander Matveev Date: Wed, 15 Apr 2026 20:07:18 +0000 Subject: [PATCH] Sparse MLA on Hopper: Use SGLang's kernel for the sparse mla low latency runs Signed-off-by: Alexander Matveev --- CMakeLists.txt | 1 + cmake/external_projects/cutlass_fa3.cmake | 163 +++ csrc/cutlass_fa3_extension.cc | 72 + csrc/sgl_flash_kernel_ops.h | 45 + csrc/sgl_kernel_torch_shim.h | 121 ++ docs/design/attention_backends.md | 11 +- .../attention/test_cutlass_fa3_sparse.py | 1225 +++++++++++++++++ .../test_cutlass_fa3_sparse_backend.py | 968 +++++++++++++ .../layers/attention/mla_attention.py | 10 + vllm/platforms/cuda.py | 16 + .../backends/mla/cutlass_fa3_sparse.py | 633 +++++++++ vllm/v1/attention/backends/registry.py | 3 + vllm/v1/attention/ops/cutlass_fa3.py | 162 +++ 13 files changed, 3420 insertions(+), 10 deletions(-) create mode 100644 cmake/external_projects/cutlass_fa3.cmake create mode 100644 csrc/cutlass_fa3_extension.cc create mode 100644 csrc/sgl_flash_kernel_ops.h create mode 100644 csrc/sgl_kernel_torch_shim.h create mode 100644 tests/kernels/attention/test_cutlass_fa3_sparse.py create mode 100644 tests/v1/attention/test_cutlass_fa3_sparse_backend.py create mode 100644 vllm/v1/attention/backends/mla/cutlass_fa3_sparse.py create mode 100644 vllm/v1/attention/ops/cutlass_fa3.py diff --git a/CMakeLists.txt b/CMakeLists.txt index fd4314bec52..bbda9cd97ee 100644 --- a/CMakeLists.txt +++ b/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 diff --git a/cmake/external_projects/cutlass_fa3.cmake b/cmake/external_projects/cutlass_fa3.cmake new file mode 100644 index 00000000000..45cd9be2b79 --- /dev/null +++ b/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 + $<$:-UPy_LIMITED_API> + $<$:-UPy_LIMITED_API> + $<$:-std=c++17> + $<$:-std=c++17> + $<$:--use_fast_math> + $<$:--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})") diff --git a/csrc/cutlass_fa3_extension.cc b/csrc/cutlass_fa3_extension.cc new file mode 100644 index 00000000000..a1b1996a0f4 --- /dev/null +++ b/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 +#include +#include +#include + +#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); +} diff --git a/csrc/sgl_flash_kernel_ops.h b/csrc/sgl_flash_kernel_ops.h new file mode 100644 index 00000000000..2ba942f47d0 --- /dev/null +++ b/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 +#include +#include +#include + +#include + +#include "sgl_kernel_torch_shim.h" + +/* + * From flash-attention (sgl-attn fork) + */ +std::tuple 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 k_new_, std::optional v_new_, + std::optional q_v_, // MLA value projection query + std::optional out_, std::optional cu_seqlens_q_, + std::optional cu_seqlens_k_, + std::optional cu_seqlens_k_new_, + std::optional seqused_q_, std::optional seqused_k_, + std::optional max_seqlen_q_, std::optional max_seqlen_k_, + std::optional page_table_, + std::optional kv_batch_idx_, + std::optional leftpad_k_, std::optional rotary_cos_, + std::optional rotary_sin_, + std::optional seqlens_rotary_, + std::optional q_descale_, std::optional k_descale_, + std::optional v_descale_, std::optional 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 scheduler_metadata_, int64_t num_splits, + std::optional pack_gqa_, int64_t sm_margin, + std::optional& sinks_); diff --git a/csrc/sgl_kernel_torch_shim.h b/csrc/sgl_kernel_torch_shim.h new file mode 100644 index 00000000000..c15aee4c28e --- /dev/null +++ b/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 + +/** + * 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 &` + * - `std::optional &` + * So we convert them to (respectively): + * - `int64_t` + * - `double` + * - `const std::optional&` + * - `const std::optional&` + */ + +template +struct pytorch_library_compatible_type { + using type = T; + static T convert_from_type(T arg) { return arg; } +}; + +template +using pytorch_library_compatible_type_t = + typename pytorch_library_compatible_type::type; + +template +T convert_from_pytorch_compatible_type( + pytorch_library_compatible_type_t arg) { + return pytorch_library_compatible_type::convert_from_type(arg); +} + +// Map `c10::optional &` -> `const c10::optional&` +// (NOTE: this is bit unsafe but non of the ops in flash_attn mutate +// the optional container) +template +struct pytorch_library_compatible_type&> { + using type = const c10::optional&; + static c10::optional& convert_from_type(const c10::optional& arg) { + return const_cast&>(arg); + } +}; + +// Map `c10::optional` -> +// `c10::optional>` +// (NOTE: tested for `c10::optional` -> `c10::optional`) +template +struct pytorch_library_compatible_type> { + using type = c10::optional>; + static c10::optional> convert_from_type( + c10::optional arg) { + return arg; + } +}; + +// Map `c10::optional&` -> `const c10::optional&` +template <> +struct pytorch_library_compatible_type&> { + using type = const c10::optional&; + static c10::optional& convert_from_type( + const c10::optional& arg) { + return const_cast&>( + reinterpret_cast&>(arg)); + } +}; + +// Map `int` -> `int64_t` +template <> +struct pytorch_library_compatible_type { + using type = int64_t; + static int convert_from_type(int64_t arg) { + TORCH_CHECK(arg <= std::numeric_limits::max(), + "int64_t value is too large to be converted to int"); + TORCH_CHECK(arg >= std::numeric_limits::min(), + "int64_t value is too small to be converted to int"); + return arg; + } +}; + +// Map `float` -> `double` +template <> +struct pytorch_library_compatible_type { + using type = double; + static float convert_from_type(double arg) { + TORCH_CHECK(std::abs(arg) <= std::numeric_limits::max(), + "double value is too large to be converted to float"); + return arg; + } +}; + +// +// Shim Utils +// + +template +auto make_pytorch_shim(Ret (*fun)(Args... args)) { + return [fun](pytorch_library_compatible_type_t... args) { + return fun(convert_from_pytorch_compatible_type(args)...); + }; +} diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index d10c23038a6..e6f9f4b08a3 100644 --- a/docs/design/attention_backends.md +++ b/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 | diff --git a/tests/kernels/attention/test_cutlass_fa3_sparse.py b/tests/kernels/attention/test_cutlass_fa3_sparse.py new file mode 100644 index 00000000000..f6511ba03ae --- /dev/null +++ b/tests/kernels/attention/test_cutlass_fa3_sparse.py @@ -0,0 +1,1225 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the vendored CUTLASS FA3 MLA sparse attention kernel. + +These tests verify the kernel wrapper works correctly in isolation, +following the pattern of test_flashmla_sparse.py smoke tests. + +Tests cover: + 1. Extension availability + 2. Basic MLA decode smoke test + 3. Batched decode + 4. Correctness vs SDPA reference + 5. Invalid indices handling + 6. Variable sequence lengths + 7. num_splits variants + 8. softcap parameter +""" + +import pytest +import torch +import torch.nn.functional as F + +from vllm.v1.attention.ops.cutlass_fa3 import ( + flash_attn_with_kvcache, + is_cutlass_fa3_available, +) + +# Skip all tests if FA3 is not available (non-SM90 or CUDA < 12.4) +pytestmark = pytest.mark.skipif( + not is_cutlass_fa3_available(), + reason="CUTLASS FA3 not available (requires CUDA >= 12.4, SM90)", +) + +# Common MLA dimensions for DeepSeek-V3.2 +NUM_HEADS = 16 +HEADDIM_QK = 64 # RoPE component +HEADDIM_V = 512 # kv_lora_rank (NoPE component) +HEAD_SIZE = HEADDIM_V + HEADDIM_QK # 576 +SOFTMAX_SCALE = 192 ** (-0.5) # qk_head_dim = 192 + + +def _make_kv_cache(kv_pool_size: int, dtype=torch.bfloat16, device="cuda"): + """Create mock KV cache with page_size=1 format. + + FA3 expects paged KV as 4D: [num_pages, page_size=1, num_kv_heads=1, dim] + """ + k_rope_cache = torch.randn( + kv_pool_size, 1, 1, HEADDIM_QK, dtype=dtype, device=device + ) + v_cache = torch.randn(kv_pool_size, 1, 1, HEADDIM_V, dtype=dtype, device=device) + return k_rope_cache, v_cache + + +def _make_page_table(T, topk, kv_pool_size, device="cuda"): + """Create page table with unique random indices per token.""" + page_table = torch.stack( + [torch.randperm(kv_pool_size, device=device)[:topk] for _ in range(T)] + ).to(torch.int32) + return page_table + + +def _make_cu_seqlens(cache_seqlens, device="cuda"): + """Create cumulative sequence length tensors for FA3.""" + T = cache_seqlens.shape[0] + cu_seqlens_q = torch.arange(0, T + 1, dtype=torch.int32, device=device) + cu_seqlens_k = torch.cat( + [ + torch.zeros(1, dtype=torch.int32, device=device), + torch.cumsum(cache_seqlens, dim=0), + ] + ) + return cu_seqlens_q, cu_seqlens_k + + +# ─── TEST 1.1: Availability ────────────────────────────────────────── + + +def test_cutlass_fa3_availability(): + """Verify the _cutlass_fa3_C extension loads successfully.""" + assert is_cutlass_fa3_available() + assert hasattr(torch.ops, "_cutlass_fa3_C") + assert hasattr(torch.ops._cutlass_fa3_C, "fwd") + + +# ─── TEST 1.2: Basic Smoke Test ────────────────────────────────────── + + +def test_fa3_mla_decode_smoke(): + """Basic smoke test: 1 token, 1 request, small topk.""" + device = "cuda" + T = 1 + topk = 128 + kv_pool_size = 256 + + q_rope = torch.randn(T, NUM_HEADS, HEADDIM_QK, dtype=torch.bfloat16, device=device) + q_nope = torch.randn(T, NUM_HEADS, HEADDIM_V, dtype=torch.bfloat16, device=device) + k_rope_cache, v_cache = _make_kv_cache(kv_pool_size, device=device) + page_table = _make_page_table(T, topk, kv_pool_size, device=device) + cache_seqlens = torch.tensor([topk], dtype=torch.int32, device=device) + cu_seqlens_q, cu_seqlens_k = _make_cu_seqlens(cache_seqlens, device=device) + + out = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope_cache, + v_cache=v_cache, + qv=q_nope, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k_new=cu_seqlens_k, + max_seqlen_q=1, + softmax_scale=SOFTMAX_SCALE, + causal=True, + window_size=(-1, -1), + softcap=0.0, + num_splits=0, + ) + + assert out.shape == (T, NUM_HEADS, HEADDIM_V) + assert out.dtype == torch.bfloat16 + assert not out.isnan().any(), "Output contains NaN values" + assert not out.isinf().any(), "Output contains Inf values" + + +# ─── TEST 1.3: Batched Decode ──────────────────────────────────────── + + +@pytest.mark.parametrize("batch_size", [1, 2, 4, 8, 16, 32]) +def test_fa3_mla_decode_batched(batch_size): + """Verify batched decode with multiple tokens.""" + device = "cuda" + T = batch_size + topk = 256 + kv_pool_size = 4096 + + q_rope = torch.randn(T, NUM_HEADS, HEADDIM_QK, dtype=torch.bfloat16, device=device) + q_nope = torch.randn(T, NUM_HEADS, HEADDIM_V, dtype=torch.bfloat16, device=device) + k_rope_cache, v_cache = _make_kv_cache(kv_pool_size, device=device) + page_table = _make_page_table(T, topk, kv_pool_size, device=device) + cache_seqlens = torch.full((T,), topk, dtype=torch.int32, device=device) + cu_seqlens_q, cu_seqlens_k = _make_cu_seqlens(cache_seqlens, device=device) + + out = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope_cache, + v_cache=v_cache, + qv=q_nope, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k_new=cu_seqlens_k, + max_seqlen_q=1, + softmax_scale=SOFTMAX_SCALE, + causal=True, + num_splits=0, + ) + + assert out.shape == (T, NUM_HEADS, HEADDIM_V) + assert not out.isnan().any() + + +# ─── TEST 1.4: Correctness vs SDPA Reference ───────────────────────── + + +def test_fa3_mla_correctness_vs_reference(): + """Verify numerical correctness against PyTorch SDPA reference. + + For each batch element, we manually compute attention using the same + Q, K, V data and compare against FA3 output. + """ + device = "cuda" + T = 4 + topk = 64 # small for reference tractability + kv_pool_size = 256 + + q_rope = torch.randn(T, NUM_HEADS, HEADDIM_QK, dtype=torch.bfloat16, device=device) + q_nope = torch.randn(T, NUM_HEADS, HEADDIM_V, dtype=torch.bfloat16, device=device) + k_rope_cache, v_cache = _make_kv_cache(kv_pool_size, device=device) + page_table = _make_page_table(T, topk, kv_pool_size, device=device) + cache_seqlens = torch.full((T,), topk, dtype=torch.int32, device=device) + cu_seqlens_q, cu_seqlens_k = _make_cu_seqlens(cache_seqlens, device=device) + + # FA3 output + fa3_out = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope_cache, + v_cache=v_cache, + qv=q_nope, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k_new=cu_seqlens_k, + max_seqlen_q=1, + softmax_scale=SOFTMAX_SCALE, + causal=True, + num_splits=1, # deterministic + ) + # fa3_out is already [T, N, 512] in varlen mode + + # Reference: per-token SDPA + for b in range(T): + indices = page_table[b] # [topk] + # Gather KV from cache (4D: [topk,1,1,dim] -> squeeze to [topk,dim]) + k_rope_gathered = k_rope_cache[indices].squeeze(1).squeeze(1) # [topk, 64] + v_gathered = v_cache[indices].squeeze(1).squeeze(1) # [topk, 512] + + # Full Q: [q_nope | q_rope] concatenated + q_r = q_rope[b] # [N, 64] + q_n = q_nope[b] # [N, 512] + + # Full K: [k_nope_from_v | k_rope] — for MLA, K_nope = V (latent) + k_full = torch.cat([v_gathered, k_rope_gathered], dim=-1) # [topk, 576] + + # Expand for MQA: [topk, 1, dim] -> [topk, N, dim] + k_expanded = k_full.unsqueeze(1).expand(-1, NUM_HEADS, -1) # [topk, N, 576] + v_expanded = v_gathered.unsqueeze(1).expand(-1, NUM_HEADS, -1) # [topk, N, 512] + + # Q full: [q_nope | q_rope] + q_full = torch.cat([q_n, q_r], dim=-1) # [N, 576] + + # SDPA expects: Q[batch, heads, seq_q, dim], K[batch, heads, seq_k, dim] + # Q: [1, N, 1, 576], K: [1, N, topk, 576], V: [1, N, topk, 512] + ref_out = F.scaled_dot_product_attention( + q_full.unsqueeze(0).unsqueeze(2), # [1, N, 1, 576] + k_expanded.transpose(0, 1).unsqueeze(0), # [1, N, topk, 576] + v_expanded.transpose(0, 1).unsqueeze(0), # [1, N, topk, 512] + scale=SOFTMAX_SCALE, + ) # -> [1, N, 1, 512] + ref_out = ref_out.squeeze(0).squeeze(1) # [N, 512] + + # Compare (BF16 tolerances) + torch.testing.assert_close( + fa3_out[b].float(), ref_out.float(), rtol=0.02, atol=0.02 + ) + + +# ─── TEST 1.5: Invalid Indices ─────────────────────────────────────── + + +def test_fa3_mla_with_invalid_indices(): + """Verify handling of -1 (padding) entries in page_table.""" + device = "cuda" + T = 1 + topk_valid = 64 + topk_total = 256 + kv_pool_size = 512 + + q_rope = torch.randn(T, NUM_HEADS, HEADDIM_QK, dtype=torch.bfloat16, device=device) + q_nope = torch.randn(T, NUM_HEADS, HEADDIM_V, dtype=torch.bfloat16, device=device) + k_rope_cache, v_cache = _make_kv_cache(kv_pool_size, device=device) + + # Mix of valid and -1 entries + page_table = torch.full((T, topk_total), -1, dtype=torch.int32, device=device) + page_table[0, :topk_valid] = torch.randperm(kv_pool_size, device=device)[ + :topk_valid + ].to(torch.int32) + + cache_seqlens = torch.tensor([topk_valid], dtype=torch.int32, device=device) + cu_seqlens_q, cu_seqlens_k = _make_cu_seqlens(cache_seqlens, device=device) + + out = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope_cache, + v_cache=v_cache, + qv=q_nope, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k_new=cu_seqlens_k, + max_seqlen_q=1, + softmax_scale=SOFTMAX_SCALE, + causal=True, + num_splits=0, + ) + + assert out.shape == (T, NUM_HEADS, HEADDIM_V) + assert not out.isnan().any() + assert not out.isinf().any() + + +# ─── TEST 1.6: Variable Sequence Lengths ───────────────────────────── + + +@pytest.mark.parametrize( + "seq_lens", + [ + [100, 200, 500, 1000], + [1, 1, 1, 1], + [2048, 2048, 2048, 2048], + [10, 2048, 50, 1500], + ], +) +def test_fa3_mla_variable_seqlens(seq_lens): + """Verify with variable cache_seqlens per batch element.""" + device = "cuda" + T = len(seq_lens) + topk = 2048 + kv_pool_size = 8192 + + q_rope = torch.randn(T, NUM_HEADS, HEADDIM_QK, dtype=torch.bfloat16, device=device) + q_nope = torch.randn(T, NUM_HEADS, HEADDIM_V, dtype=torch.bfloat16, device=device) + k_rope_cache, v_cache = _make_kv_cache(kv_pool_size, device=device) + + page_table = torch.full((T, topk), 0, dtype=torch.int32, device=device) + actual_seqlens = [] + for i, sl in enumerate(seq_lens): + actual_topk = min(sl, topk) + actual_seqlens.append(actual_topk) + page_table[i, :actual_topk] = torch.randperm(kv_pool_size, device=device)[ + :actual_topk + ].to(torch.int32) + + cache_seqlens = torch.tensor(actual_seqlens, dtype=torch.int32, device=device) + cu_seqlens_q, cu_seqlens_k = _make_cu_seqlens(cache_seqlens, device=device) + + out = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope_cache, + v_cache=v_cache, + qv=q_nope, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k_new=cu_seqlens_k, + max_seqlen_q=1, + softmax_scale=SOFTMAX_SCALE, + causal=True, + num_splits=0, + ) + + assert out.shape == (T, NUM_HEADS, HEADDIM_V) + assert not out.isnan().any() + + +# ─── TEST 1.7: num_splits Variants ─────────────────────────────────── + + +@pytest.mark.parametrize("num_splits", [0, 1, 2, 4]) +def test_fa3_mla_num_splits(num_splits): + """Verify different num_splits values produce valid results.""" + device = "cuda" + T = 4 + topk = 256 + kv_pool_size = 4096 + + q_rope = torch.randn(T, NUM_HEADS, HEADDIM_QK, dtype=torch.bfloat16, device=device) + q_nope = torch.randn(T, NUM_HEADS, HEADDIM_V, dtype=torch.bfloat16, device=device) + k_rope_cache, v_cache = _make_kv_cache(kv_pool_size, device=device) + page_table = _make_page_table(T, topk, kv_pool_size, device=device) + cache_seqlens = torch.full((T,), topk, dtype=torch.int32, device=device) + cu_seqlens_q, cu_seqlens_k = _make_cu_seqlens(cache_seqlens, device=device) + + out = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope_cache, + v_cache=v_cache, + qv=q_nope, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k_new=cu_seqlens_k, + max_seqlen_q=1, + softmax_scale=SOFTMAX_SCALE, + causal=True, + num_splits=num_splits, + ) + + assert out.shape == (T, NUM_HEADS, HEADDIM_V) + assert not out.isnan().any() + + +# ─── TEST 1.8: Softcap Parameter ───────────────────────────────────── + + +@pytest.mark.parametrize("softcap", [0.0, 30.0, 50.0]) +def test_fa3_mla_softcap(softcap): + """Verify logits_soft_cap parameter works.""" + device = "cuda" + T = 2 + topk = 128 + kv_pool_size = 512 + + q_rope = torch.randn(T, NUM_HEADS, HEADDIM_QK, dtype=torch.bfloat16, device=device) + q_nope = torch.randn(T, NUM_HEADS, HEADDIM_V, dtype=torch.bfloat16, device=device) + k_rope_cache, v_cache = _make_kv_cache(kv_pool_size, device=device) + page_table = _make_page_table(T, topk, kv_pool_size, device=device) + cache_seqlens = torch.full((T,), topk, dtype=torch.int32, device=device) + cu_seqlens_q, cu_seqlens_k = _make_cu_seqlens(cache_seqlens, device=device) + + out = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope_cache, + v_cache=v_cache, + qv=q_nope, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k_new=cu_seqlens_k, + max_seqlen_q=1, + softmax_scale=SOFTMAX_SCALE, + causal=True, + softcap=softcap, + num_splits=0, + ) + + assert out.shape == (T, NUM_HEADS, HEADDIM_V) + assert not out.isnan().any(), "Softcap test output contains NaN" + + +# ─── TEST 1.9: Prefill Short Sequence (REGRESSION for Issue #1) ───── + + +@pytest.mark.parametrize("seq_len", [1, 2, 3, 4, 8]) +def test_fa3_mla_prefill_short_sequence(seq_len): + """REGRESSION TEST for Issue #1: illegal memory access on short prefill. + + Root cause: cache_seqlens was set to min(seq_len, topk) for ALL tokens, + but for prefill tokens at the start of a sequence, the indexer produces + fewer valid topk entries than cache_seqlens (due to causal masking). + FA3 then reads -1 (invalid) entries from page_table -> crash. + + Fix: Use valid_counts from triton_convert_req_index_to_global_index + as cache_seqlens, and clamp page_table -1 entries to 0. + + This test directly reproduces the bug scenario: + - T tokens, each with different numbers of valid topk entries + - Token 0 has 1 valid entry, token 1 has 2, etc. + - page_table has 0-padding beyond valid entries (simulating clamp) + - cache_seqlens = actual valid count per token (not min(seq_len, topk)) + """ + device = "cuda" + T = seq_len + topk = 2048 + kv_pool_size = 4096 + + q_rope = torch.randn(T, NUM_HEADS, HEADDIM_QK, dtype=torch.bfloat16, device=device) + q_nope = torch.randn(T, NUM_HEADS, HEADDIM_V, dtype=torch.bfloat16, device=device) + k_rope_cache, v_cache = _make_kv_cache(kv_pool_size, device=device) + + # Simulate prefill: token i has (i+1) valid topk entries + page_table = torch.zeros((T, topk), dtype=torch.int32, device=device) + valid_counts = [] + for i in range(T): + num_valid = i + 1 # causal: token i sees positions 0..i + valid_counts.append(num_valid) + page_table[i, :num_valid] = torch.randperm(kv_pool_size, device=device)[ + :num_valid + ].to(torch.int32) + + cache_seqlens = torch.tensor(valid_counts, dtype=torch.int32, device=device) + cu_seqlens_q, _ = _make_cu_seqlens(cache_seqlens, device=device) + + out = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope_cache, + v_cache=v_cache, + qv=q_nope, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + softmax_scale=SOFTMAX_SCALE, + causal=True, + num_splits=0, + ) + + assert out.shape == (T, NUM_HEADS, HEADDIM_V) + assert not out.isnan().any(), f"NaN in output for seq_len={seq_len}" + assert not out.isinf().any(), f"Inf in output for seq_len={seq_len}" + + +# ─── TEST 1.10: Page Table with Clamped -1 Entries ────────────────── + + +def test_fa3_mla_page_table_minus1_clamped(): + """Verify that page_table with -1 entries clamped to 0 doesn't crash. + + Tests the fix pattern: after converting topk indices to global cache + slot IDs, -1 entries are replaced with 0 (a valid page index), and + cache_seqlens is set to the actual valid count. + """ + device = "cuda" + T = 4 + topk = 256 + kv_pool_size = 1024 + + q_rope = torch.randn(T, NUM_HEADS, HEADDIM_QK, dtype=torch.bfloat16, device=device) + q_nope = torch.randn(T, NUM_HEADS, HEADDIM_V, dtype=torch.bfloat16, device=device) + k_rope_cache, v_cache = _make_kv_cache(kv_pool_size, device=device) + + # Create page_table with -1 entries then clamp to 0 + page_table = torch.full((T, topk), -1, dtype=torch.int32, device=device) + valid_per_token = [10, 50, 100, 200] + for i in range(T): + nv = valid_per_token[i] + page_table[i, :nv] = torch.randperm(kv_pool_size, device=device)[:nv].to( + torch.int32 + ) + + page_table = page_table.clamp(min=0) + cache_seqlens = torch.tensor(valid_per_token, dtype=torch.int32, device=device) + cu_seqlens_q, _ = _make_cu_seqlens(cache_seqlens, device=device) + + out = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope_cache, + v_cache=v_cache, + qv=q_nope, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + softmax_scale=SOFTMAX_SCALE, + causal=True, + num_splits=0, + ) + + assert out.shape == (T, NUM_HEADS, HEADDIM_V) + assert not out.isnan().any() + assert not out.isinf().any() + + +# ─── TEST 1.11: Single Token (Minimum Batch) ──────────────────────── + + +def test_fa3_mla_single_token(): + """Edge case: single token with cache_seqlens=1.""" + device = "cuda" + T = 1 + topk = 1 + kv_pool_size = 64 + + q_rope = torch.randn(T, NUM_HEADS, HEADDIM_QK, dtype=torch.bfloat16, device=device) + q_nope = torch.randn(T, NUM_HEADS, HEADDIM_V, dtype=torch.bfloat16, device=device) + k_rope_cache, v_cache = _make_kv_cache(kv_pool_size, device=device) + + page_table = torch.zeros((T, topk), dtype=torch.int32, device=device) + cache_seqlens = torch.tensor([1], dtype=torch.int32, device=device) + cu_seqlens_q = torch.arange(0, T + 1, dtype=torch.int32, device=device) + + out = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope_cache, + v_cache=v_cache, + qv=q_nope, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + softmax_scale=SOFTMAX_SCALE, + causal=True, + num_splits=0, + ) + + assert out.shape == (T, NUM_HEADS, HEADDIM_V) + assert not out.isnan().any() + + +# ─── TEST 1.12: cu_seqlens_k_new=None ─────────────────────────────── + + +def test_fa3_mla_cu_seqlens_k_none(): + """Verify FA3 works correctly when cu_seqlens_k_new=None. + + The fix passes None for cu_seqlens_k_new since it is unused when + k_new is None. This test verifies the kernel accepts None. + """ + device = "cuda" + T = 4 + topk = 128 + kv_pool_size = 512 + + q_rope = torch.randn(T, NUM_HEADS, HEADDIM_QK, dtype=torch.bfloat16, device=device) + q_nope = torch.randn(T, NUM_HEADS, HEADDIM_V, dtype=torch.bfloat16, device=device) + k_rope_cache, v_cache = _make_kv_cache(kv_pool_size, device=device) + page_table = _make_page_table(T, topk, kv_pool_size, device=device) + cache_seqlens = torch.full((T,), topk, dtype=torch.int32, device=device) + cu_seqlens_q = torch.arange(0, T + 1, dtype=torch.int32, device=device) + + out = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope_cache, + v_cache=v_cache, + qv=q_nope, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k_new=None, + max_seqlen_q=1, + softmax_scale=SOFTMAX_SCALE, + causal=True, + num_splits=0, + ) + + assert out.shape == (T, NUM_HEADS, HEADDIM_V) + assert not out.isnan().any() + + +# ─── TEST 1.13: Valid Counts Match Cache Seqlens ───────────────────── + + +def test_fa3_valid_counts_correctness(): + """Verify that using valid_counts as cache_seqlens produces correct results. + + Tests the complete fix flow: + 1. Create page_table with varying valid entries per token + 2. Use valid_counts (= number of non-(-1) entries) as cache_seqlens + 3. Clamp page_table -1 to 0 + 4. Run FA3 and verify correctness vs reference + """ + device = "cuda" + T = 4 + topk = 128 + kv_pool_size = 256 + + q_rope = torch.randn(T, NUM_HEADS, HEADDIM_QK, dtype=torch.bfloat16, device=device) + q_nope = torch.randn(T, NUM_HEADS, HEADDIM_V, dtype=torch.bfloat16, device=device) + k_rope_cache, v_cache = _make_kv_cache(kv_pool_size, device=device) + + valid_per_token = [1, 8, 32, 64] + page_table_raw = torch.full((T, topk), -1, dtype=torch.int32, device=device) + for i in range(T): + nv = valid_per_token[i] + page_table_raw[i, :nv] = torch.randperm(kv_pool_size, device=device)[:nv].to( + torch.int32 + ) + + page_table = page_table_raw.clamp(min=0) + cache_seqlens = torch.tensor(valid_per_token, dtype=torch.int32, device=device) + cu_seqlens_q = torch.arange(0, T + 1, dtype=torch.int32, device=device) + + out = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope_cache, + v_cache=v_cache, + qv=q_nope, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + softmax_scale=SOFTMAX_SCALE, + causal=True, + num_splits=1, + ) + + for b in range(T): + nv = valid_per_token[b] + indices_raw = page_table_raw[b, :nv] + k_gathered = k_rope_cache[indices_raw].squeeze(1).squeeze(1) + v_gathered = v_cache[indices_raw].squeeze(1).squeeze(1) + + q_r = q_rope[b] + q_n = q_nope[b] + k_full = torch.cat([v_gathered, k_gathered], dim=-1) + + k_exp = k_full.unsqueeze(1).expand(-1, NUM_HEADS, -1) + v_exp = v_gathered.unsqueeze(1).expand(-1, NUM_HEADS, -1) + q_full = torch.cat([q_n, q_r], dim=-1) + + ref_out = ( + F.scaled_dot_product_attention( + q_full.unsqueeze(0).unsqueeze(2), + k_exp.transpose(0, 1).unsqueeze(0), + v_exp.transpose(0, 1).unsqueeze(0), + scale=SOFTMAX_SCALE, + ) + .squeeze(0) + .squeeze(1) + ) + + torch.testing.assert_close( + out[b].float(), + ref_out.float(), + rtol=0.02, + atol=0.02, + msg=f"Token {b} (valid={nv}) mismatch", + ) + + +# ─── TEST 1.14: Max topk=2048 ─────────────────────────────────────── + + +def test_fa3_mla_max_topk(): + """Verify with maximum topk=2048 entries.""" + device = "cuda" + T = 2 + topk = 2048 + kv_pool_size = 4096 + + q_rope = torch.randn(T, NUM_HEADS, HEADDIM_QK, dtype=torch.bfloat16, device=device) + q_nope = torch.randn(T, NUM_HEADS, HEADDIM_V, dtype=torch.bfloat16, device=device) + k_rope_cache, v_cache = _make_kv_cache(kv_pool_size, device=device) + page_table = _make_page_table(T, topk, kv_pool_size, device=device) + cache_seqlens = torch.full((T,), topk, dtype=torch.int32, device=device) + cu_seqlens_q = torch.arange(0, T + 1, dtype=torch.int32, device=device) + + out = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope_cache, + v_cache=v_cache, + qv=q_nope, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + softmax_scale=SOFTMAX_SCALE, + causal=True, + num_splits=0, + ) + + assert out.shape == (T, NUM_HEADS, HEADDIM_V) + assert not out.isnan().any() + + +# ─── TEST 1.15: Mixed Batch (Prefill + Decode) ────────────────────── + + +def test_fa3_mla_mixed_batch_prefill_decode(): + """Mixed batch: 3 prefill tokens with causal masking + 1 decode token. + + This tests the scenario where a batch contains tokens from different + requests with DIFFERENT numbers of valid topk entries: + - Tokens 0-2: prefill with ascending valid counts [1, 2, 3] + - Token 3: decode with full topk valid entries + + The fix must handle each token's valid_count independently. + This test was added in V2-FIXED to close the mixed-batch test gap. + """ + device = "cuda" + T = 4 + topk = 256 + kv_pool_size = 4096 + + q_rope = torch.randn(T, NUM_HEADS, HEADDIM_QK, dtype=torch.bfloat16, device=device) + q_nope = torch.randn(T, NUM_HEADS, HEADDIM_V, dtype=torch.bfloat16, device=device) + k_rope_cache, v_cache = _make_kv_cache(kv_pool_size, device=device) + + # Mixed batch: 3 prefill tokens (ascending valid) + 1 decode (full topk) + valid_per_token = [1, 2, 3, topk] + page_table = torch.zeros((T, topk), dtype=torch.int32, device=device) + for i in range(T): + nv = valid_per_token[i] + page_table[i, :nv] = torch.randperm(kv_pool_size, device=device)[:nv].to( + torch.int32 + ) + + cache_seqlens = torch.tensor(valid_per_token, dtype=torch.int32, device=device) + cu_seqlens_q = torch.arange(0, T + 1, dtype=torch.int32, device=device) + + out = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope_cache, + v_cache=v_cache, + qv=q_nope, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + softmax_scale=SOFTMAX_SCALE, + causal=True, + num_splits=0, + ) + + assert out.shape == (T, NUM_HEADS, HEADDIM_V) + assert not out.isnan().any(), "Mixed batch output contains NaN" + assert not out.isinf().any(), "Mixed batch output contains Inf" + + +# ─── TEST 1.16: End-to-end KV Cache Write → FA3 Read ──────────────── + + +def test_fa3_mla_kv_cache_write_then_read(): + """End-to-end test: write KV to cache via concat_and_cache_mla, then + read via FA3 with page_table and valid_counts. + + This tests the complete flow from KV cache population through FA3 + kernel invocation, verifying that the written data is correctly read + back by FA3 when using the valid_counts fix with varying valid entries. + Added in V3 to close the end-to-end gap. + """ + device = "cuda" + T = 4 # 4 tokens (simulating causal prefill) + num_blocks = 4 + block_size = 64 + kv_lora_rank = 512 + qk_rope_head_dim = 64 + head_size = kv_lora_rank + qk_rope_head_dim # 576 + + # 1) Create BF16 KV cache and write known values + cache = torch.zeros( + num_blocks, + block_size, + head_size, + dtype=torch.bfloat16, + device=device, + ) + 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) + # Write to slots 0,1,2,3 + slot_mapping = torch.arange(T, 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, + ) + + # 2) Reshape cache for FA3 (page_size=1 format) + S = num_blocks * block_size + kv_flat = cache.reshape(S, head_size) + c_kv = kv_flat[:, :kv_lora_rank].reshape(S, 1, 1, kv_lora_rank) + k_rope = kv_flat[:, kv_lora_rank:].reshape(S, 1, 1, qk_rope_head_dim) + + # 3) Create Q and page_table simulating causal prefill + q_rope = torch.randn( + T, NUM_HEADS, qk_rope_head_dim, dtype=torch.bfloat16, device=device + ) + q_nope = torch.randn( + T, NUM_HEADS, kv_lora_rank, dtype=torch.bfloat16, device=device + ) + + # Causal page_table: token i sees slots 0..i, rest padded with 0 + topk = 128 + page_table = torch.zeros((T, topk), dtype=torch.int32, device=device) + valid_per_token = [] + for i in range(T): + nv = i + 1 # token i can attend to i+1 positions + valid_per_token.append(nv) + page_table[i, :nv] = torch.arange(nv, dtype=torch.int32, device=device) + + cache_seqlens = torch.tensor(valid_per_token, dtype=torch.int32, device=device) + cu_seqlens_q = torch.arange(0, T + 1, dtype=torch.int32, device=device) + + # 4) Run FA3 kernel + out = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope, + v_cache=c_kv, + qv=q_nope, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + softmax_scale=SOFTMAX_SCALE, + causal=True, + num_splits=1, # deterministic + ) + + # 5) Verify output shape, no NaN/Inf + assert out.shape == (T, NUM_HEADS, kv_lora_rank) + assert not out.isnan().any(), "E2E KV write+read output contains NaN" + assert not out.isinf().any(), "E2E KV write+read output contains Inf" + + # 6) Verify vs reference: each token's output should match SDPA + # using the SAME KV data that was written to cache + for b in range(T): + nv = valid_per_token[b] + # Gather the actual written KV from cache slots 0..nv-1 + k_gathered = k_rope[:nv].squeeze(1).squeeze(1) # [nv, 64] + v_gathered = c_kv[:nv].squeeze(1).squeeze(1) # [nv, 512] + + q_r = q_rope[b] # [N, 64] + q_n = q_nope[b] # [N, 512] + k_full = torch.cat([v_gathered, k_gathered], dim=-1) # [nv, 576] + + k_exp = k_full.unsqueeze(1).expand(-1, NUM_HEADS, -1) + v_exp = v_gathered.unsqueeze(1).expand(-1, NUM_HEADS, -1) + q_full = torch.cat([q_n, q_r], dim=-1) + + ref_out = ( + F.scaled_dot_product_attention( + q_full.unsqueeze(0).unsqueeze(2), + k_exp.transpose(0, 1).unsqueeze(0), + v_exp.transpose(0, 1).unsqueeze(0), + scale=SOFTMAX_SCALE, + ) + .squeeze(0) + .squeeze(1) + ) + + torch.testing.assert_close( + out[b].float(), + ref_out.float(), + rtol=0.02, + atol=0.02, + msg=f"E2E token {b} (valid={nv}): FA3 vs SDPA mismatch", + ) + + +# ─── TEST 1.17: Batch Size Gating — FA3 vs FlashMLA Routing ────────── + + +def test_batch_size_gating_constant(): + """Verify that MAX_BATCH_SIZE_FOR_FA3 is defined and is 16.""" + from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import ( + MAX_BATCH_SIZE_FOR_FA3, + ) + + assert MAX_BATCH_SIZE_FOR_FA3 == 16, ( + f"MAX_BATCH_SIZE_FOR_FA3 should be 16, got {MAX_BATCH_SIZE_FOR_FA3}" + ) + + +@pytest.mark.parametrize("batch_size", [1, 4, 8, 16]) +def test_fa3_small_batch_correctness(batch_size): + """Verify FA3 produces correct results for small batch sizes (<=16). + + These batch sizes should use the FA3 kernel path, which is faster + for small batches due to lower kernel launch overhead. + """ + device = "cuda" + T = batch_size + topk = 128 + kv_pool_size = 1024 + + q_rope = torch.randn(T, NUM_HEADS, HEADDIM_QK, dtype=torch.bfloat16, device=device) + q_nope = torch.randn(T, NUM_HEADS, HEADDIM_V, dtype=torch.bfloat16, device=device) + k_rope_cache, v_cache = _make_kv_cache(kv_pool_size, device=device) + page_table = _make_page_table(T, topk, kv_pool_size, device=device) + cache_seqlens = torch.full((T,), topk, dtype=torch.int32, device=device) + cu_seqlens_q, cu_seqlens_k = _make_cu_seqlens(cache_seqlens, device=device) + + out = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope_cache, + v_cache=v_cache, + qv=q_nope, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + softmax_scale=SOFTMAX_SCALE, + causal=True, + num_splits=1, + ) + + assert out.shape == (T, NUM_HEADS, HEADDIM_V) + assert not out.isnan().any(), f"NaN in FA3 output for batch_size={batch_size}" + assert not out.isinf().any(), f"Inf in FA3 output for batch_size={batch_size}" + + # Verify correctness against SDPA reference for first token + indices = page_table[0] + k_gathered = k_rope_cache[indices].squeeze(1).squeeze(1) + v_gathered = v_cache[indices].squeeze(1).squeeze(1) + q_full = torch.cat([q_nope[0], q_rope[0]], dim=-1) + k_full = torch.cat([v_gathered, k_gathered], dim=-1) + k_exp = k_full.unsqueeze(1).expand(-1, NUM_HEADS, -1) + v_exp = v_gathered.unsqueeze(1).expand(-1, NUM_HEADS, -1) + ref_out = ( + F.scaled_dot_product_attention( + q_full.unsqueeze(0).unsqueeze(2), + k_exp.transpose(0, 1).unsqueeze(0), + v_exp.transpose(0, 1).unsqueeze(0), + scale=SOFTMAX_SCALE, + ) + .squeeze(0) + .squeeze(1) + ) + torch.testing.assert_close( + out[0].float(), + ref_out.float(), + rtol=0.02, + atol=0.02, + msg=f"FA3 small batch (bs={batch_size}) token 0 mismatch", + ) + + +# ─── TEST 1.18: FlashMLA Fallback Available ────────────────────────── + + +def test_flashmla_sparse_fallback_available(): + """Verify that the FlashMLA BF16 sparse fallback is available on SM90.""" + from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import ( + _flashmla_sparse_available, + ) + + # On SM90 with FlashMLA compiled, the fallback should be available + # This test may skip if FlashMLA is not compiled (non-standard build) + if not _flashmla_sparse_available: + pytest.skip("FlashMLA sparse not available for fallback testing") + assert _flashmla_sparse_available + + +# ─── TEST 1.19: FlashMLA BF16 Fallback Correctness ────────────────── + + +@pytest.mark.parametrize("batch_size", [17, 32, 64]) +def test_flashmla_bf16_fallback_correctness(batch_size): + """End-to-end correctness test for FlashMLA BF16 sparse fallback. + + Verifies that the FlashMLA BF16 sparse prefill kernel (used as fallback + for batch sizes > MAX_BATCH_SIZE_FOR_FA3) produces correct results by + comparing against per-token SDPA reference. + + This tests the complete fallback path: + 1. Q concatenation [ql_nope | q_pe] -> [T, N, 576] + 2. Head padding N -> 64 + 3. KV reshape to [S, 1, 576] + 4. Index reshape to [T, 1, topk] + 5. flash_mla_sparse_fwd with topk_length + 6. Output unpadding + + NOTE: The FlashMLA sparse_prefill kernel requires topk % 128 == 0 + (assertion: topk % (2*B_TOPK) == 0 where B_TOPK=64). The real system + uses topk=2048 which satisfies this. We use topk=256 for tractability. + """ + from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import ( + _FLASHMLA_SM90_HEAD_PADDING, + _flashmla_sparse_available, + ) + + if not _flashmla_sparse_available: + pytest.skip("FlashMLA sparse not available for fallback testing") + + from vllm.v1.attention.ops.flashmla import flash_mla_sparse_fwd + + device = "cuda" + T = batch_size + # topk must be divisible by 128 (kernel constraint: topk % (2*B_TOPK) == 0) + topk = 256 + kv_pool_size = 512 + + q_rope = torch.randn(T, NUM_HEADS, HEADDIM_QK, dtype=torch.bfloat16, device=device) + q_nope = torch.randn(T, NUM_HEADS, HEADDIM_V, dtype=torch.bfloat16, device=device) + + # Create combined KV cache: [kv_pool_size, 1, 576] + # Layout: [kv_c_normed(512) | k_pe(64)] + kv_combined = torch.randn( + kv_pool_size, 1, HEADDIM_V + HEADDIM_QK, dtype=torch.bfloat16, device=device + ) + + # Create page table with unique random indices per token + page_table = torch.stack( + [torch.randperm(kv_pool_size, device=device)[:topk] for _ in range(T)] + ).to(torch.int32) + + # Use valid_counts = topk for all tokens to test full topk path. + # In the real system, valid_counts may be < topk for prefill tokens. + valid_counts = torch.full((T,), topk, dtype=torch.int32, device=device) + + # 1) Concatenate Q: [ql_nope(512) | q_pe(64)] -> [T, N, 576] + q_concat = torch.cat([q_nope, q_rope], dim=-1) # [T, N, 576] + + # 2) Pad heads to 64 + padded_heads = _FLASHMLA_SM90_HEAD_PADDING + if padded_heads > NUM_HEADS: + q_padded = q_concat.new_zeros((T, padded_heads, q_concat.shape[-1])) + q_padded[:, :NUM_HEADS, :] = q_concat + q_concat = q_padded + + # 3) Reshape indices for MQA: (T, topk) -> (T, 1, topk) + indices = page_table.unsqueeze(1) + + # 4) Call FlashMLA BF16 sparse prefill kernel + output = flash_mla_sparse_fwd( + q_concat, # [T, padded_heads, 576] + kv_combined, # [kv_pool_size, 1, 576] + indices, # [T, 1, topk] + SOFTMAX_SCALE, # 192**-0.5 + d_v=HEADDIM_V, # 512 + topk_length=valid_counts, # [T] + )[0] + + # 5) Unpad heads + output = output[:, :NUM_HEADS, :] # [T, N, 512] + + assert output.shape == (T, NUM_HEADS, HEADDIM_V), ( + f"Output shape mismatch: {output.shape} vs expected ({T}, {NUM_HEADS}, {HEADDIM_V})" + ) + assert not output.isnan().any(), ( + f"NaN in FlashMLA fallback output for bs={batch_size}" + ) + assert not output.isinf().any(), ( + f"Inf in FlashMLA fallback output for bs={batch_size}" + ) + + # 6) Verify correctness against SDPA reference for first 4 tokens + for b in range(min(4, T)): + idx = page_table[b] # [topk] + # Gather KV from combined cache: [topk, 1, 576] -> split + kv_gathered = kv_combined[idx].squeeze(1) # [topk, 576] + v_gathered = kv_gathered[:, :HEADDIM_V] # [topk, 512] (kv_c_normed) + k_gathered = kv_gathered # [topk, 576] (full key) + + # Full Q for this token + q_full = torch.cat([q_nope[b], q_rope[b]], dim=-1) # [N, 576] + + # Expand for MQA: (topk, 1, dim) -> (topk, N, dim) + k_expanded = k_gathered.unsqueeze(1).expand(-1, NUM_HEADS, -1) # [topk, N, 576] + v_expanded = v_gathered.unsqueeze(1).expand(-1, NUM_HEADS, -1) # [topk, N, 512] + + ref_out = ( + F.scaled_dot_product_attention( + q_full.unsqueeze(0).unsqueeze(2), # [1, N, 1, 576] + k_expanded.transpose(0, 1).unsqueeze(0), # [1, N, topk, 576] + v_expanded.transpose(0, 1).unsqueeze(0), # [1, N, topk, 512] + scale=SOFTMAX_SCALE, + ) + .squeeze(0) + .squeeze(1) + ) # [N, 512] + + torch.testing.assert_close( + output[b].float(), + ref_out.float(), + rtol=0.02, + atol=0.02, + msg=f"FlashMLA fallback (bs={batch_size}) token {b} mismatch", + ) + + +# ─── TEST 1.20: FA3 vs FlashMLA Cross-Kernel Consistency ───────────── + + +@pytest.mark.parametrize("batch_size", [1, 4, 8, 16]) +def test_fa3_vs_flashmla_cross_kernel_consistency(batch_size): + """Verify that FA3 and FlashMLA BF16 fallback produce consistent results. + + This is the ultimate correctness test: given identical inputs, both + the CUTLASS FA3 kernel and the FlashMLA BF16 sparse prefill kernel + should produce the same output (within numerical tolerance). + + Uses small batch sizes where both kernels can run, with topk=256 + (must be divisible by 128 for FlashMLA's kernel constraint). + + Includes T=16 (the MAX_BATCH_SIZE_FOR_FA3 boundary) to verify both + kernels agree at the exact gating threshold. + """ + from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import ( + _FLASHMLA_SM90_HEAD_PADDING, + _flashmla_sparse_available, + ) + + if not _flashmla_sparse_available: + pytest.skip("FlashMLA sparse not available for cross-kernel testing") + + from vllm.v1.attention.ops.flashmla import flash_mla_sparse_fwd + + device = "cuda" + T = batch_size + # topk must be divisible by 128 for FlashMLA kernel constraint + topk = 256 + kv_pool_size = 1024 + + q_rope = torch.randn(T, NUM_HEADS, HEADDIM_QK, dtype=torch.bfloat16, device=device) + q_nope = torch.randn(T, NUM_HEADS, HEADDIM_V, dtype=torch.bfloat16, device=device) + + # Create combined KV cache in BF16: [kv_pool_size, 576] + # This flat format is what both kernels see after reshaping + kv_flat = torch.randn( + kv_pool_size, HEADDIM_V + HEADDIM_QK, dtype=torch.bfloat16, device=device + ) + + # Create page table with unique random indices per token + page_table = torch.stack( + [torch.randperm(kv_pool_size, device=device)[:topk] for _ in range(T)] + ).to(torch.int32) + + cache_seqlens = torch.full((T,), topk, dtype=torch.int32, device=device) + cu_seqlens_q = torch.arange(0, T + 1, dtype=torch.int32, device=device) + + # === FA3 PATH === + # FA3 expects separate RoPE (k_cache) and NoPE (v_cache) in paged format + c_kv = kv_flat[:, :HEADDIM_V].reshape(kv_pool_size, 1, 1, HEADDIM_V) + k_rope = kv_flat[:, HEADDIM_V:].reshape(kv_pool_size, 1, 1, HEADDIM_QK) + + fa3_out = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope, + v_cache=c_kv, + qv=q_nope, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + softmax_scale=SOFTMAX_SCALE, + causal=True, + num_splits=1, # deterministic for comparison + ) + + # === FlashMLA BF16 PATH === + # FlashMLA expects concatenated Q and KV in [S, 1, 576] format + q_concat = torch.cat([q_nope, q_rope], dim=-1) # [T, N, 576] + + padded_heads = _FLASHMLA_SM90_HEAD_PADDING + if padded_heads > NUM_HEADS: + q_padded = q_concat.new_zeros((T, padded_heads, q_concat.shape[-1])) + q_padded[:, :NUM_HEADS, :] = q_concat + q_concat = q_padded + + kv_for_flashmla = kv_flat.reshape(kv_pool_size, 1, HEADDIM_V + HEADDIM_QK) + indices = page_table.unsqueeze(1) # [T, 1, topk] + + flashmla_out = flash_mla_sparse_fwd( + q_concat, + kv_for_flashmla, + indices, + SOFTMAX_SCALE, + d_v=HEADDIM_V, + topk_length=cache_seqlens, + )[0] + flashmla_out = flashmla_out[:, :NUM_HEADS, :] + + # === COMPARISON === + assert fa3_out.shape == flashmla_out.shape == (T, NUM_HEADS, HEADDIM_V), ( + f"Shape mismatch: FA3={fa3_out.shape}, FlashMLA={flashmla_out.shape}" + ) + assert not fa3_out.isnan().any(), "FA3 output contains NaN" + assert not flashmla_out.isnan().any(), "FlashMLA output contains NaN" + + # Cross-kernel comparison with slightly relaxed tolerance + # (different kernels may have different accumulation order) + torch.testing.assert_close( + fa3_out.float(), + flashmla_out.float(), + rtol=0.03, + atol=0.03, + msg=f"FA3 vs FlashMLA mismatch at bs={batch_size}", + ) diff --git a/tests/v1/attention/test_cutlass_fa3_sparse_backend.py b/tests/v1/attention/test_cutlass_fa3_sparse_backend.py new file mode 100644 index 00000000000..16ccf8d7ca0 --- /dev/null +++ b/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" + ) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index cbbf5f3c3ca..f54c5dcf26b 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/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 diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index d79d3191820..c8d57c54ff8 100644 --- a/vllm/platforms/cuda.py +++ b/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, diff --git a/vllm/v1/attention/backends/mla/cutlass_fa3_sparse.py b/vllm/v1/attention/backends/mla/cutlass_fa3_sparse.py new file mode 100644 index 00000000000..e6be8cfbebc --- /dev/null +++ b/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, + ) diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index f31edfafc38..86a078ec02f 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/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" diff --git a/vllm/v1/attention/ops/cutlass_fa3.py b/vllm/v1/attention/ops/cutlass_fa3.py new file mode 100644 index 00000000000..e3918522e58 --- /dev/null +++ b/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