diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 81fa9c5f1d..3644990674 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.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 diff --git a/tests/kernels/attention/test_flashmla.py b/tests/kernels/attention/test_flashmla.py index 657b256f46..84744a81e7 100644 --- a/tests/kernels/attention/test_flashmla.py +++ b/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): diff --git a/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py b/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py index 48a236a015..ac46073559 100644 --- a/tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py +++ b/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 )