Add FlashMLA H100 tests to CI, fix them after #32810 (#50322)
Jane (Yuan) Xu committed
Jul 30, 2026 at 01:06 UTC
a7a204cc6ec99b375985e564f4f271e68ee8f5b8
3 files changed
+44
-6
.buildkite/test_areas/kernels.yaml
+18
@@ -143,6 +143,24 @@ steps:
143
commands:
144
- pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py
145
146
+- label: Kernels FlashMLA Test (H100)
147
+ key: kernels-flashmla-test-h100
148
+ timeout_in_minutes: 25
149
+ device: h100
150
+ num_devices: 1
151
+ source_file_dependencies:
152
+ - cmake/external_projects/flashmla.cmake
153
+ - vllm/v1/attention/ops/flashmla.py
154
+ - vllm/v1/attention/backends/mla/flashmla.py
155
+ - vllm/v1/attention/backends/mla/flashmla_sparse.py
156
+ - tests/kernels/attention/test_flashmla.py
157
+ - tests/kernels/attention/test_flashmla_sparse.py
158
+ - tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py
159
+ commands:
160
+ - pytest -v -s kernels/attention/test_flashmla.py
161
+ - pytest -v -s kernels/attention/test_flashmla_sparse.py
162
+ - pytest -v -s kernels/attention/test_mla_cross_layer_kernel_equivalence.py
163
+
164
- label: Kernels Quantization Test %N
165
key: kernels-quantization-test
166
timeout_in_minutes: 60
tests/kernels/attention/test_flashmla.py
+23
-5
@@ -10,7 +10,9 @@ import torch
10
from vllm.triton_utils import triton
11
from vllm.v1.attention.ops.flashmla import (
12
flash_mla_with_kvcache,
13
+ flash_mla_with_kvcache_fp8,
14
get_mla_metadata,
15
+ get_mla_metadata_dense_fp8,
16
is_flashmla_dense_supported,
17
)
18
@@ -86,9 +88,14 @@ def test_flash_mla(
88
)
89
blocked_v = blocked_k[..., :dv]
90
89
- tile_scheduler_metadata, num_splits = get_mla_metadata(
90
- cache_seqlens, s_q * h_q // h_kv, h_kv
91
- )
91
+ if use_fp8:
92
+ tile_scheduler_metadata, num_splits = get_mla_metadata_dense_fp8(
93
+ cache_seqlens, s_q * h_q // h_kv, h_kv
94
+ )
95
+ else:
96
+ tile_scheduler_metadata, num_splits = get_mla_metadata(
97
+ cache_seqlens, s_q * h_q // h_kv, h_kv
98
+ )
99
100
init_dtype = q.dtype
101
if use_fp8:
@@ -104,6 +111,19 @@ def test_flash_mla(
111
descale_k = None
112
113
def flash_mla():
114
+ if use_fp8:
115
+ return flash_mla_with_kvcache_fp8(
116
+ q,
117
+ blocked_k,
118
+ block_table,
119
+ cache_seqlens,
120
+ dv,
121
+ tile_scheduler_metadata,
122
+ num_splits,
123
+ causal=causal,
124
+ descale_q=descale_q,
125
+ descale_k=descale_k,
126
+ )
127
return flash_mla_with_kvcache(
128
q,
129
blocked_k,
@@ -113,8 +133,6 @@ def test_flash_mla(
133
tile_scheduler_metadata,
134
num_splits,
135
causal=causal,
116
- descale_q=descale_q,
117
- descale_k=descale_k,
136
)
137
138
def scaled_dot_product_attention(query, key, value, is_causal=False):
tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py
+3
-1
@@ -442,7 +442,9 @@ def test_flashmla_dense_fp8_decode_unified_slot_view():
442
n_layers = 3
443
layer = 1
444
445
- q = torch.randn(bs, 1, h_q, head_dim, device=dev, dtype=torch.bfloat16) * 0.1
445
+ # With a quantized KV cache the decode query is quantized to fp8 as well;
446
+ # the fp8 dense MLA kernel requires q already in fp8.
447
+ q = (torch.randn(bs, 1, h_q, head_dim, device=dev) * 0.1).to(torch.float8_e4m3fn)
448
kv_data = (torch.randn(num_blocks, page, head_dim, device=dev) * 0.1).to(
449
torch.float8_e4m3fn
450
)