Add FlashMLA H100 tests to CI, fix them after #32810 (#50322)

.buildkite/test_areas/kernels.yaml

@@ -143,6 +143,24 @@ steps:
   commands:
     - pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py
 
+- label: Kernels FlashMLA Test (H100)
+  key: kernels-flashmla-test-h100
+  timeout_in_minutes: 25
+  device: h100
+  num_devices: 1
+  source_file_dependencies:
+  - cmake/external_projects/flashmla.cmake
+  - vllm/v1/attention/ops/flashmla.py
+  - vllm/v1/attention/backends/mla/flashmla.py
+  - vllm/v1/attention/backends/mla/flashmla_sparse.py
+  - tests/kernels/attention/test_flashmla.py
+  - tests/kernels/attention/test_flashmla_sparse.py
+  - tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py
+  commands:
+    - pytest -v -s kernels/attention/test_flashmla.py
+    - pytest -v -s kernels/attention/test_flashmla_sparse.py
+    - pytest -v -s kernels/attention/test_mla_cross_layer_kernel_equivalence.py
+
 - label: Kernels Quantization Test %N
   key: kernels-quantization-test
   timeout_in_minutes: 60

tests/kernels/attention/test_flashmla.py

@@ -10,7 +10,9 @@ import torch
 from vllm.triton_utils import triton
 from vllm.v1.attention.ops.flashmla import (
     flash_mla_with_kvcache,
+    flash_mla_with_kvcache_fp8,
     get_mla_metadata,
+    get_mla_metadata_dense_fp8,
     is_flashmla_dense_supported,
 )
 
@@ -86,9 +88,14 @@ def test_flash_mla(
         )
     blocked_v = blocked_k[..., :dv]
 
-    tile_scheduler_metadata, num_splits = get_mla_metadata(
-        cache_seqlens, s_q * h_q // h_kv, h_kv
-    )
+    if use_fp8:
+        tile_scheduler_metadata, num_splits = get_mla_metadata_dense_fp8(
+            cache_seqlens, s_q * h_q // h_kv, h_kv
+        )
+    else:
+        tile_scheduler_metadata, num_splits = get_mla_metadata(
+            cache_seqlens, s_q * h_q // h_kv, h_kv
+        )
 
     init_dtype = q.dtype
     if use_fp8:
@@ -104,6 +111,19 @@ def test_flash_mla(
         descale_k = None
 
     def flash_mla():
+        if use_fp8:
+            return flash_mla_with_kvcache_fp8(
+                q,
+                blocked_k,
+                block_table,
+                cache_seqlens,
+                dv,
+                tile_scheduler_metadata,
+                num_splits,
+                causal=causal,
+                descale_q=descale_q,
+                descale_k=descale_k,
+            )
         return flash_mla_with_kvcache(
             q,
             blocked_k,
@@ -113,8 +133,6 @@ def test_flash_mla(
             tile_scheduler_metadata,
             num_splits,
             causal=causal,
-            descale_q=descale_q,
-            descale_k=descale_k,
         )
 
     def scaled_dot_product_attention(query, key, value, is_causal=False):

tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py

@@ -442,7 +442,9 @@ def test_flashmla_dense_fp8_decode_unified_slot_view():
     n_layers = 3
     layer = 1
 
-    q = torch.randn(bs, 1, h_q, head_dim, device=dev, dtype=torch.bfloat16) * 0.1
+    # With a quantized KV cache the decode query is quantized to fp8 as well;
+    # the fp8 dense MLA kernel requires q already in fp8.
+    q = (torch.randn(bs, 1, h_q, head_dim, device=dev) * 0.1).to(torch.float8_e4m3fn)
     kv_data = (torch.randn(num_blocks, page, head_dim, device=dev) * 0.1).to(
         torch.float8_e4m3fn
     )