@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / eb04c769d3
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

feat: MLA prefill enable FA4 fp8 output (#43050)

Signed-off-by: Carl You <4531192+carlyou@users.noreply.github.com>

Carl Y committed Jun 16, 2026 at 07:10 UTC eb04c769d38d0f39160bd1a2c5279359519de802
13 files changed +489 -20
benchmarks/attention_benchmarks/benchmark.py
+82 -1
@@ -53,6 +53,16 @@ from common import (
53 from vllm.v1.worker.workspace import init_workspace_manager
54
55
56 +def _str2bool(v) -> bool:
57 + if isinstance(v, bool):
58 + return v
59 + if v.lower() in ("true", "1", "yes", "t"):
60 + return True
61 + if v.lower() in ("false", "0", "no", "f"):
62 + return False
63 + raise argparse.ArgumentTypeError(f"expected a boolean, got {v!r}")
64 +
65 +
66 def run_standard_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
67 """Run standard attention benchmark (Flash/Triton/FlashInfer)."""
68 from runner import run_attention_benchmark
@@ -485,6 +495,20 @@ def main():
495 help="Prefill backends to compare (fa2, fa3, fa4). "
496 "Uses the first decode backend for impl construction.",
497 )
498 + parser.add_argument(
499 + "--fp8-output-scale",
500 + type=float,
501 + help="Static per-tensor scale enabling the MLA prefill FP8-output "
502 + "comparison on FA4 (fused write vs standalone post-quant).",
503 + )
504 + parser.add_argument(
505 + "--fuse-quant-op",
506 + nargs="+",
507 + type=_str2bool,
508 + help="FP8-output write path(s) to run: false = bf16 attention + "
509 + "standalone static-FP8 quant, true = FA4 writes FP8 directly. "
510 + "Default: both.",
511 + )
512
513 # Batch specifications
514 parser.add_argument(
@@ -618,6 +642,12 @@ def main():
642 # Prefill backends (e.g., ["fa3", "fa4"])
643 args.prefill_backends = yaml_config.get("prefill_backends", None)
644
645 + # FP8 output benchmark knobs; CLI wins.
646 + if args.fp8_output_scale is None:
647 + args.fp8_output_scale = yaml_config.get("fp8_output_scale", None)
648 + if args.fuse_quant_op is None:
649 + args.fuse_quant_op = yaml_config.get("fuse_quant_op", None)
650 +
651 # Check for special modes
652 args.mode = yaml_config.get("mode", None)
653
@@ -787,8 +817,59 @@ def main():
817 "skipped (timings are placeholder zeros).[/]"
818 )
819
820 + # FA4 fused FP8 output vs standalone post-quant, on the same fa4 kernel:
821 + # the delta is the post-quant kernel the fused path removes.
822 + fp8_output_scale = getattr(args, "fp8_output_scale", None)
823 + if fp8_output_scale is not None:
824 + decode_backend = backends[0]
825 + fuse_variants = args.fuse_quant_op or [False, True]
826 + label_of = {False: "post_quant", True: "fused"}
827 + console.print(
828 + f"[yellow]FP8 output comparison @ scale={fp8_output_scale} "
829 + f"(prefill=fa4, decode impl={decode_backend})[/]"
830 + )
831 + fp8_results = []
832 + total = len(fuse_variants) * len(args.batch_specs)
833 + with tqdm(total=total, desc="FP8 output benchmarking") as pbar:
834 + for spec in args.batch_specs:
835 + for fuse in fuse_variants:
836 + config = BenchmarkConfig(
837 + backend=decode_backend,
838 + batch_spec=spec,
839 + num_layers=args.num_layers,
840 + head_dim=args.head_dim,
841 + num_q_heads=args.num_q_heads,
842 + num_kv_heads=args.num_kv_heads,
843 + block_size=args.block_size,
844 + device=args.device,
845 + repeats=args.repeats,
846 + warmup_iters=args.warmup_iters,
847 + profile_memory=args.profile_memory,
848 + kv_cache_dtype=args.kv_cache_dtype,
849 + use_cuda_graphs=args.cuda_graphs,
850 + prefill_backend="fa4",
851 + )
852 + result = run_benchmark(
853 + config, output_scale=fp8_output_scale, fuse_quant_op=fuse
854 + )
855 + label = label_of[fuse]
856 + labeled_config = replace(result.config, backend=label)
857 + result = replace(result, config=labeled_config)
858 + fp8_results.append(result)
859 +
860 + if not result.success:
861 + console.print(f"[red]Error {label} {spec}: {result.error}[/]")
862 +
863 + pbar.update(1)
864 +
865 + console.print("\n[bold green]FP8 Output Results:[/]")
866 + formatter = ResultsFormatter(console)
867 + labels = [label_of[f] for f in fuse_variants]
868 + formatter.print_table(fp8_results, labels, compare_to_fastest=True)
869 + all_results = fp8_results
870 +
871 # Handle special mode: decode_vs_prefill comparison
791 - if hasattr(args, "mode") and args.mode == "decode_vs_prefill":
872 + elif hasattr(args, "mode") and args.mode == "decode_vs_prefill":
873 console.print("[yellow]Mode: Decode vs Prefill pipeline comparison[/]")
874 console.print(
875 "[dim]For each query length, testing both decode and prefill pipelines[/]"
benchmarks/attention_benchmarks/configs/mla_fa4_fp8_output.yaml new
+44
@@ -0,0 +1,44 @@
1 +# MLA prefill FP8-output microbenchmark (FA4).
2 +# Compares the fused FP8 write against bf16 attention + a standalone static-FP8
3 +# quant; the delta is the post-quant kernel the fused path removes.
4 +# DeepSeek-Coder-V2-Lite dims; FA4 needs SM100/110.
5 +#
6 +# Usage:
7 +# python benchmark.py --config configs/mla_fa4_fp8_output.yaml
8 +
9 +description: "MLA prefill FA4 fused-FP8 output vs post-quant"
10 +
11 +model:
12 + name: "deepseek-v2-lite"
13 + num_layers: 27
14 + num_q_heads: 16
15 + num_kv_heads: 1
16 + head_dim: 576
17 + kv_lora_rank: 512
18 + qk_nope_head_dim: 128
19 + qk_rope_head_dim: 64
20 + v_head_dim: 128
21 + block_size: 128
22 +
23 +# Pure prefill (q_len == kv_len) so every token goes through forward_mha.
24 +batch_specs:
25 + - "q512"
26 + - "q1k"
27 + - "q2k"
28 + - "q4k"
29 + - "q8k"
30 + - "2q4k"
31 + - "4q4k"
32 + - "8q4k"
33 +
34 +# Only used to construct the MLA impl; the pure-prefill specs skip decode.
35 +decode_backends:
36 + - CUTLASS_MLA
37 +
38 +# Sweep the two FP8 write paths (prefill backend is fixed to fa4).
39 +fp8_output_scale: 0.1
40 +fuse_quant_op: [false, true]
41 +
42 +device: "cuda:0"
43 +repeats: 50
44 +warmup_iters: 10
benchmarks/attention_benchmarks/mla_runner.py
+64 -11
@@ -708,6 +708,8 @@ def _run_single_benchmark(
708 device: torch.device,
709 indexer=None,
710 kv_cache_dtype: str | None = None,
711 + output_scale: float | None = None,
712 + fuse_quant_op: bool = False,
713 ) -> BenchmarkResult:
714 """
715 Run a single benchmark iteration.
@@ -721,6 +723,11 @@ def _run_single_benchmark(
723 mla_dims: MLA dimension configuration
724 device: Target device
725 indexer: Optional MockIndexer for sparse backends
726 + output_scale: Static per-tensor FP8 scale for prefill output. None
727 + keeps the plain bf16 output (no quantization).
728 + fuse_quant_op: With output_scale set, True lets the prefill kernel write
729 + FP8 directly; False runs bf16 attention then a standalone static-FP8
730 + quant. The delta isolates the saved post-quant kernel.
731
732 Returns:
733 BenchmarkResult with timing statistics
@@ -824,23 +831,55 @@ def _run_single_benchmark(
831 num_prefill, mla_dims, query_fmt, device, torch.bfloat16
832 )
833
834 + # Prefill FP8 output: fused (kernel writes e4m3) vs separate post-quant.
835 + prefill_fp8_output = None
836 + prefill_output_scale = None
837 + prefill_quant_op = None
838 + if has_prefill and output_scale is not None:
839 + from vllm.platforms import current_platform
840 +
841 + prefill_output_scale = torch.tensor(
842 + [output_scale], device=device, dtype=torch.float32
843 + )
844 + if fuse_quant_op:
845 + prefill_fp8_output = torch.empty_like(
846 + prefill_inputs["output"], dtype=current_platform.fp8_dtype()
847 + )
848 + else:
849 + from vllm.model_executor.layers.quantization.input_quant_fp8 import (
850 + QuantFP8,
851 + )
852 + from vllm.model_executor.layers.quantization.utils.quant_utils import (
853 + GroupShape,
854 + )
855 +
856 + prefill_quant_op = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
857 +
858 + fused_output = output_scale is not None and fuse_quant_op
859 +
860 # Build forward function (runs a single decode/prefill pass)
861 def forward_fn():
862 results = []
863 if has_decode:
864 results.append(impl.forward_mqa(decode_inputs, kv_cache, metadata, layer))
865 if has_prefill:
833 - results.append(
834 - impl.forward_mha(
835 - prefill_inputs["q"],
836 - prefill_inputs["k_c_normed"],
837 - prefill_inputs["k_pe"],
838 - kv_cache,
839 - metadata,
840 - prefill_inputs["k_scale"],
841 - prefill_inputs["output"],
842 - )
866 + out = impl.forward_mha(
867 + prefill_inputs["q"],
868 + prefill_inputs["k_c_normed"],
869 + prefill_inputs["k_pe"],
870 + kv_cache,
871 + metadata,
872 + prefill_inputs["k_scale"],
873 + prefill_fp8_output if fused_output else prefill_inputs["output"],
874 + prefill_output_scale if fused_output else None,
875 )
876 + if fused_output:
877 + out = prefill_fp8_output
878 + elif prefill_quant_op is not None:
879 + out, _ = prefill_quant_op(
880 + prefill_inputs["output"], prefill_output_scale
881 + )
882 + results.append(out)
883 return results[0] if len(results) == 1 else tuple(results)
884
885 def benchmark_fn():
@@ -881,6 +920,8 @@ def _run_mla_benchmark_batched(
920 configs_with_params: list[tuple], # [(config, threshold, num_splits), ...]
921 index_topk: int = 2048,
922 prefill_backend: str | None = None,
923 + output_scale: float | None = None,
924 + fuse_quant_op: bool = False,
925 ) -> list[BenchmarkResult]:
926 """
927 Unified batched MLA benchmark runner for all backends.
@@ -1020,6 +1061,8 @@ def _run_mla_benchmark_batched(
1061 device,
1062 indexer=indexer,
1063 kv_cache_dtype=kv_cache_dtype,
1064 + output_scale=output_scale,
1065 + fuse_quant_op=fuse_quant_op,
1066 )
1067 results.append(result)
1068
@@ -1047,6 +1090,8 @@ def run_mla_benchmark(
1090 num_kv_splits: int | None = None,
1091 index_topk: int = 2048,
1092 prefill_backend: str | None = None,
1093 + output_scale: float | None = None,
1094 + fuse_quant_op: bool = False,
1095 ) -> BenchmarkResult | list[BenchmarkResult]:
1096 """
1097 Unified MLA benchmark runner for all backends.
@@ -1066,6 +1111,9 @@ def run_mla_benchmark(
1111 index_topk: Topk value for sparse MLA backends (default 2048)
1112 prefill_backend: Prefill backend name (e.g., "fa3", "fa4").
1113 When set, forces the specified FlashAttention version for prefill.
1114 + output_scale: Static per-tensor FP8 scale for prefill output (None = bf16).
1115 + fuse_quant_op: With output_scale set, fuse the FP8 write into the prefill
1116 + kernel vs a standalone post-quant kernel. See _run_single_benchmark.
1117
1118 Returns:
1119 BenchmarkResult (single mode) or list of BenchmarkResult (batched mode)
@@ -1090,7 +1138,12 @@ def run_mla_benchmark(
1138
1139 # Use unified batched execution
1140 results = _run_mla_benchmark_batched(
1093 - backend, configs_with_params, index_topk, prefill_backend=prefill_backend
1141 + backend,
1142 + configs_with_params,
1143 + index_topk,
1144 + prefill_backend=prefill_backend,
1145 + output_scale=output_scale,
1146 + fuse_quant_op=fuse_quant_op,
1147 )
1148
1149 # Return single result or list based on input
tests/v1/attention/test_mla_prefill_quant_output.py new
+207
@@ -0,0 +1,207 @@
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 unittest.mock import patch
14 +
15 +import pytest
16 +import torch
17 +
18 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
19 + kFp8Dynamic128Sym,
20 + kFp8StaticTensorSym,
21 + kNvfp4Dynamic,
22 +)
23 +from vllm.platforms.interface import DeviceCapability
24 +from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
25 +from vllm.v1.attention.backends.mla.prefill.flash_attn import (
26 + FlashAttnPrefillBackend,
27 +)
28 +
29 +_FA_MODULE = "vllm.v1.attention.backends.mla.prefill.flash_attn"
30 +
31 +
32 +class _DummyPrefillBackend(MLAPrefillBackend):
33 + """Concrete backend that does NOT override supports_quant_output."""
34 +
35 + @staticmethod
36 + def get_name() -> str:
37 + return "DUMMY"
38 +
39 + def run_prefill_new_tokens(self, *args, **kwargs): # pragma: no cover
40 + raise NotImplementedError
41 +
42 + def run_prefill_context_chunk(self, *args, **kwargs): # pragma: no cover
43 + raise NotImplementedError
44 +
45 +
46 +@pytest.mark.parametrize(
47 + "quant_key", [kFp8StaticTensorSym, kFp8Dynamic128Sym, kNvfp4Dynamic, None]
48 +)
49 +def test_base_backend_never_supports_quant_output(quant_key):
50 + """The base default opts every backend out unless it overrides."""
51 + backend = object.__new__(_DummyPrefillBackend)
52 + assert backend.supports_quant_output(quant_key) is False
53 +
54 +
55 +def _make_fa_backend(version: int | None, is_vllm_fa: bool):
56 + """Build a FlashAttnPrefillBackend without running its heavy __init__."""
57 + backend = object.__new__(FlashAttnPrefillBackend)
58 + backend.vllm_flash_attn_version = version
59 + backend._is_vllm_fa = is_vllm_fa
60 + return backend
61 +
62 +
63 +@pytest.mark.parametrize(
64 + ("version", "is_vllm_fa", "dc_major", "quant_key", "expected"),
65 + [
66 + # FA4 + vLLM-FA + Blackwell SM100/SM110 + static FP8 -> fused.
67 + (4, True, 10, kFp8StaticTensorSym, True),
68 + (4, True, 11, kFp8StaticTensorSym, True),
69 + # Wrong compute capability (SM90 / SM120) -> not supported (#135).
70 + (4, True, 9, kFp8StaticTensorSym, False),
71 + (4, True, 12, kFp8StaticTensorSym, False),
72 + # Not FA4.
73 + (3, True, 10, kFp8StaticTensorSym, False),
74 + (2, True, 10, kFp8StaticTensorSym, False),
75 + (None, True, 10, kFp8StaticTensorSym, False),
76 + # Upstream (ROCm) flash-attn, not vLLM-FA.
77 + (4, False, 10, kFp8StaticTensorSym, False),
78 + # Quant keys not wired through FA4 yet.
79 + (4, True, 10, kFp8Dynamic128Sym, False),
80 + (4, True, 10, kNvfp4Dynamic, False),
81 + ],
82 +)
83 +def test_flash_attn_supports_quant_output(
84 + version, is_vllm_fa, dc_major, quant_key, expected
85 +):
86 + backend = _make_fa_backend(version, is_vllm_fa)
87 + with patch(f"{_FA_MODULE}.current_platform") as plat:
88 + plat.get_device_capability.return_value = DeviceCapability(
89 + major=dc_major, minor=0
90 + )
91 + assert backend.supports_quant_output(quant_key) is expected
92 +
93 +
94 +def test_flash_attn_supports_quant_output_unknown_device():
95 + """A None device capability (e.g. capability probe failed) is safe."""
96 + backend = _make_fa_backend(version=4, is_vllm_fa=True)
97 + with patch(f"{_FA_MODULE}.current_platform") as plat:
98 + plat.get_device_capability.return_value = None
99 + assert backend.supports_quant_output(kFp8StaticTensorSym) is False
100 +
101 +
102 +def test_flash_attn_prefill_backend_signature_accepts_fused_kwargs():
103 + """run_prefill_new_tokens must accept out/output_scale so the direct
104 + (non-**kwargs) call in forward_mha type- and runtime-checks."""
105 + import inspect
106 +
107 + params = inspect.signature(
108 + FlashAttnPrefillBackend.run_prefill_new_tokens
109 + ).parameters
110 + assert "out" in params
111 + assert "output_scale" in params
112 + # The base contract must expose them too (Liskov / direct call site).
113 + base_params = inspect.signature(MLAPrefillBackend.run_prefill_new_tokens).parameters
114 + assert "out" in base_params
115 + assert "output_scale" in base_params
116 +
117 +
118 +def test_mla_impl_forward_mha_accepts_output_scale():
119 + """The abstract MLA impl forward_mha must carry output_scale so every
120 + override (and the unconditional forward_impl call) stays compatible."""
121 + import inspect
122 +
123 + from vllm.v1.attention.backend import MLAAttentionImpl
124 +
125 + params = inspect.signature(MLAAttentionImpl.forward_mha).parameters
126 + assert "output_scale" in params
127 + assert params["output_scale"].default is None
128 +
129 +
130 +def _fused_fp8_skip_reason() -> str | None:
131 + """FA4 fused FP8 output needs a real Blackwell SM100/SM110 GPU."""
132 + if not torch.cuda.is_available():
133 + return "requires CUDA"
134 + major = torch.cuda.get_device_capability()[0]
135 + if major not in (10, 11):
136 + return f"FA4 fused FP8 output requires SM100/SM110, got SM{major}x"
137 + return None
138 +
139 +
140 +_FUSED_FP8_SKIP = _fused_fp8_skip_reason()
141 +
142 +
143 +@pytest.mark.skipif(_FUSED_FP8_SKIP is not None, reason=_FUSED_FP8_SKIP or "")
144 +def test_fa4_fused_fp8_output_matches_post_quant(default_vllm_config):
145 + """FA4's fused FP8 write (output_scale, flash-attention#135) must match the
146 + bf16-attention + standalone static-FP8-quant path it replaces, since
147 + production uses the same output_scale for both."""
148 + from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
149 + from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
150 + from vllm.platforms import current_platform
151 + from vllm.vllm_flash_attn import flash_attn_varlen_func
152 +
153 + torch.manual_seed(0)
154 + device = torch.device("cuda")
155 + fp8_dtype = current_platform.fp8_dtype()
156 +
157 + # MLA prefill head dims (post kv_b_proj): q/k = qk_nope(128)+qk_rope(64),
158 + # v = v_head_dim(128); DeepSeek-V2-Lite has 16 query heads.
159 + num_heads, qk_head_dim, v_head_dim, seqlen = 16, 192, 128, 512
160 + cu_seqlens = torch.tensor([0, seqlen], dtype=torch.int32, device=device)
161 + q = torch.randn(seqlen, num_heads, qk_head_dim, dtype=torch.bfloat16, device=device)
162 + k = torch.randn(seqlen, num_heads, qk_head_dim, dtype=torch.bfloat16, device=device)
163 + v = torch.randn(seqlen, num_heads, v_head_dim, dtype=torch.bfloat16, device=device)
164 +
165 + fa_kwargs = dict(
166 + cu_seqlens_q=cu_seqlens,
167 + cu_seqlens_k=cu_seqlens,
168 + max_seqlen_q=seqlen,
169 + max_seqlen_k=seqlen,
170 + causal=True,
171 + fa_version=4,
172 + )
173 +
174 + # Reference: bf16 attention, then standalone static per-tensor FP8 quant.
175 + out_bf16 = flash_attn_varlen_func(q=q, k=k, v=v, **fa_kwargs)
176 + out_2d = out_bf16.reshape(seqlen, num_heads * v_head_dim)
177 + # Scale the amax near e4m3 max so the check uses the representable range.
178 + finfo = torch.finfo(fp8_dtype)
179 + scale = (out_2d.abs().max() / finfo.max).to(torch.float32).reshape(1)
180 + quant_op = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
181 + ref_fp8, _ = quant_op(out_2d, scale)
182 +
183 + # Feature: FA4 writes e4m3 into the (tokens, heads*dim) buffer directly.
184 + fused_fp8 = torch.empty(
185 + seqlen, num_heads * v_head_dim, dtype=fp8_dtype, device=device
186 + )
187 + flash_attn_varlen_func(
188 + q=q,
189 + k=k,
190 + v=v,
191 + out=fused_fp8.view(seqlen, num_heads, v_head_dim),
192 + output_scale=scale,
193 + **fa_kwargs,
194 + )
195 +
196 + # Non-degenerate (catches a no-op / all-zero write).
197 + assert torch.isfinite(fused_fp8.float()).all()
198 + assert fused_fp8.float().abs().any()
199 +
200 + # e4m3 has 3 mantissa bits, so allow ~1 mantissa step of rounding slack.
201 + ref = ref_fp8.float() * scale
202 + got = fused_fp8.float() * scale
203 + torch.testing.assert_close(got, ref, rtol=0.125, atol=float(scale) * 2)
204 +
205 + # ...and most elements land in the exact same fp8 bucket.
206 + exact = (fused_fp8.view(torch.uint8) == ref_fp8.view(torch.uint8)).float().mean()
207 + assert exact > 0.9, f"only {exact:.1%} of fused FP8 outputs matched the baseline"
vllm/model_executor/layers/attention/mla_attention.py
+38 -8
@@ -687,7 +687,23 @@ class MLAAttention(nn.Module, AttentionLayerBase):
687 num_mqa_tokens = attn_metadata.num_decode_tokens
688 num_mha_tokens = q.size(0) - num_mqa_tokens
689
690 + mha_use_quant_output = (
691 + quant_key is not None
692 + and self.prefill_backend.supports_quant_output(quant_key)
693 + and attn_metadata is not None
694 + and attn_metadata.prefill is not None
695 + and attn_metadata.prefill.chunked_context is None
696 + and self.impl.dcp_world_size <= 1
697 + )
698 +
699 if num_mha_tokens > 0:
700 + if mha_use_quant_output:
701 + mha_output = quant_output
702 + mha_output_scale = output_scale
703 + else:
704 + mha_output = output
705 + mha_output_scale = None
706 +
707 self.impl.forward_mha( # type: ignore[attr-defined]
708 q[num_mqa_tokens:],
709 k_c_normed[num_mqa_tokens:],
@@ -695,7 +711,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):
711 kv_cache,
712 attn_metadata,
713 self._k_scale,
698 - output=output[num_mqa_tokens:],
714 + output=mha_output[num_mqa_tokens:num_actual_toks],
715 + output_scale=mha_output_scale,
716 )
717
718 if num_mqa_tokens > 0:
@@ -794,13 +811,15 @@ class MLAAttention(nn.Module, AttentionLayerBase):
811 self._v_up_proj(attn_out, out=mqa_output_slice)
812
813 if quant_key is not None:
797 - # Quantize the BF16 computation result into the quantized output
798 - actual = output[:num_actual_toks]
814 + quant_idx = num_mqa_tokens if mha_use_quant_output else num_actual_toks
815 + if quant_idx == 0:
816 + return quant_output
817 + actual = output[:quant_idx]
818 if quant_key == kNvfp4Dynamic:
819 # NVFP4: two FP4 values packed into one uint8
820 assert output_block_scale is not None
821 fp4_data, fp4_scales = ops.scaled_fp4_quant(actual, output_scale)
803 - quant_output[:num_actual_toks].copy_(fp4_data)
822 + quant_output[:quant_idx].copy_(fp4_data)
823 output_block_scale[: fp4_scales.shape[0]].copy_(fp4_scales)
824 elif quant_key in (kFp8Dynamic128Sym, kFp8Dynamic64Sym):
825 # Per-group FP8
@@ -812,8 +831,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):
831 finfo = torch.finfo(_FP8_DTYPE)
832 torch.ops._C.per_token_group_fp8_quant(
833 actual,
815 - quant_output[:num_actual_toks],
816 - output_block_scale[:num_actual_toks],
834 + quant_output[:quant_idx],
835 + output_block_scale[:quant_idx],
836 quant_group_size,
837 1e-10, # eps
838 finfo.min,
@@ -825,7 +844,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
844 elif quant_key == kFp8StaticTensorSym:
845 # Static FP8 quantization
846 fp8_data, _ = self._quant_fp8_op(actual, output_scale)
828 - quant_output[:num_actual_toks].copy_(fp8_data)
847 + quant_output[:quant_idx].copy_(fp8_data)
848 else:
849 raise ValueError(f"Unsupported quant_key: {quant_key}")
850 return quant_output
@@ -2254,6 +2273,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
2273 attn_metadata: MLACommonMetadata,
2274 k_scale: torch.Tensor,
2275 output: torch.Tensor,
2276 + output_scale: torch.Tensor | None = None,
2277 ) -> None:
2278 assert attn_metadata.prefill is not None
2279 assert self.dcp_world_size != -1
@@ -2267,6 +2287,9 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
2287 q = q.to(prefill_metadata.q_data_type)
2288
2289 has_context = prefill_metadata.chunked_context is not None
2290 + assert output_scale is None or not has_context, (
2291 + "Fused FP8 output is only wired for the non-chunked-context path"
2292 + )
2293
2294 kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
2295 -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim
@@ -2283,6 +2306,12 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
2306 k=k,
2307 v=v,
2308 return_softmax_lse=has_context,
2309 + out=(
2310 + output.view(-1, self.num_heads, self.v_head_dim)
2311 + if output_scale is not None
2312 + else None
2313 + ),
2314 + output_scale=output_scale,
2315 )
2316
2317 if has_context:
@@ -2312,7 +2341,8 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
2341 suffix_lse=suffix_lse,
2342 prefill_tokens_with_context=prefill_metadata.chunked_context.prefill_tokens_with_context,
2343 )
2315 - else:
2344 + elif output_scale is None:
2345 + # With output_scale set, backend already wrote into `output` in place.
2346 assert isinstance(output_prefill, torch.Tensor)
2347 output_prefill = output_prefill.flatten(start_dim=-2)
2348 output.copy_(output_prefill)
vllm/v1/attention/backend.py
+1
@@ -899,6 +899,7 @@ class MLAAttentionImpl(AttentionImplBase[T], Generic[T]):
899 attn_metadata: T,
900 k_scale: torch.Tensor,
901 output: torch.Tensor,
902 + output_scale: torch.Tensor | None = None,
903 ) -> None:
904 """MHA-style prefill forward pass."""
905 raise NotImplementedError
vllm/v1/attention/backends/mla/prefill/base.py
+9
@@ -12,6 +12,7 @@ if TYPE_CHECKING:
12 from vllm.model_executor.layers.attention.mla_attention import (
13 MLACommonPrefillMetadata,
14 )
15 + from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
16 from vllm.platforms.interface import DeviceCapability
17 from vllm.v1.attention.backends.mla.prefill.selector import (
18 MLAPrefillSelectorConfig,
@@ -44,6 +45,12 @@ class MLAPrefillBackend(ABC):
45 def is_available(cls) -> bool:
46 return True
47
48 + def supports_quant_output(self, quant_key: "QuantKey") -> bool:
49 + """Whether `run_prefill_new_tokens` can write quantized output
50 + directly (fused) for the given quant key, skipping the post-quant
51 + pass. Overridden by backends that support it."""
52 + return False
53 +
54 @classmethod
55 def validate_configuration(
56 cls,
@@ -107,6 +114,8 @@ class MLAPrefillBackend(ABC):
114 k: torch.Tensor,
115 v: torch.Tensor,
116 return_softmax_lse: bool,
117 + out: torch.Tensor | None = None,
118 + output_scale: torch.Tensor | None = None,
119 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
120 raise NotImplementedError
121
vllm/v1/attention/backends/mla/prefill/flash_attn.py
+23
@@ -8,6 +8,9 @@ from typing import TYPE_CHECKING
8 import torch
9
10 import vllm.envs as envs
11 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
12 + kFp8StaticTensorSym,
13 +)
14 from vllm.platforms import current_platform
15 from vllm.v1.attention.backends.fa_utils import (
16 get_flash_attn_version,
@@ -17,6 +20,7 @@ from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
20
21 if TYPE_CHECKING:
22 from vllm.config import VllmConfig
23 + from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
24
25 if is_flash_attn_varlen_func_available():
26 from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func
@@ -87,6 +91,16 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
91 # Track whether we're using vllm's FA or upstream (for ROCm)
92 self._is_vllm_fa = current_platform.is_cuda() or current_platform.is_xpu()
93
94 + def supports_quant_output(self, quant_key: "QuantKey") -> bool:
95 + device_capability = current_platform.get_device_capability()
96 + return (
97 + self.vllm_flash_attn_version == 4
98 + and self._is_vllm_fa
99 + and device_capability is not None
100 + and device_capability[0] in (10, 11)
101 + and quant_key == kFp8StaticTensorSym
102 + )
103 +
104 def _flash_attn_varlen_diff_headdims(
105 self,
106 q: torch.Tensor,
@@ -94,6 +108,8 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
108 v: torch.Tensor,
109 return_softmax_lse: bool = False,
110 softmax_scale: float | None = None,
111 + out: torch.Tensor | None = None,
112 + output_scale: torch.Tensor | None = None,
113 **kwargs,
114 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
115 maybe_padded_v = v
@@ -104,10 +120,13 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
120
121 if self._is_vllm_fa:
122 kwargs["return_softmax_lse"] = return_softmax_lse
123 + kwargs["out"] = out
124 + kwargs["output_scale"] = output_scale
125 else:
126 # ROCm leverages the upstream flash_attn, which takes a parameter
127 # called "return_attn_probs" instead of return_softmax_lse
128 kwargs["return_attn_probs"] = return_softmax_lse
129 + assert out is None and output_scale is None
130 if envs.VLLM_BATCH_INVARIANT:
131 kwargs["num_splits"] = 1
132
@@ -140,6 +159,8 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
159 k: torch.Tensor,
160 v: torch.Tensor,
161 return_softmax_lse: bool,
162 + out: torch.Tensor | None = None,
163 + output_scale: torch.Tensor | None = None,
164 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
165 return self._flash_attn_varlen_diff_headdims(
166 q=q,
@@ -152,6 +173,8 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
173 softmax_scale=self.scale,
174 causal=True,
175 return_softmax_lse=return_softmax_lse,
176 + out=out,
177 + output_scale=output_scale,
178 )
179
180 def run_prefill_context_chunk(
vllm/v1/attention/backends/mla/prefill/flashinfer.py
+2
@@ -188,6 +188,8 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
188 k: torch.Tensor,
189 v: torch.Tensor,
190 return_softmax_lse: bool,
191 + out: torch.Tensor | None = None,
192 + output_scale: torch.Tensor | None = None,
193 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
194 assert self._prefill_main is not None
195
vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py
+2
@@ -115,6 +115,8 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend):
115 k: torch.Tensor,
116 v: torch.Tensor,
117 return_softmax_lse: bool,
118 + out: torch.Tensor | None = None,
119 + output_scale: torch.Tensor | None = None,
120 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
121 from tokenspeed_mla import tokenspeed_mla_prefill
122
vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py
+2
@@ -83,6 +83,8 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
83 k: torch.Tensor,
84 v: torch.Tensor,
85 return_softmax_lse: bool,
86 + out: torch.Tensor | None = None,
87 + output_scale: torch.Tensor | None = None,
88 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
89 from flashinfer.prefill import trtllm_ragged_attention_deepseek
90
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+7
@@ -841,6 +841,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
841 attn_metadata: MLACommonMetadata,
842 k_scale: torch.Tensor,
843 output: torch.Tensor,
844 + output_scale: torch.Tensor | None = None,
845 ) -> None:
846 """Dispatch prefill to the FP8 ASM kernel when available.
847
@@ -866,6 +867,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
867 attn_metadata,
868 k_scale,
869 output,
870 + output_scale,
871 )
872
873 assert attn_metadata.prefill is not None
@@ -881,8 +883,13 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
883 attn_metadata,
884 k_scale,
885 output,
886 + output_scale,
887 )
888
889 + assert output_scale is None, (
890 + "fused FP8 output not supported by the AITER FP8 MLA prefill path"
891 + )
892 +
893 kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
894 -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim
895 )
vllm/vllm_flash_attn/flash_attn_interface.py
+8
@@ -200,6 +200,8 @@ def flash_attn_varlen_func(
200 k_descale=None,
201 v_descale=None,
202 num_splits: int = 0,
203 + # FA4 Only
204 + output_scale=None,
205 # Version selector
206 fa_version: int = DEFAULT_FA_VERSION,
207 s_aux=None,
@@ -273,6 +275,11 @@ def flash_attn_varlen_func(
275 "seqused_k must be provided if block_table is provided"
276 )
277
278 + assert output_scale is None or fa_version == 4, (
279 + f"Fused FP8 output (output_scale) is only supported by FA4, "
280 + f"got fa_version={fa_version}"
281 + )
282 +
283 if softmax_scale is None:
284 softmax_scale = q.shape[-1] ** (-0.5)
285 # custom op does not support non-tuple input
@@ -403,6 +410,7 @@ def flash_attn_varlen_func(
410 learnable_sink=s_aux,
411 mask_mod=mask_mod,
412 aux_tensors=aux_tensors,
413 + output_scale=output_scale,
414 )
415 else:
416 raise ValueError(f"Unsupported FA version: {fa_version}")