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 14, 05:30 UTC
[Kernel] Flashinfer MLA (trtllm-gen) decode kernel integration (#21078)
Signed-off-by: hjjq <hanjieq@nvidia.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Hanjie Qiu committed
Sep 10, 2025 at 18:31 UTC
dcb28a332bfb7c7b1991ab8c6694233956a4eafe
8 files changed
+255
-1
.buildkite/test-pipeline.yaml
+2
-1
@@ -729,7 +729,8 @@ steps:
729
# num_heads2 broken by https://github.com/flashinfer-ai/flashinfer/issues/1353
730
- pytest -v -s tests/kernels/attention/test_flashinfer.py -k 'not num_heads2'
731
- pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py
732
- - pytest -v -s tests/kernels/test_cutlass_mla_decode.py
732
+ - pytest -v -s tests/kernels/attention/test_cutlass_mla_decode.py
733
+ - pytest -v -s tests/kernels/attention/test_flashinfer_mla_decode.py
734
# Quantization
735
- pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8'
736
- pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py
tests/kernels/attention/test_cutlass_mla_decode.py
renamed
tests/kernels/attention/test_flashinfer_mla_decode.py
new
+123
@@ -0,0 +1,123 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pytest
4
+import torch
5
+import torch.nn.functional as F
6
+from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
7
+from torch import Tensor
8
+
9
+from vllm.platforms import current_platform
10
+
11
+FLASHINFER_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
12
+
13
+if not current_platform.has_device_capability(100):
14
+ pytest.skip(
15
+ reason="FlashInfer MLA Requires compute capability of 10 or above.",
16
+ allow_module_level=True)
17
+
18
+
19
+def ref_mla(
20
+ out: Tensor, # (bs, num_heads, v_head_dim)
21
+ query: Tensor, # (bs, num_heads, head_dim)
22
+ kv_cache: Tensor, # (num_blocks, block_size, head_dim)
23
+ scale: float,
24
+ block_tables: Tensor, # (bs, max_num_blocks)
25
+ seq_lens: Tensor, # (bs,)
26
+):
27
+ bs, num_heads, v_head_dim = out.shape
28
+ head_dim = query.shape[2]
29
+
30
+ for i in range(bs):
31
+ # gather and flatten KV-cache
32
+ kv = kv_cache[
33
+ block_tables[i]] # (max_num_blocks, block_size, head_dim)
34
+ kv = kv.view(1, -1,
35
+ head_dim)[:, :seq_lens[i]] # (1, seq_len, head_dim)
36
+ v = kv[:, :, :v_head_dim]
37
+
38
+ q = query[i].view(num_heads, 1, head_dim)
39
+ o = F.scaled_dot_product_attention(q,
40
+ kv,
41
+ v,
42
+ scale=scale,
43
+ enable_gqa=True)
44
+ out[i] = o.view(num_heads, v_head_dim)
45
+
46
+ return out
47
+
48
+
49
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
50
+@pytest.mark.parametrize("bs", [1, 2, 4, 16])
51
+@pytest.mark.parametrize("block_size", [32, 64])
52
+def test_flashinfer_mla_decode(dtype: torch.dtype, bs: int, block_size: int):
53
+ torch.set_default_device('cuda')
54
+ torch.manual_seed(42)
55
+
56
+ # Deepseek R1 config
57
+ num_heads = 128
58
+ kv_lora_rank = 512
59
+ qk_nope_head_dim = 128
60
+ qk_rope_head_dim = 64
61
+ qk_head_dim = kv_lora_rank + qk_rope_head_dim
62
+ scale = (qk_nope_head_dim + qk_rope_head_dim)**-0.5
63
+
64
+ MAX_SEQ_LEN = 1024
65
+
66
+ seq_lens = [torch.randint(2, MAX_SEQ_LEN, (1, )).item() for _ in range(bs)]
67
+ seq_lens[-1] = MAX_SEQ_LEN
68
+ max_seq_len = max(seq_lens)
69
+ seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32)
70
+
71
+ # Generate block tables with random but unique block IDs
72
+ # From https://github.com/flashinfer-ai/flashinfer/pull/1222
73
+ blocks_per_seq = (seq_lens_tensor + block_size - 1) // block_size
74
+ max_num_blocks_per_seq = max(blocks_per_seq.max().item(), 4)
75
+ total_blocks_needed = sum(blocks_per_seq)
76
+ # Get random unique IDs for all blocks
77
+ all_block_ids = torch.randperm(total_blocks_needed)
78
+
79
+ block_id = 0
80
+ block_tables = torch.zeros(
81
+ (bs, max_num_blocks_per_seq),
82
+ dtype=torch.int32,
83
+ )
84
+
85
+ # Populate block tables and track block assignments
86
+ block_id = 0
87
+ for i in range(bs):
88
+ num_blocks_needed = blocks_per_seq[i]
89
+ block_tables[i, :num_blocks_needed] = all_block_ids[block_id:block_id +
90
+ num_blocks_needed]
91
+ block_id += num_blocks_needed
92
+
93
+ kv_cache = torch.randn(block_tables.numel(), block_size,
94
+ qk_head_dim).to(dtype)
95
+ q = torch.randn(bs, num_heads, qk_head_dim).to(dtype)
96
+
97
+ out_ref = q.new_zeros(bs, num_heads, kv_lora_rank)
98
+ ref_mla(out_ref, q, kv_cache, scale, block_tables, seq_lens_tensor)
99
+
100
+ workspace_buffer = torch.zeros(
101
+ FLASHINFER_WORKSPACE_BUFFER_SIZE,
102
+ dtype=torch.uint8,
103
+ device=q.device,
104
+ )
105
+ # Flashinfer MLA expects the query to be of shape
106
+ # (bs, q_len_per_request, num_heads, qk_head_dim),
107
+ # where q_len_per_request is the MTP query length (=1 without MTP)
108
+ q = q.unsqueeze(1)
109
+
110
+ out_ans = trtllm_batch_decode_with_kv_cache_mla(
111
+ query=q,
112
+ kv_cache=kv_cache.unsqueeze(1),
113
+ workspace_buffer=workspace_buffer,
114
+ qk_nope_head_dim=qk_nope_head_dim,
115
+ kv_lora_rank=kv_lora_rank,
116
+ qk_rope_head_dim=qk_rope_head_dim,
117
+ block_tables=block_tables,
118
+ seq_lens=seq_lens_tensor,
119
+ max_seq_len=max_seq_len,
120
+ bmm1_scale=scale,
121
+ )
122
+ out_ans = out_ans.squeeze(1)
123
+ torch.testing.assert_close(out_ans, out_ref, atol=1e-2, rtol=1e-2)
vllm/engine/arg_utils.py
+1
@@ -1504,6 +1504,7 @@ class EngineArgs:
1504
"FLASH_ATTN_MLA",
1505
"FLASHINFER",
1506
"FLASHINFER_VLLM_V1",
1507
+ "FLASHINFER_MLA",
1508
"ROCM_AITER_MLA",
1509
"TORCH_SDPA_VLLM_V1",
1510
"FLEX_ATTENTION",
vllm/platforms/cuda.py
+15
@@ -228,6 +228,8 @@ class CudaPlatformBase(Platform):
228
use_cutlassmla = selected_backend == _Backend.CUTLASS_MLA or (
229
selected_backend is None and cls.is_device_capability(100)
230
and block_size == 128)
231
+ use_flashinfermla = (selected_backend == _Backend.FLASHINFER_MLA
232
+ and cls.has_device_capability(100))
233
use_flashmla = selected_backend in [
234
_Backend.FLASHMLA, _Backend.FLASHMLA_VLLM_V1
235
] or (selected_backend is None and is_flashmla_supported()[0])
@@ -252,6 +254,19 @@ class CudaPlatformBase(Platform):
254
else:
255
logger.warning(
256
"Cutlass MLA backend is only supported on V1 engine")
257
+ if use_flashinfermla:
258
+ if use_v1:
259
+ from vllm.v1.attention.backends.utils import (
260
+ set_kv_cache_layout)
261
+ set_kv_cache_layout("HND")
262
+ logger.info_once(
263
+ "Using FlashInfer MLA backend on V1 engine.")
264
+ return ("vllm.v1.attention.backends.mla."
265
+ "flashinfer_mla.FlashInferMLABackend")
266
+ else:
267
+ logger.warning(
268
+ "FlashInfer MLA backend is only supported on V1 engine"
269
+ )
270
if use_flashmla:
271
if block_size != 64:
272
logger.warning(
vllm/platforms/interface.py
+1
@@ -51,6 +51,7 @@ class _Backend(enum.Enum):
51
TORCH_SDPA_VLLM_V1 = enum.auto()
52
FLASHINFER = enum.auto()
53
FLASHINFER_VLLM_V1 = enum.auto()
54
+ FLASHINFER_MLA = enum.auto()
55
TRITON_MLA = enum.auto() # Supported by V1
56
TRITON_MLA_VLLM_V1 = enum.auto()
57
CUTLASS_MLA = enum.auto()
vllm/v1/attention/backends/mla/common.py
+3
@@ -381,6 +381,7 @@ class MLACommonMetadata(Generic[D]):
381
382
num_reqs: int
383
max_query_len: int
384
+ max_seq_len: int
385
386
num_actual_tokens: int # Number of tokens excluding padding.
387
query_start_loc: torch.Tensor
@@ -644,6 +645,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
645
num_reqs = common_attn_metadata.num_reqs
646
num_tokens = common_attn_metadata.num_actual_tokens
647
max_query_len = common_attn_metadata.max_query_len
648
+ max_seq_len = common_attn_metadata.max_seq_len
649
650
# Note(simon): be careful about the CPU <> GPU memory movement in this
651
# function. We should avoid GPU -> CPU sync as much as possible because
@@ -830,6 +832,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
832
attn_metadata = self.metadata_cls(
833
num_reqs=common_attn_metadata.num_reqs,
834
max_query_len=common_attn_metadata.max_query_len,
835
+ max_seq_len=max_seq_len,
836
num_actual_tokens=num_tokens,
837
query_start_loc=query_start_loc,
838
slot_mapping=slot_mapping,
vllm/v1/attention/backends/mla/flashinfer_mla.py
new
+110
@@ -0,0 +1,110 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from typing import Optional, Union
5
+
6
+import torch
7
+from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
8
+
9
+from vllm.attention.backends.abstract import (AttentionLayer, AttentionType,
10
+ is_quantized_kv_cache)
11
+from vllm.logger import init_logger
12
+from vllm.v1.attention.backends.mla.common import (MLACommonBackend,
13
+ MLACommonImpl,
14
+ MLACommonMetadata)
15
+
16
+logger = init_logger(__name__)
17
+
18
+FLASHINFER_MLA_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
19
+
20
+
21
+class FlashInferMLABackend(MLACommonBackend):
22
+
23
+ @staticmethod
24
+ def get_name() -> str:
25
+ return "FLASHINFER_MLA"
26
+
27
+ @staticmethod
28
+ def get_impl_cls() -> type["FlashInferMLAImpl"]:
29
+ return FlashInferMLAImpl
30
+
31
+
32
+g_fi_workspace = torch.zeros(
33
+ FLASHINFER_MLA_WORKSPACE_BUFFER_SIZE,
34
+ dtype=torch.uint8,
35
+ device="cuda",
36
+)
37
+
38
+
39
+class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
40
+
41
+ def __init__(
42
+ self,
43
+ num_heads: int,
44
+ head_size: int,
45
+ scale: float,
46
+ num_kv_heads: int,
47
+ alibi_slopes: Optional[list[float]],
48
+ sliding_window: Optional[int],
49
+ kv_cache_dtype: str,
50
+ logits_soft_cap: Optional[float],
51
+ attn_type: str,
52
+ kv_sharing_target_layer_name: Optional[str],
53
+ # MLA Specific Arguments
54
+ **mla_args) -> None:
55
+ super().__init__(num_heads, head_size, scale, num_kv_heads,
56
+ alibi_slopes, sliding_window, kv_cache_dtype,
57
+ logits_soft_cap, attn_type,
58
+ kv_sharing_target_layer_name, **mla_args)
59
+
60
+ unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
61
+ if any(unsupported_features):
62
+ raise NotImplementedError(
63
+ "FlashInferMLAImpl does not support one of the following: "
64
+ "alibi_slopes, sliding_window, logits_soft_cap")
65
+
66
+ if attn_type != AttentionType.DECODER:
67
+ raise NotImplementedError("Encoder self-attention and "
68
+ "encoder/decoder cross-attention "
69
+ "are not implemented for "
70
+ "FlashInferMLAImpl")
71
+
72
+ if is_quantized_kv_cache(self.kv_cache_dtype):
73
+ raise NotImplementedError(
74
+ "FlashInferMLA V1 with FP8 KV cache not yet supported")
75
+
76
+ self._workspace_buffer = g_fi_workspace
77
+
78
+ def _forward_decode(
79
+ self,
80
+ q: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],
81
+ kv_c_and_k_pe_cache: torch.Tensor,
82
+ attn_metadata: MLACommonMetadata,
83
+ layer: AttentionLayer,
84
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
85
+ assert kv_c_and_k_pe_cache.numel() > 0
86
+ assert attn_metadata.decode is not None
87
+
88
+ if isinstance(q, tuple):
89
+ q_nope, q_pe = q
90
+ q = torch.cat([q_nope, q_pe], dim=-1)
91
+
92
+ # trtllm API requires extra dimension q_len_per_request for MTP
93
+ q = q.unsqueeze(1)
94
+
95
+ o = trtllm_batch_decode_with_kv_cache_mla(
96
+ query=q,
97
+ kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
98
+ workspace_buffer=self._workspace_buffer,
99
+ qk_nope_head_dim=self.qk_nope_head_dim,
100
+ kv_lora_rank=self.kv_lora_rank,
101
+ qk_rope_head_dim=self.qk_rope_head_dim,
102
+ block_tables=attn_metadata.decode.block_table,
103
+ seq_lens=attn_metadata.decode.seq_lens,
104
+ max_seq_len=attn_metadata.max_seq_len,
105
+ bmm1_scale=self.scale,
106
+ )
107
+
108
+ # TODO: Return LSE pending support from Flashinfer API:
109
+ # https://github.com/flashinfer-ai/flashinfer/pull/1566
110
+ return o, None