@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7bd4c37ae7
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 15, 02:15 UTC

[Core] Add Flashinfer TRTLLM Backend for Flashinfer decode path (SM100). (#19825)

Signed-off-by: Pavani Majety <pmajety@nvidia.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: shuw <shuw@nvidia.com> Co-authored-by: mgoin <mgoin64@gmail.com>

Pavani Majety committed Jul 11, 2025 at 02:23 UTC 7bd4c37ae7c6f2223c1a031bbdd2e3435d53da94
8 files changed +667 -56
benchmarks/kernels/benchmark_trtllm_attention.py new
+240
@@ -0,0 +1,240 @@
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_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,
39 +):
40 + torch.set_default_device("cuda")
41 + device = "cuda"
42 + torch.manual_seed(0)
43 +
44 + # Currently only HEAD_GRP_SIZE == 8 is supported
45 + HEAD_GRP_SIZE = 8
46 + MAX_SEQ_LEN = max_seq_len
47 +
48 + # large number to reduce kv_cache reuse
49 + NUM_BLOCKS = int(256000 / page_size)
50 +
51 + workspace_buffer = torch.empty(1024 * 1024 * 1024, dtype=torch.int8, device=device)
52 +
53 + # For decode, batch_size is num_decode_token
54 + num_qo_heads = num_kv_heads * HEAD_GRP_SIZE
55 + sm_scale = float(1.0 / (head_dim**0.5))
56 + q = torch.randn(num_seqs, num_qo_heads, head_dim, device=device, dtype=dtype)
57 + kv_lens = [random.randint(1, MAX_SEQ_LEN) for _ in range(num_seqs)]
58 +
59 + max_kv_len = max(kv_lens)
60 + kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int, device=device)
61 + max_num_blocks_per_seq = (max_kv_len + page_size - 1) // page_size
62 +
63 + block_tables = torch.randint(
64 + 0, NUM_BLOCKS, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
65 + )
66 +
67 + kv_cache_shape = (NUM_BLOCKS, 2, num_kv_heads, page_size, head_dim)
68 + kv_cache = torch.randn(size=kv_cache_shape, device=device, dtype=dtype)
69 + k_scale = v_scale = 1.0
70 +
71 + if kv_cache_dtype.startswith("fp8"):
72 + kv_cache, _ = to_float8(kv_cache)
73 +
74 + # Benchmark TRT decode
75 + def trt_decode():
76 + return flashinfer.decode.trtllm_batch_decode_with_kv_cache(
77 + q,
78 + kv_cache,
79 + workspace_buffer,
80 + num_qo_heads,
81 + num_kv_heads,
82 + sm_scale,
83 + block_tables,
84 + kv_lens_tensor,
85 + page_size,
86 + max_kv_len,
87 + kv_cache_dtype,
88 + k_scale,
89 + v_scale,
90 + )
91 +
92 + def time_fn(fn, warmup=10, trials=20):
93 + torch.cuda.synchronize()
94 + start = torch.cuda.Event(enable_timing=True)
95 + end = torch.cuda.Event(enable_timing=True)
96 + times = []
97 + for i in range(warmup):
98 + fn()
99 + for i in range(trials):
100 + start.record()
101 + fn()
102 + end.record()
103 + torch.cuda.synchronize()
104 + times.append(start.elapsed_time(end)) # ms
105 + return sum(times) / len(times), torch.std(torch.tensor(times))
106 +
107 + # TRT Decode
108 + trt_mean, trt_std = time_fn(trt_decode)
109 +
110 + kv_indptr = [0]
111 + kv_indices = []
112 + kv_last_page_lens = []
113 + for i in range(num_seqs):
114 + seq_len = kv_lens[i]
115 + assert seq_len > 0
116 + num_blocks = (seq_len + page_size - 1) // page_size
117 + kv_indices.extend(block_tables[i, :num_blocks])
118 + kv_indptr.append(kv_indptr[-1] + num_blocks)
119 + kv_last_page_len = seq_len % page_size
120 + if kv_last_page_len == 0:
121 + kv_last_page_len = page_size
122 + kv_last_page_lens.append(kv_last_page_len)
123 +
124 + kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32)
125 + kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
126 + kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
127 +
128 + wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
129 + workspace_buffer,
130 + kv_layout,
131 + use_tensor_cores=((num_qo_heads // num_kv_heads) > 4),
132 + )
133 +
134 + wrapper.plan(
135 + kv_indptr,
136 + kv_indices,
137 + kv_last_page_lens,
138 + num_qo_heads,
139 + num_kv_heads,
140 + head_dim,
141 + page_size,
142 + "NONE",
143 + q_data_type=dtype,
144 + kv_data_type=torch.float8_e4m3fn if kv_cache_dtype.startswith("fp8") else dtype,
145 + )
146 +
147 + def baseline_decode():
148 + return wrapper.run(q, kv_cache, sm_scale, k_scale, v_scale)
149 +
150 + baseline_mean, baseline_std = time_fn(baseline_decode)
151 +
152 + # Calculate percentage speedup (positive means TRT is faster)
153 + speedup_percent = (baseline_mean - trt_mean) / baseline_mean
154 +
155 + print(
156 + f"\t{num_seqs}\t{max_seq_len}\t{trt_mean:.3f}\t{trt_std.item():.3f}"
157 + f"\t{baseline_mean:.3f}\t{baseline_std.item():.3f}\t{speedup_percent:.3f}"
158 + )
159 +
160 + # Return results for CSV writing
161 + return {
162 + "num_seqs": num_seqs,
163 + "trt_mean": trt_mean,
164 + "trt_std": trt_std.item(),
165 + "baseline_mean": baseline_mean,
166 + "baseline_std": baseline_std.item(),
167 + "speedup_percent": speedup_percent,
168 + "q_dtype": str(dtype),
169 + "kv_cache_dtype": kv_cache_dtype,
170 + "page_size": page_size,
171 + "num_kv_heads": num_kv_heads,
172 + "head_dim": head_dim,
173 + "max_seq_len": max_seq_len,
174 + }
175 +
176 +
177 +def write_results_to_csv(results, filename=None):
178 + """Write benchmark results to CSV file."""
179 + if filename is None:
180 + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
181 + filename = f"flashinfer_trtllm_benchmark_{timestamp}.csv"
182 +
183 + fieldnames = [
184 + "num_seqs",
185 + "trt_mean",
186 + "trt_std",
187 + "baseline_mean",
188 + "baseline_std",
189 + "speedup_percent",
190 + "q_dtype",
191 + "kv_cache_dtype",
192 + "page_size",
193 + "num_kv_heads",
194 + "head_dim",
195 + "max_seq_len",
196 + ]
197 +
198 + file_exists = os.path.exists(filename)
199 +
200 + with open(filename, "a", newline="") as csvfile:
201 + writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
202 +
203 + if not file_exists:
204 + writer.writeheader()
205 +
206 + for result in results:
207 + writer.writerow(result)
208 +
209 + print(f"Results written to {filename}")
210 +
211 +
212 +if __name__ == "__main__":
213 + num_seqs = [1, 4, 8, 16, 32, 64, 128, 256]
214 + max_seq_lens = [1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072]
215 + all_results = []
216 +
217 + print("Running benchmark for kv_cache_dtype: bfloat16")
218 + print(
219 + "\tnum_seqs\tmax_seq_len\ttrt_mean\ttrt_std\tbaseline_mean\tbaseline_std\tspeedup_percent"
220 + )
221 + for max_seq_len in max_seq_lens:
222 + for bs in num_seqs:
223 + result = benchmark_decode(
224 + bs, max_seq_len, dtype=torch.bfloat16, kv_cache_dtype="auto"
225 + )
226 + all_results.append(result)
227 +
228 + print("Running benchmark for q_dtype = bfloat16, kv_cache_dtype: fp8")
229 + print(
230 + "\tnum_seqs\tmax_seq_len\ttrt_mean\ttrt_std\tbaseline_mean\tbaseline_std\tspeedup_percent"
231 + )
232 + for max_seq_len in max_seq_lens:
233 + for bs in num_seqs:
234 + result = benchmark_decode(
235 + bs, max_seq_len, dtype=torch.bfloat16, kv_cache_dtype="fp8"
236 + )
237 + all_results.append(result)
238 +
239 + # Write all results to CSV
240 + write_results_to_csv(all_results)
tests/kernels/attention/test_flashinfer_trtllm_decode_attention.py new
+140
@@ -0,0 +1,140 @@
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 = wrapper.run(query, key_value_cache, scale)
117 +
118 + # TRTLLM Decode
119 + max_kv_len = max(kv_lens)
120 + kv_lens_tensor = torch.tensor(kv_lens,
121 + dtype=torch.int,
122 + device=query.device)
123 + output_trtllm = flashinfer.decode.trtllm_batch_decode_with_kv_cache(
124 + query.contiguous(),
125 + key_value_cache,
126 + workspace_buffer,
127 + num_query_heads,
128 + num_kv_heads,
129 + scale,
130 + block_tables,
131 + kv_lens_tensor,
132 + block_size,
133 + max_kv_len,
134 + "auto",
135 + k_scale,
136 + v_scale,
137 + )
138 +
139 + torch.testing.assert_close(output, output_trtllm, atol=1e-2, rtol=1e-2), \
140 + f"{torch.max(torch.abs(output - output_trtllm))}"
vllm/attention/backends/flashinfer.py
+107 -16
@@ -11,7 +11,8 @@ from vllm.multimodal import MultiModalPlaceholderMap
11
12 try:
13 from flashinfer import BatchDecodeWithPagedKVCacheWrapper
14 - from flashinfer.decode import CUDAGraphBatchDecodeWithPagedKVCacheWrapper
14 + from flashinfer.decode import (CUDAGraphBatchDecodeWithPagedKVCacheWrapper,
15 + trtllm_batch_decode_with_kv_cache)
16 from flashinfer.prefill import BatchPrefillWithPagedKVCacheWrapper
17
18 from vllm.vllm_flash_attn import flash_attn_varlen_func
@@ -22,7 +23,10 @@ except ImportError:
23 BatchDecodeWithPagedKVCacheWrapper = None
24 CUDAGraphBatchDecodeWithPagedKVCacheWrapper = None
25 BatchPrefillWithPagedKVCacheWrapper = None
26 + trtllm_batch_decode_with_kv_cache = None
27 FLASHINFER_WORKSPACE_BUFFER_SIZE = 0
28 + raise ImportError("FlashInfer is not installed. Please install it from "
29 + "https://github.com/flashinfer-ai/flashinfer") from None
30
31 import torch
32
@@ -40,6 +44,7 @@ from vllm.attention.layer import Attention
44 from vllm.attention.ops.paged_attn import PagedAttention
45 from vllm.config import VllmConfig, get_layers_from_vllm_config
46 from vllm.logger import init_logger
47 +from vllm.platforms import current_platform
48 from vllm.utils import (async_tensor_h2d, get_kv_cache_torch_dtype,
49 make_tensor_with_pad)
50
@@ -49,10 +54,9 @@ if TYPE_CHECKING:
54 from vllm.worker.model_runner import (ModelInputForGPUBuilder,
55 ModelInputForGPUWithSamplingMetadata)
56
52 -FLASHINFER_KV_CACHE_LAYOUT: str = envs.VLLM_KV_CACHE_LAYOUT or "NHD"
53 -
57
58 class FlashInferBackend(AttentionBackend):
59 + cached_sm100a_supported: Optional[bool] = None
60
61 @staticmethod
62 def get_name() -> str:
@@ -85,7 +89,7 @@ class FlashInferBackend(AttentionBackend):
89
90 @staticmethod
91 def get_kv_cache_stride_order() -> Tuple[int, ...]:
88 - cache_layout = FLASHINFER_KV_CACHE_LAYOUT
92 + cache_layout = FlashInferState.get_kv_cache_layout()
93 assert (cache_layout in ("NHD", "HND"))
94 stride_order = (0, 1, 2, 3, 4) if cache_layout == "NHD" else (0, 1, 3,
95 2, 4)
@@ -119,6 +123,47 @@ class FlashInferBackend(AttentionBackend):
123 else:
124 raise ValueError(f"Unrecognized FP8 dtype: {kv_cache_dtype}")
125
126 + @staticmethod
127 + def use_trtllm_decode_attention(
128 + batch_size: int,
129 + max_seq_len: int,
130 + kv_cache_dtype: str,
131 + num_qo_heads: Optional[int],
132 + num_kv_heads: Optional[int],
133 + attn_head_size: Optional[int],
134 + ) -> bool:
135 + if FlashInferBackend.cached_sm100a_supported is None:
136 + FlashInferBackend.cached_sm100a_supported = (
137 + current_platform.has_device_capability(100))
138 + if not FlashInferBackend.cached_sm100a_supported:
139 + return False
140 + # Check if the dimensions are supported by TRTLLM decode attention
141 + if (attn_head_size is None or num_qo_heads is None
142 + or num_kv_heads is None or num_qo_heads // num_kv_heads > 8
143 + or num_qo_heads % num_kv_heads != 0 or attn_head_size != 128):
144 + return False
145 + env_value = envs.VLLM_USE_TRTLLM_DECODE_ATTENTION
146 + if env_value is not None:
147 + logger.info_once("VLLM_USE_TRTLLM_DECODE_ATTENTION is set to %s",
148 + env_value)
149 + # Environment variable is set - respect it
150 + # Making the conditional check for zero because
151 + # the path is automatically enabled if the batch size condition
152 + # is satisfied.
153 + no_use_trtllm = (env_value == "0")
154 + if not no_use_trtllm:
155 + logger.info_once("Using TRTLLM decode attention.")
156 + return not no_use_trtllm
157 + else:
158 + # Environment variable not set - use auto-detection
159 + use_trtllm = (FlashInferBackend.cached_sm100a_supported
160 + and batch_size <= 256 and max_seq_len < 131072
161 + and kv_cache_dtype == "auto")
162 + if use_trtllm:
163 + logger.warning_once(
164 + "Using TRTLLM decode attention (auto-detected).")
165 + return use_trtllm
166 +
167
168 @dataclass
169 class PerLayerParameters:
@@ -207,10 +252,19 @@ class FlashInferState(AttentionState):
252 device=self.runner.device)
253 return self._workspace_buffer
254
210 - def get_kv_cache_layout(self):
211 - if self._kv_cache_layout is None:
212 - self._kv_cache_layout = FLASHINFER_KV_CACHE_LAYOUT
213 - return self._kv_cache_layout
255 + @staticmethod
256 + def get_kv_cache_layout():
257 + from vllm.v1.attention.backends.utils import _KV_CACHE_LAYOUT_OVERRIDE
258 + if _KV_CACHE_LAYOUT_OVERRIDE is not None:
259 + logger.info_once("Using KV cache layout %s",
260 + _KV_CACHE_LAYOUT_OVERRIDE)
261 + return _KV_CACHE_LAYOUT_OVERRIDE
262 + cache_layout = envs.VLLM_KV_CACHE_LAYOUT
263 + if cache_layout is None:
264 + logger.info_once("Using default KV cache layout NHD")
265 + return "NHD"
266 + logger.info_once("Using KV cache layout %s", cache_layout)
267 + return cache_layout
268
269 def _get_prefill_wrapper(self):
270 if self._prefill_wrapper is None:
@@ -323,6 +377,8 @@ class FlashInferState(AttentionState):
377 num_prefill_tokens=0,
378 num_decode_tokens=batch_size,
379 max_prefill_seq_len=0,
380 + max_decode_seq_len=0,
381 + seq_lens_tensor=self._graph_seq_lens,
382 block_tables=self._graph_block_tables,
383 paged_kv_indptr=paged_kv_indptr_tensor_host,
384 paged_kv_indices=paged_kv_indices_tensor_host,
@@ -348,6 +404,8 @@ class FlashInferState(AttentionState):
404 attn_metadata,
405 is_encoder_decoder_model: bool = False):
406 return {
407 + "block_tables": attn_metadata.block_tables,
408 + "seq_lens_tensor": attn_metadata.seq_lens_tensor,
409 "slot_mapping": attn_metadata.slot_mapping,
410 }
411
@@ -355,7 +413,13 @@ class FlashInferState(AttentionState):
413 input_buffers,
414 attn_metadata,
415 is_encoder_decoder_model: bool = False):
358 - return
416 + # FlashInfer-specific logic: copy additional tensors
417 + num_total_blocks = attn_metadata.decode_metadata.seq_lens_tensor.shape[
418 + 0]
419 + input_buffers["seq_lens_tensor"][:num_total_blocks].copy_(
420 + attn_metadata.seq_lens_tensor, non_blocking=True)
421 + input_buffers["block_tables"][:num_total_blocks].copy_(
422 + attn_metadata.block_tables, non_blocking=True)
423
424 def begin_forward(self, model_input):
425 assert not self._is_graph_capturing
@@ -388,6 +452,8 @@ class FlashInferMetadata(AttentionMetadata):
452 # Maximum sequence length among prefill batch. 0 if there are decoding
453 # requests only.
454 max_prefill_seq_len: int
455 + max_decode_seq_len: int
456 +
457 # Number of query tokens for each request in the batch.
458 # Currently, we require that all requests have the same number of query
459 # tokens during the decoding phase. When speculavie decoding is enabled,
@@ -790,6 +856,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
856 use_captured_graph = cuda_graph_pad_size != -1
857
858 max_prefill_seq_len = max(self.prefill_seq_lens, default=0)
859 + max_decode_seq_len = max(self.curr_seq_lens, default=0)
860 num_decode_tokens = self.num_decode_tokens
861 decode_query_len = max(query_lens[self.num_prefills:], default=1)
862
@@ -895,6 +962,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
962 num_prefill_tokens=self.num_prefill_tokens,
963 num_decode_tokens=num_decode_tokens,
964 max_prefill_seq_len=max_prefill_seq_len,
965 + max_decode_seq_len=max_decode_seq_len,
966 block_tables=block_tables,
967 paged_kv_indptr=paged_kv_indptr_tensor,
968 paged_kv_indices=paged_kv_indices_tensor,
@@ -1081,13 +1149,36 @@ class FlashInferImpl(AttentionImpl):
1149 assert decode_meta.decode_wrapper._logits_soft_cap == (
1150 logits_soft_cap or 0.0)
1151 assert decode_meta.decode_wrapper._sm_scale == softmax_scale
1084 -
1085 - decode_output = decode_meta.decode_wrapper.run(
1086 - decode_query,
1087 - kv_cache.permute(*stride_order),
1088 - k_scale=layer._k_scale_float,
1089 - v_scale=layer._v_scale_float,
1090 - )
1152 + # TODO: @pavanimajety Remove this once the switch happens
1153 + # inside flashinfer.
1154 + if not FlashInferBackend.use_trtllm_decode_attention(
1155 + num_decode_tokens, attn_metadata.max_decode_seq_len,
1156 + kv_cache_dtype, attn_metadata.num_qo_heads,
1157 + attn_metadata.num_kv_heads, attn_metadata.head_dim):
1158 + decode_output = decode_meta.decode_wrapper.run(
1159 + decode_query,
1160 + kv_cache.permute(*stride_order),
1161 + k_scale=layer._k_scale_float,
1162 + v_scale=layer._v_scale_float,
1163 + )
1164 + else:
1165 + workspace_buffer = (
1166 + decode_meta.decode_wrapper._int_workspace_buffer)
1167 + assert FlashInferState.get_kv_cache_layout() == "HND"
1168 + decode_output = trtllm_batch_decode_with_kv_cache(
1169 + query=decode_query,
1170 + kv_cache=kv_cache.permute(*stride_order),
1171 + workspace_buffer=workspace_buffer,
1172 + num_heads=num_heads,
1173 + num_kv_heads=num_kv_heads,
1174 + scale=softmax_scale,
1175 + block_tables=attn_metadata.block_tables,
1176 + seq_lens=decode_meta.seq_lens_tensor,
1177 + block_size=attn_metadata.page_size,
1178 + max_seq_len=attn_metadata.max_decode_seq_len,
1179 + kv_cache_dtype=kv_cache_dtype,
1180 + k_scale=layer._k_scale_float,
1181 + v_scale=layer._v_scale_float)
1182
1183 if prefill_output is None and decode_output is not None:
1184 # Decode only batch.
vllm/engine/arg_utils.py
+2
@@ -1424,6 +1424,8 @@ class EngineArgs:
1424 from vllm.attention.utils.fa_utils import (
1425 flash_attn_supports_fp8)
1426 supported = flash_attn_supports_fp8()
1427 + elif envs.VLLM_USE_TRTLLM_DECODE_ATTENTION:
1428 + supported = True
1429 if not supported:
1430 _raise_or_fallback(feature_name="--kv-cache-dtype",
1431 recommend_to_remove=False)
vllm/envs.py
+5 -1
@@ -959,7 +959,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
959 # consumer. This is only applicable when using NixlConnector in a
960 # disaggregated decode-prefill setup.
961 "VLLM_NIXL_ABORT_REQUEST_TIMEOUT":
962 - lambda: int(os.getenv("VLLM_NIXL_ABORT_REQUEST_TIMEOUT", "120"))
962 + lambda: int(os.getenv("VLLM_NIXL_ABORT_REQUEST_TIMEOUT", "120")),
963 +
964 + # If set to 1, use the TRTLLM Decode Attention backend in flashinfer.
965 + "VLLM_USE_TRTLLM_DECODE_ATTENTION":
966 + lambda: os.getenv("VLLM_USE_TRTLLM_DECODE_ATTENTION", None),
967 }
968
969 # --8<-- [end:env-vars-definition]
vllm/platforms/cuda.py
+17 -2
@@ -244,6 +244,10 @@ class CudaPlatformBase(Platform):
244
245 if selected_backend == _Backend.FLASHINFER:
246 logger.info_once("Using FlashInfer backend on V1 engine.")
247 + if cls.has_device_capability(100):
248 + from vllm.v1.attention.backends.utils import (
249 + set_kv_cache_layout)
250 + set_kv_cache_layout("HND")
251 return FLASHINFER_V1
252 elif selected_backend == _Backend.FLEX_ATTENTION:
253 logger.info_once("Using FlexAttention backend on V1 engine.")
@@ -271,9 +275,13 @@ class CudaPlatformBase(Platform):
275 supports_head_size(FLASHINFER_V1, head_size):
276 try:
277 import flashinfer # noqa: F401
278 +
279 + from vllm.v1.attention.backends.utils import (
280 + set_kv_cache_layout)
281 logger.info_once(
275 - "Using FlashInfer backend on V1 engine by default for "
276 - "Blackwell (SM 10.0) GPUs.")
282 + "Using FlashInfer backend with HND KV cache layout on "
283 + "V1 engine by default for Blackwell (SM 10.0) GPUs.")
284 + set_kv_cache_layout("HND")
285 return FLASHINFER_V1
286 except ImportError:
287 logger.info_once(
@@ -293,6 +301,13 @@ class CudaPlatformBase(Platform):
301 # Backends for V0 engine
302 if selected_backend == _Backend.FLASHINFER:
303 logger.info("Using FlashInfer backend.")
304 + if cls.has_device_capability(100):
305 + from vllm.v1.attention.backends.utils import (
306 + set_kv_cache_layout)
307 + logger.info_once(
308 + "Using HND KV cache layout on V1 engine by default for "
309 + "Blackwell (SM 10.0) GPUs.")
310 + set_kv_cache_layout("HND")
311 return "vllm.attention.backends.flashinfer.FlashInferBackend"
312 elif selected_backend == _Backend.XFORMERS:
313 logger.info("Using XFormers backend.")
vllm/v1/attention/backends/flashinfer.py
+147 -36
@@ -10,11 +10,13 @@ import torch
10 from flashinfer import (BatchDecodeWithPagedKVCacheWrapper,
11 BatchPrefillWithPagedKVCacheWrapper,
12 MultiLevelCascadeAttentionWrapper)
13 +from flashinfer.decode import trtllm_batch_decode_with_kv_cache
14
15 import vllm.envs as envs
16 from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
17 AttentionType)
18 from vllm.logger import init_logger
19 +from vllm.platforms import current_platform
20 from vllm.v1.attention.backends.flash_attn import use_cascade_attention
21 from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
22 CommonAttentionMetadata,
@@ -38,6 +40,7 @@ logger = init_logger(__name__)
40 class FlashInferBackend(AttentionBackend):
41
42 accept_output_buffer: bool = True
43 + cached_sm100a_supported: Optional[bool] = None
44
45 @classmethod
46 def get_supported_head_sizes(cls) -> list[int]:
@@ -93,6 +96,57 @@ class FlashInferBackend(AttentionBackend):
96 raise ValueError(f"Unknown cache layout format {cache_layout}.")
97 return stride_order
98
99 + @staticmethod
100 + def use_trtllm_decode_attention(
101 + batch_size: int,
102 + max_seq_len: int,
103 + kv_cache_dtype: str,
104 + num_qo_heads: int,
105 + num_kv_heads: int,
106 + attn_head_size: int,
107 + ) -> bool:
108 + if FlashInferBackend.cached_sm100a_supported is None:
109 + FlashInferBackend.cached_sm100a_supported = (
110 + current_platform.has_device_capability(100))
111 + if not FlashInferBackend.cached_sm100a_supported:
112 + return False
113 + if (num_qo_heads // num_kv_heads > 8
114 + or num_qo_heads % num_kv_heads != 0 or attn_head_size != 128):
115 + return False
116 + env_value = envs.VLLM_USE_TRTLLM_DECODE_ATTENTION
117 + if env_value is not None:
118 + logger.info_once("VLLM_USE_TRTLLM_DECODE_ATTENTION is set to %s",
119 + env_value)
120 + # Environment variable is set - respect it
121 + # Making the conditional check for zero because
122 + # the path is automatically enabled if the batch size condition
123 + # is satisfied.
124 + no_use_trtllm = env_value == "0"
125 + if not no_use_trtllm:
126 + logger.info_once(
127 + "VLLM_USE_TRTLLM_DECODE_ATTENTION is set to 1, "
128 + "using TRTLLM decode attention.")
129 + return not no_use_trtllm
130 + else:
131 + # Environment variable not set - use auto-detection
132 + # Only supports attention head size of 128
133 + use_trtllm = (FlashInferBackend.cached_sm100a_supported
134 + and batch_size <= 256 and max_seq_len < 131072
135 + and kv_cache_dtype == "auto")
136 + if use_trtllm:
137 + logger.warning_once(
138 + "Using TRTLLM decode attention (auto-detected).")
139 + return use_trtllm
140 +
141 + @staticmethod
142 + def get_fp8_dtype_for_flashinfer(kv_cache_dtype: str) -> torch.dtype:
143 + if kv_cache_dtype in ("fp8", "fp8_e4m3"):
144 + return torch.float8_e4m3fn
145 + elif kv_cache_dtype == "fp8_e5m2":
146 + return torch.float8_e5m2
147 + else:
148 + raise ValueError(f"Unrecognized FP8 dtype: {kv_cache_dtype}")
149 +
150
151 @dataclass
152 class FlashInferMetadata:
@@ -127,12 +181,18 @@ class FlashInferMetadata:
181 # Block size of vllm
182 page_size: int
183 # The data type of the paged kv cache
130 - data_type: torch.dtype
184 + kv_data_type: torch.dtype
185 # The data type of the query
186 q_data_type: torch.dtype
187
188 slot_mapping: torch.Tensor
189
190 + # For flashinfer trtllm batch decode
191 + max_seq_len: int
192 + seq_lens: torch.Tensor
193 + block_table_tensor: torch.Tensor
194 + workspace_buffer: torch.Tensor
195 +
196 # For handling prefill decode split
197 num_decodes: int
198 num_decode_tokens: int
@@ -299,6 +359,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
359 window_left=self.global_hyperparameters.window_left,
360 logits_soft_cap=self.global_hyperparameters.logits_soft_cap,
361 q_data_type=attn_metadata.q_data_type,
362 + kv_data_type=attn_metadata.kv_data_type,
363 )
364 else:
365 # Regular attention (common case).
@@ -334,28 +395,33 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
395 logits_soft_cap=self.global_hyperparameters.
396 logits_soft_cap,
397 q_data_type=attn_metadata.q_data_type,
337 - kv_data_type=attn_metadata.data_type,
398 + kv_data_type=attn_metadata.kv_data_type,
399 )
400
401 if self._num_decodes > 0:
402 attn_metadata.decode_wrapper = self._get_decode_wrapper()
342 - attn_metadata.decode_wrapper.plan(
343 - attn_metadata.paged_kv_indptr[:self._num_decodes + 1],
344 - attn_metadata.paged_kv_indices,
345 - attn_metadata.paged_kv_last_page_len[:self._num_decodes],
346 - attn_metadata.num_qo_heads,
347 - attn_metadata.num_kv_heads,
348 - attn_metadata.head_dim,
349 - attn_metadata.page_size,
350 - # Disable flashinfer's pos encoding and use vllm's rope.
351 - pos_encoding_mode="NONE",
352 - sm_scale=self.global_hyperparameters.sm_scale,
353 - window_left=self.global_hyperparameters.window_left,
354 - logits_soft_cap=self.global_hyperparameters.
355 - logits_soft_cap,
356 - q_data_type=attn_metadata.q_data_type,
357 - kv_data_type=attn_metadata.data_type,
358 - )
403 + if not FlashInferBackend.use_trtllm_decode_attention(
404 + self._num_decodes, attn_metadata.max_seq_len,
405 + attn_metadata.kv_data_type, attn_metadata.num_qo_heads,
406 + attn_metadata.num_kv_heads, attn_metadata.head_dim):
407 + attn_metadata.decode_wrapper.plan(
408 + attn_metadata.paged_kv_indptr[:self._num_decodes + 1],
409 + attn_metadata.paged_kv_indices,
410 + attn_metadata.paged_kv_last_page_len[:self.
411 + _num_decodes],
412 + attn_metadata.num_qo_heads,
413 + attn_metadata.num_kv_heads,
414 + attn_metadata.head_dim,
415 + attn_metadata.page_size,
416 + # Disable flashinfer's pos encoding and use vllm's rope.
417 + pos_encoding_mode="NONE",
418 + sm_scale=self.global_hyperparameters.sm_scale,
419 + window_left=self.global_hyperparameters.window_left,
420 + logits_soft_cap=self.global_hyperparameters.
421 + logits_soft_cap,
422 + q_data_type=attn_metadata.q_data_type,
423 + kv_data_type=attn_metadata.kv_data_type,
424 + )
425
426 def build(self, common_prefix_len: int,
427 common_attn_metadata: CommonAttentionMetadata):
@@ -368,6 +434,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
434 page_size = self.kv_cache_spec.block_size
435 device = self.runner.device
436 qo_indptr = common_attn_metadata.query_start_loc
437 + max_seq_len = int(self.runner.seq_lens_np[:num_reqs].max())
438 seq_lens = common_attn_metadata.seq_lens
439 block_table_tensor = self.block_table.get_device_tensor()[:num_reqs]
440 slot_mapping = self.block_table.slot_mapping_cpu[:num_actual_tokens].to(
@@ -416,7 +483,12 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
483 paged_kv_last_page_len = seq_lens % page_size
484 paged_kv_last_page_len = torch.where(paged_kv_last_page_len == 0,
485 page_size, paged_kv_last_page_len)
419 -
486 + cache_dtype = self.runner.cache_config.cache_dtype
487 + if cache_dtype.startswith("fp8"):
488 + kv_cache_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(
489 + cache_dtype)
490 + else:
491 + kv_cache_dtype = self.kv_cache_spec.dtype
492 attn_metadata = FlashInferMetadata(
493 num_actual_tokens=num_actual_tokens,
494 qo_indptr=qo_indptr,
@@ -427,7 +499,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
499 num_kv_heads=self.kv_cache_spec.num_kv_heads,
500 head_dim=self.kv_cache_spec.head_size,
501 page_size=page_size,
430 - data_type=self.kv_cache_spec.dtype,
502 + kv_data_type=kv_cache_dtype,
503 q_data_type=self.runner.dtype,
504 slot_mapping=slot_mapping,
505 num_decodes=self._num_decodes,
@@ -439,6 +511,10 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
511 shared_kv_page_indptr=shared_kv_page_indptr,
512 shared_kv_page_indices=shared_kv_page_indices,
513 shared_kv_last_page_len=shared_kv_last_page_len,
514 + max_seq_len=max_seq_len,
515 + seq_lens=seq_lens,
516 + block_table_tensor=block_table_tensor,
517 + workspace_buffer=self._workspace_buffer,
518 )
519
520 self._plan(attn_metadata)
@@ -514,7 +590,11 @@ class FlashInferImpl(AttentionImpl):
590 query: shape = [num_tokens, num_heads, head_size]
591 key: shape = [num_tokens, num_kv_heads, head_size]
592 value: shape = [num_tokens, num_kv_heads, head_size]
517 - kv_cache = [num_blocks, 2, block_size, num_kv_heads, head_size]
593 + kv_cache: shape -
594 + # NHD: [num_blocks, 2, block_size, num_kv_heads, head_size]
595 + # HND: [num_blocks, 2, num_kv_heads, block_size, head_size]
596 +
597 +
598 attn_metadata: Metadata for attention.
599 Returns:
600 shape = [num_tokens, num_heads * head_size]
@@ -560,6 +640,13 @@ class FlashInferImpl(AttentionImpl):
640 layer._v_scale,
641 )
642
643 + # The FlashInfer api requires data to be in fp8_e4m3 or fp8_e5m2
644 + # to process the cache when the kv_cache_dtype is fp8
645 + if self.kv_cache_dtype.startswith("fp8"):
646 + torch_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(
647 + self.kv_cache_dtype)
648 + kv_cache = kv_cache.view(torch_dtype)
649 +
650 window_left = (self.sliding_window[0]
651 if self.sliding_window is not None else -1)
652
@@ -597,21 +684,45 @@ class FlashInferImpl(AttentionImpl):
684 v_scale=layer._v_scale_float,
685 out=output[num_decode_tokens:],
686 )
600 -
687 if decode_wrapper := attn_metadata.decode_wrapper:
688 decode_query = query[:num_decode_tokens]
689 assert decode_query.shape[0] == num_decode_tokens
604 - assert decode_wrapper is not None
605 - assert decode_wrapper._window_left == window_left
606 - assert decode_wrapper._logits_soft_cap == (self.logits_soft_cap
607 - or 0.0)
608 - assert decode_wrapper._sm_scale == self.scale
609 - decode_wrapper.run(
610 - decode_query,
611 - kv_cache.permute(*stride_order),
612 - k_scale=layer._k_scale_float,
613 - v_scale=layer._v_scale_float,
614 - out=output[:num_decode_tokens],
615 - )
616 -
690 + if not FlashInferBackend.use_trtllm_decode_attention(
691 + attn_metadata.num_decodes, attn_metadata.max_seq_len,
692 + self.kv_cache_dtype, attn_metadata.num_qo_heads,
693 + attn_metadata.num_kv_heads, attn_metadata.head_dim):
694 + assert decode_wrapper is not None
695 + assert decode_wrapper._window_left == window_left
696 + assert decode_wrapper._logits_soft_cap == (self.logits_soft_cap
697 + or 0.0)
698 + assert decode_wrapper._sm_scale == self.scale
699 + decode_wrapper.run(
700 + decode_query,
701 + kv_cache.permute(*stride_order),
702 + k_scale=layer._k_scale_float,
703 + v_scale=layer._v_scale_float,
704 + out=output[:num_decode_tokens],
705 + )
706 + else:
707 + # This path needs to be enabled with VLLM_KV_CACHE_LAYOUT = HND
708 + if num_decode_tokens > 0:
709 + assert get_kv_cache_layout() == "HND"
710 + output[:num_decode_tokens] = (
711 + trtllm_batch_decode_with_kv_cache(
712 + query=decode_query,
713 + kv_cache=kv_cache.permute(*stride_order),
714 + workspace_buffer=attn_metadata.workspace_buffer,
715 + num_heads=self.num_heads,
716 + num_kv_heads=self.num_kv_heads,
717 + scale=self.scale,
718 + block_tables=attn_metadata.
719 + block_table_tensor[:num_decode_tokens],
720 + seq_lens=attn_metadata.
721 + seq_lens[:num_decode_tokens],
722 + block_size=attn_metadata.page_size,
723 + max_seq_len=attn_metadata.max_seq_len,
724 + kv_cache_dtype=self.kv_cache_dtype,
725 + k_scale=layer._k_scale_float,
726 + v_scale=layer._v_scale_float,
727 + ))
728 return output_padded
vllm/v1/attention/backends/utils.py
+9 -1
@@ -24,6 +24,7 @@ from vllm.distributed.kv_transfer.kv_connector.utils import (
24 from vllm.logger import init_logger
25
26 logger = init_logger(__name__)
27 +_KV_CACHE_LAYOUT_OVERRIDE = None
28
29
30 @dataclass
@@ -103,6 +104,7 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):
104
105 @functools.lru_cache
106 def get_kv_cache_layout():
107 + global _KV_CACHE_LAYOUT_OVERRIDE
108 # Override with format specified by the user.
109 cache_layout = envs.VLLM_KV_CACHE_LAYOUT
110 if cache_layout is None:
@@ -110,10 +112,16 @@ def get_kv_cache_layout():
112 else:
113 logger.info_once("`VLLM_KV_CACHE_LAYOUT` environment variable " \
114 "detected. Setting KV cache layout to %s.", cache_layout)
113 -
115 + if _KV_CACHE_LAYOUT_OVERRIDE is not None:
116 + cache_layout = _KV_CACHE_LAYOUT_OVERRIDE
117 return cache_layout
118
119
120 +def set_kv_cache_layout(cache_layout: str):
121 + global _KV_CACHE_LAYOUT_OVERRIDE
122 + _KV_CACHE_LAYOUT_OVERRIDE = cache_layout
123 +
124 +
125 @dataclass
126 class PerLayerParameters:
127 """