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, 23:45 UTC
[Bugfix] Fix workspace buffer None issue for Flashinfer TRTLLM Backend (#21525)
Signed-off-by: elvischenv <219235043+elvischenv@users.noreply.github.com>
elvischenv committed
Jul 29, 2025 at 22:34 UTC
58b11b24a69f0d5fc48f3a6ce8291e8d92af26e2
4 files changed
+61
-42
benchmarks/kernels/benchmark_trtllm_attention.py
+28
-14
@@ -71,22 +71,20 @@ def benchmark_decode(
71
if kv_cache_dtype.startswith("fp8"):
72
kv_cache, _ = to_float8(kv_cache)
73
74
+ output_trtllm = torch.empty(q.shape, dtype=dtype)
75
+
76
# Benchmark TRT decode
77
def trt_decode():
78
return flashinfer.decode.trtllm_batch_decode_with_kv_cache(
79
q,
80
kv_cache,
81
workspace_buffer,
80
- num_qo_heads,
81
- num_kv_heads,
82
- sm_scale,
82
block_tables,
83
kv_lens_tensor,
85
- page_size,
84
max_kv_len,
87
- kv_cache_dtype,
88
- k_scale,
89
- v_scale,
85
+ bmm1_scale=k_scale * sm_scale,
86
+ bmm2_scale=v_scale,
87
+ out=output_trtllm,
88
)
89
90
def time_fn(fn, warmup=10, trials=20):
@@ -125,6 +123,8 @@ def benchmark_decode(
123
kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
124
kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
125
126
+ output_baseline = torch.empty(q.shape, dtype=dtype)
127
+
128
wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
129
workspace_buffer,
130
kv_layout,
@@ -145,7 +145,7 @@ def benchmark_decode(
145
)
146
147
def baseline_decode():
148
- return wrapper.run(q, kv_cache, sm_scale, k_scale, v_scale)
148
+ return wrapper.run(q, kv_cache, sm_scale, k_scale, v_scale, output_baseline)
149
150
baseline_mean, baseline_std = time_fn(baseline_decode)
151
@@ -214,25 +214,39 @@ if __name__ == "__main__":
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")
217
print(
219
- "\tnum_seqs\tmax_seq_len\ttrt_mean\ttrt_std\tbaseline_mean\tbaseline_std\tspeedup_percent"
218
+ "Running benchmark for q_dtype = bfloat16, kv_cache_dtype: bfloat16, "
219
+ "output_dtype: bfloat16"
220
+ )
221
+ print(
222
+ "\tnum_seqs\tmax_seq_len\ttrt_mean\ttrt_std\tbaseline_mean\t"
223
+ "baseline_std\tspeedup_percent"
224
)
225
for max_seq_len in max_seq_lens:
226
for bs in num_seqs:
227
result = benchmark_decode(
224
- bs, max_seq_len, dtype=torch.bfloat16, kv_cache_dtype="auto"
228
+ bs,
229
+ max_seq_len,
230
+ dtype=torch.bfloat16,
231
+ kv_cache_dtype="auto",
232
)
233
all_results.append(result)
234
228
- print("Running benchmark for q_dtype = bfloat16, kv_cache_dtype: fp8")
235
print(
230
- "\tnum_seqs\tmax_seq_len\ttrt_mean\ttrt_std\tbaseline_mean\tbaseline_std\tspeedup_percent"
236
+ "Running benchmark for q_dtype = bfloat16, kv_cache_dtype: fp8, "
237
+ "output_dtype: bfloat16"
238
+ )
239
+ print(
240
+ "\tnum_seqs\tmax_seq_len\ttrt_mean\ttrt_std\tbaseline_mean\t"
241
+ "baseline_std\tspeedup_percent"
242
)
243
for max_seq_len in max_seq_lens:
244
for bs in num_seqs:
245
result = benchmark_decode(
235
- bs, max_seq_len, dtype=torch.bfloat16, kv_cache_dtype="fp8"
246
+ bs,
247
+ max_seq_len,
248
+ dtype=torch.bfloat16,
249
+ kv_cache_dtype="fp8",
250
)
251
all_results.append(result)
252
tests/kernels/attention/test_flashinfer_trtllm_decode_attention.py
+7
-9
@@ -113,27 +113,25 @@ def test_flashinfer_trtllm_decode_with_baseline(
113
kv_data_type=dtype,
114
logits_soft_cap=soft_cap)
115
116
- output = wrapper.run(query, key_value_cache, scale)
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)
123
- output_trtllm = flashinfer.decode.trtllm_batch_decode_with_kv_cache(
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,
127
- num_query_heads,
128
- num_kv_heads,
129
- scale,
129
block_tables,
130
kv_lens_tensor,
132
- block_size,
131
max_kv_len,
134
- "auto",
135
- k_scale,
136
- v_scale,
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), \
vllm/attention/backends/flashinfer.py
+11
-4
@@ -1104,7 +1104,12 @@ class FlashInferImpl(AttentionImpl):
1104
window_left = window_size[0] if window_size is not None else -1
1105
1106
prefill_output: Optional[torch.Tensor] = None
1107
- decode_output: Optional[torch.Tensor] = None
1107
+ if num_decode_tokens > 0:
1108
+ decode_output = torch.empty(decode_query.shape,
1109
+ dtype=decode_query.dtype,
1110
+ device=decode_query.device)
1111
+ else:
1112
+ decode_output = None
1113
stride_order = FlashInferBackend.get_kv_cache_stride_order()
1114
if prefill_meta := attn_metadata.prefill_metadata:
1115
# We will use flash attention for prefill
@@ -1155,17 +1160,18 @@ class FlashInferImpl(AttentionImpl):
1160
num_decode_tokens, attn_metadata.max_decode_seq_len,
1161
kv_cache_dtype, attn_metadata.num_qo_heads,
1162
attn_metadata.num_kv_heads, attn_metadata.head_dim):
1158
- decode_output = decode_meta.decode_wrapper.run(
1163
+ decode_meta.decode_wrapper.run(
1164
decode_query,
1165
kv_cache.permute(*stride_order),
1166
k_scale=layer._k_scale_float,
1167
v_scale=layer._v_scale_float,
1168
+ out=decode_output,
1169
)
1170
else:
1171
workspace_buffer = (
1166
- decode_meta.decode_wrapper._int_workspace_buffer)
1172
+ decode_meta.decode_wrapper._float_workspace_buffer)
1173
assert FlashInferState.get_kv_cache_layout() == "HND"
1168
- decode_output = trtllm_batch_decode_with_kv_cache(
1174
+ trtllm_batch_decode_with_kv_cache(
1175
query=decode_query,
1176
kv_cache=kv_cache.permute(*stride_order),
1177
workspace_buffer=workspace_buffer,
@@ -1174,6 +1180,7 @@ class FlashInferImpl(AttentionImpl):
1180
max_seq_len=attn_metadata.max_decode_seq_len,
1181
bmm1_scale=layer._k_scale_float * softmax_scale,
1182
bmm2_scale=layer._v_scale_float,
1183
+ out=decode_output,
1184
)
1185
1186
if prefill_output is None and decode_output is not None:
vllm/v1/attention/backends/flashinfer.py
+15
-15
@@ -194,7 +194,6 @@ class FlashInferMetadata:
194
max_seq_len: int
195
seq_lens: torch.Tensor
196
block_table_tensor: torch.Tensor
197
- workspace_buffer: torch.Tensor
197
198
# For handling prefill decode split
199
num_decodes: int
@@ -473,7 +472,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
472
max_seq_len=max_seq_len,
473
seq_lens=seq_lens,
474
block_table_tensor=block_table_tensor,
476
- workspace_buffer=self._get_workspace_buffer(),
475
)
476
477
self._plan(num_prefills, num_decodes, attn_metadata)
@@ -641,11 +639,11 @@ class FlashInferImpl(AttentionImpl):
639
if decode_wrapper := attn_metadata.decode_wrapper:
640
decode_query = query[:num_decode_tokens]
641
assert decode_query.shape[0] == num_decode_tokens
642
+ assert decode_wrapper is not None
643
if not FlashInferBackend.use_trtllm_decode_attention(
644
attn_metadata.num_decodes, attn_metadata.max_seq_len,
645
self.kv_cache_dtype, attn_metadata.num_qo_heads,
646
attn_metadata.num_kv_heads, attn_metadata.head_dim):
648
- assert decode_wrapper is not None
647
assert decode_wrapper._window_left == window_left
648
assert decode_wrapper._logits_soft_cap == (self.logits_soft_cap
649
or 0.0)
@@ -666,22 +664,24 @@ class FlashInferImpl(AttentionImpl):
664
num_decode_tokens]
665
seq_lens_decode = attn_metadata.seq_lens[:
666
num_decode_tokens]
667
+ workspace_buffer = decode_wrapper._float_workspace_buffer
668
669
assert get_kv_cache_layout() == "HND"
670
assert decode_query.is_contiguous()
671
assert kv_cache_permute.is_contiguous()
672
assert block_tables_decode.is_contiguous()
673
assert seq_lens_decode.is_contiguous()
675
-
676
- output[:num_decode_tokens] = (
677
- trtllm_batch_decode_with_kv_cache(
678
- query=decode_query,
679
- kv_cache=kv_cache_permute,
680
- workspace_buffer=attn_metadata.workspace_buffer,
681
- block_tables=block_tables_decode,
682
- seq_lens=seq_lens_decode,
683
- max_seq_len=attn_metadata.max_seq_len,
684
- bmm1_scale=layer._k_scale_float * self.scale,
685
- bmm2_scale=layer._v_scale_float,
686
- ))
674
+ assert workspace_buffer.is_contiguous()
675
+
676
+ trtllm_batch_decode_with_kv_cache(
677
+ query=decode_query,
678
+ kv_cache=kv_cache_permute,
679
+ workspace_buffer=workspace_buffer,
680
+ block_tables=block_tables_decode,
681
+ seq_lens=seq_lens_decode,
682
+ max_seq_len=attn_metadata.max_seq_len,
683
+ bmm1_scale=layer._k_scale_float * self.scale,
684
+ bmm2_scale=layer._v_scale_float,
685
+ out=output[:num_decode_tokens],
686
+ )
687
return output_padded