@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f95903909f
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 22:15 UTC

[Kernel] FlashMLA integration (#13747)

Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Lucas Wilkinson committed Feb 26, 2025 at 21:35 UTC f95903909f07999ef77a9ce81290ad6b848fa08c
11 files changed +733 -86
CMakeLists.txt
+4 -73
@@ -575,77 +575,8 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
575 WITH_SOABI)
576 endif()
577
578 -# vllm-flash-attn currently only supported on CUDA
579 -if (NOT VLLM_GPU_LANG STREQUAL "CUDA")
580 - return()
578 +# For CUDA we also build and ship some external projects.
579 +if (VLLM_GPU_LANG STREQUAL "CUDA")
580 + include(cmake/external_projects/flashmla.cmake)
581 + include(cmake/external_projects/vllm_flash_attn.cmake)
582 endif ()
582 -
583 -# vLLM flash attention requires VLLM_GPU_ARCHES to contain the set of target
584 -# arches in the CMake syntax (75-real, 89-virtual, etc), since we clear the
585 -# arches in the CUDA case (and instead set the gencodes on a per file basis)
586 -# we need to manually set VLLM_GPU_ARCHES here.
587 -if(VLLM_GPU_LANG STREQUAL "CUDA")
588 - foreach(_ARCH ${CUDA_ARCHS})
589 - string(REPLACE "." "" _ARCH "${_ARCH}")
590 - list(APPEND VLLM_GPU_ARCHES "${_ARCH}-real")
591 - endforeach()
592 -endif()
593 -
594 -#
595 -# Build vLLM flash attention from source
596 -#
597 -# IMPORTANT: This has to be the last thing we do, because vllm-flash-attn uses the same macros/functions as vLLM.
598 -# Because functions all belong to the global scope, vllm-flash-attn's functions overwrite vLLMs.
599 -# They should be identical but if they aren't, this is a massive footgun.
600 -#
601 -# The vllm-flash-attn install rules are nested under vllm to make sure the library gets installed in the correct place.
602 -# To only install vllm-flash-attn, use --component _vllm_fa2_C (for FA2) or --component _vllm_fa3_C (for FA3).
603 -# If no component is specified, vllm-flash-attn is still installed.
604 -
605 -# If VLLM_FLASH_ATTN_SRC_DIR is set, vllm-flash-attn is installed from that directory instead of downloading.
606 -# This is to enable local development of vllm-flash-attn within vLLM.
607 -# It can be set as an environment variable or passed as a cmake argument.
608 -# The environment variable takes precedence.
609 -if (DEFINED ENV{VLLM_FLASH_ATTN_SRC_DIR})
610 - set(VLLM_FLASH_ATTN_SRC_DIR $ENV{VLLM_FLASH_ATTN_SRC_DIR})
611 -endif()
612 -
613 -if(VLLM_FLASH_ATTN_SRC_DIR)
614 - FetchContent_Declare(
615 - vllm-flash-attn SOURCE_DIR
616 - ${VLLM_FLASH_ATTN_SRC_DIR}
617 - BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
618 - )
619 -else()
620 - FetchContent_Declare(
621 - vllm-flash-attn
622 - GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
623 - GIT_TAG 720c94869cf2e0ff5a706e9c7f1dce0939686ade
624 - GIT_PROGRESS TRUE
625 - # Don't share the vllm-flash-attn build between build types
626 - BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
627 - )
628 -endif()
629 -
630 -
631 -# Fetch the vllm-flash-attn library
632 -FetchContent_MakeAvailable(vllm-flash-attn)
633 -message(STATUS "vllm-flash-attn is available at ${vllm-flash-attn_SOURCE_DIR}")
634 -
635 -# Copy over the vllm-flash-attn python files (duplicated for fa2 and fa3, in
636 -# case only one is built, in the case both are built redundant work is done)
637 -install(
638 - DIRECTORY ${vllm-flash-attn_SOURCE_DIR}/vllm_flash_attn/
639 - DESTINATION vllm_flash_attn
640 - COMPONENT _vllm_fa2_C
641 - FILES_MATCHING PATTERN "*.py"
642 -)
643 -
644 -install(
645 - DIRECTORY ${vllm-flash-attn_SOURCE_DIR}/vllm_flash_attn/
646 - DESTINATION vllm_flash_attn
647 - COMPONENT _vllm_fa3_C
648 - FILES_MATCHING PATTERN "*.py"
649 -)
650 -
651 -# Nothing after vllm-flash-attn, see comment about macros above
cmake/external_projects/flashmla.cmake new
+66
@@ -0,0 +1,66 @@
1 +include(FetchContent)
2 +
3 +# If FLASH_MLA_SRC_DIR is set, flash-mla is installed from that directory
4 +# instead of downloading.
5 +# It can be set as an environment variable or passed as a cmake argument.
6 +# The environment variable takes precedence.
7 +if (DEFINED ENV{FLASH_MLA_SRC_DIR})
8 + set(FLASH_MLA_SRC_DIR $ENV{FLASH_MLA_SRC_DIR})
9 +endif()
10 +
11 +if(FLASH_MLA_SRC_DIR)
12 + FetchContent_Declare(
13 + flashmla
14 + SOURCE_DIR ${FLASH_MLA_SRC_DIR}
15 + CONFIGURE_COMMAND ""
16 + BUILD_COMMAND ""
17 + )
18 +else()
19 + FetchContent_Declare(
20 + flashmla
21 + GIT_REPOSITORY https://github.com/vllm-project/FlashMLA.git
22 + GIT_TAG 575f7724b9762f265bbee5889df9c7d630801845
23 + GIT_PROGRESS TRUE
24 + CONFIGURE_COMMAND ""
25 + BUILD_COMMAND ""
26 + )
27 +endif()
28 +
29 +
30 +FetchContent_MakeAvailable(flashmla)
31 +message(STATUS "FlashMLA is available at ${flashmla_SOURCE_DIR}")
32 +
33 +# The FlashMLA kernels only work on hopper and require CUDA 12.3 or later.
34 +# Only build FlashMLA kernels if we are building for something compatible with
35 +# sm90a
36 +cuda_archs_loose_intersection(FLASH_MLA_ARCHS "9.0a" "${CUDA_ARCHS}")
37 +if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.3 AND FLASH_MLA_ARCHS)
38 + set(FlashMLA_SOURCES
39 + ${flashmla_SOURCE_DIR}/csrc/flash_api.cpp
40 + ${flashmla_SOURCE_DIR}/csrc/flash_fwd_mla_bf16_sm90.cu
41 + ${flashmla_SOURCE_DIR}/csrc/flash_fwd_mla_fp16_sm90.cu
42 + ${flashmla_SOURCE_DIR}/csrc/flash_fwd_mla_metadata.cu)
43 +
44 + set(FlashMLA_INCLUDES
45 + ${flashmla_SOURCE_DIR}/csrc/cutlass/include
46 + ${flashmla_SOURCE_DIR}/csrc/include)
47 +
48 + set_gencode_flags_for_srcs(
49 + SRCS "${FlashMLA_SOURCES}"
50 + CUDA_ARCHS "${FLASH_MLA_ARCHS}")
51 +
52 + define_gpu_extension_target(
53 + _flashmla_C
54 + DESTINATION vllm
55 + LANGUAGE ${VLLM_GPU_LANG}
56 + SOURCES ${FlashMLA_SOURCES}
57 + COMPILE_FLAGS ${VLLM_GPU_FLAGS}
58 + ARCHITECTURES ${VLLM_GPU_ARCHES}
59 + INCLUDE_DIRECTORIES ${FlashMLA_INCLUDES}
60 + USE_SABI 3
61 + WITH_SOABI)
62 +else()
63 + # Create an empty target for setup.py when not targeting sm90a systems
64 + add_custom_target(_flashmla_C)
65 +endif()
66 +
cmake/external_projects/vllm_flash_attn.cmake new
+67
@@ -0,0 +1,67 @@
1 +# vLLM flash attention requires VLLM_GPU_ARCHES to contain the set of target
2 +# arches in the CMake syntax (75-real, 89-virtual, etc), since we clear the
3 +# arches in the CUDA case (and instead set the gencodes on a per file basis)
4 +# we need to manually set VLLM_GPU_ARCHES here.
5 +if(VLLM_GPU_LANG STREQUAL "CUDA")
6 + foreach(_ARCH ${CUDA_ARCHS})
7 + string(REPLACE "." "" _ARCH "${_ARCH}")
8 + list(APPEND VLLM_GPU_ARCHES "${_ARCH}-real")
9 + endforeach()
10 +endif()
11 +
12 +#
13 +# Build vLLM flash attention from source
14 +#
15 +# IMPORTANT: This has to be the last thing we do, because vllm-flash-attn uses the same macros/functions as vLLM.
16 +# Because functions all belong to the global scope, vllm-flash-attn's functions overwrite vLLMs.
17 +# They should be identical but if they aren't, this is a massive footgun.
18 +#
19 +# The vllm-flash-attn install rules are nested under vllm to make sure the library gets installed in the correct place.
20 +# To only install vllm-flash-attn, use --component _vllm_fa2_C (for FA2) or --component _vllm_fa3_C (for FA3).
21 +# If no component is specified, vllm-flash-attn is still installed.
22 +
23 +# If VLLM_FLASH_ATTN_SRC_DIR is set, vllm-flash-attn is installed from that directory instead of downloading.
24 +# This is to enable local development of vllm-flash-attn within vLLM.
25 +# It can be set as an environment variable or passed as a cmake argument.
26 +# The environment variable takes precedence.
27 +if (DEFINED ENV{VLLM_FLASH_ATTN_SRC_DIR})
28 + set(VLLM_FLASH_ATTN_SRC_DIR $ENV{VLLM_FLASH_ATTN_SRC_DIR})
29 +endif()
30 +
31 +if(VLLM_FLASH_ATTN_SRC_DIR)
32 + FetchContent_Declare(
33 + vllm-flash-attn SOURCE_DIR
34 + ${VLLM_FLASH_ATTN_SRC_DIR}
35 + BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
36 + )
37 +else()
38 + FetchContent_Declare(
39 + vllm-flash-attn
40 + GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
41 + GIT_TAG 720c94869cf2e0ff5a706e9c7f1dce0939686ade
42 + GIT_PROGRESS TRUE
43 + # Don't share the vllm-flash-attn build between build types
44 + BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
45 + )
46 +endif()
47 +
48 +
49 +# Fetch the vllm-flash-attn library
50 +FetchContent_MakeAvailable(vllm-flash-attn)
51 +message(STATUS "vllm-flash-attn is available at ${vllm-flash-attn_SOURCE_DIR}")
52 +
53 +# Copy over the vllm-flash-attn python files (duplicated for fa2 and fa3, in
54 +# case only one is built, in the case both are built redundant work is done)
55 +install(
56 + DIRECTORY ${vllm-flash-attn_SOURCE_DIR}/vllm_flash_attn/
57 + DESTINATION vllm_flash_attn
58 + COMPONENT _vllm_fa2_C
59 + FILES_MATCHING PATTERN "*.py"
60 +)
61 +
62 +install(
63 + DIRECTORY ${vllm-flash-attn_SOURCE_DIR}/vllm_flash_attn/
64 + DESTINATION vllm_flash_attn
65 + COMPONENT _vllm_fa3_C
66 + FILES_MATCHING PATTERN "*.py"
67 +)
\ No newline at end of file
setup.py
+6
@@ -328,6 +328,7 @@ class repackage_wheel(build_ext):
328 files_to_copy = [
329 "vllm/_C.abi3.so",
330 "vllm/_moe_C.abi3.so",
331 + "vllm/_flashmla_C.abi3.so",
332 "vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
333 "vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
334 "vllm/vllm_flash_attn/flash_attn_interface.py",
@@ -612,6 +613,11 @@ if _is_cuda():
613 # FA3 requires CUDA 12.0 or later
614 ext_modules.append(
615 CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa3_C"))
616 + if envs.VLLM_USE_PRECOMPILED or get_nvcc_cuda_version() >= Version("12.3"):
617 + # Optional since this doesn't get built (produce an .so file) when
618 + # not targeting a hopper system
619 + ext_modules.append(
620 + CMakeExtension(name="vllm._flashmla_C", optional=True))
621 ext_modules.append(CMakeExtension(name="vllm.cumem_allocator"))
622
623 if _build_custom_ops():
tests/kernels/test_flashmla.py new
+132
@@ -0,0 +1,132 @@
1 +# Adapted from: https://github.com/deepseek-ai/FlashMLA/blob/main/tests/test_flash_mla.py
2 +# SPDX-License-Identifier: Apache-2.0
3 +import math
4 +import random
5 +
6 +import pytest
7 +import torch
8 +import triton
9 +
10 +from vllm.attention.ops.flashmla import (flash_mla_with_kvcache,
11 + get_mla_metadata,
12 + is_flashmla_supported)
13 +
14 +
15 +def cal_diff(x: torch.Tensor, y: torch.Tensor, name: str) -> None:
16 + x, y = x.double(), y.double()
17 + cos_diff = 1 - 2 * (x * y).sum().item() / max(
18 + (x * x + y * y).sum().item(), 1e-12)
19 + assert cos_diff < 1e-5
20 +
21 +FLASH_MLA_UNSUPPORTED_REASON = is_flashmla_supported()[1] \
22 + if not is_flashmla_supported()[0] else "FlashMLA is supported"
23 +
24 +
25 +@pytest.mark.skipif(not is_flashmla_supported()[0],
26 + reason=FLASH_MLA_UNSUPPORTED_REASON)
27 +@pytest.mark.parametrize("b", [128])
28 +@pytest.mark.parametrize("s_q", [1, 2])
29 +@pytest.mark.parametrize("mean_sk", [4096, 8192])
30 +@pytest.mark.parametrize("h_q", [16, 32, 64, 128])
31 +@pytest.mark.parametrize("h_kv", [1])
32 +@pytest.mark.parametrize("d", [576])
33 +@pytest.mark.parametrize("dv", [512])
34 +@pytest.mark.parametrize("block_size", [64])
35 +@pytest.mark.parametrize("causal", [True])
36 +@pytest.mark.parametrize("varlen", [False, True])
37 +@torch.inference_mode()
38 +def test_flash_mla(b, s_q, mean_sk, h_q, h_kv, d, dv, block_size, causal,
39 + varlen):
40 + # TODO: parametrize using pytest
41 + dtype = torch.bfloat16
42 + device = torch.device("cuda:0")
43 + torch.set_default_dtype(dtype)
44 + torch.set_default_device(device)
45 + torch.cuda.set_device(device)
46 + torch.manual_seed(0)
47 + random.seed(0)
48 +
49 + print(f"{b=}, {s_q=}, {mean_sk=}, {h_q=}, {h_kv=}, "
50 + f"{d=}, {dv=}, {causal=}, {varlen=}")
51 +
52 + cache_seqlens = torch.full((b, ), mean_sk, dtype=torch.int32)
53 + if varlen:
54 + for i in range(b):
55 + cache_seqlens[i] = max(random.normalvariate(mean_sk, mean_sk / 2),
56 + s_q)
57 + total_seqlens = cache_seqlens.sum().item()
58 + max_seqlen = cache_seqlens.max().item()
59 + max_seqlen_pad = triton.cdiv(max_seqlen, 256) * 256
60 +
61 + q = torch.randn(b, s_q, h_q, d)
62 + block_table = torch.arange(b * max_seqlen_pad // block_size,
63 + dtype=torch.int32).view(
64 + b, max_seqlen_pad // block_size)
65 + blocked_k = torch.randn(block_table.numel(), block_size, h_kv, d)
66 + for i in range(b):
67 + blocked_k.view(b, max_seqlen_pad, h_kv,
68 + d)[i, cache_seqlens[i].item():] = float("nan")
69 + blocked_v = blocked_k[..., :dv]
70 +
71 + tile_scheduler_metadata, num_splits = get_mla_metadata(
72 + cache_seqlens, s_q * h_q // h_kv, h_kv)
73 +
74 + def flash_mla():
75 + return flash_mla_with_kvcache(
76 + q,
77 + blocked_k,
78 + block_table,
79 + cache_seqlens,
80 + dv,
81 + tile_scheduler_metadata,
82 + num_splits,
83 + causal=causal,
84 + )
85 +
86 + def scaled_dot_product_attention(query, key, value, is_causal=False):
87 + query = query.float()
88 + key = key.float()
89 + value = value.float()
90 + key = key.repeat_interleave(h_q // h_kv, dim=0)
91 + value = value.repeat_interleave(h_q // h_kv, dim=0)
92 + attn_weight = query @ key.transpose(-2, -1) / math.sqrt(query.size(-1))
93 + if is_causal:
94 + s_q = query.shape[-2]
95 + s_k = key.shape[-2]
96 + attn_bias = torch.zeros(s_q, s_k, dtype=query.dtype)
97 + temp_mask = torch.ones(s_q, s_k,
98 + dtype=torch.bool).tril(diagonal=s_k - s_q)
99 + attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
100 + attn_bias.to(query.dtype)
101 + attn_weight += attn_bias
102 + lse = attn_weight.logsumexp(dim=-1)
103 + attn_weight = torch.softmax(attn_weight, dim=-1, dtype=torch.float32)
104 + return attn_weight @ value, lse
105 +
106 + def ref_mla():
107 + out = torch.empty(b, s_q, h_q, dv, dtype=torch.float32)
108 + lse = torch.empty(b, h_q, s_q, dtype=torch.float32)
109 + for i in range(b):
110 + begin = i * max_seqlen_pad
111 + end = begin + cache_seqlens[i]
112 + ref_O, LSE = scaled_dot_product_attention(
113 + q[i].transpose(0, 1),
114 + blocked_k.view(-1, h_kv, d)[begin:end].transpose(0, 1),
115 + blocked_v.view(-1, h_kv, dv)[begin:end].transpose(0, 1),
116 + is_causal=causal,
117 + )
118 + out[i] = ref_O.transpose(0, 1)
119 + lse[i] = LSE
120 + return out, lse
121 +
122 + out_flash, lse_flash = flash_mla()
123 + out_torch, lse_torch = ref_mla()
124 + cal_diff(out_flash, out_torch, "out")
125 + cal_diff(lse_flash, lse_torch, "lse")
126 +
127 + t = triton.testing.do_bench(flash_mla, fast_flush=False)
128 + FLOPS = s_q * total_seqlens * h_q * (d + dv) * 2
129 + bytes = (total_seqlens * h_kv * d + b * s_q * h_q * d +
130 + b * s_q * h_q * dv) * (torch.finfo(dtype).bits // 8)
131 + print(f"{t:.3f} ms, {FLOPS / 10 ** 9 / t:.0f} "
132 + f"TFLOPS, {bytes / 10 ** 6 / t:.0f} GB/s")
vllm/_custom_ops.py
+64
@@ -1163,3 +1163,67 @@ def get_graph_buffer_ipc_meta(fa: int) -> Tuple[List[int], List[int]]:
1163 def register_graph_buffers(fa: int, handles: List[List[int]],
1164 offsets: List[List[int]]) -> None:
1165 torch.ops._C_custom_ar.register_graph_buffers(fa, handles, offsets)
1166 +
1167 +
1168 +def get_flash_mla_metadata(
1169 + cache_seqlens: torch.Tensor,
1170 + num_heads_per_head_k: int,
1171 + num_heads_k: int,
1172 +) -> Tuple[torch.Tensor, torch.Tensor]:
1173 + """
1174 + Arguments:
1175 + cache_seqlens: (batch_size), dtype torch.int32.
1176 + num_heads_per_head_k: Equals to seq_len_q * num_heads_q // num_heads_k.
1177 + num_heads_k: num_heads_k.
1178 +
1179 + Return:
1180 + tile_scheduler_metadata: (num_sm_parts, TileSchedulerMetaDataSize), dtype torch.int32.
1181 + num_splits: (batch_size + 1), dtype torch.int32.
1182 + """
1183 + return torch.ops._C.get_flash_mla_metadata(cache_seqlens,
1184 + num_heads_per_head_k,
1185 + num_heads_k)
1186 +
1187 +
1188 +def flash_mla_with_kvcache(
1189 + q: torch.Tensor,
1190 + k_cache: torch.Tensor,
1191 + block_table: torch.Tensor,
1192 + cache_seqlens: torch.Tensor,
1193 + head_dim_v: int,
1194 + tile_scheduler_metadata: torch.Tensor,
1195 + num_splits: torch.Tensor,
1196 + softmax_scale: Optional[float] = None,
1197 + causal: bool = False,
1198 +) -> Tuple[torch.Tensor, torch.Tensor]:
1199 + """
1200 + Arguments:
1201 + q: (batch_size, seq_len_q, num_heads_q, head_dim).
1202 + k_cache: (num_blocks, page_block_size, num_heads_k, head_dim).
1203 + block_table: (batch_size, max_num_blocks_per_seq), torch.int32.
1204 + cache_seqlens: (batch_size), torch.int32.
1205 + head_dim_v: Head_dim of v.
1206 + tile_scheduler_metadata: (num_sm_parts, TileSchedulerMetaDataSize), torch.int32, return by get_mla_metadata.
1207 + num_splits: (batch_size + 1), torch.int32, return by get_mla_metadata.
1208 + softmax_scale: float. The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim).
1209 + causal: bool. Whether to apply causal attention mask.
1210 +
1211 + Return:
1212 + out: (batch_size, seq_len_q, num_heads_q, head_dim_v).
1213 + softmax_lse: (batch_size, num_heads_q, seq_len_q), torch.float32.
1214 + """
1215 + if softmax_scale is None:
1216 + softmax_scale = q.shape[-1]**(-0.5)
1217 + out, softmax_lse = torch.ops._C.flash_mla_fwd_kvcache(
1218 + q,
1219 + k_cache,
1220 + None,
1221 + head_dim_v,
1222 + cache_seqlens,
1223 + block_table,
1224 + softmax_scale,
1225 + causal,
1226 + tile_scheduler_metadata,
1227 + num_splits,
1228 + )
1229 + return out, softmax_lse
vllm/attention/backends/flashmla.py new
+239
@@ -0,0 +1,239 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from contextlib import contextmanager
4 +from dataclasses import dataclass
5 +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type
6 +
7 +import torch
8 +
9 +from vllm.attention.backends.abstract import AttentionType
10 +from vllm.attention.backends.mla.common import (MLACommonBackend,
11 + MLACommonImpl,
12 + MLACommonMetadata,
13 + MLACommonMetadataBuilder,
14 + MLACommonState)
15 +from vllm.attention.ops.flashmla import (flash_mla_with_kvcache,
16 + get_mla_metadata,
17 + is_flashmla_supported)
18 +
19 +if TYPE_CHECKING:
20 + from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata
21 +
22 +
23 +class FlashMLABackend(MLACommonBackend):
24 +
25 + @staticmethod
26 + def get_name() -> str:
27 + return "FLASHMLA"
28 +
29 + @staticmethod
30 + def get_impl_cls() -> Type["FlashMLAImpl"]:
31 + return FlashMLAImpl
32 +
33 + @staticmethod
34 + def get_metadata_cls() -> Type["FlashMLAMetadata"]:
35 + return FlashMLAMetadata
36 +
37 + @staticmethod
38 + def get_builder_cls() -> Type["FlashMLAMetadataBuilder"]:
39 + return FlashMLAMetadataBuilder
40 +
41 + @staticmethod
42 + def get_state_cls() -> Type["FlashMLAState"]:
43 + return FlashMLAState
44 +
45 +
46 +@dataclass
47 +class FlashMLAMetadata(MLACommonMetadata):
48 + decode_tile_scheduler_metadata: Optional[Tuple[torch.Tensor,
49 + torch.Tensor]] = None
50 + decode_num_splits: Optional[torch.Tensor] = None
51 +
52 + @property
53 + def decode_metadata(self):
54 + decode_metadata = super().decode_metadata
55 + # TODO: cache assignment?
56 + if decode_metadata is not None:
57 + decode_metadata.decode_tile_scheduler_metadata=\
58 + self.decode_tile_scheduler_metadata
59 + decode_metadata.decode_num_splits=\
60 + self.decode_num_splits
61 + return decode_metadata
62 +
63 + def advance_step(self,
64 + model_input: "ModelInputForGPUWithSamplingMetadata",
65 + sampled_token_ids: Optional[torch.Tensor],
66 + block_size: int,
67 + num_seqs: int,
68 + num_queries: int,
69 + turn_prefills_into_decodes: bool = False):
70 + raise NotImplementedError(
71 + "advance_step is not implemented for FlashMLA")
72 +
73 +
74 +class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
75 +
76 + def __init__(self, *args, **kwargs):
77 + super().__init__(*args, **kwargs)
78 +
79 + self.num_q_heads = self.runner.model_config.get_num_attention_heads(
80 + self.runner.parallel_config)
81 +
82 + def build(self, seq_lens: List[int], query_lens: List[int],
83 + cuda_graph_pad_size: int, batch_size: int):
84 + m = super().build(seq_lens, query_lens, cuda_graph_pad_size,
85 + batch_size)
86 +
87 + if m.num_decode_tokens > 0:
88 + m.decode_tile_scheduler_metadata, m.decode_num_splits = \
89 + get_mla_metadata(
90 + m.seq_lens_tensor[m.num_prefills:],
91 + self.num_q_heads,
92 + 1, # MQA for the decode path
93 + )
94 +
95 + return m
96 +
97 +
98 +class FlashMLAState(MLACommonState[FlashMLAMetadata]):
99 +
100 + def __init__(self, *args, **kwds):
101 + super().__init__(*args, **kwds)
102 +
103 + self.num_q_heads = self.runner.model_config.get_num_attention_heads(
104 + self.runner.parallel_config)
105 +
106 + @contextmanager
107 + def graph_capture(self, max_batch_size: int):
108 + # Run a dummy `get_mla_metadata` so we can get the right shapes
109 + self._graph_decoder_tile_scheduler_metadata, \
110 + self._graph_decode_num_splits = get_mla_metadata(
111 + torch.ones(
112 + max_batch_size, dtype=torch.int32, device=self.runner.device),
113 + self.num_q_heads,
114 + 1, # MQA for the decode path
115 + )
116 +
117 + with super().graph_capture(max_batch_size):
118 + yield
119 +
120 + del self._graph_decoder_tile_scheduler_metadata
121 + del self._graph_decode_num_splits
122 +
123 + def graph_capture_get_metadata_for_batch(
124 + self, batch_size: int, is_encoder_decoder_model: bool = False):
125 + metadata = super().graph_capture_get_metadata_for_batch(
126 + batch_size, is_encoder_decoder_model)
127 + assert metadata.num_decode_tokens > 0
128 +
129 + decoder_tile_scheduler_metadata, decode_num_splits = get_mla_metadata(
130 + self._graph_seq_lens[:batch_size],
131 + self.num_q_heads,
132 + 1, # MQA for the decode path
133 + )
134 +
135 + self._graph_decoder_tile_scheduler_metadata.copy_(
136 + decoder_tile_scheduler_metadata)
137 + self._graph_decode_num_splits[:batch_size + 1].copy_(decode_num_splits)
138 +
139 + metadata.decode_tile_scheduler_metadata=\
140 + self._graph_decoder_tile_scheduler_metadata
141 + metadata.decode_num_splits=\
142 + self._graph_decode_num_splits[:batch_size + 1]
143 +
144 + return metadata
145 +
146 + def get_graph_input_buffers(self,
147 + attn_metadata,
148 + is_encoder_decoder_model: bool = False):
149 + input_buffers = super().get_graph_input_buffers(
150 + attn_metadata, is_encoder_decoder_model)
151 + input_buffers["decode_tile_scheduler_metadata"] = \
152 + attn_metadata.decode_metadata.decode_tile_scheduler_metadata
153 + input_buffers["decode_num_splits"] = \
154 + attn_metadata.decode_metadata.decode_num_splits
155 +
156 + return input_buffers
157 +
158 + def prepare_graph_input_buffers(self,
159 + input_buffers,
160 + attn_metadata,
161 + is_encoder_decoder_model: bool = False):
162 + super().prepare_graph_input_buffers(input_buffers, attn_metadata,
163 + is_encoder_decoder_model)
164 +
165 + input_buffers["decode_tile_scheduler_metadata"].copy_(
166 + attn_metadata.decode_metadata.decode_tile_scheduler_metadata)
167 + input_buffers["decode_num_splits"].copy_(
168 + attn_metadata.decode_metadata.decode_num_splits)
169 +
170 +
171 +class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
172 +
173 + def __init__(
174 + self,
175 + num_heads: int,
176 + head_size: int,
177 + scale: float,
178 + num_kv_heads: int,
179 + alibi_slopes: Optional[List[float]],
180 + sliding_window: Optional[int],
181 + kv_cache_dtype: str,
182 + blocksparse_params: Optional[Dict[str, Any]],
183 + logits_soft_cap: Optional[float],
184 + attn_type: str,
185 + # MLA Specific Arguments
186 + **mla_args) -> None:
187 + super().__init__(num_heads, head_size, scale, num_kv_heads,
188 + alibi_slopes, sliding_window, kv_cache_dtype,
189 + blocksparse_params, logits_soft_cap, attn_type,
190 + **mla_args)
191 +
192 + assert is_flashmla_supported(), \
193 + "FlashMLA is not supported on this device"
194 +
195 + unsupported_features = [
196 + alibi_slopes, sliding_window, blocksparse_params, logits_soft_cap
197 + ]
198 + if any(unsupported_features):
199 + raise NotImplementedError(
200 + "FlashMLAImpl does not support one of the following: "
201 + "alibi_slopes, sliding_window, blocksparse_params, "
202 + "logits_soft_cap")
203 +
204 + if attn_type != AttentionType.DECODER:
205 + raise NotImplementedError("Encoder self-attention and "
206 + "encoder/decoder cross-attention "
207 + "are not implemented for "
208 + "FlashMLAImpl")
209 +
210 + def _forward_decode(
211 + self,
212 + q_nope: torch.Tensor,
213 + q_pe: torch.Tensor,
214 + kv_c_and_k_pe_cache: torch.Tensor,
215 + attn_metadata: FlashMLAMetadata,
216 + ) -> torch.Tensor:
217 + assert kv_c_and_k_pe_cache.numel() > 0
218 + if self.kv_cache_dtype.startswith("fp8"):
219 + raise NotImplementedError("FP8 FlashMLA not yet supported")
220 +
221 + decode_meta = attn_metadata.decode_metadata
222 + assert decode_meta is not None
223 +
224 + q = torch.cat([q_nope, q_pe], dim=-1)\
225 + .unsqueeze(1) # Add seqlen dim of 1 (decode)
226 +
227 + o, _ = flash_mla_with_kvcache(
228 + q=q,
229 + k_cache=kv_c_and_k_pe_cache.unsqueeze(-2), # Add head dim of 1
230 + block_table=decode_meta.block_tables,
231 + cache_seqlens=decode_meta.seq_lens_tensor,
232 + head_dim_v=self.kv_lora_rank,
233 + tile_scheduler_metadata=decode_meta.decode_tile_scheduler_metadata,
234 + num_splits=decode_meta.decode_num_splits,
235 + softmax_scale=self.scale,
236 + causal=True,
237 + )
238 +
239 + return self._v_up_proj_and_o_proj(o)
vllm/attention/backends/mla/common.py
+15 -13
@@ -293,7 +293,10 @@ class MLACommonBackend(AttentionBackend):
293 return [576]
294
295
296 -class MLACommonState(AttentionState):
296 +T = TypeVar("T", bound="MLACommonMetadata")
297 +
298 +
299 +class MLACommonState(AttentionState, Generic[T]):
300
301 def __init__(self, runner):
302 self.runner = runner
@@ -355,7 +358,9 @@ class MLACommonState(AttentionState):
358 return self.__class__(self.runner)
359
360 def graph_capture_get_metadata_for_batch(
358 - self, batch_size: int, is_encoder_decoder_model: bool = False):
361 + self,
362 + batch_size: int,
363 + is_encoder_decoder_model: bool = False) -> T:
364 assert self._is_graph_capturing
365
366 attn_metadata = self.runner.attn_backend.make_metadata(
@@ -507,8 +512,8 @@ class MLACommonMetadata(AttentionMetadata):
512 # [4, 6], it is [0, 4, 10].
513 seq_start_loc: Optional[torch.Tensor] = None
514
510 - _cached_prefill_metadata: Optional["MLACommonMetadata"] = None
511 - _cached_decode_metadata: Optional["MLACommonMetadata"] = None
515 + _cached_prefill_metadata: Optional[Any] = None
516 + _cached_decode_metadata: Optional[Any] = None
517
518 num_prefill_tokens: int
519
@@ -537,7 +542,7 @@ class MLACommonMetadata(AttentionMetadata):
542 f" received {self.head_dim}.")
543
544 @property
540 - def prefill_metadata(self) -> Optional["MLACommonMetadata"]:
545 + def prefill_metadata(self):
546 if self.num_prefills == 0:
547 return None
548
@@ -565,7 +570,7 @@ class MLACommonMetadata(AttentionMetadata):
570 input_positions = (None if self.input_positions is None else
571 self.input_positions[:self.num_prefill_tokens])
572
568 - self._cached_prefill_metadata = MLACommonMetadata(
573 + self._cached_prefill_metadata = self.__class__(
574 # Required by ModelRunner
575 use_cuda_graph=False, # Not Attention Related
576 # Required by Attention Metadata
@@ -599,7 +604,7 @@ class MLACommonMetadata(AttentionMetadata):
604 return self._cached_prefill_metadata
605
606 @property
602 - def decode_metadata(self) -> Optional["MLACommonMetadata"]:
607 + def decode_metadata(self):
608 if self.num_decode_tokens == 0:
609 return None
610
@@ -617,7 +622,7 @@ class MLACommonMetadata(AttentionMetadata):
622 input_positions = (None if self.input_positions is None else
623 self.input_positions[self.num_prefill_tokens:])
624
620 - self._cached_decode_metadata = MLACommonMetadata(
625 + self._cached_decode_metadata = self.__class__(
626 # Required by ModelRunner
627 use_cuda_graph=self.use_cuda_graph, # Not Attention Related
628 # Required by Attention Metadata
@@ -723,10 +728,7 @@ class MLACommonMetadata(AttentionMetadata):
728 block_tables=self.block_tables)
729
730
726 -T = TypeVar("T", bound=MLACommonMetadata)
727 -
728 -
729 -class MLACommonMetadataBuilder(AttentionMetadataBuilder[MLACommonMetadata]):
731 +class MLACommonMetadataBuilder(AttentionMetadataBuilder[T], Generic[T]):
732 """
733 NOTE: Please read the comment at the top of the file before trying to
734 understand this class
@@ -959,7 +961,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[MLACommonMetadata]):
961 assert max(context_chunk_seq_tot) <= \
962 self.chunked_prefill_workspace_size
963
962 - return MLACommonMetadata(
964 + return self.runner.attn_backend.make_metadata(
965 # Required by ModelRunner
966 use_cuda_graph=use_captured_graph, # Not Attention Related
967 # Required by Attention Metadata
vllm/attention/ops/flashmla.py new
+115
@@ -0,0 +1,115 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# adapted from: https://github.com/deepseek-ai/FlashMLA/blob/main/flash_mla/flash_mla_interface.py
3 +from typing import Optional, Tuple
4 +
5 +import torch
6 +
7 +from vllm.logger import init_logger
8 +from vllm.platforms import current_platform
9 +
10 +logger = init_logger(__name__)
11 +
12 +if current_platform.is_cuda():
13 + try:
14 + import vllm._flashmla_C # noqa: F401
15 + _flashmla_C_AVAILABLE = True
16 + except ImportError:
17 + _flashmla_C_AVAILABLE = False
18 +else:
19 + _flashmla_C_AVAILABLE = False
20 +
21 +
22 +def is_flashmla_supported() -> Tuple[bool, Optional[str]]:
23 + """
24 + Return: is_supported_flag, unsupported_reason (optional).
25 + """
26 + if not current_platform.is_cuda():
27 + return False, "FlashMLA is only supported on CUDA devices."
28 + if current_platform.get_device_capability()[0] != 9:
29 + return False, "FlashMLA is only supported on Hopper devices."
30 + if not _flashmla_C_AVAILABLE:
31 + return False, "vllm._flashmla_C is not available, likely was not "\
32 + "compiled due to insufficient nvcc version or a supported arch "\
33 + "(only sm90a currently) was not in the list of target arches to "\
34 + "compile for."
35 + return True, None
36 +
37 +
38 +def get_mla_metadata(
39 + cache_seqlens: torch.Tensor,
40 + num_heads_per_head_k: int,
41 + num_heads_k: int,
42 +) -> Tuple[torch.Tensor, torch.Tensor]:
43 + """
44 + Arguments:
45 + cache_seqlens: (batch_size), dtype torch.int32.
46 + num_heads_per_head_k: Equals to seq_len_q * num_heads_q // num_heads_k.
47 + num_heads_k: num_heads_k.
48 +
49 + Return:
50 + tile_scheduler_metadata: (num_sm_parts, TileSchedulerMetaDataSize),
51 + dtype torch.int32.
52 + num_splits: (batch_size + 1), dtype torch.int32.
53 + """
54 + return torch.ops._flashmla_C.get_mla_metadata(cache_seqlens,
55 + num_heads_per_head_k,
56 + num_heads_k)
57 +
58 +
59 +def flash_mla_with_kvcache(
60 + q: torch.Tensor,
61 + k_cache: torch.Tensor,
62 + block_table: torch.Tensor,
63 + cache_seqlens: torch.Tensor,
64 + head_dim_v: int,
65 + tile_scheduler_metadata: torch.Tensor,
66 + num_splits: torch.Tensor,
67 + softmax_scale: Optional[float] = None,
68 + causal: bool = False,
69 +) -> Tuple[torch.Tensor, torch.Tensor]:
70 + """
71 + Arguments:
72 + q: (batch_size, seq_len_q, num_heads_q, head_dim).
73 + k_cache: (num_blocks, page_block_size, num_heads_k, head_dim).
74 + block_table: (batch_size, max_num_blocks_per_seq), torch.int32.
75 + cache_seqlens: (batch_size), torch.int32.
76 + head_dim_v: Head_dim of v.
77 + tile_scheduler_metadata: (num_sm_parts, TileSchedulerMetaDataSize),
78 + torch.int32, return by get_mla_metadata.
79 + num_splits: (batch_size + 1), torch.int32, return by get_mla_metadata.
80 + softmax_scale: float. The scaling of QK^T before applying softmax.
81 + Default to 1 / sqrt(head_dim).
82 + causal: bool. Whether to apply causal attention mask.
83 +
84 + Return:
85 + out: (batch_size, seq_len_q, num_heads_q, head_dim_v).
86 + softmax_lse: (batch_size, num_heads_q, seq_len_q), torch.float32.
87 + """
88 + if softmax_scale is None:
89 + softmax_scale = q.shape[-1]**(-0.5)
90 + out, softmax_lse = torch.ops._flashmla_C.fwd_kvcache_mla(
91 + q,
92 + k_cache,
93 + None,
94 + head_dim_v,
95 + cache_seqlens,
96 + block_table,
97 + softmax_scale,
98 + causal,
99 + tile_scheduler_metadata,
100 + num_splits,
101 + )
102 + return out, softmax_lse
103 +
104 +
105 +#
106 +# TODO: Add fake functions
107 +#
108 +# @register_fake("_flashmla_C::get_mla_metadata")
109 +# def _get_mla_metadata_fake(....) -> Tuple[torch.Tensor, torch.Tensor]:
110 +# return ....
111 +#
112 +# @register_fake("_flashmla_C::fwd_kvcache_mla")
113 +# def _fwd_kvcache_mla_fake(....) -> Tuple[torch.Tensor, torch.Tensor]:
114 +# return ....
115 +#
vllm/platforms/cuda.py
+24
@@ -141,6 +141,14 @@ class CudaPlatformBase(Platform):
141 cache_config = vllm_config.cache_config
142 if cache_config and cache_config.block_size is None:
143 cache_config.block_size = 16
144 + # TODO(lucas): handle this more gracefully
145 + if envs.VLLM_ATTENTION_BACKEND is not None \
146 + and envs.VLLM_ATTENTION_BACKEND == "FLASHMLA" \
147 + and cache_config.block_size != 64:
148 + cache_config.block_size = 64
149 + logger.info(
150 + "FlashMLA: Forcing kv cache block size to 64 since this"
151 + " is currently the only block size supported by the kernel.")
152
153 @classmethod
154 def get_current_memory_usage(cls,
@@ -157,6 +165,22 @@ class CudaPlatformBase(Platform):
165 logger.info("Using Flash Attention backend on V1 engine.")
166 return "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend"
167 if use_mla:
168 + if selected_backend == _Backend.FLASHMLA:
169 + from vllm.attention.backends.flashmla import (
170 + is_flashmla_supported)
171 + if not is_flashmla_supported()[0]:
172 + logger.warning(
173 + "FlashMLA backend is not supported due to %s",
174 + is_flashmla_supported()[1])
175 + elif block_size != 64:
176 + logger.warning(
177 + "FlashMLA backend is not supported for block size %d"
178 + " (currently only supports block size 64).",
179 + block_size)
180 + else:
181 + logger.info("Using FlashMLA backend.")
182 + return "vllm.attention.backends.flashmla.FlashMLABackend"
183 +
184 logger.info("Using Triton MLA backend.")
185 return "vllm.attention.backends.triton_mla.TritonMLABackend"
186 if selected_backend == _Backend.FLASHINFER:
vllm/platforms/interface.py
+1
@@ -35,6 +35,7 @@ class _Backend(enum.Enum):
35 OPENVINO = enum.auto()
36 FLASHINFER = enum.auto()
37 TRITON_MLA = enum.auto()
38 + FLASHMLA = enum.auto()
39 HPU_ATTN = enum.auto()
40 PALLAS = enum.auto()
41 PALLAS_VLLM_V1 = enum.auto()