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 9, 21:48 UTC
Update FlashMLA (#32491)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Jan 20, 2026 at 22:03 UTC
b4f64e5b02a949c8856c9f81990b77ca56296cdc
4 files changed
+169
-42
cmake/external_projects/flashmla.cmake
+40
-10
@@ -19,7 +19,7 @@ else()
19
FetchContent_Declare(
20
flashmla
21
GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
22
- GIT_TAG 46d64a8ebef03fa50b4ae74937276a5c940e3f95
22
+ GIT_TAG 526781394b33d9888e4c41952e692266267dd8bf
23
GIT_PROGRESS TRUE
24
CONFIGURE_COMMAND ""
25
BUILD_COMMAND ""
@@ -55,16 +55,43 @@ if(FLASH_MLA_ARCHS)
55
56
set(FlashMLA_SOURCES
57
${flashmla_SOURCE_DIR}/csrc/torch_api.cpp
58
- ${flashmla_SOURCE_DIR}/csrc/pybind.cpp
59
- ${flashmla_SOURCE_DIR}/csrc/smxx/get_mla_metadata.cu
60
- ${flashmla_SOURCE_DIR}/csrc/smxx/mla_combine.cu
61
- ${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/splitkv_mla.cu
62
- ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/splitkv_mla.cu
58
+
59
+ # Misc kernels for decoding
60
+ ${flashmla_SOURCE_DIR}/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu
61
+ ${flashmla_SOURCE_DIR}/csrc/smxx/decode/combine/combine.cu
62
+
63
+ # sm90 dense decode
64
+ ${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/fp16.cu
65
+ ${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/bf16.cu
66
+
67
+ # sm90 sparse decode
68
+ ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h64.cu
69
+ ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h128.cu
70
+ ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h64.cu
71
+ ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h128.cu
72
+
73
+ # sm90 sparse prefill
74
${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/fwd.cu
64
- ${flashmla_SOURCE_DIR}/csrc/sm100/decode/sparse_fp8/splitkv_mla.cu
75
+ ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512.cu
76
+ ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu
77
+ ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576.cu
78
+ ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu
79
+
80
+ # sm100 dense prefill & backward
81
${flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu
82
${flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu
67
- ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd.cu
83
+
84
+ # sm100 sparse prefill
85
+ ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k512.cu
86
+ ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k576.cu
87
+ ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu
88
+ ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu
89
+ ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_prefill_k512.cu
90
+
91
+ # sm100 sparse decode
92
+ ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32.cu
93
+ ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/model1.cu
94
+ ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu
95
)
96
97
set(FlashMLA_Extension_SOURCES
@@ -76,6 +103,7 @@ if(FLASH_MLA_ARCHS)
103
104
set(FlashMLA_INCLUDES
105
${flashmla_SOURCE_DIR}/csrc
106
+ ${flashmla_SOURCE_DIR}/csrc/kerutils/include
107
${flashmla_SOURCE_DIR}/csrc/sm90
108
${flashmla_SOURCE_DIR}/csrc/cutlass/include
109
${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include
@@ -83,7 +111,6 @@ if(FLASH_MLA_ARCHS)
111
112
set(FlashMLA_Extension_INCLUDES
113
${flashmla_SOURCE_DIR}/csrc
86
- ${flashmla_SOURCE_DIR}/csrc/sm90
114
${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/
115
${flashmla_SOURCE_DIR}/csrc/cutlass/include
116
${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include
@@ -110,9 +137,12 @@ if(FLASH_MLA_ARCHS)
137
138
# Keep Stable ABI for the module, but *not* for CUDA/C++ files.
139
# This prevents Py_LIMITED_API from affecting nvcc and C++ compiles.
140
+ # Also enable C++20 for the FlashMLA sources (required for std::span, requires, etc.)
141
target_compile_options(_flashmla_C PRIVATE
142
$<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API>
115
- $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>)
143
+ $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>
144
+ $<$<COMPILE_LANGUAGE:CXX>:-std=c++20>
145
+ $<$<COMPILE_LANGUAGE:CUDA>:-std=c++20>)
146
147
define_extension_target(
148
_flashmla_extension_C
tests/kernels/attention/test_flashmla_sparse.py
+1
-1
@@ -43,7 +43,7 @@ def test_sparse_flashmla_decode_smoke():
43
device = torch.device("cuda")
44
batch_size = 1
45
seqlen_q = 1
46
- num_heads_q = 1
46
+ num_heads_q = 64
47
head_dim_k = 576
48
head_dim_v = 512
49
num_heads_k = 1
tests/v1/attention/test_sparse_mla_backends.py
+64
-11
@@ -51,10 +51,34 @@ SPARSE_BACKEND_BATCH_SPECS["large_q_pure_prefill"] = BatchSpec(
51
)
52
53
54
+def _float_to_e8m0_truncate(f: float) -> float:
55
+ """Simulate SM100's float -> e8m0 -> bf16 scale conversion.
56
+
57
+ e8m0 format only stores the exponent (power of 2).
58
+ cudaRoundZero truncates toward zero, meaning we round down to the
59
+ nearest power of 2.
60
+ """
61
+ if f <= 0:
62
+ return 0.0
63
+ # e8m0 = floor(log2(f)), then 2^(e8m0)
64
+ # This is equivalent to truncating to the nearest power of 2 below f
65
+ exp = math.floor(math.log2(f))
66
+ return 2.0**exp
67
+
68
+
69
def _dequantize_fp8_ds_mla_entry(
55
- cache_slice: torch.Tensor, kv_lora_rank: int, rope_dim: int, dtype: torch.dtype
70
+ cache_slice: torch.Tensor,
71
+ kv_lora_rank: int,
72
+ rope_dim: int,
73
+ dtype: torch.dtype,
74
+ simulate_sm100_e8m0_scales: bool = False,
75
) -> tuple[torch.Tensor, torch.Tensor]:
57
- """Dequantize a single fp8_ds_mla cache entry back to latent + rope."""
76
+ """Dequantize a single fp8_ds_mla cache entry back to latent + rope.
77
+
78
+ Args:
79
+ simulate_sm100_e8m0_scales: If True, simulate the SM100 kernel's
80
+ float -> e8m0 -> bf16 scale conversion path.
81
+ """
82
83
# The first kv_lora_rank bytes store FP8 latent values with one scale per
84
# 128 element tile written as float32 right after the latent payload.
@@ -63,10 +87,14 @@ def _dequantize_fp8_ds_mla_entry(
87
for tile_idx in range(4):
88
tile_start = tile_idx * 128
89
tile_end = tile_start + 128
90
+ scale_val = float(scales[tile_idx].item())
91
+ if simulate_sm100_e8m0_scales:
92
+ # Simulate the lossy float -> e8m0 -> bf16 conversion
93
+ scale_val = _float_to_e8m0_truncate(scale_val)
94
ops.convert_fp8(
95
latent[tile_start:tile_end],
96
cache_slice[tile_start:tile_end],
69
- float(scales[tile_idx].item()),
97
+ scale_val,
98
kv_dtype="fp8",
99
)
100
latent = latent.to(dtype)
@@ -77,9 +105,18 @@ def _dequantize_fp8_ds_mla_entry(
105
106
107
def _quantize_dequantize_fp8_ds_mla(
80
- kv_c: torch.Tensor, k_pe: torch.Tensor, block_size: int, scale: torch.Tensor
108
+ kv_c: torch.Tensor,
109
+ k_pe: torch.Tensor,
110
+ block_size: int,
111
+ scale: torch.Tensor,
112
+ simulate_sm100_e8m0_scales: bool = False,
113
) -> tuple[torch.Tensor, torch.Tensor]:
82
- """Round-trip kv_c/k_pe though the fp8_ds_mla cache layout."""
114
+ """Round-trip kv_c/k_pe though the fp8_ds_mla cache layout.
115
+
116
+ Args:
117
+ simulate_sm100_e8m0_scales: If True, simulate the SM100 kernel's
118
+ float -> e8m0 -> bf16 scale conversion in dequantization.
119
+ """
120
121
if kv_c.numel() == 0:
122
return kv_c.clone(), k_pe.clone()
@@ -108,7 +145,11 @@ def _quantize_dequantize_fp8_ds_mla(
145
block_offset = slot % block_size
146
cache_slice = tmp_cache[block_idx, block_offset]
147
latent, rope_vals = _dequantize_fp8_ds_mla_entry(
111
- cache_slice, kv_lora_rank, rope_dim, kv_c.dtype
148
+ cache_slice,
149
+ kv_lora_rank,
150
+ rope_dim,
151
+ kv_c.dtype,
152
+ simulate_sm100_e8m0_scales=simulate_sm100_e8m0_scales,
153
)
154
dequant_kv_c[token_idx] = latent
155
dequant_k_pe[token_idx] = rope_vals
@@ -143,7 +184,10 @@ def test_sparse_backend_decode_correctness(
184
batch_spec = SPARSE_BACKEND_BATCH_SPECS[batch_name]
185
186
# Model hyper-parameters (kept intentionally small for the unit test)
146
- num_heads = 128
187
+ total_num_heads = 128
188
+ # Compute per-rank heads for simulated TP
189
+ num_heads = max(1, total_num_heads // tensor_parallel_size)
190
+
191
kv_lora_rank = 512
192
qk_nope_head_dim = 128
193
qk_rope_head_dim = 64
@@ -179,7 +223,7 @@ def test_sparse_backend_decode_correctness(
223
)
224
model_config.dtype = dtype
225
model_config.get_num_attention_heads = MethodType(
182
- lambda self, parallel_config: max(1, num_heads // tensor_parallel_size),
226
+ lambda self, parallel_config: num_heads,
227
model_config,
228
)
229
model_config.get_num_kv_heads = MethodType(
@@ -195,10 +239,10 @@ def test_sparse_backend_decode_correctness(
239
scale = 1.0 / math.sqrt(head_size)
240
241
# Shared MLA projection weights to keep reference and backend in sync
198
- W_UK = torch.randn(
242
+ W_UK = torch.rand(
243
kv_lora_rank, num_heads, qk_nope_head_dim, dtype=dtype, device=device
244
)
201
- W_UV = torch.randn(kv_lora_rank, num_heads, v_head_dim, dtype=dtype, device=device)
245
+ W_UV = torch.rand(kv_lora_rank, num_heads, v_head_dim, dtype=dtype, device=device)
246
247
# Build synthetic decode-only workload
248
seq_lens = batch_spec.seq_lens
@@ -225,11 +269,15 @@ def test_sparse_backend_decode_correctness(
269
kv_c_full = torch.rand(s_len, kv_lora_rank, dtype=dtype, device=device)
270
k_pe_full = torch.rand(s_len, 1, qk_rope_head_dim, dtype=dtype, device=device)
271
272
+ # SM100 (Blackwell) uses float -> e8m0 -> bf16 scale conversion
273
+ # which truncates scales to powers of 2. Simulate this in reference.
274
+ is_sm100 = torch.cuda.get_device_capability()[0] >= 10
275
kv_c_full, k_pe_full = _quantize_dequantize_fp8_ds_mla(
276
kv_c_full,
277
k_pe_full.squeeze(1),
278
block_size=vllm_config.cache_config.block_size,
279
scale=kv_cache_scale,
280
+ simulate_sm100_e8m0_scales=is_sm100,
281
)
282
283
q_nope, q_pe = q_c.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1)
@@ -381,7 +429,12 @@ def test_sparse_backend_decode_correctness(
429
assert backend_output.dtype == sdpa_reference.dtype
430
assert torch.isfinite(backend_output).all()
431
384
- torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.5, atol=0.5)
432
+ # FP8 quantization introduces some error, but should be within reasonable bounds
433
+ # BF16 (auto) should be very accurate, FP8 allows slightly more tolerance
434
+ if kv_cache_dtype == "fp8_ds_mla":
435
+ torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.05, atol=0.05)
436
+ else:
437
+ torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.01, atol=0.01)
438
439
440
def _triton_convert_reference_impl(
vllm/v1/attention/backends/mla/flashmla_sparse.py
+64
-20
@@ -17,7 +17,6 @@ from vllm.model_executor.layers.attention.mla_attention import (
17
from vllm.platforms import current_platform
18
from vllm.platforms.interface import DeviceCapability
19
from vllm.triton_utils import tl, triton
20
-from vllm.utils.math_utils import cdiv
20
from vllm.v1.attention.backend import (
21
AttentionBackend,
22
AttentionCGSupport,
@@ -397,6 +396,10 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad
396
397
self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
398
self.mla_dims = get_mla_dims(self.model_config)
399
+ # FP8 decode kernel only supports h_q = 64 or 128, so we need to pad
400
+ self.fp8_decode_padded_heads = (
401
+ FlashMLASparseImpl._compute_fp8_decode_padded_heads(self.num_heads)
402
+ )
403
404
self.topk_tokens = vllm_config.model_config.hf_config.index_topk
405
self.use_fp8_kv_cache = cache_config.cache_dtype == "fp8_ds_mla"
@@ -417,14 +420,20 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad
420
(max_num_seqs, 1), dtype=torch.int32, device=self.device
421
)
422
420
- # Equation taken from FlashMLA/csrc/pybind.cpp
421
- h_q, h_k = self.num_heads, 1
422
- s_q = 1 # inversely proportional to s_q, so s_q = 1 is the largest
423
- max_num_sm_parts = int(
424
- max((sm_count // 2) / h_k // (cdiv(h_q // h_k, 2 * 64) * s_q), 1)
425
- )
423
+ # Equation taken from FlashMLA/csrc/api/sparse_decode.h
424
+ # For sparse FP8 decode, the formula depends on architecture:
425
+ # - SM90 (Hopper): num_sm_parts = num_sms / s_q / (h_q/64)
426
+ # - SM100 (Blackwell head64/head64x2): num_sm_parts = num_sms / s_q
427
+ # - SM100 (Blackwell head128): num_sm_parts = num_sms / s_q / 2
428
+ # For max buffer size, use s_q = 1 (the case that produces largest output)
429
+ # Use padded head count since that's what will be passed to the kernel
430
+ h_q = self.fp8_decode_padded_heads
431
if current_platform.is_device_capability_family(100):
427
- max_num_sm_parts *= 2
432
+ # SM100 head64 or head64x2 uses full SM count
433
+ max_num_sm_parts = sm_count
434
+ else:
435
+ # SM90 uses h_q/64 divisor
436
+ max_num_sm_parts = sm_count // max(1, h_q // 64)
437
self.tile_scheduler_metadata_buffer = torch.empty(
438
# TileSchedulerMetaDataSize = 8
439
# see: FlashMLA/csrc/params.h
@@ -455,12 +464,15 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad
464
"""
465
num_tokens = common_attn_metadata.num_actual_tokens
466
467
+ # Use padded head count since that's what the kernel will see
468
+ padded_heads = self.fp8_decode_padded_heads
469
+
470
# Build metadata for all tokens as a single batch
471
tile_scheduler_metadata, num_splits = get_mla_metadata(
472
cache_seqlens=self.topk_tokens_tensor[:1], # Single batch
461
- num_q_tokens_per_head_k=num_tokens * self.num_heads,
473
+ num_q_tokens_per_head_k=num_tokens * padded_heads,
474
topk=self.topk_tokens,
463
- num_heads_q=self.num_heads,
475
+ num_heads_q=padded_heads,
476
num_heads_k=1,
477
is_fp8_kvcache=True,
478
)
@@ -606,11 +618,13 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad
618
query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
619
decode_query_len = (query_start_loc_cpu[1] - query_start_loc_cpu[0]).item()
620
621
+ # Use padded head count since that's what the kernel will see
622
+ padded_heads = self.fp8_decode_padded_heads
623
tile_scheduler_metadata, num_splits = get_mla_metadata(
624
cache_seqlens=self.topk_tokens_tensor[:num_decodes],
611
- num_q_tokens_per_head_k=decode_query_len * self.num_heads,
625
+ num_q_tokens_per_head_k=decode_query_len * padded_heads,
626
topk=self.topk_tokens,
613
- num_heads_q=self.num_heads,
627
+ num_heads_q=padded_heads,
628
num_heads_k=1,
629
is_fp8_kvcache=True,
630
)
@@ -689,6 +703,12 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad
703
704
705
class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):
706
+ @staticmethod
707
+ def _compute_fp8_decode_padded_heads(num_heads: int) -> int:
708
+ # FP8 decode kernel only supports h_q = 64 or 128
709
+ # Compute padded head count for decode
710
+ return 64 if num_heads <= 64 else 128
711
+
712
def __init__(
713
self,
714
num_heads: int,
@@ -722,7 +742,11 @@ class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):
742
self.softmax_scale = scale
743
assert indexer is not None
744
self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer
725
- self.padding = 128 if current_platform.is_device_capability_family(100) else 64
745
+ # Prefill BF16 kernel requires 64 on Hopper, 128 on Blackwell
746
+ self.prefill_padding = (
747
+ 128 if current_platform.is_device_capability_family(100) else 64
748
+ )
749
+ self.fp8_decode_padded_heads = self._compute_fp8_decode_padded_heads(num_heads)
750
751
if kv_cache_dtype == "fp8_ds_mla":
752
# Reserve workspace during initialization
@@ -903,8 +927,22 @@ class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):
927
kv_c_and_k_pe_cache: torch.Tensor,
928
topk_indices: torch.Tensor,
929
kernel_metadata: FlashMLASparseMetadata.FP8KernelMetadata,
906
- ) -> torch.Tensor:
907
- return flash_mla_with_kvcache(
930
+ ) -> tuple[torch.Tensor, torch.Tensor]:
931
+ # q shape: (batch, seq_len, num_heads, head_dim)
932
+ actual_num_heads = q.size(2)
933
+ padded_num_heads = self.fp8_decode_padded_heads
934
+
935
+ # Pad query if needed (kernel only supports h_q = 64 or 128)
936
+ if actual_num_heads < padded_num_heads:
937
+ logger.warning_once(
938
+ f"Padding num_heads from {actual_num_heads} to "
939
+ f"{padded_num_heads} for FP8 sparse decode kernel"
940
+ )
941
+ q_padded = q.new_zeros((q.size(0), q.size(1), padded_num_heads, q.size(3)))
942
+ q_padded[:, :, :actual_num_heads, :] = q
943
+ q = q_padded
944
+
945
+ out, lse = flash_mla_with_kvcache(
946
q=q,
947
k_cache=kv_c_and_k_pe_cache.view(torch.uint8).unsqueeze(-2),
948
block_table=kernel_metadata.dummy_block_table,
@@ -917,6 +955,12 @@ class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):
955
softmax_scale=self.softmax_scale,
956
)
957
958
+ # Slice output back to actual head count if we padded
959
+ if actual_num_heads < padded_num_heads:
960
+ out = out[:, :, :actual_num_heads, :]
961
+
962
+ return out, lse
963
+
964
def _bf16_flash_mla_kernel(
965
self,
966
q: torch.Tensor,
@@ -930,13 +974,13 @@ class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):
974
975
# NOTE(Chen): kernel requires num_local_head to be a multiple of
976
# 64 on hopper and 128 on blackwell
933
- if self.num_heads % self.padding != 0:
934
- assert self.padding % self.num_heads == 0
977
+ if self.num_heads % self.prefill_padding != 0:
978
+ assert self.prefill_padding % self.num_heads == 0
979
logger.warning_once(
936
- f"padding num_heads to {self.padding} due to sparse attn "
937
- "kernel requirement"
980
+ f"Padding num_heads from {self.num_heads} to "
981
+ f"{self.prefill_padding} for BF16 sparse prefill kernel"
982
)
939
- q_padded = q.new_empty((q.shape[0], self.padding, q.shape[2]))
983
+ q_padded = q.new_empty((q.shape[0], self.prefill_padding, q.shape[2]))
984
q_padded[:, : self.num_heads, :] = q
985
q = q_padded
986