@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 83156c7b89
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 TRT-LLM Prefill Attention Kernel (#22095)

Signed-off-by: elvischenv <219235043+elvischenv@users.noreply.github.com>

elvischenv committed Aug 5, 2025 at 17:45 UTC 83156c7b89fb880744216f3475c99f698d67a4dc
9 files changed +701 -235
.buildkite/test-pipeline.yaml
+1 -1
@@ -664,7 +664,7 @@ steps:
664 # Attention
665 # num_heads2 broken by https://github.com/flashinfer-ai/flashinfer/issues/1353
666 - pytest -v -s tests/kernels/attention/test_flashinfer.py -k 'not num_heads2'
667 - - pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_decode_attention.py
667 + - pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py
668 - pytest -v -s tests/kernels/test_cutlass_mla_decode.py
669 # Quantization
670 - pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8'
benchmarks/kernels/benchmark_trtllm_decode_attention.py renamed
-1
@@ -41,7 +41,6 @@ def benchmark_decode(
41 device = "cuda"
42 torch.manual_seed(0)
43
44 - # Currently only HEAD_GRP_SIZE == 8 is supported
44 HEAD_GRP_SIZE = 8
45 MAX_SEQ_LEN = max_seq_len
46
benchmarks/kernels/benchmark_trtllm_prefill_attention.py new
+250
@@ -0,0 +1,250 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import csv
5 +import os
6 +import random
7 +from datetime import datetime
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)
16 +
17 +
18 +def to_float8(x, dtype=torch.float8_e4m3fn):
19 + finfo = torch.finfo(dtype)
20 + min_val, max_val = x.aminmax()
21 + amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
22 + scale = finfo.max / amax * 0.1
23 + x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max)
24 + return x_scl_sat.to(dtype), scale.float().reciprocal()
25 +
26 +
27 +@torch.no_grad()
28 +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,
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
45 +
46 + # large number to reduce kv_cache reuse
47 + NUM_BLOCKS = int(256000 / page_size)
48 +
49 + workspace_buffer = torch.empty(1024 * 1024 * 1024, dtype=torch.int8)
50 +
51 + num_qo_heads = num_kv_heads * HEAD_GRP_SIZE
52 + sm_scale = float(1.0 / (head_dim**0.5))
53 +
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)
57 + q_indptr = torch.cat(
58 + [
59 + torch.tensor([0], dtype=torch.int32),
60 + torch.cumsum(
61 + torch.tensor(q_lens, dtype=torch.int32), dim=0, dtype=torch.int32
62 + ),
63 + ]
64 + )
65 + q = torch.randn(sum(q_lens), num_qo_heads, head_dim, dtype=dtype)
66 +
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 + block_tables = torch.randint(
76 + 0, NUM_BLOCKS, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
77 + )
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 +
88 + kv_indptr = [0]
89 + kv_indices = []
90 + kv_last_page_lens = []
91 + for i in range(num_seqs):
92 + seq_len = seq_lens[i]
93 + assert seq_len > 0
94 + num_blocks = (seq_len + page_size - 1) // page_size
95 + kv_indices.extend(block_tables[i, :num_blocks])
96 + kv_indptr.append(kv_indptr[-1] + num_blocks)
97 + kv_last_page_len = seq_len % page_size
98 + if kv_last_page_len == 0:
99 + kv_last_page_len = page_size
100 + kv_last_page_lens.append(kv_last_page_len)
101 +
102 + kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32)
103 + kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
104 + kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
105 +
106 + output_baseline = torch.empty(q.shape, dtype=dtype)
107 +
108 + wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(
109 + workspace_buffer, kv_layout
110 + )
111 + wrapper.plan(
112 + q_indptr,
113 + kv_indptr,
114 + kv_indices,
115 + kv_last_page_lens,
116 + num_qo_heads,
117 + num_kv_heads,
118 + head_dim,
119 + page_size,
120 + causal=True,
121 + sm_scale=sm_scale,
122 + q_data_type=dtype,
123 + kv_data_type=kv_cache.dtype,
124 + )
125 +
126 + def time_fn(fn, warmup=10, trials=20):
127 + torch.cuda.synchronize()
128 + start = torch.cuda.Event(enable_timing=True)
129 + end = torch.cuda.Event(enable_timing=True)
130 + times = []
131 + for i in range(warmup):
132 + fn()
133 + for i in range(trials):
134 + start.record()
135 + fn()
136 + end.record()
137 + torch.cuda.synchronize()
138 + times.append(start.elapsed_time(end)) # ms
139 + return sum(times) / len(times), torch.std(torch.tensor(times))
140 +
141 + def baseline_prefill():
142 + return wrapper.run(
143 + q, kv_cache, k_scale=k_scale, v_scale=v_scale, out=output_baseline
144 + )
145 +
146 + def trt_prefill():
147 + return flashinfer.prefill.trtllm_batch_context_with_kv_cache(
148 + query=q,
149 + kv_cache=kv_cache,
150 + workspace_buffer=workspace_buffer,
151 + block_tables=block_tables,
152 + seq_lens=seq_lens_tensor,
153 + max_q_len=max_q_len,
154 + max_kv_len=max_seq_len,
155 + bmm1_scale=k_scale * sm_scale,
156 + bmm2_scale=v_scale,
157 + batch_size=num_seqs,
158 + cum_seq_lens_q=q_indptr,
159 + cum_seq_lens_kv=kv_indptr,
160 + out=output_trtllm,
161 + )
162 +
163 + trt_mean, trt_std = time_fn(trt_prefill)
164 + baseline_mean, baseline_std = time_fn(baseline_prefill)
165 +
166 + # Calculate percentage speedup (positive means TRT is faster)
167 + speedup_percent = (baseline_mean - trt_mean) / baseline_mean
168 +
169 + 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}"
172 + )
173 +
174 + # Return results for CSV writing
175 + return {
176 + "num_seqs": num_seqs,
177 + "trt_mean": trt_mean,
178 + "trt_std": trt_std.item(),
179 + "baseline_mean": baseline_mean,
180 + "baseline_std": baseline_std.item(),
181 + "speedup_percent": speedup_percent,
182 + "q_dtype": str(dtype),
183 + "kv_cache_dtype": kv_cache_dtype,
184 + "page_size": page_size,
185 + "num_kv_heads": num_kv_heads,
186 + "head_dim": head_dim,
187 + "max_seq_len": max_seq_len,
188 + }
189 +
190 +
191 +def write_results_to_csv(results, filename=None):
192 + """Write benchmark results to CSV file."""
193 + if filename is None:
194 + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
195 + filename = f"flashinfer_trtllm_benchmark_{timestamp}.csv"
196 +
197 + fieldnames = [
198 + "num_seqs",
199 + "trt_mean",
200 + "trt_std",
201 + "baseline_mean",
202 + "baseline_std",
203 + "speedup_percent",
204 + "q_dtype",
205 + "kv_cache_dtype",
206 + "page_size",
207 + "num_kv_heads",
208 + "head_dim",
209 + "max_seq_len",
210 + ]
211 +
212 + file_exists = os.path.exists(filename)
213 +
214 + with open(filename, "a", newline="") as csvfile:
215 + writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
216 +
217 + if not file_exists:
218 + writer.writeheader()
219 +
220 + for result in results:
221 + writer.writerow(result)
222 +
223 + print(f"Results written to {filename}")
224 +
225 +
226 +if __name__ == "__main__":
227 + num_seqs = [1, 4, 8, 16, 32, 64, 128, 256]
228 + max_seq_lens = [1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072]
229 + all_results = []
230 +
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 +
249 + # Write all results to CSV
250 + write_results_to_csv(all_results)
tests/kernels/attention/test_flashinfer_trtllm_attention.py new
+293
@@ -0,0 +1,293 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Optional
4 +
5 +import flashinfer
6 +import pytest
7 +import torch
8 +
9 +from vllm.platforms import current_platform
10 +
11 +if not current_platform.is_device_capability(100):
12 + pytest.skip("This TRTLLM kernel requires NVIDIA Blackwell.",
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 = [(64, 8), (16, 16), (40, 8), (32, 8)]
24 +HEAD_SIZES = [128]
25 +BLOCK_SIZES = [16, 32]
26 +KV_LAYOUTS = ["HND"]
27 +DTYPES = [torch.float16, 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]
31 +
32 +
33 +def to_float8(x, dtype=torch.float8_e4m3fn):
34 + finfo = torch.finfo(dtype)
35 + min_val, max_val = x.aminmax()
36 + amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
37 + scale = finfo.max / amax * 0.1
38 + x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max)
39 + return x_scl_sat.to(dtype), scale.float().reciprocal()
40 +
41 +
42 +@pytest.mark.parametrize("batch_size", BATCH_SIZES)
43 +@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 +@torch.inference_mode
51 +def test_flashinfer_trtllm_decode_with_baseline(
52 + batch_size: int,
53 + num_heads: tuple[int, int],
54 + head_size: int,
55 + block_size: int,
56 + kv_layout: str,
57 + dtype: torch.dtype,
58 + kv_cache_dtype: Optional[torch.dtype],
59 + soft_cap: Optional[float],
60 +) -> None:
61 + kv_cache_dtype = dtype if kv_cache_dtype is None else kv_cache_dtype
62 +
63 + torch.set_default_device("cuda")
64 + current_platform.seed_everything(0)
65 +
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 +
71 + num_query_heads = num_heads[0]
72 + num_kv_heads = num_heads[1]
73 + assert num_query_heads % num_kv_heads == 0
74 +
75 + scale = head_size**-0.5
76 +
77 + query = torch.randn(num_seqs, num_query_heads, head_size, dtype=dtype)
78 +
79 + kv_cache_shape = None
80 + if kv_layout == "NHD":
81 + kv_cache_shape = (NUM_BLOCKS, 2, block_size, num_kv_heads, head_size)
82 + elif kv_layout == "HND":
83 + kv_cache_shape = (NUM_BLOCKS, 2, num_kv_heads, block_size, head_size)
84 + else:
85 + 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())
91 +
92 + max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
93 + block_tables = torch.randint(0,
94 + NUM_BLOCKS,
95 + (num_seqs, max_num_blocks_per_seq),
96 + dtype=torch.int32)
97 + k_scale = v_scale = kv_scale
98 + kv_indptr = [0]
99 + kv_indices = []
100 + kv_last_page_lens = []
101 + for i in range(num_seqs):
102 + seq_len = kv_lens[i]
103 + assert seq_len > 0
104 + num_blocks = (seq_len + block_size - 1) // block_size
105 + kv_indices.extend(block_tables[i, :num_blocks])
106 + kv_indptr.append(kv_indptr[-1] + num_blocks)
107 + kv_last_page_len = seq_len % block_size
108 + if kv_last_page_len == 0:
109 + kv_last_page_len = block_size
110 + kv_last_page_lens.append(kv_last_page_len)
111 +
112 + kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32)
113 + kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
114 + kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
115 +
116 + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8)
117 + wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
118 + workspace_buffer,
119 + kv_layout,
120 + use_tensor_cores=((num_query_heads // num_kv_heads) > 4))
121 + wrapper.plan(kv_indptr,
122 + kv_indices,
123 + kv_last_page_lens,
124 + num_query_heads,
125 + num_kv_heads,
126 + head_size,
127 + block_size,
128 + "NONE",
129 + sm_scale=scale,
130 + q_data_type=dtype,
131 + kv_data_type=kv_cache_dtype,
132 + logits_soft_cap=soft_cap)
133 +
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)
140 +
141 + # TRTLLM Decode
142 + kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
143 + output_trtllm = torch.empty(query.shape, dtype=dtype)
144 + flashinfer.decode.trtllm_batch_decode_with_kv_cache(
145 + query=query.contiguous(),
146 + kv_cache=key_value_cache,
147 + workspace_buffer=workspace_buffer,
148 + 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,
153 + out=output_trtllm,
154 + )
155 +
156 + torch.testing.assert_close(output, output_trtllm, atol=1e-2, rtol=1e-2), \
157 + f"{torch.max(torch.abs(output - output_trtllm))}"
158 +
159 +
160 +@pytest.mark.parametrize("batch_size", BATCH_SIZES)
161 +@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)
167 +@pytest.mark.parametrize("soft_cap", [None])
168 +@torch.inference_mode
169 +def test_flashinfer_trtllm_prefill_with_baseline(
170 + batch_size: int,
171 + num_heads: tuple[int, int],
172 + head_size: int,
173 + block_size: int,
174 + kv_layout: str,
175 + dtype: torch.dtype,
176 + kv_cache_dtype: Optional[torch.dtype],
177 + soft_cap: Optional[float],
178 +) -> 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 +
184 + torch.set_default_device("cuda")
185 + current_platform.seed_everything(0)
186 +
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
197 +
198 + seq_lens = kv_lens + q_lens
199 + max_seq_len = torch.max(seq_lens).item()
200 + num_seqs = len(seq_lens)
201 +
202 + num_query_heads = num_heads[0]
203 + num_kv_heads = num_heads[1]
204 + assert num_query_heads % num_kv_heads == 0
205 +
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)
212 +
213 + kv_cache_shape = None
214 + if kv_layout == "NHD":
215 + kv_cache_shape = (NUM_BLOCKS, 2, block_size, num_kv_heads, head_size)
216 + elif kv_layout == "HND":
217 + kv_cache_shape = (NUM_BLOCKS, 2, num_kv_heads, block_size, head_size)
218 + else:
219 + 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())
225 +
226 + max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size
227 + block_tables = torch.randint(0,
228 + NUM_BLOCKS,
229 + (num_seqs, max_num_blocks_per_seq),
230 + dtype=torch.int32)
231 + k_scale = v_scale = kv_scale
232 + kv_indptr = [0]
233 + kv_indices = []
234 + kv_last_page_lens = []
235 + for i in range(num_seqs):
236 + seq_len = seq_lens[i]
237 + assert seq_len > 0
238 + num_blocks = (seq_len + block_size - 1) // block_size
239 + kv_indices.extend(block_tables[i, :num_blocks])
240 + kv_indptr.append(kv_indptr[-1] + num_blocks)
241 + kv_last_page_len = seq_len % block_size
242 + if kv_last_page_len == 0:
243 + kv_last_page_len = block_size
244 + kv_last_page_lens.append(kv_last_page_len)
245 +
246 + kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32)
247 + kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
248 + kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
249 +
250 + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8)
251 + wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(
252 + workspace_buffer, kv_layout)
253 + wrapper.plan(q_indptr,
254 + kv_indptr,
255 + kv_indices,
256 + kv_last_page_lens,
257 + num_query_heads,
258 + num_kv_heads,
259 + head_size,
260 + block_size,
261 + causal=True,
262 + sm_scale=scale,
263 + q_data_type=dtype,
264 + kv_data_type=kv_cache_dtype,
265 + logits_soft_cap=soft_cap)
266 +
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)
273 +
274 + # TRTLLM Decode
275 + output_trtllm = torch.empty(query.shape, dtype=dtype)
276 + flashinfer.prefill.trtllm_batch_context_with_kv_cache(
277 + query=query.contiguous(),
278 + kv_cache=key_value_cache,
279 + workspace_buffer=workspace_buffer,
280 + block_tables=block_tables,
281 + seq_lens=seq_lens,
282 + max_q_len=max_q_len,
283 + max_kv_len=max_seq_len,
284 + bmm1_scale=k_scale * scale,
285 + bmm2_scale=v_scale,
286 + batch_size=num_seqs,
287 + cum_seq_lens_q=q_indptr,
288 + cum_seq_lens_kv=kv_indptr,
289 + out=output_trtllm,
290 + )
291 +
292 + torch.testing.assert_close(output, output_trtllm, atol=1e-2, rtol=1e-2), \
293 + f"{torch.max(torch.abs(output - output_trtllm))}"
tests/kernels/attention/test_flashinfer_trtllm_decode_attention.py deleted
-138
@@ -1,138 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import Optional
4 -
5 -import flashinfer
6 -import pytest
7 -import torch
8 -
9 -from vllm.platforms import current_platform
10 -
11 -if not current_platform.is_device_capability(100):
12 - pytest.skip("This TRTLLM kernel requires NVIDIA Blackwell.",
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 -NUM_HEADS = [(64, 8), (16, 16), (40, 8), (32, 8)]
21 -HEAD_SIZES = [128]
22 -BLOCK_SIZES = [16, 32]
23 -DTYPES = [torch.float16, torch.bfloat16]
24 -NUM_BLOCKS = 32768 # Large enough to test overflow in index calculation.
25 -SOFT_CAPS = [None, 30.0, 50.0]
26 -
27 -
28 -def to_float8(x, dtype=torch.float8_e4m3fn):
29 - finfo = torch.finfo(dtype)
30 - min_val, max_val = x.aminmax()
31 - amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
32 - scale = finfo.max / amax * 0.1
33 - x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max)
34 - return x_scl_sat.to(dtype), scale.float().reciprocal()
35 -
36 -
37 -@pytest.mark.parametrize("kv_lens", [[1328, 18, 463], [1, 54, 293, 70]])
38 -@pytest.mark.parametrize("num_heads", NUM_HEADS)
39 -@pytest.mark.parametrize("head_size", HEAD_SIZES)
40 -@pytest.mark.parametrize("block_size", BLOCK_SIZES)
41 -@pytest.mark.parametrize("kv_layout", ["HND"])
42 -@pytest.mark.parametrize("dtype", DTYPES)
43 -@pytest.mark.parametrize("soft_cap", SOFT_CAPS)
44 -@torch.inference_mode
45 -def test_flashinfer_trtllm_decode_with_baseline(
46 - kv_lens: list[int],
47 - num_heads: tuple[int, int],
48 - head_size: int,
49 - dtype: torch.dtype,
50 - block_size: int,
51 - soft_cap: Optional[float],
52 - kv_layout: str,
53 -) -> None:
54 - torch.set_default_device("cuda")
55 - current_platform.seed_everything(0)
56 - num_seqs = len(kv_lens)
57 - num_query_heads = num_heads[0]
58 - num_kv_heads = num_heads[1]
59 -
60 - assert num_query_heads % num_kv_heads == 0
61 - max_kv_len = max(kv_lens)
62 - scale = head_size**-0.5
63 -
64 - query = torch.randn(num_seqs, num_query_heads, head_size, dtype=dtype)
65 - kv_cache_shape = None
66 - if kv_layout == "NHD":
67 - kv_cache_shape = (NUM_BLOCKS, 2, block_size, num_kv_heads, head_size)
68 - elif kv_layout == "HND":
69 - kv_cache_shape = (NUM_BLOCKS, 2, num_kv_heads, block_size, head_size)
70 - else:
71 - raise ValueError(f"Invalid kv_layout: {kv_layout}")
72 - key_value_cache = torch.randn(kv_cache_shape, dtype=dtype)
73 -
74 - max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
75 - block_tables = torch.randint(0,
76 - NUM_BLOCKS,
77 - (num_seqs, max_num_blocks_per_seq),
78 - dtype=torch.int32)
79 - k_scale = v_scale = 1.0
80 - kv_indptr = [0]
81 - kv_indices = []
82 - kv_last_page_lens = []
83 - for i in range(num_seqs):
84 - seq_len = kv_lens[i]
85 - assert seq_len > 0
86 - num_blocks = (seq_len + block_size - 1) // block_size
87 - kv_indices.extend(block_tables[i, :num_blocks])
88 - kv_indptr.append(kv_indptr[-1] + num_blocks)
89 - kv_last_page_len = seq_len % block_size
90 - if kv_last_page_len == 0:
91 - kv_last_page_len = block_size
92 - kv_last_page_lens.append(kv_last_page_len)
93 -
94 - kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32)
95 - kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
96 - kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
97 -
98 - workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8)
99 - wrapper = flashinfer.\
100 - BatchDecodeWithPagedKVCacheWrapper(workspace_buffer, kv_layout,
101 - use_tensor_cores=(
102 - (num_query_heads//num_kv_heads) > 4)
103 - )
104 - wrapper.plan(kv_indptr,
105 - kv_indices,
106 - kv_last_page_lens,
107 - num_query_heads,
108 - num_kv_heads,
109 - head_size,
110 - block_size,
111 - "NONE",
112 - q_data_type=dtype,
113 - kv_data_type=dtype,
114 - logits_soft_cap=soft_cap)
115 -
116 - output = torch.empty(query.shape, dtype=dtype)
117 - wrapper.run(query, key_value_cache, scale, out=output)
118 -
119 - # TRTLLM Decode
120 - max_kv_len = max(kv_lens)
121 - kv_lens_tensor = torch.tensor(kv_lens,
122 - dtype=torch.int,
123 - device=query.device)
124 - output_trtllm = torch.empty(query.shape, dtype=dtype)
125 - flashinfer.decode.trtllm_batch_decode_with_kv_cache(
126 - query.contiguous(),
127 - key_value_cache,
128 - workspace_buffer,
129 - block_tables,
130 - kv_lens_tensor,
131 - max_kv_len,
132 - bmm1_scale=k_scale * scale,
133 - bmm2_scale=v_scale,
134 - out=output_trtllm,
135 - )
136 -
137 - torch.testing.assert_close(output, output_trtllm, atol=1e-2, rtol=1e-2), \
138 - f"{torch.max(torch.abs(output - output_trtllm))}"
vllm/attention/backends/flashinfer.py
+2 -2
@@ -46,7 +46,7 @@ from vllm.config import VllmConfig, get_layers_from_vllm_config
46 from vllm.logger import init_logger
47 from vllm.utils import (async_tensor_h2d, get_kv_cache_torch_dtype,
48 make_tensor_with_pad)
49 -from vllm.utils.flashinfer import use_trtllm_decode_attention
49 +from vllm.utils.flashinfer import use_trtllm_attention
50
51 logger = init_logger(__name__)
52
@@ -1114,7 +1114,7 @@ class FlashInferImpl(AttentionImpl):
1114 assert decode_meta.decode_wrapper._sm_scale == softmax_scale
1115 # TODO: @pavanimajety Remove this once the switch happens
1116 # inside flashinfer.
1117 - if not use_trtllm_decode_attention(
1117 + if not use_trtllm_attention(
1118 num_decode_tokens, attn_metadata.max_decode_seq_len,
1119 kv_cache_dtype, attn_metadata.num_qo_heads,
1120 attn_metadata.num_kv_heads, attn_metadata.head_dim):
vllm/envs.py
+3 -3
@@ -1027,9 +1027,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
1027 "VLLM_USE_CUDNN_PREFILL":
1028 lambda: bool(int(os.getenv("VLLM_USE_CUDNN_PREFILL", "0"))),
1029
1030 - # If set to 1, use the TRTLLM Decode Attention backend in flashinfer.
1031 - "VLLM_USE_TRTLLM_DECODE_ATTENTION":
1032 - lambda: os.getenv("VLLM_USE_TRTLLM_DECODE_ATTENTION", None),
1030 + # If set to 1, use the TRTLLM Attention backend in flashinfer.
1031 + "VLLM_USE_TRTLLM_ATTENTION":
1032 + lambda: os.getenv("VLLM_USE_TRTLLM_ATTENTION", None),
1033
1034 # Controls garbage collection during CUDA graph capture.
1035 # If set to 0 (default), enables GC freezing to speed up capture time.
vllm/utils/flashinfer.py
+7 -10
@@ -124,7 +124,7 @@ def has_flashinfer_cutlass_fused_moe() -> bool:
124 @functools.cache
125 def has_nvidia_artifactory() -> bool:
126 """Return ``True`` if NVIDIA's artifactory is accessible.
127 -
127 +
128 This checks connectivity to the kernel inference library artifactory
129 which is required for downloading certain cubin kernels like TRTLLM FHMA.
130 """
@@ -144,7 +144,7 @@ def has_nvidia_artifactory() -> bool:
144 return False
145
146
147 -def use_trtllm_decode_attention(
147 +def use_trtllm_attention(
148 num_tokens: int,
149 max_seq_len: int,
150 kv_cache_dtype: str,
@@ -159,29 +159,26 @@ def use_trtllm_decode_attention(
159
160 # Check if the dimensions are supported by TRTLLM decode attention
161 if (attn_head_size is None or num_qo_heads is None or num_kv_heads is None
162 - or num_qo_heads // num_kv_heads > 8
162 or num_qo_heads % num_kv_heads != 0 or attn_head_size != 128):
163 return False
164
166 - env_value = envs.VLLM_USE_TRTLLM_DECODE_ATTENTION
165 + env_value = envs.VLLM_USE_TRTLLM_ATTENTION
166 if env_value is not None:
168 - logger.info_once("VLLM_USE_TRTLLM_DECODE_ATTENTION is set to %s",
169 - env_value)
167 + logger.info_once("VLLM_USE_TRTLLM_ATTENTION is set to %s", env_value)
168 # Environment variable is set - respect it
169 # Making the conditional check for zero because
170 # the path is automatically enabled if the batch size condition
171 # is satisfied.
172 no_use_trtllm = (env_value == "0")
173 if not no_use_trtllm:
176 - logger.info_once("Using TRTLLM decode attention.")
174 + logger.info_once("Using TRTLLM attention.")
175 return not no_use_trtllm
176 else:
177 # Environment variable not set - use auto-detection
178 use_trtllm = (num_tokens <= 256 and max_seq_len < 131072
179 and kv_cache_dtype == "auto")
180 if use_trtllm:
183 - logger.warning_once(
184 - "Using TRTLLM decode attention (auto-detected).")
181 + logger.warning_once("Using TRTLLM attention (auto-detected).")
182 return use_trtllm
183
184
@@ -195,5 +192,5 @@ __all__ = [
192 "has_flashinfer_moe",
193 "has_flashinfer_cutlass_fused_moe",
194 "has_nvidia_artifactory",
198 - "use_trtllm_decode_attention",
195 + "use_trtllm_attention",
196 ]
vllm/v1/attention/backends/flashinfer.py
+145 -80
@@ -12,6 +12,7 @@ from flashinfer import (BatchDecodeWithPagedKVCacheWrapper,
12 MultiLevelCascadeAttentionWrapper)
13 from flashinfer.decode import (_get_range_buf, get_seq_lens,
14 trtllm_batch_decode_with_kv_cache)
15 +from flashinfer.prefill import trtllm_batch_context_with_kv_cache
16
17 import vllm.envs as envs
18 from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
@@ -19,7 +20,7 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
20 from vllm.config import VllmConfig
21 from vllm.logger import init_logger
22 from vllm.utils import cdiv, is_pin_memory_available
22 -from vllm.utils.flashinfer import use_trtllm_decode_attention
23 +from vllm.utils.flashinfer import use_trtllm_attention
24 from vllm.v1.attention.backends.flash_attn import use_cascade_attention
25 # yapf conflicts with isort for this block
26 # yapf: disable
@@ -149,9 +150,12 @@ class FlashInferMetadata:
150 slot_mapping: torch.Tensor
151
152 # For flashinfer trtllm batch decode
153 + max_q_len: int
154 max_seq_len: int
155 seq_lens: torch.Tensor
156 block_table_tensor: torch.Tensor
157 + prefill_use_trtllm: bool
158 + decode_use_trtllm: bool
159
160 # For handling prefill decode split
161 num_decodes: int
@@ -170,6 +174,9 @@ class FlashInferMetadata:
174 decode_wrapper: Optional[BatchDecodeWithPagedKVCacheWrapper] = None
175 cascade_wrapper: Optional[MultiLevelCascadeAttentionWrapper] = None
176
177 + qo_indptr_gpu: Optional[torch.Tensor] = None
178 + paged_kv_indptr_gpu: Optional[torch.Tensor] = None
179 +
180 def __post_init__(self):
181 if self.head_dim is not None:
182 FlashInferBackend.validate_head_size(self.head_dim)
@@ -305,8 +312,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
312 2, self._get_workspace_buffer(), get_kv_cache_layout())
313 return self._cascade_wrapper
314
308 - def _plan(self, num_prefills: int, num_decodes: int,
309 - attn_metadata: FlashInferMetadata):
315 + def _plan(self, attn_metadata: FlashInferMetadata):
316 if attn_metadata.use_cascade:
317 attn_metadata.cascade_wrapper = self._get_cascade_wrapper()
318 attn_metadata.cascade_wrapper.plan(
@@ -341,6 +347,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
347 # Regular attention (common case).
348 # Decodes are at the front and prefills are at the back,
349 # according to reorder_batch()
350 + num_prefills = attn_metadata.num_prefills
351 + num_decodes = attn_metadata.num_decodes
352 if num_prefills > 0:
353 # Decodes are first so prefills start after the last decode
354 prefill_start = num_decodes
@@ -356,23 +364,31 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
364 # to be relative to the start of the prefill queries.
365 qo_indptr_cpu = attn_metadata.qo_indptr_cpu[
366 prefill_start:] - attn_metadata.qo_indptr_cpu[prefill_start]
359 - attn_metadata.prefill_wrapper.plan(
360 - qo_indptr_cpu,
361 - attn_metadata.paged_kv_indptr_cpu[prefill_start:],
362 - attn_metadata.paged_kv_indices,
363 - attn_metadata.paged_kv_last_page_len_cpu[prefill_start:],
364 - attn_metadata.num_qo_heads,
365 - attn_metadata.num_kv_heads,
366 - attn_metadata.head_dim,
367 - attn_metadata.page_size,
368 - causal=True,
369 - sm_scale=self.global_hyperparameters.sm_scale,
370 - window_left=self.global_hyperparameters.window_left,
371 - logits_soft_cap=self.global_hyperparameters.
372 - logits_soft_cap,
373 - q_data_type=attn_metadata.q_data_type,
374 - kv_data_type=attn_metadata.kv_data_type,
375 - )
367 + paged_kv_indptr_cpu = attn_metadata.paged_kv_indptr_cpu[
368 + prefill_start:]
369 + if not attn_metadata.prefill_use_trtllm:
370 + attn_metadata.prefill_wrapper.plan(
371 + qo_indptr_cpu,
372 + paged_kv_indptr_cpu,
373 + attn_metadata.paged_kv_indices,
374 + attn_metadata.
375 + paged_kv_last_page_len_cpu[prefill_start:],
376 + attn_metadata.num_qo_heads,
377 + attn_metadata.num_kv_heads,
378 + attn_metadata.head_dim,
379 + attn_metadata.page_size,
380 + causal=True,
381 + sm_scale=self.global_hyperparameters.sm_scale,
382 + window_left=self.global_hyperparameters.window_left,
383 + logits_soft_cap=self.global_hyperparameters.
384 + logits_soft_cap,
385 + q_data_type=attn_metadata.q_data_type,
386 + kv_data_type=attn_metadata.kv_data_type,
387 + )
388 + else:
389 + attn_metadata.qo_indptr_gpu = qo_indptr_cpu.to(self.device)
390 + attn_metadata.paged_kv_indptr_gpu = paged_kv_indptr_cpu.to(
391 + self.device)
392
393 if num_decodes > 0:
394 pure_decode = num_prefills == 0
@@ -400,11 +416,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
416
417 attn_metadata.decode_wrapper = self._get_decode_wrapper(
418 num_input_tokens, use_cudagraph)
403 - if not use_trtllm_decode_attention(
404 - num_decodes, attn_metadata.max_seq_len,
405 - self.cache_config.cache_dtype,
406 - attn_metadata.num_qo_heads, attn_metadata.num_kv_heads,
407 - attn_metadata.head_dim):
419 + if not attn_metadata.decode_use_trtllm:
420 # Use the persistent buffer with padding length,
421 # instead of the same address but chunked version
422 # in atten_metadata when using cudagraph.
@@ -437,6 +449,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
449 split_decodes_and_prefills(common_attn_metadata)
450
451 page_size = self.kv_cache_spec.block_size
452 + max_q_len = common_attn_metadata.max_query_len
453 max_seq_len = common_attn_metadata.seq_lens_cpu.max()
454 seq_lens = common_attn_metadata.seq_lens
455 seq_lens_cpu = common_attn_metadata.seq_lens_cpu
@@ -503,6 +516,24 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
516 cache_dtype)
517 else:
518 kv_cache_dtype = self.kv_cache_spec.dtype
519 +
520 + num_qo_heads = self.vllm_config.model_config.get_num_attention_heads(
521 + self.vllm_config.parallel_config)
522 + num_kv_heads = self.kv_cache_spec.num_kv_heads
523 + head_dim = self.kv_cache_spec.head_size
524 +
525 + # currently prefill trtllm attention does not support fp8 kv cache
526 + # trtllm may not support sliding window
527 + prefill_use_trtllm = (self.global_hyperparameters.window_left == -1
528 + and not cache_dtype.startswith("fp8")
529 + and use_trtllm_attention(
530 + num_prefill_tokens, max_seq_len, cache_dtype,
531 + num_qo_heads, num_kv_heads, head_dim))
532 + decode_use_trtllm = (self.global_hyperparameters.window_left == -1
533 + and use_trtllm_attention(
534 + num_decode_tokens, max_seq_len, cache_dtype,
535 + num_qo_heads, num_kv_heads, head_dim))
536 +
537 attn_metadata = FlashInferMetadata(
538 num_actual_tokens=num_actual_tokens,
539 qo_indptr_cpu=common_attn_metadata.query_start_loc_cpu,
@@ -510,14 +541,19 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
541 paged_kv_indices=paged_kv_indices,
542 paged_kv_last_page_len_cpu=self.
543 paged_kv_last_page_len_cpu[:num_reqs],
513 - num_qo_heads=self.vllm_config.model_config.get_num_attention_heads(
514 - self.vllm_config.parallel_config),
515 - num_kv_heads=self.kv_cache_spec.num_kv_heads,
516 - head_dim=self.kv_cache_spec.head_size,
544 + num_qo_heads=num_qo_heads,
545 + num_kv_heads=num_kv_heads,
546 + head_dim=head_dim,
547 page_size=page_size,
548 kv_data_type=kv_cache_dtype,
549 q_data_type=self.vllm_config.model_config.dtype,
550 slot_mapping=common_attn_metadata.slot_mapping,
551 + max_q_len=max_q_len,
552 + max_seq_len=max_seq_len,
553 + seq_lens=seq_lens,
554 + block_table_tensor=block_table_tensor,
555 + prefill_use_trtllm=prefill_use_trtllm,
556 + decode_use_trtllm=decode_use_trtllm,
557 num_decodes=num_decodes,
558 num_decode_tokens=num_decode_tokens,
559 num_prefills=num_prefills,
@@ -527,12 +563,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
563 shared_kv_page_indptr_cpu=shared_kv_page_indptr_cpu,
564 shared_kv_page_indices_cpu=shared_kv_page_indices_cpu,
565 shared_kv_last_page_len_cpu=shared_kv_last_page_len_cpu,
530 - max_seq_len=max_seq_len,
531 - seq_lens=seq_lens,
532 - block_table_tensor=block_table_tensor,
566 )
567
535 - self._plan(num_prefills, num_decodes, attn_metadata)
568 + self._plan(attn_metadata)
569
570 return attn_metadata
571
@@ -698,30 +731,64 @@ class FlashInferImpl(AttentionImpl):
731 # Regular attention (common case).
732 # Decodes are at the front and prefills are at the back,
733 # according to reorder_batch()
701 - if prefill_wrapper := attn_metadata.prefill_wrapper:
734 + if num_prefill_tokens > 0:
735 + prefill_wrapper = attn_metadata.prefill_wrapper
736 prefill_query = query[num_decode_tokens:]
737 assert prefill_query.shape[0] == num_prefill_tokens
738 assert prefill_wrapper is not None
705 - assert prefill_wrapper._causal
706 - assert prefill_wrapper._window_left == window_left
707 - assert prefill_wrapper._logits_soft_cap == (self.logits_soft_cap
708 - or 0.0)
709 - assert prefill_wrapper._sm_scale == self.scale
710 - prefill_wrapper.run(
711 - prefill_query,
712 - kv_cache_permute,
713 - k_scale=layer._k_scale_float,
714 - v_scale=layer._v_scale_float,
715 - out=output[num_decode_tokens:],
716 - )
717 - if decode_wrapper := attn_metadata.decode_wrapper:
739 +
740 + if not attn_metadata.prefill_use_trtllm:
741 + assert prefill_wrapper._causal
742 + assert prefill_wrapper._window_left == window_left
743 + assert prefill_wrapper._logits_soft_cap == (
744 + self.logits_soft_cap or 0.0)
745 + assert prefill_wrapper._sm_scale == self.scale
746 + prefill_wrapper.run(
747 + prefill_query,
748 + kv_cache_permute,
749 + k_scale=layer._k_scale_float,
750 + v_scale=layer._v_scale_float,
751 + out=output[num_decode_tokens:],
752 + )
753 + else:
754 + # prefill_query may be non-contiguous
755 + prefill_query = prefill_query.contiguous()
756 + workspace_buffer = prefill_wrapper._float_workspace_buffer
757 + block_tables_prefill = attn_metadata.block_table_tensor[
758 + num_decode_tokens:]
759 + seq_lens_prefill = attn_metadata.seq_lens[num_decode_tokens:]
760 +
761 + # This path needs to be enabled with VLLM_KV_CACHE_LAYOUT = HND
762 + assert get_kv_cache_layout() == "HND"
763 + assert prefill_query.is_contiguous()
764 + assert kv_cache_permute.is_contiguous()
765 + assert workspace_buffer.is_contiguous()
766 + assert block_tables_prefill.is_contiguous()
767 + assert seq_lens_prefill.is_contiguous()
768 +
769 + trtllm_batch_context_with_kv_cache(
770 + query=prefill_query,
771 + kv_cache=kv_cache_permute,
772 + workspace_buffer=workspace_buffer,
773 + block_tables=block_tables_prefill,
774 + seq_lens=seq_lens_prefill,
775 + max_q_len=attn_metadata.max_q_len,
776 + max_kv_len=attn_metadata.max_seq_len,
777 + bmm1_scale=layer._k_scale_float * self.scale,
778 + bmm2_scale=layer._v_scale_float,
779 + batch_size=attn_metadata.num_prefills,
780 + cum_seq_lens_q=attn_metadata.qo_indptr_gpu,
781 + cum_seq_lens_kv=attn_metadata.paged_kv_indptr_gpu,
782 + out=output[num_decode_tokens:],
783 + )
784 +
785 + if num_decode_tokens > 0:
786 + decode_wrapper = attn_metadata.decode_wrapper
787 decode_query = query[:num_decode_tokens]
788 assert decode_query.shape[0] == num_decode_tokens
789 assert decode_wrapper is not None
721 - if not use_trtllm_decode_attention(
722 - attn_metadata.num_decodes, attn_metadata.max_seq_len,
723 - self.kv_cache_dtype, attn_metadata.num_qo_heads,
724 - attn_metadata.num_kv_heads, attn_metadata.head_dim):
790 +
791 + if not attn_metadata.decode_use_trtllm:
792 assert decode_wrapper._window_left == window_left
793 assert decode_wrapper._logits_soft_cap == (self.logits_soft_cap
794 or 0.0)
@@ -734,34 +801,32 @@ class FlashInferImpl(AttentionImpl):
801 out=output[:num_decode_tokens],
802 )
803 else:
804 + # decode_query may be non-contiguous
805 + decode_query = decode_query.contiguous()
806 + workspace_buffer = decode_wrapper._float_workspace_buffer
807 + block_tables_decode = attn_metadata.block_table_tensor[:
808 + num_decode_tokens]
809 + seq_lens_decode = attn_metadata.seq_lens[:num_decode_tokens]
810 +
811 # This path needs to be enabled with VLLM_KV_CACHE_LAYOUT = HND
738 - if num_decode_tokens > 0:
739 - # decode_query may be non-contiguous
740 - decode_query = decode_query.contiguous()
741 - block_tables_decode = attn_metadata.block_table_tensor[:
742 - num_decode_tokens]
743 - seq_lens_decode = attn_metadata.seq_lens[:
744 - num_decode_tokens]
745 - workspace_buffer = decode_wrapper._float_workspace_buffer
746 -
747 - assert get_kv_cache_layout() == "HND"
748 - assert decode_query.is_contiguous()
749 - assert kv_cache_permute.is_contiguous()
750 - assert block_tables_decode.is_contiguous()
751 - assert seq_lens_decode.is_contiguous()
752 - assert workspace_buffer.is_contiguous()
753 -
754 - trtllm_batch_decode_with_kv_cache(
755 - query=decode_query,
756 - kv_cache=kv_cache_permute,
757 - workspace_buffer=workspace_buffer,
758 - block_tables=block_tables_decode,
759 - seq_lens=seq_lens_decode,
760 - max_seq_len=attn_metadata.max_seq_len,
761 - bmm1_scale=layer._k_scale_float * self.scale,
762 - bmm2_scale=layer._v_scale_float,
763 - out=output[:num_decode_tokens],
764 - )
812 + assert get_kv_cache_layout() == "HND"
813 + assert decode_query.is_contiguous()
814 + assert kv_cache_permute.is_contiguous()
815 + assert workspace_buffer.is_contiguous()
816 + assert block_tables_decode.is_contiguous()
817 + assert seq_lens_decode.is_contiguous()
818 +
819 + trtllm_batch_decode_with_kv_cache(
820 + query=decode_query,
821 + kv_cache=kv_cache_permute,
822 + workspace_buffer=workspace_buffer,
823 + block_tables=block_tables_decode,
824 + seq_lens=seq_lens_decode,
825 + max_seq_len=attn_metadata.max_seq_len,
826 + bmm1_scale=layer._k_scale_float * self.scale,
827 + bmm2_scale=layer._v_scale_float,
828 + out=output[:num_decode_tokens],
829 + )
830 return output_padded
831
832
@@ -786,8 +851,8 @@ def fast_plan_decode(
851 non_blocking: bool = True,
852 ) -> None:
853 """
789 - A faster version of BatchDecodeWithPagedKVCacheWrapper::plan used for
790 - cudagraph capture/replay, while the no cudagraph version turns back
854 + A faster version of BatchDecodeWithPagedKVCacheWrapper::plan used for
855 + cudagraph capture/replay, while the no cudagraph version turns back
856 to the original plan.
857 using original plan after passing host-side buffers:
858 - only host-to-device copy of indptr and last_page_len buffers