| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for MLA prefill backend fused-quant-output support. |
| 4 | |
| 5 | Covers two things: |
| 6 | * `MLAPrefillBackend.supports_quant_output`, the capability gate that decides |
| 7 | whether the prefill kernel writes quantized output directly (FA4 native |
| 8 | fused FP8, see flash-attention#135) instead of the post-quant path. |
| 9 | * The numerical equivalence of that fused FP8 write versus the bf16-attention |
| 10 | + standalone static-FP8-quant path it replaces (GPU-only, SM100/SM110). |
| 11 | """ |
| 12 | |
| 13 | from types import SimpleNamespace |
| 14 | from unittest.mock import patch |
| 15 | |
| 16 | import pytest |
| 17 | import torch |
| 18 | |
| 19 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 20 | kFp8Dynamic128Sym, |
| 21 | kFp8StaticTensorSym, |
| 22 | kNvfp4Dynamic, |
| 23 | ) |
| 24 | from vllm.platforms.interface import DeviceCapability |
| 25 | from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend |
| 26 | from vllm.v1.attention.backends.mla.prefill.flash_attn import ( |
| 27 | FlashAttnPrefillBackend, |
| 28 | ) |
| 29 | |
| 30 | _FA_MODULE = "vllm.v1.attention.backends.mla.prefill.flash_attn" |
| 31 | |
| 32 | |
| 33 | class _DummyPrefillBackend(MLAPrefillBackend): |
| 34 | """Concrete backend that does NOT override supports_quant_output.""" |
| 35 | |
| 36 | @staticmethod |
| 37 | def get_name() -> str: |
| 38 | return "DUMMY" |
| 39 | |
| 40 | def run_prefill_new_tokens(self, *args, **kwargs): # pragma: no cover |
| 41 | raise NotImplementedError |
| 42 | |
| 43 | def run_prefill_context_chunk(self, *args, **kwargs): # pragma: no cover |
| 44 | raise NotImplementedError |
| 45 | |
| 46 | |
| 47 | @pytest.mark.parametrize( |
| 48 | "quant_key", [kFp8StaticTensorSym, kFp8Dynamic128Sym, kNvfp4Dynamic, None] |
| 49 | ) |
| 50 | def test_base_backend_never_supports_quant_output(quant_key): |
| 51 | """The base default opts every backend out unless it overrides.""" |
| 52 | backend = object.__new__(_DummyPrefillBackend) |
| 53 | assert backend.supports_quant_output(quant_key) is False |
| 54 | |
| 55 | |
| 56 | def _make_fa_backend(version: int | None, is_vllm_fa: bool): |
| 57 | """Build a FlashAttnPrefillBackend without running its heavy __init__.""" |
| 58 | backend = object.__new__(FlashAttnPrefillBackend) |
| 59 | backend.vllm_flash_attn_version = version |
| 60 | backend._is_vllm_fa = is_vllm_fa |
| 61 | return backend |
| 62 | |
| 63 | |
| 64 | @pytest.mark.parametrize( |
| 65 | ("version", "is_vllm_fa", "dc_major", "quant_key", "expected"), |
| 66 | [ |
| 67 | # FA4 + vLLM-FA + Blackwell SM100/SM110 + static FP8 -> fused. |
| 68 | (4, True, 10, kFp8StaticTensorSym, True), |
| 69 | (4, True, 11, kFp8StaticTensorSym, True), |
| 70 | # Wrong compute capability (SM90 / SM120) -> not supported (#135). |
| 71 | (4, True, 9, kFp8StaticTensorSym, False), |
| 72 | (4, True, 12, kFp8StaticTensorSym, False), |
| 73 | # Not FA4. |
| 74 | (3, True, 10, kFp8StaticTensorSym, False), |
| 75 | (2, True, 10, kFp8StaticTensorSym, False), |
| 76 | (None, True, 10, kFp8StaticTensorSym, False), |
| 77 | # Upstream (ROCm) flash-attn, not vLLM-FA. |
| 78 | (4, False, 10, kFp8StaticTensorSym, False), |
| 79 | # Quant keys not wired through FA4 yet. |
| 80 | (4, True, 10, kFp8Dynamic128Sym, False), |
| 81 | (4, True, 10, kNvfp4Dynamic, False), |
| 82 | ], |
| 83 | ) |
| 84 | def test_flash_attn_supports_quant_output( |
| 85 | version, is_vllm_fa, dc_major, quant_key, expected |
| 86 | ): |
| 87 | backend = _make_fa_backend(version, is_vllm_fa) |
| 88 | with patch(f"{_FA_MODULE}.current_platform") as plat: |
| 89 | plat.get_device_capability.return_value = DeviceCapability( |
| 90 | major=dc_major, minor=0 |
| 91 | ) |
| 92 | assert backend.supports_quant_output(quant_key) is expected |
| 93 | |
| 94 | |
| 95 | def test_flash_attn_supports_quant_output_unknown_device(): |
| 96 | """A None device capability (e.g. capability probe failed) is safe.""" |
| 97 | backend = _make_fa_backend(version=4, is_vllm_fa=True) |
| 98 | with patch(f"{_FA_MODULE}.current_platform") as plat: |
| 99 | plat.get_device_capability.return_value = None |
| 100 | assert backend.supports_quant_output(kFp8StaticTensorSym) is False |
| 101 | |
| 102 | |
| 103 | @pytest.mark.parametrize( |
| 104 | ("version", "enable_jit_warmup", "expected_calls"), |
| 105 | [(4, True, 1), (4, False, 0), (3, True, 0)], |
| 106 | ) |
| 107 | def test_flash_attn_registers_warmup_only_for_fa4( |
| 108 | version: int, |
| 109 | enable_jit_warmup: bool, |
| 110 | expected_calls: int, |
| 111 | ): |
| 112 | vllm_config = SimpleNamespace( |
| 113 | kernel_config=SimpleNamespace(enable_jit_warmup=enable_jit_warmup) |
| 114 | ) |
| 115 | with ( |
| 116 | patch(f"{_FA_MODULE}.flash_attn_varlen_func"), |
| 117 | patch(f"{_FA_MODULE}.get_flash_attn_version", return_value=version), |
| 118 | patch(f"{_FA_MODULE}._FA4_MLA_PREFILL_KERNEL.register_warmup") as register, |
| 119 | patch(f"{_FA_MODULE}.current_platform") as platform, |
| 120 | ): |
| 121 | platform.get_device_capability.return_value = DeviceCapability(10, 0) |
| 122 | FlashAttnPrefillBackend( |
| 123 | num_heads=16, |
| 124 | scale=1.0, |
| 125 | kv_lora_rank=512, |
| 126 | qk_nope_head_dim=128, |
| 127 | qk_rope_head_dim=64, |
| 128 | v_head_dim=128, |
| 129 | vllm_config=vllm_config, |
| 130 | ) |
| 131 | |
| 132 | assert register.call_count == expected_calls |
| 133 | |
| 134 | |
| 135 | def test_flash_attn_prefill_backend_signature_accepts_fused_kwargs(): |
| 136 | """run_prefill_new_tokens must accept out/output_scale so the direct |
| 137 | (non-**kwargs) call in forward_mha type- and runtime-checks.""" |
| 138 | import inspect |
| 139 | |
| 140 | params = inspect.signature( |
| 141 | FlashAttnPrefillBackend.run_prefill_new_tokens |
| 142 | ).parameters |
| 143 | assert "out" in params |
| 144 | assert "output_scale" in params |
| 145 | # The base contract must expose them too (Liskov / direct call site). |
| 146 | base_params = inspect.signature(MLAPrefillBackend.run_prefill_new_tokens).parameters |
| 147 | assert "out" in base_params |
| 148 | assert "output_scale" in base_params |
| 149 | |
| 150 | |
| 151 | def test_mla_impl_forward_mha_accepts_output_scale(): |
| 152 | """The abstract MLA impl forward_mha must carry output_scale so every |
| 153 | override (and the unconditional forward_impl call) stays compatible.""" |
| 154 | import inspect |
| 155 | |
| 156 | from vllm.v1.attention.backend import MLAAttentionImpl |
| 157 | |
| 158 | params = inspect.signature(MLAAttentionImpl.forward_mha).parameters |
| 159 | assert "output_scale" in params |
| 160 | assert params["output_scale"].default is None |
| 161 | |
| 162 | |
| 163 | def _fused_fp8_skip_reason() -> str | None: |
| 164 | """FA4 fused FP8 output needs a real Blackwell SM100/SM110 GPU.""" |
| 165 | if not torch.cuda.is_available(): |
| 166 | return "requires CUDA" |
| 167 | major = torch.cuda.get_device_capability()[0] |
| 168 | if major not in (10, 11): |
| 169 | return f"FA4 fused FP8 output requires SM100/SM110, got SM{major}x" |
| 170 | return None |
| 171 | |
| 172 | |
| 173 | _FUSED_FP8_SKIP = _fused_fp8_skip_reason() |
| 174 | |
| 175 | |
| 176 | @pytest.mark.skipif(_FUSED_FP8_SKIP is not None, reason=_FUSED_FP8_SKIP or "") |
| 177 | def test_fa4_fused_fp8_output_matches_post_quant(default_vllm_config): |
| 178 | """FA4's fused FP8 write (output_scale, flash-attention#135) must match the |
| 179 | bf16-attention + standalone static-FP8-quant path it replaces, since |
| 180 | production uses the same output_scale for both.""" |
| 181 | from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 |
| 182 | from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape |
| 183 | from vllm.platforms import current_platform |
| 184 | from vllm.vllm_flash_attn import flash_attn_varlen_func |
| 185 | |
| 186 | torch.manual_seed(0) |
| 187 | device = torch.device("cuda") |
| 188 | fp8_dtype = current_platform.fp8_dtype() |
| 189 | |
| 190 | # MLA prefill head dims (post kv_b_proj): q/k = qk_nope(128)+qk_rope(64), |
| 191 | # v = v_head_dim(128); DeepSeek-V2-Lite has 16 query heads. |
| 192 | num_heads, qk_head_dim, v_head_dim, seqlen = 16, 192, 128, 512 |
| 193 | cu_seqlens = torch.tensor([0, seqlen], dtype=torch.int32, device=device) |
| 194 | q = torch.randn(seqlen, num_heads, qk_head_dim, dtype=torch.bfloat16, device=device) |
| 195 | k = torch.randn(seqlen, num_heads, qk_head_dim, dtype=torch.bfloat16, device=device) |
| 196 | v = torch.randn(seqlen, num_heads, v_head_dim, dtype=torch.bfloat16, device=device) |
| 197 | |
| 198 | fa_kwargs = dict( |
| 199 | cu_seqlens_q=cu_seqlens, |
| 200 | cu_seqlens_k=cu_seqlens, |
| 201 | max_seqlen_q=seqlen, |
| 202 | max_seqlen_k=seqlen, |
| 203 | causal=True, |
| 204 | fa_version=4, |
| 205 | ) |
| 206 | |
| 207 | # Reference: bf16 attention, then standalone static per-tensor FP8 quant. |
| 208 | out_bf16 = flash_attn_varlen_func(q=q, k=k, v=v, **fa_kwargs) |
| 209 | out_2d = out_bf16.reshape(seqlen, num_heads * v_head_dim) |
| 210 | # Scale the amax near e4m3 max so the check uses the representable range. |
| 211 | finfo = torch.finfo(fp8_dtype) |
| 212 | scale = (out_2d.abs().max() / finfo.max).to(torch.float32).reshape(1) |
| 213 | quant_op = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR) |
| 214 | ref_fp8, _ = quant_op(out_2d, scale) |
| 215 | |
| 216 | # Feature: FA4 writes e4m3 into the (tokens, heads*dim) buffer directly. |
| 217 | fused_fp8 = torch.empty( |
| 218 | seqlen, num_heads * v_head_dim, dtype=fp8_dtype, device=device |
| 219 | ) |
| 220 | flash_attn_varlen_func( |
| 221 | q=q, |
| 222 | k=k, |
| 223 | v=v, |
| 224 | out=fused_fp8.view(seqlen, num_heads, v_head_dim), |
| 225 | output_scale=scale, |
| 226 | **fa_kwargs, |
| 227 | ) |
| 228 | |
| 229 | # Non-degenerate (catches a no-op / all-zero write). |
| 230 | assert torch.isfinite(fused_fp8.float()).all() |
| 231 | assert fused_fp8.float().abs().any() |
| 232 | |
| 233 | # e4m3 has 3 mantissa bits, so allow ~1 mantissa step of rounding slack. |
| 234 | ref = ref_fp8.float() * scale |
| 235 | got = fused_fp8.float() * scale |
| 236 | torch.testing.assert_close(got, ref, rtol=0.125, atol=float(scale) * 2) |
| 237 | |
| 238 | # ...and most elements land in the exact same fp8 bucket. |
| 239 | exact = (fused_fp8.view(torch.uint8) == ref_fp8.view(torch.uint8)).float().mean() |
| 240 | assert exact > 0.9, f"only {exact:.1%} of fused FP8 outputs matched the baseline" |