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

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

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 14:45 UTC

[Kernel] Fuse FP8 output quantization into merge_attn_states (#36518)

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

Carl Y committed Apr 2, 2026 at 18:47 UTC 3bc2734dd03d2073e1ba066bda042b4b085037bd
8 files changed +516 -70
benchmarks/fused_kernels/merge_attn_states_benchmarks.py new
+264
@@ -0,0 +1,264 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Benchmark: Fused FP8 output quantization in merge_attn_states
5 +
6 +Compares fused vs unfused approaches for producing FP8-quantized merged
7 +attention output:
8 + 1. Fused CUDA -- single CUDA kernel (merge + FP8 quant)
9 + 2. Fused Triton -- single Triton kernel (merge + FP8 quant)
10 + 3. Unfused CUDA -- CUDA merge + torch.compiled FP8 quant
11 + 4. Unfused Triton -- Triton merge + torch.compiled FP8 quant
12 +
13 +Usage:
14 + python benchmarks/fused_kernels/merge_attn_states_benchmarks.py
15 + python benchmarks/fused_kernels/merge_attn_states_benchmarks.py --tp 1 4 8
16 + python benchmarks/fused_kernels/merge_attn_states_benchmarks.py --dtype bfloat16
17 +"""
18 +
19 +import argparse
20 +import itertools
21 +
22 +import torch
23 +
24 +from vllm._custom_ops import merge_attn_states as merge_attn_states_cuda
25 +from vllm.benchmarks.lib.utils import default_vllm_config
26 +from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
27 +from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
28 +from vllm.platforms import current_platform
29 +from vllm.triton_utils import triton
30 +from vllm.v1.attention.ops.triton_merge_attn_states import (
31 + merge_attn_states as merge_attn_states_triton,
32 +)
33 +
34 +# ---------------------------------------------------------------------------
35 +# Configuration defaults
36 +# ---------------------------------------------------------------------------
37 +
38 +NUM_TOKENS_LIST = [1, 16, 64, 256, 1024, 4096]
39 +
40 +# (label, num_heads, head_size) — num_heads is for TP=1
41 +HEAD_CONFIGS = [
42 + ("DeepSeek-V3 MLA", 128, 128),
43 + ("Llama-70B", 64, 128),
44 + ("Llama-8B", 32, 128),
45 +]
46 +
47 +TP_SIZES = [1, 2, 4, 8]
48 +
49 +INPUT_DTYPES = [torch.float32, torch.float16, torch.bfloat16]
50 +
51 +QUANTILES = [0.5, 0.2, 0.8]
52 +
53 +
54 +# ---------------------------------------------------------------------------
55 +# Helpers
56 +# ---------------------------------------------------------------------------
57 +
58 +
59 +def short_dtype(dtype: torch.dtype) -> str:
60 + return str(dtype).removeprefix("torch.")
61 +
62 +
63 +def make_inputs(
64 + num_tokens: int,
65 + num_heads: int,
66 + head_size: int,
67 + dtype: torch.dtype,
68 +):
69 + """Create random prefix/suffix outputs and LSEs."""
70 + prefix_output = torch.randn(
71 + (num_tokens, num_heads, head_size), dtype=dtype, device="cuda"
72 + )
73 + suffix_output = torch.randn(
74 + (num_tokens, num_heads, head_size), dtype=dtype, device="cuda"
75 + )
76 + prefix_lse = torch.randn(num_heads, num_tokens, dtype=torch.float32, device="cuda")
77 + suffix_lse = torch.randn(num_heads, num_tokens, dtype=torch.float32, device="cuda")
78 + # Sprinkle some inf values to exercise edge-case paths
79 + mask = torch.rand(num_heads, num_tokens, device="cuda") < 0.05
80 + prefix_lse[mask] = float("inf")
81 + mask2 = torch.rand(num_heads, num_tokens, device="cuda") < 0.05
82 + suffix_lse[mask2] = float("inf")
83 + return prefix_output, suffix_output, prefix_lse, suffix_lse
84 +
85 +
86 +def build_configs(head_configs, num_tokens_list, input_dtypes, tp_sizes):
87 + """Build (num_tokens, num_heads, head_size, dtype_str) config tuples,
88 + applying TP division to num_heads and skipping invalid combos."""
89 + configs = []
90 + for (_, nh, hs), nt, dtype, tp in itertools.product(
91 + head_configs, num_tokens_list, input_dtypes, tp_sizes
92 + ):
93 + nh_tp = nh // tp
94 + if nh_tp >= 1:
95 + configs.append((nt, nh_tp, hs, short_dtype(dtype)))
96 + return configs
97 +
98 +
99 +def parse_args():
100 + parser = argparse.ArgumentParser(
101 + description="Benchmark merge_attn_states fused FP8 quantization"
102 + )
103 + parser.add_argument(
104 + "--num-tokens",
105 + type=int,
106 + nargs="+",
107 + default=None,
108 + help=f"Override token counts (default: {NUM_TOKENS_LIST})",
109 + )
110 + parser.add_argument(
111 + "--tp",
112 + type=int,
113 + nargs="+",
114 + default=None,
115 + help=f"TP sizes to simulate (divides num_heads) (default: {TP_SIZES})",
116 + )
117 + parser.add_argument(
118 + "--dtype",
119 + type=str,
120 + nargs="+",
121 + default=None,
122 + help="Input dtypes (e.g. bfloat16 float16 float32). "
123 + f"Default: {[short_dtype(d) for d in INPUT_DTYPES]}",
124 + )
125 + return parser.parse_args()
126 +
127 +
128 +# ---------------------------------------------------------------------------
129 +# Parse args and build configs before decorators
130 +# ---------------------------------------------------------------------------
131 +
132 +args = parse_args()
133 +
134 +num_tokens_list = args.num_tokens if args.num_tokens else NUM_TOKENS_LIST
135 +tp_sizes = args.tp if args.tp else TP_SIZES
136 +
137 +if args.dtype:
138 + from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
139 +
140 + input_dtypes = [STR_DTYPE_TO_TORCH_DTYPE[d] for d in args.dtype]
141 +else:
142 + input_dtypes = INPUT_DTYPES
143 +
144 +configs = build_configs(HEAD_CONFIGS, num_tokens_list, input_dtypes, tp_sizes)
145 +
146 +torch._dynamo.config.recompile_limit = 8888
147 +
148 +
149 +# ---------------------------------------------------------------------------
150 +# Benchmark function
151 +# ---------------------------------------------------------------------------
152 +
153 +
154 +@triton.testing.perf_report(
155 + triton.testing.Benchmark(
156 + x_names=["num_tokens", "num_heads", "head_size", "dtype_str"],
157 + x_vals=configs,
158 + line_arg="provider",
159 + line_vals=["fused_cuda", "fused_triton", "unfused_cuda", "unfused_triton"],
160 + line_names=["Fused CUDA", "Fused Triton", "Unfused CUDA", "Unfused Triton"],
161 + styles=[("blue", "-"), ("green", "-"), ("blue", "--"), ("green", "--")],
162 + ylabel="us",
163 + plot_name="merge_attn_states FP8 (fused vs unfused)",
164 + args={},
165 + )
166 +)
167 +@default_vllm_config()
168 +def benchmark(num_tokens, num_heads, head_size, dtype_str, provider):
169 + input_dtype = getattr(torch, dtype_str)
170 + fp8_dtype = current_platform.fp8_dtype()
171 + prefix_out, suffix_out, prefix_lse, suffix_lse = make_inputs(
172 + num_tokens, num_heads, head_size, input_dtype
173 + )
174 + output_scale = torch.tensor([0.1], dtype=torch.float32, device="cuda")
175 +
176 + if provider == "fused_cuda":
177 + output = torch.empty(
178 + (num_tokens, num_heads, head_size), dtype=fp8_dtype, device="cuda"
179 + )
180 + fn = lambda: merge_attn_states_cuda(
181 + output,
182 + prefix_out,
183 + prefix_lse,
184 + suffix_out,
185 + suffix_lse,
186 + output_scale=output_scale,
187 + )
188 + elif provider == "fused_triton":
189 + output = torch.empty(
190 + (num_tokens, num_heads, head_size), dtype=fp8_dtype, device="cuda"
191 + )
192 + fn = lambda: merge_attn_states_triton(
193 + output,
194 + prefix_out,
195 + prefix_lse,
196 + suffix_out,
197 + suffix_lse,
198 + output_scale=output_scale,
199 + )
200 + elif provider == "unfused_cuda":
201 + merge_buf = torch.empty(
202 + (num_tokens, num_heads, head_size), dtype=input_dtype, device="cuda"
203 + )
204 + quant_fp8 = QuantFP8(
205 + static=True,
206 + group_shape=GroupShape.PER_TENSOR,
207 + column_major_scales=False,
208 + )
209 + quant_input = merge_buf.view(-1, head_size)
210 + compiled_quant = torch.compile(
211 + quant_fp8.forward_native, fullgraph=True, dynamic=False
212 + )
213 +
214 + def unfused_fn():
215 + merge_attn_states_cuda(
216 + merge_buf, prefix_out, prefix_lse, suffix_out, suffix_lse
217 + )
218 + compiled_quant(quant_input, output_scale)
219 +
220 + fn = unfused_fn
221 + else: # unfused_triton
222 + merge_buf = torch.empty(
223 + (num_tokens, num_heads, head_size), dtype=input_dtype, device="cuda"
224 + )
225 + quant_fp8 = QuantFP8(
226 + static=True,
227 + group_shape=GroupShape.PER_TENSOR,
228 + column_major_scales=False,
229 + )
230 + quant_input = merge_buf.view(-1, head_size)
231 + compiled_quant = torch.compile(
232 + quant_fp8.forward_native, fullgraph=True, dynamic=False
233 + )
234 +
235 + def unfused_fn():
236 + merge_attn_states_triton(
237 + merge_buf, prefix_out, prefix_lse, suffix_out, suffix_lse
238 + )
239 + compiled_quant(quant_input, output_scale)
240 +
241 + fn = unfused_fn
242 +
243 + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=QUANTILES)
244 + return 1000 * ms, 1000 * max_ms, 1000 * min_ms # us
245 +
246 +
247 +# ---------------------------------------------------------------------------
248 +# Main
249 +# ---------------------------------------------------------------------------
250 +
251 +
252 +def main():
253 + device_name = current_platform.get_device_name()
254 + print(f"Device: {device_name}")
255 + print(f"Token counts: {num_tokens_list}")
256 + print(f"TP sizes: {tp_sizes}")
257 + print(f"Input dtypes: {[short_dtype(d) for d in input_dtypes]}")
258 + print(f"Head configs: {[(c[0], c[1], c[2]) for c in HEAD_CONFIGS]}")
259 + benchmark.run(print_data=True)
260 +
261 +
262 +if __name__ == "__main__":
263 + with torch.inference_mode():
264 + main()
csrc/attention/merge_attn_states.cu
+133 -41
@@ -7,19 +7,29 @@
7
8 #include "attention_dtypes.h"
9 #include "attention_utils.cuh"
10 +#include "../quantization/w8a8/fp8/common.cuh"
11 +#include "../dispatch_utils.h"
12
13 namespace vllm {
14
15 // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
16 // can be used to combine partial attention results (in the split-KV case)
15 -template <typename scalar_t, const uint NUM_THREADS>
17 +template <typename scalar_t, typename output_t, const uint NUM_THREADS,
18 + bool USE_FP8_OUTPUT>
19 __global__ void merge_attn_states_kernel(
17 - scalar_t* output, float* output_lse, const scalar_t* prefix_output,
20 + output_t* output, float* output_lse, const scalar_t* prefix_output,
21 const float* prefix_lse, const scalar_t* suffix_output,
22 const float* suffix_lse, const uint num_tokens, const uint num_heads,
23 const uint head_size, const uint prefix_head_stride,
21 - const uint output_head_stride, const uint prefix_num_tokens) {
22 - using pack_128b_t = uint4;
24 + const uint output_head_stride, const uint prefix_num_tokens,
25 + const float* output_scale) {
26 + // Inputs always load 128-bit packs (pack_size elements of scalar_t).
27 + // Outputs store pack_size elements of output_t, which is smaller for FP8.
28 + using input_pack_t = uint4;
29 + using output_pack_t =
30 + std::conditional_t<USE_FP8_OUTPUT,
31 + std::conditional_t<sizeof(scalar_t) == 4, uint, uint2>,
32 + uint4>;
33 const uint pack_size = 16 / sizeof(scalar_t);
34 const uint threads_per_head = head_size / pack_size;
35
@@ -42,15 +52,36 @@ __global__ void merge_attn_states_kernel(
52 head_idx * output_head_stride;
53 const scalar_t* prefix_head_ptr = prefix_output + src_head_offset;
54 const scalar_t* suffix_head_ptr = suffix_output + src_head_offset;
45 - scalar_t* output_head_ptr = output + dst_head_offset;
55 + output_t* output_head_ptr = output + dst_head_offset;
56 +
57 + // Pre-invert scale: multiplication is faster than division
58 + float fp8_scale_inv = 1.0f;
59 + if constexpr (USE_FP8_OUTPUT) {
60 + fp8_scale_inv = 1.0f / *output_scale;
61 + }
62
63 // If token_idx >= prefix_num_tokens, just copy from suffix
64 if (token_idx >= prefix_num_tokens) {
65 if (pack_offset < head_size) {
50 - pack_128b_t s_out_pack = reinterpret_cast<const pack_128b_t*>(
66 + input_pack_t s_out_pack = reinterpret_cast<const input_pack_t*>(
67 suffix_head_ptr)[pack_offset / pack_size];
52 - reinterpret_cast<pack_128b_t*>(output_head_ptr)[pack_offset / pack_size] =
53 - s_out_pack;
68 +
69 + if constexpr (USE_FP8_OUTPUT) {
70 + output_t o_out_pack[pack_size];
71 +#pragma unroll
72 + for (uint i = 0; i < pack_size; ++i) {
73 + const float val =
74 + vllm::to_float(reinterpret_cast<const scalar_t*>(&s_out_pack)[i]);
75 + o_out_pack[i] =
76 + vllm::scaled_fp8_conversion<true, output_t>(val, fp8_scale_inv);
77 + }
78 + reinterpret_cast<output_pack_t*>(
79 + output_head_ptr)[pack_offset / pack_size] =
80 + *reinterpret_cast<output_pack_t*>(o_out_pack);
81 + } else {
82 + reinterpret_cast<output_pack_t*>(
83 + output_head_ptr)[pack_offset / pack_size] = s_out_pack;
84 + }
85 }
86 if (output_lse != nullptr && pack_idx == 0) {
87 float s_lse = suffix_lse[head_idx * num_tokens + token_idx];
@@ -70,20 +101,34 @@ __global__ void merge_attn_states_kernel(
101 /* In certain edge cases, MLA can produce p_lse = s_lse = -inf;
102 continuing the pipeline then yields NaN. Root cause: with chunked prefill
103 a batch may be split into two chunks; if a request in that batch has no
73 - prefix hit, every LSE entry for that request’s position is -inf, and at
104 + prefix hit, every LSE entry for that request's position is -inf, and at
105 this moment we merge cross-attention at first. For now we simply emit
106 prefix_output (expected to be all zeros) and prefix_lse (-inf) to fix
107 this problem.
108 */
109 if (std::isinf(max_lse)) {
110 if (pack_offset < head_size) {
80 - // Pack 128b load
81 - pack_128b_t p_out_pack = reinterpret_cast<const pack_128b_t*>(
111 + input_pack_t p_out_pack = reinterpret_cast<const input_pack_t*>(
112 prefix_head_ptr)[pack_offset / pack_size];
113
84 - // Pack 128b storage
85 - reinterpret_cast<pack_128b_t*>(output_head_ptr)[pack_offset / pack_size] =
86 - p_out_pack;
114 + if constexpr (USE_FP8_OUTPUT) {
115 + // Convert prefix values to FP8 (since -inf means no data,
116 + // prefix_output is expected to be zeros)
117 + output_t o_out_pack[pack_size];
118 +#pragma unroll
119 + for (uint i = 0; i < pack_size; ++i) {
120 + const float val =
121 + vllm::to_float(reinterpret_cast<const scalar_t*>(&p_out_pack)[i]);
122 + o_out_pack[i] =
123 + vllm::scaled_fp8_conversion<true, output_t>(val, fp8_scale_inv);
124 + }
125 + reinterpret_cast<output_pack_t*>(
126 + output_head_ptr)[pack_offset / pack_size] =
127 + *reinterpret_cast<output_pack_t*>(o_out_pack);
128 + } else {
129 + reinterpret_cast<output_pack_t*>(
130 + output_head_ptr)[pack_offset / pack_size] = p_out_pack;
131 + }
132 }
133 // We only need to write to output_lse once per head.
134 if (output_lse != nullptr && pack_idx == 0) {
@@ -101,30 +146,43 @@ __global__ void merge_attn_states_kernel(
146 const float s_scale = s_se / out_se;
147
148 if (pack_offset < head_size) {
104 - // Pack 128b load
105 - pack_128b_t p_out_pack = reinterpret_cast<const pack_128b_t*>(
149 + input_pack_t p_out_pack = reinterpret_cast<const input_pack_t*>(
150 prefix_head_ptr)[pack_offset / pack_size];
107 - pack_128b_t s_out_pack = reinterpret_cast<const pack_128b_t*>(
151 + input_pack_t s_out_pack = reinterpret_cast<const input_pack_t*>(
152 suffix_head_ptr)[pack_offset / pack_size];
109 - pack_128b_t o_out_pack;
153
154 + // Compute merged values in float32
155 + float o_out_f[pack_size];
156 #pragma unroll
157 for (uint i = 0; i < pack_size; ++i) {
113 - // Always use float for FMA to keep high precision.
114 - // half(uint16_t), bfloat16, float -> float.
158 const float p_out_f =
159 vllm::to_float(reinterpret_cast<const scalar_t*>(&p_out_pack)[i]);
160 const float s_out_f =
161 vllm::to_float(reinterpret_cast<const scalar_t*>(&s_out_pack)[i]);
119 - // fma: a * b + c = p_out_f * p_scale + (s_out_f * s_scale)
120 - const float o_out_f = p_out_f * p_scale + (s_out_f * s_scale);
121 - // float -> half(uint16_t), bfloat16, float.
122 - vllm::from_float(reinterpret_cast<scalar_t*>(&o_out_pack)[i], o_out_f);
162 + o_out_f[i] = p_out_f * p_scale + (s_out_f * s_scale);
163 }
164
125 - // Pack 128b storage
126 - reinterpret_cast<pack_128b_t*>(output_head_ptr)[pack_offset / pack_size] =
127 - o_out_pack;
165 + // Convert and store
166 + if constexpr (USE_FP8_OUTPUT) {
167 + output_t o_out_pack[pack_size];
168 +#pragma unroll
169 + for (uint i = 0; i < pack_size; ++i) {
170 + o_out_pack[i] = vllm::scaled_fp8_conversion<true, output_t>(
171 + o_out_f[i], fp8_scale_inv);
172 + }
173 + reinterpret_cast<output_pack_t*>(
174 + output_head_ptr)[pack_offset / pack_size] =
175 + *reinterpret_cast<output_pack_t*>(o_out_pack);
176 + } else {
177 + output_pack_t o_out_pack;
178 +#pragma unroll
179 + for (uint i = 0; i < pack_size; ++i) {
180 + vllm::from_float(reinterpret_cast<scalar_t*>(&o_out_pack)[i],
181 + o_out_f[i]);
182 + }
183 + reinterpret_cast<output_pack_t*>(
184 + output_head_ptr)[pack_offset / pack_size] = o_out_pack;
185 + }
186 }
187 // We only need to write to output_lse once per head.
188 if (output_lse != nullptr && pack_idx == 0) {
@@ -151,24 +209,26 @@ __global__ void merge_attn_states_kernel(
209 } \
210 }
211
154 -#define LAUNCH_MERGE_ATTN_STATES(scalar_t, NUM_THREADS) \
212 +#define LAUNCH_MERGE_ATTN_STATES(scalar_t, output_t, NUM_THREADS, \
213 + USE_FP8_OUTPUT) \
214 { \
156 - vllm::merge_attn_states_kernel<scalar_t, NUM_THREADS> \
215 + vllm::merge_attn_states_kernel<scalar_t, output_t, NUM_THREADS, \
216 + USE_FP8_OUTPUT> \
217 <<<grid, block, 0, stream>>>( \
158 - reinterpret_cast<scalar_t*>(output.data_ptr()), output_lse_ptr, \
218 + reinterpret_cast<output_t*>(output.data_ptr()), output_lse_ptr, \
219 reinterpret_cast<scalar_t*>(prefix_output.data_ptr()), \
220 reinterpret_cast<float*>(prefix_lse.data_ptr()), \
221 reinterpret_cast<scalar_t*>(suffix_output.data_ptr()), \
222 reinterpret_cast<float*>(suffix_lse.data_ptr()), num_tokens, \
223 num_heads, head_size, prefix_head_stride, output_head_stride, \
164 - prefix_num_tokens); \
224 + prefix_num_tokens, output_scale_ptr); \
225 }
226
227 /*@brief Merges the attention states from prefix and suffix
228 * into the output tensor. NUM_TOKENS: n, NUM_HEADS: h, HEAD_SIZE: d
229 *
230 * @param output [n,h,d] The output tensor to store the merged attention states.
171 - * @param output_lse [h,d] Optional tensor to store the log-sum-exp values.
231 + * @param output_lse [h,n] Optional tensor to store the log-sum-exp values.
232 * @param prefix_output [n,h,d] The prefix attention states.
233 * @param prefix_lse [h,n] The log-sum-exp values for the prefix attention
234 * states.
@@ -180,19 +240,23 @@ __global__ void merge_attn_states_kernel(
240 * is computed by merging prefix_output and suffix_output. For remaining tokens
241 * (prefill_tokens_with_context <= token_idx < n), output is copied directly
242 * from suffix_output.
243 + * @param output_scale Optional scalar tensor for FP8 static quantization.
244 + * When provided, output must be FP8 dtype.
245 */
246 template <typename scalar_t>
247 void merge_attn_states_launcher(
248 torch::Tensor& output, std::optional<torch::Tensor> output_lse,
249 const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse,
250 const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse,
189 - const std::optional<int64_t> prefill_tokens_with_context) {
251 + const std::optional<int64_t> prefill_tokens_with_context,
252 + const std::optional<torch::Tensor>& output_scale) {
253 constexpr uint NUM_THREADS = 128;
254 const uint num_tokens = output.size(0);
255 const uint num_heads = output.size(1);
256 const uint head_size = output.size(2);
257 const uint prefix_head_stride = prefix_output.stride(1);
258 const uint output_head_stride = output.stride(1);
259 + // Thread mapping is based on input BF16 pack_size
260 const uint pack_size = 16 / sizeof(scalar_t);
261 TORCH_CHECK(head_size % pack_size == 0,
262 "headsize must be multiple of pack_size:", pack_size);
@@ -208,6 +272,10 @@ void merge_attn_states_launcher(
272 if (output_lse.has_value()) {
273 output_lse_ptr = output_lse.value().data_ptr<float>();
274 }
275 + float* output_scale_ptr = nullptr;
276 + if (output_scale.has_value()) {
277 + output_scale_ptr = output_scale.value().data_ptr<float>();
278 + }
279 // Process one pack elements per thread. for float, the
280 // pack_size is 4 for half/bf16, the pack_size is 8.
281 const uint threads_per_head = head_size / pack_size;
@@ -219,20 +287,44 @@ void merge_attn_states_launcher(
287 const c10::cuda::OptionalCUDAGuard device_guard(prefix_output.device());
288 auto stream = at::cuda::getCurrentCUDAStream();
289
222 - LAUNCH_MERGE_ATTN_STATES(scalar_t, NUM_THREADS);
290 + if (output_scale.has_value()) {
291 + // FP8 output path - dispatch on output FP8 type
292 + VLLM_DISPATCH_FP8_TYPES(output.scalar_type(), "merge_attn_states_fp8", [&] {
293 + LAUNCH_MERGE_ATTN_STATES(scalar_t, fp8_t, NUM_THREADS, true);
294 + });
295 + } else {
296 + // Original BF16/FP16/FP32 output path
297 + LAUNCH_MERGE_ATTN_STATES(scalar_t, scalar_t, NUM_THREADS, false);
298 + }
299 }
300
301 #define CALL_MERGE_ATTN_STATES_LAUNCHER(scalar_t) \
302 { \
303 merge_attn_states_launcher<scalar_t>( \
304 output, output_lse, prefix_output, prefix_lse, suffix_output, \
229 - suffix_lse, prefill_tokens_with_context); \
305 + suffix_lse, prefill_tokens_with_context, output_scale); \
306 }
307
232 -void merge_attn_states(
233 - torch::Tensor& output, std::optional<torch::Tensor> output_lse,
234 - const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse,
235 - const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse,
236 - std::optional<int64_t> prefill_tokens_with_context = std::nullopt) {
237 - DISPATCH_BY_SCALAR_DTYPE(output.dtype(), CALL_MERGE_ATTN_STATES_LAUNCHER);
308 +void merge_attn_states(torch::Tensor& output,
309 + std::optional<torch::Tensor> output_lse,
310 + const torch::Tensor& prefix_output,
311 + const torch::Tensor& prefix_lse,
312 + const torch::Tensor& suffix_output,
313 + const torch::Tensor& suffix_lse,
314 + std::optional<int64_t> prefill_tokens_with_context,
315 + const std::optional<torch::Tensor>& output_scale) {
316 + if (output_scale.has_value()) {
317 + TORCH_CHECK(output.scalar_type() == at::ScalarType::Float8_e4m3fn ||
318 + output.scalar_type() == at::ScalarType::Float8_e4m3fnuz,
319 + "output must be FP8 when output_scale is provided, got: ",
320 + output.scalar_type());
321 + } else {
322 + TORCH_CHECK(output.scalar_type() == prefix_output.scalar_type(),
323 + "output dtype (", output.scalar_type(),
324 + ") must match prefix_output dtype (",
325 + prefix_output.scalar_type(), ") when output_scale is not set");
326 + }
327 + // Always dispatch on prefix_output (input) dtype
328 + DISPATCH_BY_SCALAR_DTYPE(prefix_output.dtype(),
329 + CALL_MERGE_ATTN_STATES_LAUNCHER);
330 }
csrc/ops.h
+2 -1
@@ -57,7 +57,8 @@ void merge_attn_states(
57 torch::Tensor& output, std::optional<torch::Tensor> output_lse,
58 const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse,
59 const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse,
60 - const std::optional<int64_t> prefill_tokens_with_context);
60 + const std::optional<int64_t> prefill_tokens_with_context,
61 + const std::optional<torch::Tensor>& output_scale = std::nullopt);
62 #ifndef USE_ROCM
63 void convert_vertical_slash_indexes(
64 torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS]
csrc/torch_bindings.cpp
+2 -1
@@ -73,7 +73,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
73 " Tensor prefix_lse,"
74 " Tensor suffix_output,"
75 " Tensor suffix_lse,"
76 - " int!? prefill_tokens_with_context) -> ()");
76 + " int!? prefill_tokens_with_context,"
77 + " Tensor? output_scale=None) -> ()");
78 ops.impl("merge_attn_states", torch::kCUDA, &merge_attn_states);
79 #ifndef USE_ROCM
80 ops.def(
tests/kernels/attention/test_merge_attn_states.py
+64 -15
@@ -4,7 +4,12 @@
4 import pytest
5 import torch
6
7 -from vllm._custom_ops import merge_attn_states as merge_attn_states_cuda
7 +from vllm._custom_ops import (
8 + merge_attn_states as merge_attn_states_cuda,
9 +)
10 +from vllm._custom_ops import (
11 + scaled_fp8_quant,
12 +)
13 from vllm.platforms import current_platform
14 from vllm.v1.attention.ops.triton_merge_attn_states import (
15 merge_attn_states as merge_attn_states_triton,
@@ -21,6 +26,7 @@ def merge_attn_states_torch(
26 suffix_lse: torch.Tensor, # [NUM_HEADS, NUM_TOKENS]
27 output_lse: torch.Tensor | None = None, # [NUM_HEADS, NUM_TOKENS]
28 prefill_tokens_with_context: int | None = None,
29 + output_scale: torch.Tensor | None = None, # scalar, per-tensor FP8 scale
30 ):
31 # Apply prefill_tokens_with_context mask if needed
32 if prefill_tokens_with_context is None:
@@ -49,9 +55,13 @@ def merge_attn_states_torch(
55 s_scale = s_lse_exp / out_se # [NUM_HEADS, NUM_TOKENS]
56 p_scale = torch.transpose(p_scale, 0, 1).unsqueeze(2) # [NUM_TOKENS, NUM_HEADS, 1]
57 s_scale = torch.transpose(s_scale, 0, 1).unsqueeze(2) # [NUM_TOKENS, NUM_HEADS, 1]
52 - output.copy_(
53 - prefix_output * p_scale * mask + suffix_output * (s_scale * mask + (1 - mask))
58 + output = prefix_output * p_scale * mask + suffix_output * (
59 + s_scale * mask + (1 - mask)
60 )
61 + if output_scale is not None:
62 + shape = output.shape
63 + output, _ = scaled_fp8_quant(output.float().view(-1, shape[-1]), output_scale)
64 + output = output.view(shape)
65 return output, output_lse
66
67
@@ -102,18 +112,20 @@ def generate_markdown_table():
112 )
113
114
115 +@pytest.mark.parametrize("use_fp8", [False, True])
116 @pytest.mark.parametrize("prefill_tokens_with_context", [None, 128])
117 @pytest.mark.parametrize("num_tokens", NUM_BATCH_TOKENS)
118 @pytest.mark.parametrize("num_query_heads", NUM_QUERY_HEADS)
119 @pytest.mark.parametrize("head_size", HEAD_SIZES)
109 -@pytest.mark.parametrize("output_dtype", DTYPES)
120 +@pytest.mark.parametrize("input_dtype", DTYPES)
121 @torch.inference_mode()
122 def test_merge_attn_states(
123 prefill_tokens_with_context: int | None,
124 num_tokens: int,
125 num_query_heads: int,
126 head_size: int,
116 - output_dtype: torch.dtype,
127 + input_dtype: torch.dtype,
128 + use_fp8: bool,
129 ):
130 if not current_platform.is_cuda():
131 pytest.skip(
@@ -125,9 +137,18 @@ def test_merge_attn_states(
137 NUM_HEADS = num_query_heads
138 HEAD_SIZE = head_size
139
140 + # When use_fp8 is set, inputs stay as input_dtype (bf16/fp16/fp32)
141 + # and output becomes FP8.
142 + output_dtype = input_dtype
143 + output_scale = None
144 + if use_fp8:
145 + output_dtype = current_platform.fp8_dtype()
146 + output_scale = torch.tensor([0.05], dtype=torch.float32, device="cuda")
147 +
148 print(
149 f"\nNUM_TOKENS:{NUM_TOKENS}, NUM_HEADS:{NUM_HEADS}, "
130 - f"HEAD_SIZE:{HEAD_SIZE}, DTYPE: {output_dtype}, "
150 + f"HEAD_SIZE:{HEAD_SIZE}, input_dtype: {input_dtype}, "
151 + f"output_dtype: {output_dtype}, use_fp8: {use_fp8}, "
152 f"prefill_tokens_with_context: {prefill_tokens_with_context}, "
153 f"Device: {current_platform.get_device_name()}"
154 )
@@ -156,10 +177,10 @@ def test_merge_attn_states(
177 (NUM_HEADS, NUM_TOKENS), dtype=torch.float32, device="cuda"
178 )
179 prefix_output = torch.randn(
159 - (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=output_dtype, device="cuda"
180 + (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=input_dtype, device="cuda"
181 )
182 suffix_output = torch.randn(
162 - (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=output_dtype, device="cuda"
183 + (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=input_dtype, device="cuda"
184 )
185
186 warmup_times = 2
@@ -183,6 +204,7 @@ def test_merge_attn_states(
204 suffix_lse_torch,
205 output_lse_torch,
206 prefill_tokens_with_context,
207 + output_scale,
208 )
209 torch.accelerator.synchronize()
210
@@ -196,6 +218,7 @@ def test_merge_attn_states(
218 suffix_lse_torch,
219 output_lse_torch,
220 prefill_tokens_with_context,
221 + output_scale,
222 )
223 end.record()
224 torch.accelerator.synchronize()
@@ -220,6 +243,7 @@ def test_merge_attn_states(
243 suffix_lse,
244 output_lse_ref_triton,
245 prefill_tokens_with_context,
246 + output_scale,
247 )
248 torch.accelerator.synchronize()
249
@@ -233,6 +257,7 @@ def test_merge_attn_states(
257 suffix_lse,
258 output_lse_ref_triton,
259 prefill_tokens_with_context,
260 + output_scale,
261 )
262 end.record()
263 torch.accelerator.synchronize()
@@ -254,6 +279,7 @@ def test_merge_attn_states(
279 suffix_lse,
280 output_lse_cuda,
281 prefill_tokens_with_context,
282 + output_scale,
283 )
284 torch.accelerator.synchronize()
285
@@ -267,6 +293,7 @@ def test_merge_attn_states(
293 suffix_lse,
294 output_lse_cuda,
295 prefill_tokens_with_context,
296 + output_scale,
297 )
298 end.record()
299 torch.accelerator.synchronize()
@@ -288,7 +315,19 @@ def test_merge_attn_states(
315 # Liger Kernel: Efficient Triton Kernels for LLM Training
316 # https://arxiv.org/pdf/2410.10989, 3.3 Correctness
317 # use rtol = 1e-2 for bfloat16.
291 - rtol = 1e-2 if output_dtype == torch.bfloat16 else 1e-3
318 + if use_fp8:
319 + # Compare in dequantized space (multiply back by scale) so that
320 + # absolute differences reflect real precision, not amplified FP8
321 + # quantization steps.
322 + atol, rtol = 1e-1, 1e-1
323 + assert output_scale is not None
324 + scale = output_scale.item()
325 + elif output_dtype == torch.bfloat16:
326 + atol, rtol = 1e-3, 1e-2
327 + scale = 1.0
328 + else:
329 + atol, rtol = 1e-3, 1e-3
330 + scale = 1.0
331
332 def diff(a: torch.Tensor, b: torch.Tensor):
333 max_diff = torch.max(torch.abs(a.float() - b.float()))
@@ -300,16 +339,26 @@ def test_merge_attn_states(
339 output_ref = output_ref_triton
340 output_lse_ref = output_lse_ref_triton
341 torch.testing.assert_close(
303 - output_cuda.float(), output_ref.float(), atol=1e-3, rtol=rtol
342 + output_cuda.float() * scale,
343 + output_ref.float() * scale,
344 + atol=atol,
345 + rtol=rtol,
346 + )
347 + print(
348 + "Output all match, max abs diff (dequantized):"
349 + if use_fp8
350 + else "Output all match, max abs diff:"
351 )
305 - print("Output all match, max abs diff:")
306 - print(f"(Triton vs Torch) : {diff(output_torch, output_ref)}")
307 - print(f" (CUDA vs Torch) : {diff(output_torch, output_cuda)}")
308 - print(f" (CUDA vs Triton): {diff(output_ref, output_cuda)}")
352 + _diff = diff(output_ref.float() * scale, output_torch.float() * scale)
353 + print(f"(Triton vs Torch) : {_diff}")
354 + _diff = diff(output_torch.float() * scale, output_cuda.float() * scale)
355 + print(f" (CUDA vs Torch) : {_diff}")
356 + _diff = diff(output_ref.float() * scale, output_cuda.float() * scale)
357 + print(f" (CUDA vs Triton): {_diff}")
358 print("-" * 100)
359
360 torch.testing.assert_close(
312 - output_lse_cuda.float(), output_lse_ref.float(), atol=1e-3, rtol=rtol
361 + output_lse_cuda.float(), output_lse_ref.float(), atol=atol, rtol=rtol
362 )
363 print("Output LSE all match, max abs diff:")
364 print(f"(Triton vs Torch) : {diff(output_lse_torch, output_lse_ref)}")
vllm/_custom_ops.py
+2
@@ -265,6 +265,7 @@ def merge_attn_states(
265 suffix_lse: torch.Tensor,
266 output_lse: torch.Tensor | None = None,
267 prefill_tokens_with_context: int | None = None,
268 + output_scale: torch.Tensor | None = None,
269 ) -> None:
270 torch.ops._C.merge_attn_states(
271 output,
@@ -274,6 +275,7 @@ def merge_attn_states(
275 suffix_output,
276 suffix_lse,
277 prefill_tokens_with_context,
278 + output_scale,
279 )
280
281
vllm/v1/attention/ops/merge_attn_states.py
+25 -10
@@ -14,6 +14,7 @@ def merge_attn_states(
14 suffix_lse: torch.Tensor,
15 output_lse: torch.Tensor | None = None,
16 prefill_tokens_with_context: int | None = None,
17 + output_scale: torch.Tensor | None = None,
18 ) -> None:
19 """Merge partial attention outputs from prefix (KV cache) and suffix
20 (new tokens) into a single output tensor using the log-sum-exp (LSE)
@@ -41,27 +42,37 @@ def merge_attn_states(
42 >= this value are decode or context-free prefill tokens whose
43 output is taken directly from suffix_output. If None, all tokens
44 are treated as having context.
45 + output_scale: Optional scalar tensor for FP8 static quantization.
46 + When provided, output must be FP8 dtype.
47 """
48
49 # NOTE(DefTruth): Currently, custom merge_attn_states CUDA kernel
47 - # does not support FP8 dtype, fallback to use Triton kernel.
48 - def supported_dtypes(o: torch.Tensor) -> bool:
49 - return o.dtype in [torch.float32, torch.half, torch.bfloat16]
50 + # does not support FP8 dtype for inputs, fallback to use Triton kernel.
51 + # However, when output_scale is provided, the inputs are still BF16/FP16
52 + # and the output is FP8 — both CUDA and Triton support this.
53 + # FP8 output requires output_scale to be set.
54 + if output.dtype not in (torch.float32, torch.half, torch.bfloat16):
55 + assert output_scale is not None, (
56 + f"output_scale is required when output is {output.dtype}"
57 + )
58 +
59 + def supported_dtypes(prefix: torch.Tensor) -> bool:
60 + return prefix.dtype in [torch.float32, torch.half, torch.bfloat16]
61
62 # NOTE(DefTruth): Currently, custom merge_attn_states CUDA
63 # kernel load/store 128b(16 bytes) per memory issue within
64 # thread. Namely, the headsize(headdim) must be multiple of
54 - # pack_size (float32 -> 4, half/bfloat16 -> 8).
55 - def supported_headdim(o: torch.Tensor) -> bool:
56 - headdim = o.shape[2] # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
57 - if o.dtype == torch.float32:
65 + # pack_size based on input dtype (float32 -> 4, half/bfloat16 -> 8).
66 + def supported_headdim(prefix: torch.Tensor) -> bool:
67 + headdim = prefix.shape[2] # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
68 + if prefix.dtype == torch.float32:
69 return headdim % 4 == 0
70 return headdim % 8 == 0
71
72 if (
73 current_platform.is_cuda()
63 - and supported_dtypes(output)
64 - and supported_headdim(output)
74 + and supported_dtypes(prefix_output)
75 + and supported_headdim(prefix_output)
76 ):
77 from vllm._custom_ops import merge_attn_states
78
@@ -73,9 +84,12 @@ def merge_attn_states(
84 suffix_lse,
85 output_lse,
86 prefill_tokens_with_context,
87 + output_scale,
88 )
89 else:
78 - from vllm.v1.attention.ops.triton_merge_attn_states import merge_attn_states
90 + from vllm.v1.attention.ops.triton_merge_attn_states import (
91 + merge_attn_states,
92 + )
93
94 return merge_attn_states(
95 output,
@@ -85,4 +99,5 @@ def merge_attn_states(
99 suffix_lse,
100 output_lse,
101 prefill_tokens_with_context,
102 + output_scale,
103 )
vllm/v1/attention/ops/triton_merge_attn_states.py
+24 -2
@@ -3,8 +3,11 @@
3
4 import torch
5
6 +from vllm.platforms import current_platform
7 from vllm.triton_utils import tl, triton
8
9 +float8_info = torch.finfo(current_platform.fp8_dtype())
10 +
11
12 # Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
13 # can be used to combine partial attention results (in the split-KV case)
@@ -16,14 +19,15 @@ def merge_attn_states(
19 suffix_lse: torch.Tensor,
20 output_lse: torch.Tensor | None = None,
21 prefill_tokens_with_context: int | None = None,
22 + output_scale: torch.Tensor | None = None,
23 ) -> None:
24 num_tokens = output.shape[0]
25 num_query_heads = output.shape[1]
26 head_size = output.shape[2]
27 padded_head_size = triton.next_power_of_2(head_size)
28 # We assume the output stride on num_head is not always as same as the
25 - # `suffix_output` and `prefix_output`, as them might be padded by the attention
26 - # backend.
29 + # `suffix_output` and `prefix_output`, as them might be padded by the
30 + # attention backend.
31 prefix_head_stride = prefix_output.stride(1)
32 output_head_stride = output.stride(1)
33
@@ -41,10 +45,12 @@ def merge_attn_states(
45 suffix_lse,
46 prefix_head_stride,
47 output_head_stride,
48 + output_scale,
49 head_size,
50 padded_head_size,
51 output_lse is not None,
52 prefill_tokens_with_context,
53 + output_scale is not None,
54 )
55
56
@@ -58,10 +64,14 @@ def merge_attn_states_kernel(
64 suffix_lse, # [NUM_HEADS, NUM_TOKENS]
65 prefix_head_stride,
66 output_head_stride,
67 + output_scale, # scale tensor or None
68 HEAD_SIZE: tl.constexpr,
69 PADDED_HEAD_SIZE: tl.constexpr,
70 OUTPUT_LSE: tl.constexpr,
71 prefill_tokens_with_context: tl.constexpr,
72 + USE_FP8: tl.constexpr,
73 + FP8_MIN: tl.constexpr = float8_info.min,
74 + FP8_MAX: tl.constexpr = float8_info.max,
75 ):
76 token_idx = tl.program_id(0)
77 num_tokens = tl.num_programs(0)
@@ -87,6 +97,12 @@ def merge_attn_states_kernel(
97 + head_arange,
98 mask=head_mask,
99 )
100 +
101 + if USE_FP8:
102 + s_out = s_out * (1.0 / tl.load(output_scale))
103 + s_out = tl.clamp(s_out, FP8_MIN, FP8_MAX)
104 + s_out = s_out.to(output.dtype.element_ty)
105 +
106 tl.store(
107 output
108 + token_idx * num_heads * output_head_stride
@@ -143,6 +159,12 @@ def merge_attn_states_kernel(
159 p_scale = p_se / out_se
160 s_scale = s_se / out_se
161 out = p_out * p_scale + s_out * s_scale
162 +
163 + if USE_FP8:
164 + out = out * (1.0 / tl.load(output_scale))
165 + out = tl.clamp(out, FP8_MIN, FP8_MAX)
166 + out = out.to(output.dtype.element_ty)
167 +
168 tl.store(
169 output
170 + token_idx * num_heads * output_head_stride