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
[Flashinfer][gpt-oss] Support FP8-qkv Flashinfer TRTLLM Sinks Attention (#25674)
Signed-off-by: elvischenv <219235043+elvischenv@users.noreply.github.com>
elvischenv committed
Oct 10, 2025 at 04:13 UTC
44f633dba17f03e10fc4fc375df2e86c7d53fcc1
2 files changed
+76
-46
tests/kernels/attention/test_flashinfer_trtllm_attention.py
+76
-41
@@ -49,6 +49,7 @@ KV_LAYOUT = ["HND"] # currently only HND is supported
49
BLOCK_SIZE = [16]
50
WINDOW_LEFT = [-1, 127]
51
SOFT_CAP = [None, 50.0]
52
+HAS_SINKS = [True, False]
53
54
NUM_BLOCKS = 32768 # Large enough to test overflow in index calculation.
55
@@ -63,6 +64,7 @@ NUM_BLOCKS = 32768 # Large enough to test overflow in index calculation.
64
@pytest.mark.parametrize("block_size", BLOCK_SIZE)
65
@pytest.mark.parametrize("window_left", WINDOW_LEFT)
66
@pytest.mark.parametrize("soft_cap", SOFT_CAP)
67
+@pytest.mark.parametrize("has_sinks", HAS_SINKS)
68
@torch.inference_mode
69
def test_flashinfer_trtllm_decode_with_baseline(
70
dtype: torch.dtype,
@@ -77,9 +79,10 @@ def test_flashinfer_trtllm_decode_with_baseline(
79
block_size: int,
80
window_left: int,
81
soft_cap: Optional[float],
82
+ has_sinks: bool,
83
) -> None:
84
torch.set_default_device("cuda")
82
- current_platform.seed_everything(0)
85
+ current_platform.seed_everything(42)
86
87
q_quant_dtype, kv_quant_dtype, o_quant_dtype = quant_dtypes
88
q_quant_dtype = q_quant_dtype or dtype
@@ -101,7 +104,16 @@ def test_flashinfer_trtllm_decode_with_baseline(
104
else:
105
raise ValueError(f"Invalid kv_layout: {kv_layout}")
106
104
- query = torch.randn(batch_size, num_qo_heads, head_size, dtype=dtype)
107
+ # max_q_len = 1
108
+ q_lens = torch.ones((batch_size,), dtype=torch.int32)
109
+ q_indptr = torch.cat(
110
+ [
111
+ torch.tensor([0], dtype=torch.int32),
112
+ torch.cumsum(q_lens, dim=0, dtype=torch.int32),
113
+ ]
114
+ )
115
+
116
+ query = torch.randn(torch.sum(q_lens).item(), num_qo_heads, head_size, dtype=dtype)
117
if q_quant_dtype == FP8_DTYPE:
118
query, q_scale = to_float8(query)
119
ref_query = query.to(dtype) * q_scale
@@ -112,7 +124,7 @@ def test_flashinfer_trtllm_decode_with_baseline(
124
kv_lens = torch.randint(1, max_kv_len, (batch_size,), dtype=torch.int32)
125
kv_lens[-1] = max_kv_len
126
115
- seq_lens = kv_lens
127
+ seq_lens = kv_lens + q_lens
128
max_seq_len = torch.max(seq_lens).item()
129
130
kv_cache = torch.randn(kv_cache_shape, dtype=dtype)
@@ -148,27 +160,36 @@ def test_flashinfer_trtllm_decode_with_baseline(
160
workspace_buffer = torch.zeros(128 * 1024 * 1024, dtype=torch.int8)
161
162
# Baseline Decode
151
- wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
152
- workspace_buffer, kv_layout, use_tensor_cores=True
153
- )
163
+ if has_sinks:
164
+ sinks = torch.rand(num_qo_heads, dtype=torch.float32) * 5
165
+ wrapper = flashinfer.BatchAttentionWithAttentionSinkWrapper(
166
+ float_workspace_buffer=workspace_buffer, kv_layout=kv_layout, backend="fa2"
167
+ )
168
+ else:
169
+ sinks = None
170
+ wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(
171
+ float_workspace_buffer=workspace_buffer, kv_layout=kv_layout, backend="fa2"
172
+ )
173
+
174
wrapper.plan(
155
- kv_indptr,
156
- kv_indices,
157
- kv_last_page_lens,
158
- num_qo_heads,
159
- num_kv_heads,
160
- head_size,
161
- block_size,
162
- "NONE",
175
+ qo_indptr=q_indptr,
176
+ paged_kv_indptr=kv_indptr,
177
+ paged_kv_indices=kv_indices,
178
+ paged_kv_last_page_len=kv_last_page_lens,
179
+ num_qo_heads=num_qo_heads,
180
+ num_kv_heads=num_kv_heads,
181
+ head_dim_qk=head_size,
182
+ page_size=block_size,
183
+ causal=True,
184
sm_scale=sm_scale,
164
- q_data_type=dtype,
165
- kv_data_type=dtype,
185
window_left=window_left,
186
logits_soft_cap=soft_cap,
187
+ q_data_type=dtype,
188
+ kv_data_type=dtype,
189
)
169
-
190
output = torch.empty(ref_query.shape, dtype=dtype)
171
- wrapper.run(ref_query, ref_kv_cache, out=output)
191
+ wrapper.run(ref_query, ref_kv_cache, sinks, sm_scale, out=output)
192
+
193
o_scale = 1.0
194
o_sf_scale_float = None
195
if o_quant_dtype == FP8_DTYPE:
@@ -202,6 +223,7 @@ def test_flashinfer_trtllm_decode_with_baseline(
223
bmm1_scale=q_scale * k_scale * sm_scale,
224
bmm2_scale=v_scale / o_scale,
225
window_left=window_left,
226
+ sinks=sinks,
227
o_sf_scale=o_sf_scale_float,
228
out=output_trtllm,
229
)
@@ -217,11 +239,13 @@ def test_flashinfer_trtllm_decode_with_baseline(
239
output_trtllm = output_trtllm.reshape(-1, query.shape[1], query.shape[2])
240
241
if q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP4_DTYPE:
220
- rtol, atol = 3e-1, 1e0
242
+ rtol, atol = 7e-2, 9e-2
243
elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP8_DTYPE:
222
- rtol, atol = 5e-2, 7e-2
223
- else:
244
+ rtol, atol = 2e-2, 4e-2
245
+ elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == dtype:
246
rtol, atol = 1e-2, 2e-2
247
+ else:
248
+ rtol, atol = 1e-2, 1e-2
249
250
(
251
torch.testing.assert_close(output, output_trtllm, atol=atol, rtol=rtol),
@@ -239,6 +263,7 @@ def test_flashinfer_trtllm_decode_with_baseline(
263
@pytest.mark.parametrize("block_size", BLOCK_SIZE)
264
@pytest.mark.parametrize("window_left", WINDOW_LEFT)
265
@pytest.mark.parametrize("soft_cap", [None])
266
+@pytest.mark.parametrize("has_sinks", HAS_SINKS)
267
@torch.inference_mode
268
def test_flashinfer_trtllm_prefill_with_baseline(
269
dtype: torch.dtype,
@@ -253,9 +278,10 @@ def test_flashinfer_trtllm_prefill_with_baseline(
278
block_size: int,
279
window_left: int,
280
soft_cap: Optional[float],
281
+ has_sinks: bool,
282
) -> None:
283
torch.set_default_device("cuda")
258
- current_platform.seed_everything(0)
284
+ current_platform.seed_everything(42)
285
286
q_quant_dtype, kv_quant_dtype, o_quant_dtype = quant_dtypes
287
q_quant_dtype = q_quant_dtype or dtype
@@ -297,7 +323,7 @@ def test_flashinfer_trtllm_prefill_with_baseline(
323
q_scale = 1.0
324
ref_query = query
325
300
- kv_lens = torch.randint(0, max_kv_len, (batch_size,), dtype=torch.int32)
326
+ kv_lens = torch.randint(1, max_kv_len, (batch_size,), dtype=torch.int32)
327
kv_lens[-1] = max_kv_len
328
329
seq_lens = kv_lens + q_lens
@@ -336,28 +362,36 @@ def test_flashinfer_trtllm_prefill_with_baseline(
362
workspace_buffer = torch.zeros(128 * 1024 * 1024, dtype=torch.int8)
363
364
# Baseline Prefill
339
- wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(
340
- workspace_buffer, kv_layout
341
- )
365
+ if has_sinks:
366
+ sinks = torch.rand(num_qo_heads, dtype=torch.float32) * 5
367
+ wrapper = flashinfer.BatchAttentionWithAttentionSinkWrapper(
368
+ float_workspace_buffer=workspace_buffer, kv_layout=kv_layout, backend="fa2"
369
+ )
370
+ else:
371
+ sinks = None
372
+ wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(
373
+ float_workspace_buffer=workspace_buffer, kv_layout=kv_layout, backend="fa2"
374
+ )
375
+
376
wrapper.plan(
343
- q_indptr,
344
- kv_indptr,
345
- kv_indices,
346
- kv_last_page_lens,
347
- num_qo_heads,
348
- num_kv_heads,
349
- head_size,
350
- block_size,
377
+ qo_indptr=q_indptr,
378
+ paged_kv_indptr=kv_indptr,
379
+ paged_kv_indices=kv_indices,
380
+ paged_kv_last_page_len=kv_last_page_lens,
381
+ num_qo_heads=num_qo_heads,
382
+ num_kv_heads=num_kv_heads,
383
+ head_dim_qk=head_size,
384
+ page_size=block_size,
385
causal=True,
386
sm_scale=sm_scale,
353
- q_data_type=dtype,
354
- kv_data_type=dtype,
387
window_left=window_left,
388
logits_soft_cap=soft_cap,
389
+ q_data_type=dtype,
390
+ kv_data_type=dtype,
391
)
358
-
392
output = torch.empty(ref_query.shape, dtype=dtype)
360
- wrapper.run(ref_query, ref_kv_cache, out=output)
393
+ wrapper.run(ref_query, ref_kv_cache, sinks, sm_scale, out=output)
394
+
395
o_scale = 1.0
396
o_sf_scale_float = None
397
if o_quant_dtype == FP8_DTYPE:
@@ -395,6 +429,7 @@ def test_flashinfer_trtllm_prefill_with_baseline(
429
cum_seq_lens_q=q_indptr,
430
cum_seq_lens_kv=kv_indptr,
431
window_left=window_left,
432
+ sinks=sinks,
433
o_sf_scale=o_sf_scale_float,
434
out=output_trtllm,
435
)
@@ -410,11 +445,11 @@ def test_flashinfer_trtllm_prefill_with_baseline(
445
output_trtllm = output_trtllm.reshape(-1, query.shape[1], query.shape[2])
446
447
if q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP4_DTYPE:
413
- rtol, atol = 4e-1, 1e0
448
+ rtol, atol = 1e-1, 2e-1
449
elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP8_DTYPE:
415
- rtol, atol = 5e-2, 7e-2
416
- elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == dtype:
450
rtol, atol = 4e-2, 6e-2
451
+ elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == dtype:
452
+ rtol, atol = 2e-2, 3e-2
453
else:
454
rtol, atol = 1e-2, 1e-2
455
vllm/utils/flashinfer.py
-5
@@ -269,11 +269,6 @@ def use_trtllm_attention(
269
270
# Must use TRTLLM attention if query is FP8 quantized
271
if q_dtype == current_platform.fp8_dtype():
272
- if has_sinks:
273
- raise RuntimeError(
274
- "TRTLLM FP8-qkv kernel is not supported for attention sinks. "
275
- "Use kv_cache_dtype=auto for now."
276
- )
272
logger.info_once("Using TRTLLM attention (query is quantized).")
273
return True
274