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] Support Flashinfer TRTLLM FP8-qkv BF16/FP16-out Attention Kernel (#23647)
Signed-off-by: elvischenv <219235043+elvischenv@users.noreply.github.com>
elvischenv committed
Sep 9, 2025 at 11:53 UTC
bba1042c6fcb01a3ec12803ab979852040d3e4d6
5 files changed
+22
-11
benchmarks/kernels/benchmark_trtllm_decode_attention.py
+1
@@ -259,6 +259,7 @@ if __name__ == "__main__":
259
# (q_quant_dtype, kv_quant_dtype, o_quant_dtype)
260
(None, None, None),
261
(None, FP8_DTYPE, None),
262
+ (FP8_DTYPE, FP8_DTYPE, None),
263
(FP8_DTYPE, FP8_DTYPE, FP8_DTYPE),
264
(FP8_DTYPE, FP8_DTYPE, FP4_DTYPE),
265
]
benchmarks/kernels/benchmark_trtllm_prefill_attention.py
+1
@@ -274,6 +274,7 @@ if __name__ == "__main__":
274
quant_dtypes = [
275
# (q_quant_dtype, kv_quant_dtype, o_quant_dtype)
276
(None, None, None),
277
+ (FP8_DTYPE, FP8_DTYPE, None),
278
(FP8_DTYPE, FP8_DTYPE, FP8_DTYPE),
279
(FP8_DTYPE, FP8_DTYPE, FP4_DTYPE),
280
]
tests/kernels/attention/test_flashinfer_trtllm_attention.py
+12
@@ -35,6 +35,7 @@ QUANT_DTYPES = [
35
# (q_quant_dtype, kv_quant_dtype, o_quant_dtype)
36
(None, None, None),
37
(None, FP8_DTYPE, None),
38
+ (FP8_DTYPE, FP8_DTYPE, None),
39
(FP8_DTYPE, FP8_DTYPE, FP8_DTYPE),
40
(FP8_DTYPE, FP8_DTYPE, FP4_DTYPE),
41
]
@@ -44,6 +45,7 @@ NUM_HEADS = [(64, 8), (40, 8)]
45
HEAD_SIZE = [128]
46
KV_LAYOUT = ["HND"] # currently only HND is supported
47
BLOCK_SIZE = [16]
48
+WINDOW_LEFT = [-1, 127]
49
SOFT_CAP = [None, 50.0]
50
51
NUM_BLOCKS = 32768 # Large enough to test overflow in index calculation.
@@ -57,6 +59,7 @@ NUM_BLOCKS = 32768 # Large enough to test overflow in index calculation.
59
@pytest.mark.parametrize("head_size", HEAD_SIZE)
60
@pytest.mark.parametrize("kv_layout", KV_LAYOUT)
61
@pytest.mark.parametrize("block_size", BLOCK_SIZE)
62
+@pytest.mark.parametrize("window_left", WINDOW_LEFT)
63
@pytest.mark.parametrize("soft_cap", SOFT_CAP)
64
@torch.inference_mode
65
def test_flashinfer_trtllm_decode_with_baseline(
@@ -69,6 +72,7 @@ def test_flashinfer_trtllm_decode_with_baseline(
72
head_size: int,
73
kv_layout: str,
74
block_size: int,
75
+ window_left: int,
76
soft_cap: Optional[float],
77
) -> None:
78
torch.set_default_device("cuda")
@@ -155,6 +159,7 @@ def test_flashinfer_trtllm_decode_with_baseline(
159
sm_scale=sm_scale,
160
q_data_type=dtype,
161
kv_data_type=dtype,
162
+ window_left=window_left,
163
logits_soft_cap=soft_cap)
164
165
output = torch.empty(ref_query.shape, dtype=dtype)
@@ -188,6 +193,7 @@ def test_flashinfer_trtllm_decode_with_baseline(
193
max_seq_len=max_seq_len,
194
bmm1_scale=q_scale * k_scale * sm_scale,
195
bmm2_scale=v_scale / o_scale,
196
+ window_left=window_left,
197
o_sf_scale=o_sf_scale,
198
out=output_trtllm,
199
)
@@ -222,6 +228,7 @@ def test_flashinfer_trtllm_decode_with_baseline(
228
@pytest.mark.parametrize("head_size", HEAD_SIZE)
229
@pytest.mark.parametrize("kv_layout", KV_LAYOUT)
230
@pytest.mark.parametrize("block_size", BLOCK_SIZE)
231
+@pytest.mark.parametrize("window_left", WINDOW_LEFT)
232
@pytest.mark.parametrize("soft_cap", [None])
233
@torch.inference_mode
234
def test_flashinfer_trtllm_prefill_with_baseline(
@@ -234,6 +241,7 @@ def test_flashinfer_trtllm_prefill_with_baseline(
241
head_size: int,
242
kv_layout: str,
243
block_size: int,
244
+ window_left: int,
245
soft_cap: Optional[float],
246
) -> None:
247
torch.set_default_device("cuda")
@@ -334,6 +342,7 @@ def test_flashinfer_trtllm_prefill_with_baseline(
342
sm_scale=sm_scale,
343
q_data_type=dtype,
344
kv_data_type=dtype,
345
+ window_left=window_left,
346
logits_soft_cap=soft_cap)
347
348
output = torch.empty(ref_query.shape, dtype=dtype)
@@ -371,6 +380,7 @@ def test_flashinfer_trtllm_prefill_with_baseline(
380
batch_size=batch_size,
381
cum_seq_lens_q=q_indptr,
382
cum_seq_lens_kv=kv_indptr,
383
+ window_left=window_left,
384
o_sf_scale=o_sf_scale,
385
out=output_trtllm,
386
)
@@ -390,6 +400,8 @@ def test_flashinfer_trtllm_prefill_with_baseline(
400
rtol, atol = 4e-1, 1e0
401
elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP8_DTYPE:
402
rtol, atol = 5e-2, 7e-2
403
+ elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == dtype:
404
+ rtol, atol = 4e-2, 6e-2
405
else:
406
rtol, atol = 1e-2, 1e-2
407
vllm/compilation/fusion_attn.py
+4
-2
@@ -258,8 +258,10 @@ class AttnFusionPass(VllmInductorPass):
258
pattern_fp8 = AttentionFp8StaticQuantPattern(layer)
259
pattern_fp8.register_if_supported(self.patterns)
260
261
- pattern_nvfp4 = AttentionNvfp4QuantPattern(layer)
262
- pattern_nvfp4.register_if_supported(self.patterns)
261
+ if current_platform.is_cuda() and hasattr(torch.ops._C,
262
+ "scaled_fp4_quant"):
263
+ pattern_nvfp4 = AttentionNvfp4QuantPattern(layer)
264
+ pattern_nvfp4.register_if_supported(self.patterns)
265
266
if len(attn_layers) == 0:
267
logger.warning(
vllm/v1/attention/backends/flashinfer.py
+4
-9
@@ -194,19 +194,15 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
194
FlashInferBackend.validate_head_size(self.head_dim)
195
self.page_size = self.kv_cache_spec.block_size
196
197
- self.enable_fusion = (
198
- self.compilation_config.pass_config.enable_attn_fusion)
199
- self.q_data_type = self.model_config.dtype
197
self.cache_dtype = self.cache_config.cache_dtype
198
if self.cache_dtype.startswith("fp8"):
199
self.kv_cache_dtype = (
200
FlashInferBackend.get_fp8_dtype_for_flashinfer(
201
self.cache_dtype))
205
- # Insert FP8 quant for query if FP8 kv cache and attn fusion enabled
206
- if self.enable_fusion:
207
- self.q_data_type = self.kv_cache_dtype
202
else:
203
+ assert self.kv_cache_spec.dtype == self.model_config.dtype
204
self.kv_cache_dtype = self.kv_cache_spec.dtype
205
+ self.q_data_type = self.kv_cache_dtype
206
207
self._cascade_wrapper = None # Wrapper for cascade attention
208
@@ -668,8 +664,6 @@ class FlashInferImpl(AttentionImpl):
664
665
# The attn+quant fusion happens when output_scale is provided.
666
if output_scale is None:
671
- assert attn_metadata.q_data_type != FP8_DTYPE, \
672
- "Query can only be FP8 if output fusion happened."
667
assert output_block_scale is None, "output_block_scale "\
668
"is not supported when fusion has not happened"
669
else:
@@ -697,7 +691,8 @@ class FlashInferImpl(AttentionImpl):
691
elif output.dtype == FP4_DTYPE:
692
self.o_sf_scale = layer._o_scale_float
693
700
- # Insert FP8 quant for query
694
+ # Insert FP8 quant for query
695
+ if attn_metadata.q_data_type == FP8_DTYPE:
696
num_tokens, num_heads, head_size = query.shape
697
query, _ = ops.scaled_fp8_quant(
698
query.reshape(