@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 03752dba8f
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 14, 05:30 UTC

[NVIDIA] Support Flashinfer TRTLLM FP8-q/kv/out Attention Kernel (#21716)

Signed-off-by: elvischenv <219235043+elvischenv@users.noreply.github.com> Co-authored-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>

elvischenv committed Aug 19, 2025 at 20:22 UTC 03752dba8f1c4cd8eda99949b7d7ae90a47425d5
9 files changed +849 -433
.buildkite/test-pipeline.yaml
+2
@@ -631,6 +631,7 @@ steps:
631 - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
632 - vllm/v1/attention/backends/flashinfer.py
633 - vllm/compilation/fusion.py
634 + - vllm/compilation/fusion_attn.py
635 commands:
636 - nvidia-smi
637 - python3 examples/offline_inference/basic/chat.py
@@ -647,6 +648,7 @@ steps:
648 - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
649 # Fusion
650 - pytest -v -s tests/compile/test_fusion_all_reduce.py
651 + - pytest -v -s tests/compile/test_fusion_attn.py::test_attention_quant_pattern
652
653 ##### 1 GPU test #####
654 ##### multi gpus test #####
benchmarks/kernels/benchmark_trtllm_decode_attention.py
+149 -134
@@ -3,16 +3,14 @@
3
4 import csv
5 import os
6 -import random
6 from datetime import datetime
7 +from typing import Optional
8
9 import flashinfer
10 import torch
11
12 FLOAT32_BYTES = torch.finfo(torch.float).bits // 8
13 -
14 -# KV Cache Layout for TRT-LLM
15 -# kv_cache_shape = (num_blocks, 2, num_kv_heads, page_size, head_dim)
13 +FP8_DTYPE = torch.float8_e4m3fn
14
15
16 def to_float8(x, dtype=torch.float8_e4m3fn):
@@ -26,149 +24,168 @@ def to_float8(x, dtype=torch.float8_e4m3fn):
24
25 @torch.no_grad()
26 def benchmark_decode(
29 - num_seqs,
30 - max_seq_len,
31 - page_size=16,
32 - dtype=torch.bfloat16,
33 - kv_layout="HND",
34 - num_kv_heads=8,
35 - kv_cache_dtype="auto",
36 - head_dim=128,
37 - warmup=10,
38 - trials=20,
27 + dtype: torch.dtype,
28 + quant_dtypes: tuple[
29 + Optional[torch.dtype], Optional[torch.dtype], Optional[torch.dtype]
30 + ],
31 + batch_size: int,
32 + max_seq_len: int,
33 + num_heads: tuple[int, int] = (64, 8),
34 + head_size: int = 128,
35 + kv_layout: str = "HND",
36 + block_size: int = 16,
37 + warmup: int = 10,
38 + trials: int = 20,
39 ):
40 torch.set_default_device("cuda")
41 - device = "cuda"
41 torch.manual_seed(0)
42
44 - HEAD_GRP_SIZE = 8
45 - MAX_SEQ_LEN = max_seq_len
46 -
47 - # large number to reduce kv_cache reuse
48 - NUM_BLOCKS = int(256000 / page_size)
49 -
50 - workspace_buffer = torch.empty(1024 * 1024 * 1024, dtype=torch.int8, device=device)
43 + q_quant_dtype, kv_quant_dtype, o_quant_dtype = quant_dtypes
44 + q_quant_dtype = q_quant_dtype or dtype
45 + kv_quant_dtype = kv_quant_dtype or dtype
46 + o_quant_dtype = o_quant_dtype or dtype
47
52 - # For decode, batch_size is num_decode_token
53 - num_qo_heads = num_kv_heads * HEAD_GRP_SIZE
54 - sm_scale = float(1.0 / (head_dim**0.5))
55 - q = torch.randn(num_seqs, num_qo_heads, head_dim, device=device, dtype=dtype)
56 - kv_lens = [random.randint(1, MAX_SEQ_LEN) for _ in range(num_seqs)]
48 + num_qo_heads, num_kv_heads = num_heads
49 + assert num_qo_heads % num_kv_heads == 0
50
58 - max_kv_len = max(kv_lens)
59 - kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int, device=device)
60 - max_num_blocks_per_seq = (max_kv_len + page_size - 1) // page_size
51 + sm_scale = float(1.0 / (head_size**0.5))
52
53 + # large number to reduce kv_cache reuse
54 + NUM_BLOCKS = int(256000 / block_size)
55 +
56 + kv_cache_shape = None
57 + if kv_layout == "NHD":
58 + kv_cache_shape = (NUM_BLOCKS, 2, block_size, num_kv_heads, head_size)
59 + elif kv_layout == "HND":
60 + kv_cache_shape = (NUM_BLOCKS, 2, num_kv_heads, block_size, head_size)
61 + else:
62 + raise ValueError(f"Invalid kv_layout: {kv_layout}")
63 +
64 + query = torch.randn(batch_size, num_qo_heads, head_size, dtype=dtype)
65 + if q_quant_dtype == FP8_DTYPE:
66 + query, q_scale = to_float8(query)
67 + ref_query = query.to(dtype) * q_scale
68 + else:
69 + q_scale = 1.0
70 + ref_query = query
71 +
72 + kv_lens = torch.randint(1, max_seq_len, (batch_size,), dtype=torch.int32)
73 + kv_lens[-1] = max_seq_len
74 +
75 + seq_lens = kv_lens
76 + max_seq_len = torch.max(seq_lens).item()
77 +
78 + kv_cache = torch.randn(kv_cache_shape, dtype=dtype)
79 + if kv_quant_dtype == FP8_DTYPE:
80 + kv_cache, kv_scale = to_float8(kv_cache)
81 + ref_kv_cache = kv_cache.to(dtype) * kv_scale
82 + else:
83 + kv_scale = 1.0
84 + ref_kv_cache = kv_cache
85 + k_scale = v_scale = kv_scale
86 +
87 + max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size
88 block_tables = torch.randint(
63 - 0, NUM_BLOCKS, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
89 + 0, NUM_BLOCKS, (batch_size, max_num_blocks_per_seq), dtype=torch.int32
90 )
65 -
66 - kv_cache_shape = (NUM_BLOCKS, 2, num_kv_heads, page_size, head_dim)
67 - kv_cache = torch.randn(size=kv_cache_shape, device=device, dtype=dtype)
68 - k_scale = v_scale = 1.0
69 -
70 - if kv_cache_dtype.startswith("fp8"):
71 - kv_cache, _ = to_float8(kv_cache)
72 -
73 - output_trtllm = torch.empty(q.shape, dtype=dtype)
74 -
75 - # Benchmark TRT decode
76 - def trt_decode():
77 - return flashinfer.decode.trtllm_batch_decode_with_kv_cache(
78 - q,
79 - kv_cache,
80 - workspace_buffer,
81 - block_tables,
82 - kv_lens_tensor,
83 - max_kv_len,
84 - bmm1_scale=k_scale * sm_scale,
85 - bmm2_scale=v_scale,
86 - out=output_trtllm,
87 - )
88 -
89 - def time_fn(fn, warmup=10, trials=20):
90 - torch.cuda.synchronize()
91 - start = torch.cuda.Event(enable_timing=True)
92 - end = torch.cuda.Event(enable_timing=True)
93 - times = []
94 - for i in range(warmup):
95 - fn()
96 - for i in range(trials):
97 - start.record()
98 - fn()
99 - end.record()
100 - torch.cuda.synchronize()
101 - times.append(start.elapsed_time(end)) # ms
102 - return sum(times) / len(times), torch.std(torch.tensor(times))
103 -
104 - # TRT Decode
105 - trt_mean, trt_std = time_fn(trt_decode)
106 -
91 kv_indptr = [0]
92 kv_indices = []
93 kv_last_page_lens = []
110 - for i in range(num_seqs):
111 - seq_len = kv_lens[i]
94 + for i in range(batch_size):
95 + seq_len = seq_lens[i]
96 assert seq_len > 0
113 - num_blocks = (seq_len + page_size - 1) // page_size
97 + num_blocks = (seq_len + block_size - 1) // block_size
98 kv_indices.extend(block_tables[i, :num_blocks])
99 kv_indptr.append(kv_indptr[-1] + num_blocks)
116 - kv_last_page_len = seq_len % page_size
100 + kv_last_page_len = seq_len % block_size
101 if kv_last_page_len == 0:
118 - kv_last_page_len = page_size
102 + kv_last_page_len = block_size
103 kv_last_page_lens.append(kv_last_page_len)
104
105 kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32)
106 kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
107 kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
124 -
125 - output_baseline = torch.empty(q.shape, dtype=dtype)
108 + workspace_buffer = torch.zeros(1024 * 1024 * 1024, dtype=torch.int8)
109
110 wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
111 workspace_buffer,
112 kv_layout,
113 use_tensor_cores=((num_qo_heads // num_kv_heads) > 4),
114 )
132 -
115 wrapper.plan(
116 kv_indptr,
117 kv_indices,
118 kv_last_page_lens,
119 num_qo_heads,
120 num_kv_heads,
139 - head_dim,
140 - page_size,
121 + head_size,
122 + block_size,
123 "NONE",
124 + sm_scale=sm_scale,
125 q_data_type=dtype,
143 - kv_data_type=torch.float8_e4m3fn if kv_cache_dtype.startswith("fp8") else dtype,
126 + kv_data_type=dtype,
127 )
128
129 + def time_fn(fn, warmup=10, trials=20):
130 + torch.cuda.synchronize()
131 + start = torch.cuda.Event(enable_timing=True)
132 + end = torch.cuda.Event(enable_timing=True)
133 + times = []
134 + for i in range(warmup):
135 + fn()
136 + for i in range(trials):
137 + start.record()
138 + fn()
139 + end.record()
140 + torch.cuda.synchronize()
141 + times.append(start.elapsed_time(end)) # ms
142 + return sum(times) / len(times), torch.std(torch.tensor(times))
143 +
144 + o_scale = 1.0
145 + output_baseline = torch.empty(ref_query.shape, dtype=dtype)
146 + output_trtllm = torch.empty(query.shape, dtype=o_quant_dtype)
147 +
148 def baseline_decode():
147 - return wrapper.run(q, kv_cache, sm_scale, k_scale, v_scale, output_baseline)
149 + return wrapper.run(ref_query, ref_kv_cache, out=output_baseline)
150 +
151 + def trtllm_decode():
152 + return flashinfer.decode.trtllm_batch_decode_with_kv_cache(
153 + query=query,
154 + kv_cache=kv_cache,
155 + workspace_buffer=workspace_buffer,
156 + block_tables=block_tables,
157 + seq_lens=seq_lens,
158 + max_seq_len=max_seq_len,
159 + bmm1_scale=q_scale * k_scale * sm_scale,
160 + bmm2_scale=v_scale / o_scale,
161 + out=output_trtllm,
162 + )
163
164 baseline_mean, baseline_std = time_fn(baseline_decode)
165 + trtllm_mean, trtllm_std = time_fn(trtllm_decode)
166
167 # Calculate percentage speedup (positive means TRT is faster)
152 - speedup_percent = (baseline_mean - trt_mean) / baseline_mean
168 + speedup_percent = (baseline_mean - trtllm_mean) / baseline_mean
169
170 print(
155 - f"\t{num_seqs}\t{max_seq_len}\t{trt_mean:.3f}\t{trt_std.item():.3f}"
171 + f"\t{batch_size}\t{max_seq_len}\t{trtllm_mean:.3f}\t{trtllm_std.item():.3f}"
172 f"\t{baseline_mean:.3f}\t{baseline_std.item():.3f}\t{speedup_percent:.3f}"
173 )
174
175 # Return results for CSV writing
176 return {
161 - "num_seqs": num_seqs,
162 - "trt_mean": trt_mean,
163 - "trt_std": trt_std.item(),
177 + "batch_size": batch_size,
178 + "trtllm_mean": trtllm_mean,
179 + "trtllm_std": trtllm_std.item(),
180 "baseline_mean": baseline_mean,
181 "baseline_std": baseline_std.item(),
182 "speedup_percent": speedup_percent,
167 - "q_dtype": str(dtype),
168 - "kv_cache_dtype": kv_cache_dtype,
169 - "page_size": page_size,
183 + "q_dtype": str(q_quant_dtype),
184 + "kv_cache_dtype": str(kv_quant_dtype),
185 + "output_dtype": str(o_quant_dtype),
186 + "block_size": block_size,
187 "num_kv_heads": num_kv_heads,
171 - "head_dim": head_dim,
188 + "head_size": head_size,
189 "max_seq_len": max_seq_len,
190 }
191
@@ -180,17 +197,18 @@ def write_results_to_csv(results, filename=None):
197 filename = f"flashinfer_trtllm_benchmark_{timestamp}.csv"
198
199 fieldnames = [
183 - "num_seqs",
184 - "trt_mean",
185 - "trt_std",
200 + "batch_size",
201 + "trtllm_mean",
202 + "trtllm_std",
203 "baseline_mean",
204 "baseline_std",
205 "speedup_percent",
206 "q_dtype",
207 "kv_cache_dtype",
191 - "page_size",
208 + "output_dtype",
209 + "block_size",
210 "num_kv_heads",
193 - "head_dim",
211 + "head_size",
212 "max_seq_len",
213 ]
214
@@ -209,45 +227,42 @@ def write_results_to_csv(results, filename=None):
227
228
229 if __name__ == "__main__":
212 - num_seqs = [1, 4, 8, 16, 32, 64, 128, 256]
230 + batch_sizes = [1, 4, 8, 16, 32, 64, 128, 256]
231 max_seq_lens = [1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072]
232 all_results = []
233
216 - print(
217 - "Running benchmark for q_dtype = bfloat16, kv_cache_dtype: bfloat16, "
218 - "output_dtype: bfloat16"
219 - )
220 - print(
221 - "\tnum_seqs\tmax_seq_len\ttrt_mean\ttrt_std\tbaseline_mean\t"
222 - "baseline_std\tspeedup_percent"
223 - )
224 - for max_seq_len in max_seq_lens:
225 - for bs in num_seqs:
226 - result = benchmark_decode(
227 - bs,
228 - max_seq_len,
229 - dtype=torch.bfloat16,
230 - kv_cache_dtype="auto",
231 - )
232 - all_results.append(result)
234 + dtype = torch.bfloat16
235 + quant_dtypes = [
236 + # (q_quant_dtype, kv_quant_dtype, o_quant_dtype)
237 + (None, None, None),
238 + (None, FP8_DTYPE, None),
239 + (FP8_DTYPE, FP8_DTYPE, FP8_DTYPE),
240 + ]
241
234 - print(
235 - "Running benchmark for q_dtype = bfloat16, kv_cache_dtype: fp8, "
236 - "output_dtype: bfloat16"
237 - )
238 - print(
239 - "\tnum_seqs\tmax_seq_len\ttrt_mean\ttrt_std\tbaseline_mean\t"
240 - "baseline_std\tspeedup_percent"
241 - )
242 - for max_seq_len in max_seq_lens:
243 - for bs in num_seqs:
244 - result = benchmark_decode(
245 - bs,
246 - max_seq_len,
247 - dtype=torch.bfloat16,
248 - kv_cache_dtype="fp8",
249 - )
250 - all_results.append(result)
242 + for quant_dtype in quant_dtypes:
243 + q_quant_dtype, kv_quant_dtype, o_quant_dtype = quant_dtype
244 + q_quant_dtype = q_quant_dtype or dtype
245 + kv_quant_dtype = kv_quant_dtype or dtype
246 + o_quant_dtype = o_quant_dtype or dtype
247 +
248 + print(
249 + f"Running benchmark for q_dtype = {q_quant_dtype}, "
250 + f"kv_cache_dtype: {kv_quant_dtype}, "
251 + f"output_dtype: {o_quant_dtype}"
252 + )
253 + print(
254 + "\tbatch_size\tmax_seq_len\ttrtllm_mean\ttrtllm_std\tbaseline_mean\t"
255 + "baseline_std\tspeedup_percent"
256 + )
257 + for max_seq_len in max_seq_lens:
258 + for bs in batch_sizes:
259 + result = benchmark_decode(
260 + dtype=dtype,
261 + quant_dtypes=quant_dtype,
262 + batch_size=bs,
263 + max_seq_len=max_seq_len,
264 + )
265 + all_results.append(result)
266
267 # Write all results to CSV
268 write_results_to_csv(all_results)
benchmarks/kernels/benchmark_trtllm_prefill_attention.py
+130 -99
@@ -3,16 +3,14 @@
3
4 import csv
5 import os
6 -import random
6 from datetime import datetime
7 +from typing import Optional
8
9 import flashinfer
10 import torch
11
12 FLOAT32_BYTES = torch.finfo(torch.float).bits // 8
13 -
14 -# KV Cache Layout for TRT-LLM
15 -# kv_cache_shape = (num_blocks, 2, num_kv_heads, page_size, head_dim)
13 +FP8_DTYPE = torch.float8_e4m3fn
14
15
16 def to_float8(x, dtype=torch.float8_e4m3fn):
@@ -26,84 +24,99 @@ def to_float8(x, dtype=torch.float8_e4m3fn):
24
25 @torch.no_grad()
26 def benchmark_prefill(
29 - num_seqs,
30 - max_seq_len,
31 - page_size=16,
32 - dtype=torch.bfloat16,
33 - kv_layout="HND",
34 - num_kv_heads=8,
35 - kv_cache_dtype="auto",
36 - head_dim=128,
37 - warmup=10,
38 - trials=20,
27 + dtype: torch.dtype,
28 + quant_dtypes: tuple[
29 + Optional[torch.dtype], Optional[torch.dtype], Optional[torch.dtype]
30 + ],
31 + batch_size: int,
32 + max_seq_len: int,
33 + num_heads: tuple[int, int] = (64, 8),
34 + head_size: int = 128,
35 + kv_layout: str = "HND",
36 + block_size: int = 16,
37 + warmup: int = 10,
38 + trials: int = 20,
39 ):
40 torch.set_default_device("cuda")
41 torch.manual_seed(0)
42
43 - HEAD_GRP_SIZE = 8
44 - MAX_SEQ_LEN = max_seq_len
43 + q_quant_dtype, kv_quant_dtype, o_quant_dtype = quant_dtypes
44 + q_quant_dtype = q_quant_dtype or dtype
45 + kv_quant_dtype = kv_quant_dtype or dtype
46 + o_quant_dtype = o_quant_dtype or dtype
47
46 - # large number to reduce kv_cache reuse
47 - NUM_BLOCKS = int(256000 / page_size)
48 + max_q_len = max_kv_len = max_seq_len
49
49 - workspace_buffer = torch.empty(1024 * 1024 * 1024, dtype=torch.int8)
50 + num_qo_heads, num_kv_heads = num_heads
51 + assert num_qo_heads % num_kv_heads == 0
52
51 - num_qo_heads = num_kv_heads * HEAD_GRP_SIZE
52 - sm_scale = float(1.0 / (head_dim**0.5))
53 + sm_scale = float(1.0 / (head_size**0.5))
54
54 - q_lens = [random.randint(1, MAX_SEQ_LEN) for _ in range(num_seqs)]
55 - q_lens[-1] = MAX_SEQ_LEN
56 - max_q_len = max(q_lens)
55 + # large number to reduce kv_cache reuse
56 + NUM_BLOCKS = int(256000 / block_size)
57 +
58 + kv_cache_shape = None
59 + if kv_layout == "NHD":
60 + kv_cache_shape = (NUM_BLOCKS, 2, block_size, num_kv_heads, head_size)
61 + elif kv_layout == "HND":
62 + kv_cache_shape = (NUM_BLOCKS, 2, num_kv_heads, block_size, head_size)
63 + else:
64 + raise ValueError(f"Invalid kv_layout: {kv_layout}")
65 +
66 + q_lens = torch.randint(1, max_q_len, (batch_size,), dtype=torch.int32)
67 + q_lens[-1] = max_q_len
68 q_indptr = torch.cat(
69 [
70 torch.tensor([0], dtype=torch.int32),
60 - torch.cumsum(
61 - torch.tensor(q_lens, dtype=torch.int32), dim=0, dtype=torch.int32
62 - ),
71 + torch.cumsum(q_lens, dim=0, dtype=torch.int32),
72 ]
73 )
65 - q = torch.randn(sum(q_lens), num_qo_heads, head_dim, dtype=dtype)
74
67 - kv_lens = [random.randint(0, MAX_SEQ_LEN) for _ in range(num_seqs)]
68 - kv_lens[-1] = MAX_SEQ_LEN
69 -
70 - seq_lens = [q_len + kv_len for q_len, kv_len in zip(q_lens, kv_lens)]
71 - max_seq_len = max(seq_lens)
72 - seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32)
73 -
74 - max_num_blocks_per_seq = (max_seq_len + page_size - 1) // page_size
75 + query = torch.randn(torch.sum(q_lens).item(), num_qo_heads, head_size, dtype=dtype)
76 + if q_quant_dtype == FP8_DTYPE:
77 + query, q_scale = to_float8(query)
78 + ref_query = query.to(dtype) * q_scale
79 + else:
80 + q_scale = 1.0
81 + ref_query = query
82 +
83 + kv_lens = torch.randint(0, max_kv_len, (batch_size,), dtype=torch.int32)
84 + kv_lens[-1] = max_kv_len
85 +
86 + seq_lens = kv_lens + q_lens
87 + max_seq_len = torch.max(seq_lens).item()
88 +
89 + kv_cache = torch.randn(kv_cache_shape, dtype=dtype)
90 + if kv_quant_dtype == FP8_DTYPE:
91 + kv_cache, kv_scale = to_float8(kv_cache)
92 + ref_kv_cache = kv_cache.to(dtype) * kv_scale
93 + else:
94 + kv_scale = 1.0
95 + ref_kv_cache = kv_cache
96 + k_scale = v_scale = kv_scale
97 +
98 + max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size
99 block_tables = torch.randint(
76 - 0, NUM_BLOCKS, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
100 + 0, NUM_BLOCKS, (batch_size, max_num_blocks_per_seq), dtype=torch.int32
101 )
78 -
79 - kv_cache_shape = (NUM_BLOCKS, 2, num_kv_heads, page_size, head_dim)
80 - kv_cache = torch.randn(size=kv_cache_shape, dtype=dtype)
81 - k_scale = v_scale = 1.0
82 -
83 - if kv_cache_dtype.startswith("fp8"):
84 - kv_cache, _ = to_float8(kv_cache)
85 -
86 - output_trtllm = torch.empty(q.shape, dtype=dtype)
87 -
102 kv_indptr = [0]
103 kv_indices = []
104 kv_last_page_lens = []
91 - for i in range(num_seqs):
105 + for i in range(batch_size):
106 seq_len = seq_lens[i]
107 assert seq_len > 0
94 - num_blocks = (seq_len + page_size - 1) // page_size
108 + num_blocks = (seq_len + block_size - 1) // block_size
109 kv_indices.extend(block_tables[i, :num_blocks])
110 kv_indptr.append(kv_indptr[-1] + num_blocks)
97 - kv_last_page_len = seq_len % page_size
111 + kv_last_page_len = seq_len % block_size
112 if kv_last_page_len == 0:
99 - kv_last_page_len = page_size
113 + kv_last_page_len = block_size
114 kv_last_page_lens.append(kv_last_page_len)
115
116 kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32)
117 kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
118 kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
105 -
106 - output_baseline = torch.empty(q.shape, dtype=dtype)
119 + workspace_buffer = torch.zeros(1024 * 1024 * 1024, dtype=torch.int8)
120
121 wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(
122 workspace_buffer, kv_layout
@@ -115,12 +128,12 @@ def benchmark_prefill(
128 kv_last_page_lens,
129 num_qo_heads,
130 num_kv_heads,
118 - head_dim,
119 - page_size,
131 + head_size,
132 + block_size,
133 causal=True,
134 sm_scale=sm_scale,
135 q_data_type=dtype,
123 - kv_data_type=kv_cache.dtype,
136 + kv_data_type=dtype,
137 )
138
139 def time_fn(fn, warmup=10, trials=20):
@@ -138,52 +151,55 @@ def benchmark_prefill(
151 times.append(start.elapsed_time(end)) # ms
152 return sum(times) / len(times), torch.std(torch.tensor(times))
153
154 + o_scale = 1.0
155 + output_baseline = torch.empty(ref_query.shape, dtype=dtype)
156 + output_trtllm = torch.empty(query.shape, dtype=o_quant_dtype)
157 +
158 def baseline_prefill():
142 - return wrapper.run(
143 - q, kv_cache, k_scale=k_scale, v_scale=v_scale, out=output_baseline
144 - )
159 + return wrapper.run(ref_query, ref_kv_cache, out=output_baseline)
160
146 - def trt_prefill():
161 + def trtllm_prefill():
162 return flashinfer.prefill.trtllm_batch_context_with_kv_cache(
148 - query=q,
163 + query=query,
164 kv_cache=kv_cache,
165 workspace_buffer=workspace_buffer,
166 block_tables=block_tables,
152 - seq_lens=seq_lens_tensor,
167 + seq_lens=seq_lens,
168 max_q_len=max_q_len,
169 max_kv_len=max_seq_len,
155 - bmm1_scale=k_scale * sm_scale,
156 - bmm2_scale=v_scale,
157 - batch_size=num_seqs,
170 + bmm1_scale=q_scale * k_scale * sm_scale,
171 + bmm2_scale=v_scale / o_scale,
172 + batch_size=batch_size,
173 cum_seq_lens_q=q_indptr,
174 cum_seq_lens_kv=kv_indptr,
175 out=output_trtllm,
176 )
177
163 - trt_mean, trt_std = time_fn(trt_prefill)
178 baseline_mean, baseline_std = time_fn(baseline_prefill)
179 + trtllm_mean, trtllm_std = time_fn(trtllm_prefill)
180
181 # Calculate percentage speedup (positive means TRT is faster)
167 - speedup_percent = (baseline_mean - trt_mean) / baseline_mean
182 + speedup_percent = (baseline_mean - trtllm_mean) / baseline_mean
183
184 print(
170 - f"\t{num_seqs}\t{max_seq_len}\t{trt_mean:.5f}\t{trt_std.item():.5f}"
171 - f"\t{baseline_mean:.5f}\t{baseline_std.item():.5f}\t{speedup_percent:.5f}"
185 + f"\t{batch_size}\t{max_seq_len}\t{trtllm_mean:8.3f}\t{trtllm_std.item():8.3f}"
186 + f"\t{baseline_mean:8.3f}\t{baseline_std.item():8.3f}\t{speedup_percent:8.3f}"
187 )
188
189 # Return results for CSV writing
190 return {
176 - "num_seqs": num_seqs,
177 - "trt_mean": trt_mean,
178 - "trt_std": trt_std.item(),
191 + "batch_size": batch_size,
192 + "trtllm_mean": trtllm_mean,
193 + "trtllm_std": trtllm_std.item(),
194 "baseline_mean": baseline_mean,
195 "baseline_std": baseline_std.item(),
196 "speedup_percent": speedup_percent,
182 - "q_dtype": str(dtype),
183 - "kv_cache_dtype": kv_cache_dtype,
184 - "page_size": page_size,
197 + "q_dtype": str(q_quant_dtype),
198 + "kv_cache_dtype": str(kv_quant_dtype),
199 + "output_dtype": str(o_quant_dtype),
200 + "block_size": block_size,
201 "num_kv_heads": num_kv_heads,
186 - "head_dim": head_dim,
202 + "head_size": head_size,
203 "max_seq_len": max_seq_len,
204 }
205
@@ -195,17 +211,18 @@ def write_results_to_csv(results, filename=None):
211 filename = f"flashinfer_trtllm_benchmark_{timestamp}.csv"
212
213 fieldnames = [
198 - "num_seqs",
199 - "trt_mean",
200 - "trt_std",
214 + "batch_size",
215 + "trtllm_mean",
216 + "trtllm_std",
217 "baseline_mean",
218 "baseline_std",
219 "speedup_percent",
220 "q_dtype",
221 "kv_cache_dtype",
206 - "page_size",
222 + "output_dtype",
223 + "block_size",
224 "num_kv_heads",
208 - "head_dim",
225 + "head_size",
226 "max_seq_len",
227 ]
228
@@ -224,27 +241,41 @@ def write_results_to_csv(results, filename=None):
241
242
243 if __name__ == "__main__":
227 - num_seqs = [1, 4, 8, 16, 32, 64, 128, 256]
244 + batch_sizes = [1, 4, 8, 16, 32, 64, 128, 256]
245 max_seq_lens = [1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072]
246 all_results = []
247
231 - print(
232 - "Running benchmark for q_dtype = bfloat16, kv_cache_dtype: bfloat16, "
233 - "output_dtype: bfloat16"
234 - )
235 - print(
236 - "\tnum_seqs\tmax_seq_len\ttrt_mean\ttrt_std\tbaseline_mean\t"
237 - "baseline_std\tspeedup_percent"
238 - )
239 - for max_seq_len in max_seq_lens:
240 - for bs in num_seqs:
241 - result = benchmark_prefill(
242 - bs,
243 - max_seq_len,
244 - dtype=torch.bfloat16,
245 - kv_cache_dtype="auto",
246 - )
247 - all_results.append(result)
248 + dtype = torch.bfloat16
249 + quant_dtypes = [
250 + # (q_quant_dtype, kv_quant_dtype, o_quant_dtype)
251 + (None, None, None),
252 + (FP8_DTYPE, FP8_DTYPE, FP8_DTYPE),
253 + ]
254 +
255 + for quant_dtype in quant_dtypes:
256 + q_quant_dtype, kv_quant_dtype, o_quant_dtype = quant_dtype
257 + q_quant_dtype = q_quant_dtype or dtype
258 + kv_quant_dtype = kv_quant_dtype or dtype
259 + o_quant_dtype = o_quant_dtype or dtype
260 +
261 + print(
262 + f"Running benchmark for q_dtype = {q_quant_dtype}, "
263 + f"kv_cache_dtype: {kv_quant_dtype}, "
264 + f"output_dtype: {o_quant_dtype}"
265 + )
266 + print(
267 + "\tbatch_size\tmax_seq_len\ttrtllm_mean\ttrtllm_std\tbaseline_mean\t"
268 + "baseline_std\tspeedup_percent"
269 + )
270 + for max_seq_len in max_seq_lens:
271 + for bs in batch_sizes:
272 + result = benchmark_prefill(
273 + dtype=dtype,
274 + quant_dtypes=quant_dtype,
275 + batch_size=bs,
276 + max_seq_len=max_seq_len,
277 + )
278 + all_results.append(result)
279
280 # Write all results to CSV
281 write_results_to_csv(all_results)
tests/compile/test_fusion_attn.py
+248 -1
@@ -1,5 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import copy
4 from typing import Optional
5
6 import pytest
@@ -7,13 +8,27 @@ import torch._dynamo
8
9 from tests.compile.backend import TestBackend
10 from tests.models.utils import check_outputs_equal
11 +from tests.v1.attention.utils import (BatchSpec, _Backend,
12 + create_common_attn_metadata)
13 from vllm import LLM, SamplingParams
14 +from vllm.attention import Attention
15 +from vllm.attention.selector import global_force_attn_backend_context_manager
16 from vllm.compilation.fusion import QUANT_OPS, QuantKey, kFp8StaticTensorSym
17 from vllm.compilation.fusion_attn import ATTN_OP, AttnFusionPass
18 from vllm.compilation.fx_utils import find_op_nodes
19 from vllm.compilation.noop_elimination import NoOpEliminationPass
15 -from vllm.config import CompilationConfig, CompilationLevel, VllmConfig
20 +from vllm.config import (CacheConfig, CompilationConfig, CompilationLevel,
21 + ModelConfig, PassConfig, SchedulerConfig, VllmConfig,
22 + set_current_vllm_config)
23 +from vllm.forward_context import get_forward_context, set_forward_context
24 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
25 + GroupShape)
26 +from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
27 + Fp8LinearOp)
28 from vllm.platforms import current_platform
29 +from vllm.v1.kv_cache_interface import AttentionSpec
30 +
31 +FP8_DTYPE = current_platform.fp8_dtype()
32
33 # globals needed for string-import custom Dynamo backend field
34 backend: Optional[TestBackend] = None
@@ -132,3 +147,235 @@ def test_attention_fusion(example_prompts, monkeypatch, model: str,
147
148 # Reset backend to make sure llm2 gets released
149 backend = None
150 +
151 +
152 +class TestAttentionStaticQuantPatternModel(torch.nn.Module):
153 + """Test model for AttentionStaticQuantPattern fusion."""
154 +
155 + def __init__(self, num_qo_heads: int, num_kv_heads: int, head_size: int,
156 + kv_cache_dtype: torch.dtype, device: torch.device,
157 + vllm_config: VllmConfig):
158 + super().__init__()
159 + self.num_qo_heads = num_qo_heads
160 + self.num_kv_heads = num_kv_heads
161 + self.head_size = head_size
162 + self.kv_cache_dtype = kv_cache_dtype
163 + self.device = device
164 + self.vllm_config = vllm_config
165 +
166 + self.attn = Attention(
167 + num_heads=self.num_qo_heads,
168 + head_size=self.head_size,
169 + scale=1.0 / (self.head_size**0.5),
170 + num_kv_heads=self.num_kv_heads,
171 + cache_config=vllm_config.cache_config,
172 + prefix="model.layers.0.self_attn.attn",
173 + )
174 +
175 + self.fp8_linear = Fp8LinearOp(
176 + act_quant_static=True, act_quant_group_shape=GroupShape.PER_TENSOR)
177 + self.wscale = torch.tensor([1.0], dtype=torch.float32)
178 + self.scale = torch.tensor([1.0], dtype=torch.float32)
179 +
180 + self.block_size = 16
181 +
182 + # Initialize attn MetadataBuilder
183 + self.builder = self.attn.attn_backend.get_builder_cls()(
184 + kv_cache_spec=AttentionSpec(
185 + block_size=self.block_size,
186 + num_kv_heads=self.num_kv_heads,
187 + head_size=self.head_size,
188 + dtype=self.kv_cache_dtype,
189 + use_mla=False,
190 + ),
191 + layer_names=[self.attn.layer_name],
192 + vllm_config=self.vllm_config,
193 + device=self.device,
194 + )
195 +
196 + def build_attn_metadata(self, batch_size: int):
197 + """Initialize attention metadata."""
198 +
199 + # Create common attn metadata
200 + batch_spec = BatchSpec(seq_lens=[1] * batch_size,
201 + query_lens=[1] * batch_size)
202 + common_attn_metadata = create_common_attn_metadata(
203 + batch_spec,
204 + self.block_size,
205 + self.device,
206 + arange_block_indices=True)
207 +
208 + max_blocks = (max(batch_spec.seq_lens) + self.block_size -
209 + 1) // self.block_size
210 + num_blocks = batch_size * max_blocks
211 +
212 + # Create dummy KV cache for FlashInfer TRTLLM
213 + # - NHD: [num_blocks, 2, block_size, num_kv_heads, head_size]
214 + # - HND: [num_blocks, 2, num_kv_heads, block_size, head_size]
215 + # Create kv_cache in HND layout and permute to NHD layout
216 + # (later will be permuted back to HND layout in forward pass)
217 + kv_cache = torch.zeros(num_blocks,
218 + 2,
219 + self.num_kv_heads,
220 + self.block_size,
221 + self.head_size,
222 + dtype=self.kv_cache_dtype,
223 + device=self.device)
224 + kv_cache = kv_cache.permute(0, 1, 3, 2, 4)
225 + self.attn.kv_cache = [kv_cache]
226 +
227 + # Build attn metadata
228 + self.attn_metadata = self.builder.build(
229 + common_prefix_len=0, common_attn_metadata=common_attn_metadata)
230 +
231 + return self.attn_metadata
232 +
233 + def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
234 + w: torch.Tensor):
235 + """Forward pass that creates the pattern to be fused."""
236 + attn_output = self.attn(q, k, v)
237 + return self.fp8_linear.apply(input=attn_output,
238 + weight=w,
239 + weight_scale=self.wscale,
240 + input_scale=self.scale)
241 +
242 +
243 +@pytest.mark.parametrize("num_qo_heads, num_kv_heads", [(64, 8), (40, 8)])
244 +@pytest.mark.parametrize("head_size", [128])
245 +@pytest.mark.parametrize("batch_size", [7, 256, 533])
246 +@pytest.mark.parametrize("dtype", [torch.bfloat16])
247 +@pytest.mark.parametrize(
248 + "model_name, quant_key",
249 + [("nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", kFp8StaticTensorSym)])
250 +@pytest.mark.parametrize("backend", [_Backend.FLASHINFER])
251 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")
252 +@pytest.mark.skipif(not current_platform.supports_fp8(), reason="Need FP8")
253 +@pytest.mark.skipif(not current_platform.is_device_capability((10, 0)),
254 + reason="Only test on SM100(Blackwell)")
255 +def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,
256 + head_size: int, batch_size: int,
257 + dtype: torch.dtype, model_name: str,
258 + quant_key: QuantKey, backend: _Backend,
259 + monkeypatch, dist_init):
260 + """Test AttentionStaticQuantPattern fusion pass"""
261 +
262 + monkeypatch.setenv("VLLM_USE_V1", "1")
263 +
264 + device = torch.device("cuda:0")
265 + torch.manual_seed(42)
266 +
267 + vllm_config = VllmConfig(
268 + model_config=ModelConfig(
269 + model=model_name,
270 + max_model_len=2048,
271 + ),
272 + scheduler_config=SchedulerConfig(max_num_seqs=1024),
273 + compilation_config=CompilationConfig(
274 + level=CompilationLevel.PIECEWISE,
275 + custom_ops=["+quant_fp8"],
276 + ),
277 + cache_config=CacheConfig(cache_dtype="fp8"))
278 +
279 + # Create test inputs
280 + hidden_size = num_qo_heads * head_size
281 + q = torch.randn(batch_size, hidden_size, dtype=dtype, device=device)
282 + k = torch.randn(batch_size,
283 + num_kv_heads * head_size,
284 + dtype=dtype,
285 + device=device)
286 + v = torch.randn(batch_size,
287 + num_kv_heads * head_size,
288 + dtype=dtype,
289 + device=device)
290 + linear_w = torch.randn(hidden_size, hidden_size).to(FP8_DTYPE).t()
291 +
292 + # Mark first dimension as dynamic for realistic testing
293 + torch._dynamo.mark_dynamic(q, 0)
294 + torch._dynamo.mark_dynamic(k, 0)
295 + torch._dynamo.mark_dynamic(v, 0)
296 +
297 + # Run model directly without compilation and fusion
298 + vllm_config_unfused = copy.deepcopy(vllm_config)
299 + with set_current_vllm_config(vllm_config_unfused), set_forward_context(
300 + attn_metadata=None, vllm_config=vllm_config_unfused
301 + ), global_force_attn_backend_context_manager(backend):
302 + model_unfused = TestAttentionStaticQuantPatternModel(
303 + num_qo_heads, num_kv_heads, head_size, FP8_DTYPE, device,
304 + vllm_config_unfused)
305 + model_unfused = model_unfused.to(device)
306 +
307 + forward_ctx = get_forward_context()
308 + forward_ctx.attn_metadata = model_unfused.build_attn_metadata(
309 + batch_size)
310 +
311 + # Run model directly without compilation and fusion
312 + result_unfused = model_unfused(q, k, v, linear_w)
313 +
314 + # Run model with attn fusion enabled
315 + vllm_config.compilation_config.pass_config = PassConfig(
316 + enable_attn_fusion=True, enable_noop=True)
317 + with set_current_vllm_config(vllm_config), set_forward_context(
318 + attn_metadata=None, vllm_config=vllm_config
319 + ), global_force_attn_backend_context_manager(backend):
320 + model_fused = TestAttentionStaticQuantPatternModel(
321 + num_qo_heads, num_kv_heads, head_size, FP8_DTYPE, device,
322 + vllm_config)
323 + model_fused = model_fused.to(device)
324 +
325 + forward_ctx = get_forward_context()
326 + forward_ctx.attn_metadata = model_fused.build_attn_metadata(batch_size)
327 +
328 + # Create test backend with fusion passes enabled
329 + noop_pass = NoOpEliminationPass(vllm_config)
330 + attn_pass = lambda *args, **kw: AttnFusionPass(vllm_config)(*args, **kw
331 + )
332 + test_backend = TestBackend(noop_pass, attn_pass)
333 +
334 + # Compile model with fusion enabled
335 + model_compiled = torch.compile(model_fused,
336 + backend=test_backend,
337 + fullgraph=True)
338 + assert model_compiled.attn._o_scale_float is None
339 + result_fused_1 = model_compiled(q, k, v, linear_w)
340 +
341 + # After the 1st round of the forward pass, output quant scale should be
342 + # loaded into the attn layer's _o_scale_float, the 2nd round should
343 + # reuse the loaded _o_scale_float
344 + assert model_compiled.attn._o_scale_float is not None
345 + result_fused_2 = model_compiled(q, k, v, linear_w)
346 + assert model_compiled.attn._o_scale_float is not None
347 +
348 + # Check attn fusion support
349 + attn_fusion_supported = [
350 + layer.impl.fused_output_quant_supported(quant_key.dtype,
351 + quant_key.static,
352 + quant_key.group_shape) for key,
353 + layer in vllm_config.compilation_config.static_forward_context.items()
354 + ]
355 + if any(attn_fusion_supported):
356 + # Check quantization ops in the graph before and after fusion
357 + test_backend.check_before_ops([QUANT_OPS[quant_key]],
358 + fully_replaced=True)
359 +
360 + # Check attention ops in the graph before and after fusion
361 + attn_nodes_pre = list(find_op_nodes(ATTN_OP, test_backend.graph_pre_pass))
362 + attn_nodes_post = list(find_op_nodes(ATTN_OP,
363 + test_backend.graph_post_pass))
364 +
365 + assert len(attn_nodes_pre) > 0, "Should have attention nodes before fusion"
366 + assert len(attn_nodes_pre) == len(attn_nodes_post), \
367 + "Should have same number of attention nodes before and after fusion"
368 + assert attn_nodes_pre[0].kwargs.get("output_scale") is None, \
369 + "Attention should not have output_scale before fusion"
370 + assert attn_nodes_post[0].kwargs.get("output_scale") is not None, \
371 + "Attention should have output_scale after fusion"
372 +
373 + # Check that results are closed
374 + torch.testing.assert_close(result_unfused,
375 + result_fused_1,
376 + atol=1e-2,
377 + rtol=1e-2)
378 + torch.testing.assert_close(result_unfused,
379 + result_fused_2,
380 + atol=1e-2,
381 + rtol=1e-2)
tests/kernels/attention/test_flashinfer_trtllm_attention.py
+172 -128
@@ -13,21 +13,7 @@ if not current_platform.is_device_capability(100):
13 allow_module_level=True)
14
15 FLOAT32_BYTES = torch.finfo(torch.float).bits // 8
16 -
17 -# KV Cache Layout for TRT-LLM
18 -# kv_cache_shape = (num_blocks, 2, num_kv_heads, page_size, head_dim)
19 -
20 -MAX_Q_LEN = 1024
21 -MAX_KV_LEN = 4096
22 -BATCH_SIZES = [4, 12]
23 -NUM_HEADS = [(16, 16), (40, 8)]
24 -HEAD_SIZES = [128]
25 -BLOCK_SIZES = [16]
26 -KV_LAYOUTS = ["HND"]
27 -DTYPES = [torch.bfloat16]
28 -KV_CACHE_DTYPES = [None, current_platform.fp8_dtype()]
29 -NUM_BLOCKS = 32768 # Large enough to test overflow in index calculation.
30 -SOFT_CAPS = [None, 50.0]
16 +FP8_DTYPE = current_platform.fp8_dtype()
17
18
19 def to_float8(x, dtype=torch.float8_e4m3fn):
@@ -39,42 +25,59 @@ def to_float8(x, dtype=torch.float8_e4m3fn):
25 return x_scl_sat.to(dtype), scale.float().reciprocal()
26
27
42 -@pytest.mark.parametrize("batch_size", BATCH_SIZES)
28 +DTYPE = [torch.bfloat16]
29 +QUANT_DTYPES = [
30 + # (q_quant_dtype, kv_quant_dtype, o_quant_dtype)
31 + (None, None, None),
32 + (FP8_DTYPE, FP8_DTYPE, FP8_DTYPE),
33 +]
34 +BATCH_SIZE = [4, 12]
35 +MAX_SEQ_LENS = [(1024, 4096)]
36 +NUM_HEADS = [(64, 8), (40, 8)]
37 +HEAD_SIZE = [128]
38 +KV_LAYOUT = ["HND"] # currently only HND is supported
39 +BLOCK_SIZE = [16]
40 +SOFT_CAP = [None, 50.0]
41 +
42 +NUM_BLOCKS = 32768 # Large enough to test overflow in index calculation.
43 +
44 +
45 +@pytest.mark.parametrize("dtype", DTYPE)
46 +@pytest.mark.parametrize("quant_dtypes", QUANT_DTYPES)
47 +@pytest.mark.parametrize("batch_size", BATCH_SIZE)
48 +@pytest.mark.parametrize("max_seq_lens", MAX_SEQ_LENS)
49 @pytest.mark.parametrize("num_heads", NUM_HEADS)
44 -@pytest.mark.parametrize("head_size", HEAD_SIZES)
45 -@pytest.mark.parametrize("block_size", BLOCK_SIZES)
46 -@pytest.mark.parametrize("kv_layout", KV_LAYOUTS)
47 -@pytest.mark.parametrize("dtype", DTYPES)
48 -@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
49 -@pytest.mark.parametrize("soft_cap", SOFT_CAPS)
50 +@pytest.mark.parametrize("head_size", HEAD_SIZE)
51 +@pytest.mark.parametrize("kv_layout", KV_LAYOUT)
52 +@pytest.mark.parametrize("block_size", BLOCK_SIZE)
53 +@pytest.mark.parametrize("soft_cap", SOFT_CAP)
54 @torch.inference_mode
55 def test_flashinfer_trtllm_decode_with_baseline(
56 + dtype: torch.dtype,
57 + quant_dtypes: tuple[Optional[torch.dtype], Optional[torch.dtype],
58 + Optional[torch.dtype]],
59 batch_size: int,
60 + max_seq_lens: tuple[int, int],
61 num_heads: tuple[int, int],
62 head_size: int,
55 - block_size: int,
63 kv_layout: str,
57 - dtype: torch.dtype,
58 - kv_cache_dtype: Optional[torch.dtype],
64 + block_size: int,
65 soft_cap: Optional[float],
66 ) -> None:
61 - kv_cache_dtype = dtype if kv_cache_dtype is None else kv_cache_dtype
62 -
67 torch.set_default_device("cuda")
68 current_platform.seed_everything(0)
69
66 - kv_lens = torch.randint(1, MAX_KV_LEN, (batch_size, ), dtype=torch.int32)
67 - kv_lens[-1] = MAX_KV_LEN
68 - max_kv_len = torch.max(kv_lens).item()
69 - num_seqs = len(kv_lens)
70 + q_quant_dtype, kv_quant_dtype, o_quant_dtype = quant_dtypes
71 + q_quant_dtype = q_quant_dtype or dtype
72 + kv_quant_dtype = kv_quant_dtype or dtype
73 + o_quant_dtype = o_quant_dtype or dtype
74
71 - num_query_heads = num_heads[0]
72 - num_kv_heads = num_heads[1]
73 - assert num_query_heads % num_kv_heads == 0
75 + _, max_kv_len = max_seq_lens
76
75 - scale = head_size**-0.5
77 + num_qo_heads, num_kv_heads = num_heads
78 + assert num_qo_heads % num_kv_heads == 0
79
77 - query = torch.randn(num_seqs, num_query_heads, head_size, dtype=dtype)
80 + sm_scale = float(1.0 / (head_size**0.5))
81
82 kv_cache_shape = None
83 if kv_layout == "NHD":
@@ -83,23 +86,40 @@ def test_flashinfer_trtllm_decode_with_baseline(
86 kv_cache_shape = (NUM_BLOCKS, 2, num_kv_heads, block_size, head_size)
87 else:
88 raise ValueError(f"Invalid kv_layout: {kv_layout}")
86 - key_value_cache = torch.randn(kv_cache_shape, dtype=dtype)
87 - kv_scale = 1.0
88 - if kv_cache_dtype is current_platform.fp8_dtype():
89 - key_value_cache, kv_scale = to_float8(key_value_cache,
90 - current_platform.fp8_dtype())
89
92 - max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
90 + query = torch.randn(batch_size, num_qo_heads, head_size, dtype=dtype)
91 + if q_quant_dtype == FP8_DTYPE:
92 + query, q_scale = to_float8(query)
93 + ref_query = query.to(dtype) * q_scale
94 + else:
95 + q_scale = 1.0
96 + ref_query = query
97 +
98 + kv_lens = torch.randint(1, max_kv_len, (batch_size, ), dtype=torch.int32)
99 + kv_lens[-1] = max_kv_len
100 +
101 + seq_lens = kv_lens
102 + max_seq_len = torch.max(seq_lens).item()
103 +
104 + kv_cache = torch.randn(kv_cache_shape, dtype=dtype)
105 + if kv_quant_dtype == FP8_DTYPE:
106 + kv_cache, kv_scale = to_float8(kv_cache)
107 + ref_kv_cache = kv_cache.to(dtype) * kv_scale
108 + else:
109 + kv_scale = 1.0
110 + ref_kv_cache = kv_cache
111 + k_scale = v_scale = kv_scale
112 +
113 + max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size
114 block_tables = torch.randint(0,
115 NUM_BLOCKS,
95 - (num_seqs, max_num_blocks_per_seq),
116 + (batch_size, max_num_blocks_per_seq),
117 dtype=torch.int32)
97 - k_scale = v_scale = kv_scale
118 kv_indptr = [0]
119 kv_indices = []
120 kv_last_page_lens = []
101 - for i in range(num_seqs):
102 - seq_len = kv_lens[i]
121 + for i in range(batch_size):
122 + seq_len = seq_lens[i]
123 assert seq_len > 0
124 num_blocks = (seq_len + block_size - 1) // block_size
125 kv_indices.extend(block_tables[i, :num_blocks])
@@ -112,103 +132,93 @@ def test_flashinfer_trtllm_decode_with_baseline(
132 kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32)
133 kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
134 kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
115 -
135 workspace_buffer = torch.zeros(128 * 1024 * 1024, dtype=torch.int8)
136 +
137 + # Baseline Decode
138 wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
139 workspace_buffer,
140 kv_layout,
120 - use_tensor_cores=((num_query_heads // num_kv_heads) > 4))
141 + use_tensor_cores=((num_qo_heads // num_kv_heads) > 4))
142 wrapper.plan(kv_indptr,
143 kv_indices,
144 kv_last_page_lens,
124 - num_query_heads,
145 + num_qo_heads,
146 num_kv_heads,
147 head_size,
148 block_size,
149 "NONE",
129 - sm_scale=scale,
150 + sm_scale=sm_scale,
151 q_data_type=dtype,
131 - kv_data_type=kv_cache_dtype,
152 + kv_data_type=dtype,
153 logits_soft_cap=soft_cap)
154
134 - output = torch.empty(query.shape, dtype=dtype)
135 - wrapper.run(query,
136 - key_value_cache,
137 - k_scale=k_scale,
138 - v_scale=v_scale,
139 - out=output)
155 + output = torch.empty(ref_query.shape, dtype=dtype)
156 + wrapper.run(ref_query, ref_kv_cache, out=output)
157 + o_scale = 1.0
158 + if o_quant_dtype == FP8_DTYPE:
159 + _, o_scale = to_float8(output)
160
161 # TRTLLM Decode
142 - kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
143 - output_trtllm = torch.empty(query.shape, dtype=dtype)
162 + output_trtllm = torch.empty(query.shape, dtype=o_quant_dtype)
163 flashinfer.decode.trtllm_batch_decode_with_kv_cache(
145 - query=query.contiguous(),
146 - kv_cache=key_value_cache,
164 + query=query,
165 + kv_cache=kv_cache,
166 workspace_buffer=workspace_buffer,
167 block_tables=block_tables,
149 - seq_lens=kv_lens_tensor,
150 - max_seq_len=max_kv_len,
151 - bmm1_scale=k_scale * scale,
152 - bmm2_scale=v_scale,
168 + seq_lens=seq_lens,
169 + max_seq_len=max_seq_len,
170 + bmm1_scale=q_scale * k_scale * sm_scale,
171 + bmm2_scale=v_scale / o_scale,
172 out=output_trtllm,
173 )
174 + if o_quant_dtype == FP8_DTYPE:
175 + output_trtllm = output_trtllm.to(dtype) * o_scale
176
156 - torch.testing.assert_close(output, output_trtllm, atol=1e-2, rtol=1e-2), \
177 + if q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP8_DTYPE:
178 + rtol, atol = 5e-2, 7e-2
179 + else:
180 + rtol, atol = 1e-2, 1e-2
181 +
182 + torch.testing.assert_close(output, output_trtllm, atol=atol, rtol=rtol), \
183 f"{torch.max(torch.abs(output - output_trtllm))}"
184
185
160 -@pytest.mark.parametrize("batch_size", BATCH_SIZES)
186 +@pytest.mark.parametrize("dtype", DTYPE)
187 +@pytest.mark.parametrize("quant_dtypes", QUANT_DTYPES)
188 +@pytest.mark.parametrize("batch_size", BATCH_SIZE)
189 +@pytest.mark.parametrize("max_seq_lens", MAX_SEQ_LENS)
190 @pytest.mark.parametrize("num_heads", NUM_HEADS)
162 -@pytest.mark.parametrize("head_size", HEAD_SIZES)
163 -@pytest.mark.parametrize("block_size", BLOCK_SIZES)
164 -@pytest.mark.parametrize("kv_layout", KV_LAYOUTS)
165 -@pytest.mark.parametrize("dtype", DTYPES)
166 -@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
191 +@pytest.mark.parametrize("head_size", HEAD_SIZE)
192 +@pytest.mark.parametrize("kv_layout", KV_LAYOUT)
193 +@pytest.mark.parametrize("block_size", BLOCK_SIZE)
194 @pytest.mark.parametrize("soft_cap", [None])
195 @torch.inference_mode
196 def test_flashinfer_trtllm_prefill_with_baseline(
197 + dtype: torch.dtype,
198 + quant_dtypes: tuple[Optional[torch.dtype], Optional[torch.dtype],
199 + Optional[torch.dtype]],
200 batch_size: int,
201 + max_seq_lens: tuple[int, int],
202 num_heads: tuple[int, int],
203 head_size: int,
173 - block_size: int,
204 kv_layout: str,
175 - dtype: torch.dtype,
176 - kv_cache_dtype: Optional[torch.dtype],
205 + block_size: int,
206 soft_cap: Optional[float],
207 ) -> None:
179 - kv_cache_dtype = dtype if kv_cache_dtype is None else kv_cache_dtype
180 - if dtype != kv_cache_dtype:
181 - pytest.skip(f"Not supported dtype({dtype}) with "
182 - "kv_cache_dtype({kv_cache_dtype})")
183 -
208 torch.set_default_device("cuda")
209 current_platform.seed_everything(0)
210
187 - q_lens = torch.randint(1, MAX_Q_LEN, (batch_size, ), dtype=torch.int32)
188 - q_lens[-1] = MAX_Q_LEN
189 - max_q_len = torch.max(q_lens).item()
190 - q_indptr = torch.cat([
191 - torch.tensor([0], dtype=torch.int32),
192 - torch.cumsum(q_lens, dim=0, dtype=torch.int32),
193 - ])
194 -
195 - kv_lens = torch.randint(0, MAX_KV_LEN, (batch_size, ), dtype=torch.int32)
196 - kv_lens[-1] = MAX_KV_LEN
211 + q_quant_dtype, kv_quant_dtype, o_quant_dtype = quant_dtypes
212 + q_quant_dtype = q_quant_dtype or dtype
213 + kv_quant_dtype = kv_quant_dtype or dtype
214 + o_quant_dtype = o_quant_dtype or dtype
215
198 - seq_lens = kv_lens + q_lens
199 - max_seq_len = torch.max(seq_lens).item()
200 - num_seqs = len(seq_lens)
216 + max_q_len, max_kv_len = max_seq_lens
217
202 - num_query_heads = num_heads[0]
203 - num_kv_heads = num_heads[1]
204 - assert num_query_heads % num_kv_heads == 0
218 + num_qo_heads, num_kv_heads = num_heads
219 + assert num_qo_heads % num_kv_heads == 0
220
206 - scale = head_size**-0.5
207 -
208 - query = torch.randn(torch.sum(q_lens).item(),
209 - num_query_heads,
210 - head_size,
211 - dtype=dtype)
221 + sm_scale = float(1.0 / (head_size**0.5))
222
223 kv_cache_shape = None
224 if kv_layout == "NHD":
@@ -217,22 +227,49 @@ def test_flashinfer_trtllm_prefill_with_baseline(
227 kv_cache_shape = (NUM_BLOCKS, 2, num_kv_heads, block_size, head_size)
228 else:
229 raise ValueError(f"Invalid kv_layout: {kv_layout}")
220 - key_value_cache = torch.randn(kv_cache_shape, dtype=dtype)
221 - kv_scale = 1.0
222 - if kv_cache_dtype is current_platform.fp8_dtype():
223 - key_value_cache, kv_scale = to_float8(key_value_cache,
224 - current_platform.fp8_dtype())
230 +
231 + q_lens = torch.randint(1, max_q_len, (batch_size, ), dtype=torch.int32)
232 + q_lens[-1] = max_q_len
233 + q_indptr = torch.cat([
234 + torch.tensor([0], dtype=torch.int32),
235 + torch.cumsum(q_lens, dim=0, dtype=torch.int32),
236 + ])
237 +
238 + query = torch.randn(torch.sum(q_lens).item(),
239 + num_qo_heads,
240 + head_size,
241 + dtype=dtype)
242 + if q_quant_dtype == FP8_DTYPE:
243 + query, q_scale = to_float8(query)
244 + ref_query = query.to(dtype) * q_scale
245 + else:
246 + q_scale = 1.0
247 + ref_query = query
248 +
249 + kv_lens = torch.randint(0, max_kv_len, (batch_size, ), dtype=torch.int32)
250 + kv_lens[-1] = max_kv_len
251 +
252 + seq_lens = kv_lens + q_lens
253 + max_seq_len = torch.max(seq_lens).item()
254 +
255 + kv_cache = torch.randn(kv_cache_shape, dtype=dtype)
256 + if kv_quant_dtype == FP8_DTYPE:
257 + kv_cache, kv_scale = to_float8(kv_cache)
258 + ref_kv_cache = kv_cache.to(dtype) * kv_scale
259 + else:
260 + kv_scale = 1.0
261 + ref_kv_cache = kv_cache
262 + k_scale = v_scale = kv_scale
263
264 max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size
265 block_tables = torch.randint(0,
266 NUM_BLOCKS,
229 - (num_seqs, max_num_blocks_per_seq),
267 + (batch_size, max_num_blocks_per_seq),
268 dtype=torch.int32)
231 - k_scale = v_scale = kv_scale
269 kv_indptr = [0]
270 kv_indices = []
271 kv_last_page_lens = []
235 - for i in range(num_seqs):
272 + for i in range(batch_size):
273 seq_len = seq_lens[i]
274 assert seq_len > 0
275 num_blocks = (seq_len + block_size - 1) // block_size
@@ -246,48 +283,55 @@ def test_flashinfer_trtllm_prefill_with_baseline(
283 kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32)
284 kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
285 kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
249 -
286 workspace_buffer = torch.zeros(128 * 1024 * 1024, dtype=torch.int8)
287 +
288 + # Baseline Prefill
289 wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(
290 workspace_buffer, kv_layout)
291 wrapper.plan(q_indptr,
292 kv_indptr,
293 kv_indices,
294 kv_last_page_lens,
257 - num_query_heads,
295 + num_qo_heads,
296 num_kv_heads,
297 head_size,
298 block_size,
299 causal=True,
262 - sm_scale=scale,
300 + sm_scale=sm_scale,
301 q_data_type=dtype,
264 - kv_data_type=kv_cache_dtype,
302 + kv_data_type=dtype,
303 logits_soft_cap=soft_cap)
304
267 - output = torch.empty(query.shape, dtype=dtype)
268 - wrapper.run(query,
269 - key_value_cache,
270 - k_scale=k_scale,
271 - v_scale=v_scale,
272 - out=output)
305 + output = torch.empty(ref_query.shape, dtype=dtype)
306 + wrapper.run(ref_query, ref_kv_cache, out=output)
307 + o_scale = 1.0
308 + if o_quant_dtype == FP8_DTYPE:
309 + _, o_scale = to_float8(output)
310
274 - # TRTLLM Decode
275 - output_trtllm = torch.empty(query.shape, dtype=dtype)
311 + # TRTLLM Prefill
312 + output_trtllm = torch.empty(query.shape, dtype=o_quant_dtype)
313 flashinfer.prefill.trtllm_batch_context_with_kv_cache(
277 - query=query.contiguous(),
278 - kv_cache=key_value_cache,
314 + query=query,
315 + kv_cache=kv_cache,
316 workspace_buffer=workspace_buffer,
317 block_tables=block_tables,
318 seq_lens=seq_lens,
319 max_q_len=max_q_len,
320 max_kv_len=max_seq_len,
284 - bmm1_scale=k_scale * scale,
285 - bmm2_scale=v_scale,
286 - batch_size=num_seqs,
321 + bmm1_scale=q_scale * k_scale * sm_scale,
322 + bmm2_scale=v_scale / o_scale,
323 + batch_size=batch_size,
324 cum_seq_lens_q=q_indptr,
325 cum_seq_lens_kv=kv_indptr,
326 out=output_trtllm,
327 )
328 + if o_quant_dtype == FP8_DTYPE:
329 + output_trtllm = output_trtllm.to(dtype) * o_scale
330 +
331 + if q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP8_DTYPE:
332 + rtol, atol = 5e-2, 7e-2
333 + else:
334 + rtol, atol = 1e-2, 1e-2
335
292 - torch.testing.assert_close(output, output_trtllm, atol=1e-2, rtol=1e-2), \
336 + torch.testing.assert_close(output, output_trtllm, atol=atol, rtol=rtol), \
337 f"{torch.max(torch.abs(output - output_trtllm))}"
vllm/attention/layer.py
+9 -2
@@ -128,11 +128,17 @@ class Attention(nn.Module):
128 self._q_scale = torch.tensor(1.0, dtype=torch.float32)
129 self._prob_scale = torch.tensor(1.0, dtype=torch.float32)
130
131 - # We also keep the float32 versions of k/v_scale for attention
132 - # backends that don't support tensors (Flashinfer)
131 + # We also keep q/k/v_scale on host (cpu) memory for attention
132 + # backends that require the scales to be on host instead of on device.
133 + # e.g. Flashinfer
134 + self._q_scale_float = 1.0
135 self._k_scale_float = 1.0
136 self._v_scale_float = 1.0
137
138 + # The output scale on host memory. This should be the input scale of
139 + # the quant op after this attention layer.
140 + self._o_scale_float: Optional[float] = None
141 +
142 self.use_mla = use_mla
143 self.num_heads = num_heads
144 self.head_size = head_size
@@ -291,6 +297,7 @@ class Attention(nn.Module):
297 self._q_scale.copy_(torch.abs(query).max() / self.q_range)
298 self._k_scale.copy_(torch.abs(key).max() / self.k_range)
299 self._v_scale.copy_(torch.abs(value).max() / self.v_range)
300 + self._q_scale_float = self._q_scale.item()
301 self._k_scale_float = self._k_scale.item()
302 self._v_scale_float = self._v_scale.item()
303 # We only calculate the scales once
vllm/compilation/fusion_attn.py
+43 -34
@@ -9,7 +9,7 @@ from torch._subclasses.fake_tensor import (FakeTensorMode,
9 unset_fake_temporarily)
10
11 from vllm.attention import Attention
12 -from vllm.config import VllmConfig
12 +from vllm.config import VllmConfig, get_layers_from_vllm_config
13 from vllm.logger import init_logger
14 from vllm.platforms import current_platform
15
@@ -18,23 +18,32 @@ from .vllm_inductor_pass import VllmInductorPass
18
19 logger = init_logger(__name__)
20
21 +FP8_DTYPE = current_platform.fp8_dtype()
22 +
23 ATTN_OP = torch.ops.vllm.unified_attention_with_output.default
24 RESHAPE_OP = torch.ops.aten.reshape.default
25
26
27 class AttentionStaticQuantPattern:
28 + """
29 + Fusion for Attention+StaticQuant.
30 +
31 + Only triggers when the attention implementation returns True in
32 + `fused_output_quant_supported()`. If the pattern is found, the StaticQuant
33 + op will be removed from the graph, and its scale will be passed into
34 + Attention op as the `output_scale` argument.
35 + """
36
37 def __init__(
38 self,
29 - layer_name: str,
30 - num_heads: int,
31 - head_size: int,
39 + layer: Attention,
40 quant_dtype: torch.dtype,
41 symmetric=True,
42 ):
35 - self.layer_name = layer_name
36 - self.num_heads = num_heads
37 - self.head_size = head_size
43 + self.layer = layer
44 + self.layer_name = layer.layer_name
45 + self.num_heads = layer.num_heads
46 + self.head_size = layer.head_size
47 self.quant_dtype = quant_dtype
48 self.quant_key = QuantKey(dtype=quant_dtype,
49 static=True,
@@ -48,11 +57,10 @@ class AttentionStaticQuantPattern:
57 kwargs = {'dtype': self.quant_dtype, 'device': "cuda", **kwargs}
58 return torch.empty(*args, **kwargs)
59
51 - def register_if_supported(self, pm_pass: PatternMatcherPass,
52 - layer: Attention):
53 - if layer.impl.fused_output_quant_supported(self.quant_dtype,
54 - self.quant_key.static,
55 - self.quant_key.group_shape):
60 + def register_if_supported(self, pm_pass: PatternMatcherPass):
61 + if self.layer.impl.fused_output_quant_supported(
62 + self.quant_dtype, self.quant_key.static,
63 + self.quant_key.group_shape):
64 self._register(pm_pass)
65
66 def _register(self, pm_pass: PatternMatcherPass):
@@ -60,19 +68,15 @@ class AttentionStaticQuantPattern:
68 def pattern(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
69 output_attn: torch.Tensor, output_quant: torch.Tensor,
70 scale: torch.Tensor):
63 - view_7 = RESHAPE_OP(output_attn,
64 - [-1, self.num_heads, self.head_size])
65 -
71 at1 = auto_functionalized(ATTN_OP,
72 query=q,
73 key=k,
74 value=v,
70 - output=view_7,
75 + output=output_attn,
76 layer_name=self.layer_name,
77 output_scale=None)
78 attn_out_view = RESHAPE_OP(at1[1],
79 [-1, self.num_heads * self.head_size])
75 -
80 at2 = auto_functionalized(self.QUANT_OP,
81 result=output_quant,
82 input=attn_out_view,
@@ -82,17 +86,19 @@ class AttentionStaticQuantPattern:
86 def replacement(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
87 output_attn: torch.Tensor, output_quant: torch.Tensor,
88 scale: torch.Tensor):
85 - view_7 = RESHAPE_OP(output_quant,
86 - [-1, self.num_heads, self.head_size])
87 -
89 + # attn output in quant_dtype
90 + output_attn = torch.ops.aten.full.default(
91 + [q.shape[0], self.num_heads, self.head_size],
92 + 0.0,
93 + dtype=self.quant_dtype,
94 + device=q.device)
95 at1 = auto_functionalized(ATTN_OP,
96 query=q,
97 key=k,
98 value=v,
92 - output=view_7,
99 + output=output_attn,
100 layer_name=self.layer_name,
101 output_scale=scale)
95 -
102 return RESHAPE_OP(at1[1], [-1, self.num_heads * self.head_size])
103
104 # Need custom fake mode, otherwise tracing happens with real tensors.
@@ -102,7 +108,7 @@ class AttentionStaticQuantPattern:
108 empty_bf16(5, self.num_heads, self.head_size), # q
109 empty_bf16(5, self.num_heads, self.head_size), # k
110 empty_bf16(5, self.num_heads, self.head_size), # v
105 - empty_bf16(5, self.num_heads * self.head_size), # attn_output
111 + empty_bf16(5, self.num_heads, self.head_size), # attn_output
112 self.empty_quant(5, self.num_heads *
113 self.head_size), # quant_output
114 empty_fp32(1, 1) # scale
@@ -140,27 +146,30 @@ class AttnFusionPass(VllmInductorPass):
146
147 def __init__(self, config: VllmConfig):
148 super().__init__(config)
143 - self.static_fwd_ctx = config.compilation_config.static_forward_context
149
150 self.patterns = PatternMatcherPass(pass_name="attn_fusion_pass")
151
147 - for key, layer in self.static_fwd_ctx.items():
148 - pattern = AttentionStaticQuantPattern(key, layer.num_heads,
149 - layer.head_size,
150 - current_platform.fp8_dtype())
151 - pattern.register_if_supported(self.patterns, layer)
152 - if len(self.static_fwd_ctx) == 0:
152 + attn_layers = get_layers_from_vllm_config(config, Attention)
153 + for layer_name, layer in attn_layers.items():
154 + pattern = AttentionStaticQuantPattern(layer, FP8_DTYPE)
155 + pattern.register_if_supported(self.patterns)
156 + if len(attn_layers) == 0:
157 logger.warning(
154 - "Attention + quant fusion is enabled, but "
155 - "CompilationConfig.static_forward_context is empty. "
156 - "Cannot access attention layers so no fusion "
157 - "patterns were registered.")
158 + "Attention + quant fusion is enabled, but no attention layers "
159 + "were found in CompilationConfig.static_forward_context "
160 + "so no fusion patterns were registered.")
161
162 def __call__(self, graph: torch.fx.graph.Graph) -> None:
163 self.begin()
164 self.dump_graph(graph, "before_attn_fusion")
165
166 count = self.patterns.apply(graph)
167 +
168 + # TODO: Move this to pass_manager.py after the fx graph broken issue
169 + # has been resolved.
170 + # see https://github.com/vllm-project/vllm/issues/23091
171 + graph.eliminate_dead_code()
172 +
173 logger.debug("Fused quantization onto %s attention nodes", count)
174 self.dump_graph(graph, "after_attn_fusion")
175 self.end_and_log()
vllm/utils/flashinfer.py
+16 -6
@@ -174,21 +174,30 @@ def supports_trtllm_attention() -> tuple[bool, Optional[str]]:
174
175
176 def use_trtllm_attention(
177 + num_qo_heads: int,
178 + num_kv_heads: int,
179 num_tokens: int,
180 max_seq_len: int,
181 kv_cache_dtype: str,
180 - num_qo_heads: Optional[int],
181 - num_kv_heads: Optional[int],
182 - attn_head_size: Optional[int],
182 + q_dtype: torch.dtype,
183 + is_prefill: bool,
184 has_sinks: bool = False,
185 ) -> bool:
186 use_trtllm, env_value = supports_trtllm_attention()
187 if not use_trtllm:
188 return False
189
189 - # Check if the dimensions are supported by TRTLLM decode attention
190 - if (attn_head_size is None or num_qo_heads is None or num_kv_heads is None
191 - or num_qo_heads % num_kv_heads != 0):
190 + if num_qo_heads % num_kv_heads != 0:
191 + return False
192 +
193 + # Must use TRTLLM attention if query is FP8 quantized
194 + if q_dtype == current_platform.fp8_dtype():
195 + logger.info_once("Using TRTLLM attention (query is quantized).")
196 + return True
197 +
198 + # TRTLLM prefill attention does not support FP8 kv cache with
199 + # non-quantized query
200 + if is_prefill and kv_cache_dtype.startswith("fp8"):
201 return False
202
203 # If sinks are being used, we must use TRTLLM attention as it's
@@ -290,6 +299,7 @@ __all__ = [
299 "has_flashinfer_moe",
300 "has_flashinfer_cutlass_fused_moe",
301 "has_nvidia_artifactory",
302 + "supports_trtllm_attention",
303 "use_trtllm_attention",
304 "flashinfer_scaled_fp4_mm",
305 ]
vllm/v1/attention/backends/flashinfer.py
+80 -29
@@ -15,12 +15,17 @@ from flashinfer.decode import (_get_range_buf, get_seq_lens,
15 from flashinfer.prefill import trtllm_batch_context_with_kv_cache
16
17 import vllm.envs as envs
18 +from vllm import _custom_ops as ops
19 from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
20 AttentionType)
21 from vllm.config import CUDAGraphMode, VllmConfig
22 from vllm.logger import init_logger
23 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
24 + GroupShape)
25 +from vllm.platforms import current_platform
26 from vllm.utils import cdiv, is_pin_memory_available
23 -from vllm.utils.flashinfer import use_trtllm_attention
27 +from vllm.utils.flashinfer import (supports_trtllm_attention,
28 + use_trtllm_attention)
29 from vllm.v1.attention.backends.flash_attn import use_cascade_attention
30 # yapf conflicts with isort for this block
31 # yapf: disable
@@ -35,6 +40,8 @@ from vllm.v1.kv_cache_interface import AttentionSpec
40
41 FLASHINFER_WORKSPACE_BUFFER_SIZE = 256 * 1024 * 1024
42
43 +FP8_DTYPE = current_platform.fp8_dtype()
44 +
45 logger = init_logger(__name__)
46
47
@@ -519,22 +526,27 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
526 else:
527 kv_cache_dtype = self.kv_cache_spec.dtype
528
522 - num_qo_heads = self.vllm_config.model_config.get_num_attention_heads(
523 - self.vllm_config.parallel_config)
529 + config = self.vllm_config
530 + num_qo_heads = config.model_config.get_num_attention_heads(
531 + config.parallel_config)
532 num_kv_heads = self.kv_cache_spec.num_kv_heads
533 head_dim = self.kv_cache_spec.head_size
534
535 # Check if any layer uses sinks (requires TRTLLM attention)
536 has_sinks = self.global_hyperparameters.has_sinks
537
530 - # currently prefill trtllm attention does not support fp8 kv cache
531 - prefill_use_trtllm = not cache_dtype.startswith("fp8") \
532 - and use_trtllm_attention(
533 - num_prefill_tokens, max_seq_len, cache_dtype,
534 - num_qo_heads, num_kv_heads, head_dim, has_sinks)
538 + # Insert FP8 quant for query if FP8 kv cache and attn fusion enabled
539 + q_dtype = config.model_config.dtype
540 + enable_fusion = config.compilation_config.pass_config.enable_attn_fusion
541 + if cache_dtype.startswith("fp8") and enable_fusion:
542 + q_dtype = kv_cache_dtype
543 +
544 + prefill_use_trtllm = use_trtllm_attention(
545 + num_qo_heads, num_kv_heads, num_prefill_tokens, max_seq_len,
546 + cache_dtype, q_dtype, is_prefill=True, has_sinks=has_sinks)
547 decode_use_trtllm = use_trtllm_attention(
536 - num_decode_tokens, max_seq_len, cache_dtype,
537 - num_qo_heads, num_kv_heads, head_dim, has_sinks)
548 + num_qo_heads, num_kv_heads, num_decode_tokens, max_seq_len,
549 + cache_dtype, q_dtype, is_prefill=False, has_sinks=has_sinks)
550
551 attn_metadata = FlashInferMetadata(
552 num_actual_tokens=num_actual_tokens,
@@ -548,7 +560,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
560 head_dim=head_dim,
561 page_size=page_size,
562 kv_data_type=kv_cache_dtype,
551 - q_data_type=self.vllm_config.model_config.dtype,
563 + q_data_type=q_dtype,
564 slot_mapping=common_attn_metadata.slot_mapping,
565 max_q_len=max_q_len,
566 max_seq_len=max_seq_len,
@@ -622,6 +634,8 @@ class FlashInferImpl(AttentionImpl):
634 self.sliding_window = (-1, -1)
635 else:
636 self.sliding_window = (sliding_window - 1, 0)
637 + self.window_left = (self.sliding_window[0]
638 + if self.sliding_window is not None else -1)
639 self.kv_cache_dtype = kv_cache_dtype
640 self.logits_soft_cap = logits_soft_cap
641 self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
@@ -644,6 +658,19 @@ class FlashInferImpl(AttentionImpl):
658 )
659 self.sinks = sinks
660
661 + self.support_trtllm_attn = (supports_trtllm_attention() and
662 + num_heads % num_kv_heads == 0)
663 + self.bmm1_scale: Optional[float] = None
664 + self.bmm2_scale: Optional[float] = None
665 +
666 + def fused_output_quant_supported(self, dtype: torch.dtype, static: bool,
667 + group_shape: GroupShape):
668 + supported_quant_type = (dtype == FP8_DTYPE and static and
669 + group_shape == GroupShape.PER_TENSOR)
670 + return (self.support_trtllm_attn
671 + and self.kv_cache_dtype.startswith("fp8")
672 + and supported_quant_type)
673 +
674 def forward(
675 self,
676 layer: torch.nn.Module,
@@ -672,15 +699,42 @@ class FlashInferImpl(AttentionImpl):
699 """
700 assert output is not None, "Output tensor must be provided."
701
675 - if output_scale is not None:
676 - raise NotImplementedError(
677 - "fused output quantization is not yet supported"
678 - " for FlashInferImpl")
679 -
702 if attn_metadata is None:
703 # Profiling run.
704 return output
705
706 + if self.bmm1_scale is None:
707 + self.bmm1_scale = (layer._q_scale_float * layer._k_scale_float *
708 + self.scale)
709 +
710 + if self.bmm2_scale is None:
711 + self.bmm2_scale = layer._v_scale_float
712 +
713 + # The attn+quant fusion happens when output_scale is provided.
714 + if output_scale is None:
715 + assert attn_metadata.q_data_type != FP8_DTYPE, \
716 + "Query can only be FP8 if output fusion happened."
717 + else:
718 + assert attn_metadata.q_data_type == FP8_DTYPE, \
719 + "Query must be FP8 when attn+quant fusion happened."
720 + assert (attn_metadata.prefill_use_trtllm and
721 + attn_metadata.decode_use_trtllm), "Must use TRT-LLM attn"
722 + assert output.dtype == FP8_DTYPE, \
723 + "Output must be FP8 when attn+quant fusion happened."
724 +
725 + # TRTLLM attn kernel requires o scale as a host scalar, store the
726 + # o scale to host scalar in warmup run with cuda graph not enabled
727 + if layer._o_scale_float is None:
728 + layer._o_scale_float = output_scale.cpu().item()
729 + self.bmm2_scale = self.bmm2_scale / layer._o_scale_float
730 +
731 + # Insert FP8 quant for query
732 + num_tokens, num_heads, head_size = query.shape
733 + query, _ = ops.scaled_fp8_quant(
734 + query.reshape((num_tokens, num_heads * head_size)).contiguous(),
735 + layer._q_scale)
736 + query = query.reshape((num_tokens, num_heads, head_size))
737 +
738 # IMPORTANT!
739 # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in
740 # eager-mode PyTorch. Thus, we need to be careful about any CPU overhead
@@ -718,9 +772,6 @@ class FlashInferImpl(AttentionImpl):
772 self.kv_cache_dtype)
773 kv_cache = kv_cache.view(torch_dtype)
774
721 - window_left = (self.sliding_window[0]
722 - if self.sliding_window is not None else -1)
723 -
775 # Inputs and outputs may be padded for CUDA graphs
776 query = query[:num_actual_tokens]
777 output_padded = output
@@ -748,7 +799,7 @@ class FlashInferImpl(AttentionImpl):
799
800 if not attn_metadata.prefill_use_trtllm:
801 assert prefill_wrapper._causal
751 - assert prefill_wrapper._window_left == window_left
802 + assert prefill_wrapper._window_left == self.window_left
803 assert prefill_wrapper._logits_soft_cap == (
804 self.logits_soft_cap or 0.0)
805 assert prefill_wrapper._sm_scale == self.scale
@@ -783,12 +834,12 @@ class FlashInferImpl(AttentionImpl):
834 seq_lens=seq_lens_prefill,
835 max_q_len=attn_metadata.max_q_len,
836 max_kv_len=attn_metadata.max_seq_len,
786 - bmm1_scale=layer._k_scale_float * self.scale,
787 - bmm2_scale=layer._v_scale_float,
837 + bmm1_scale=self.bmm1_scale,
838 + bmm2_scale=self.bmm2_scale,
839 batch_size=attn_metadata.num_prefills,
840 cum_seq_lens_q=attn_metadata.qo_indptr_gpu,
841 cum_seq_lens_kv=attn_metadata.paged_kv_indptr_gpu,
791 - window_left=window_left,
842 + window_left=self.window_left,
843 sinks=self.sinks,
844 out=output[num_decode_tokens:],
845 )
@@ -800,7 +851,7 @@ class FlashInferImpl(AttentionImpl):
851 assert decode_wrapper is not None
852
853 if not attn_metadata.decode_use_trtllm:
803 - assert decode_wrapper._window_left == window_left
854 + assert decode_wrapper._window_left == self.window_left
855 assert decode_wrapper._logits_soft_cap == (self.logits_soft_cap
856 or 0.0)
857 assert decode_wrapper._sm_scale == self.scale
@@ -815,8 +866,8 @@ class FlashInferImpl(AttentionImpl):
866 # decode_query may be non-contiguous
867 decode_query = decode_query.contiguous()
868 workspace_buffer = decode_wrapper._float_workspace_buffer
818 - block_tables_decode = attn_metadata.block_table_tensor[:
819 - num_decode_tokens]
869 + block_tables_decode = attn_metadata.\
870 + block_table_tensor[:num_decode_tokens]
871 seq_lens_decode = attn_metadata.seq_lens[:num_decode_tokens]
872
873 # This path needs to be enabled with VLLM_KV_CACHE_LAYOUT = HND
@@ -834,9 +885,9 @@ class FlashInferImpl(AttentionImpl):
885 block_tables=block_tables_decode,
886 seq_lens=seq_lens_decode,
887 max_seq_len=attn_metadata.max_seq_len,
837 - bmm1_scale=layer._k_scale_float * self.scale,
838 - bmm2_scale=layer._v_scale_float,
839 - window_left=window_left,
888 + bmm1_scale=self.bmm1_scale,
889 + bmm2_scale=self.bmm2_scale,
890 + window_left=self.window_left,
891 sinks=self.sinks,
892 out=output[:num_decode_tokens],
893 )