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
[Misc] Enable yapf for FlashInfer backend (#23193)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Aug 19, 2025 at 10:33 UTC
5b5f350d67a1e1efb7dbe8b18fe2353ad94857a1
1 file changed
+24
-13
vllm/v1/attention/backends/flashinfer.py
+24
-13
@@ -36,6 +36,7 @@ from vllm.v1.attention.backends.utils import (AttentionCGSupport,
36
get_per_layer_parameters,
37
infer_global_hyperparameters,
38
split_decodes_and_prefills)
39
+# yapf: enable
40
from vllm.v1.kv_cache_interface import AttentionSpec
41
42
FLASHINFER_WORKSPACE_BUFFER_SIZE = 256 * 1024 * 1024
@@ -541,12 +542,22 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
542
if cache_dtype.startswith("fp8") and enable_fusion:
543
q_dtype = kv_cache_dtype
544
544
- prefill_use_trtllm = use_trtllm_attention(
545
- num_qo_heads, num_kv_heads, num_prefill_tokens, max_seq_len,
546
- cache_dtype, q_dtype, is_prefill=True, has_sinks=has_sinks)
547
- decode_use_trtllm = use_trtllm_attention(
548
- num_qo_heads, num_kv_heads, num_decode_tokens, max_seq_len,
549
- cache_dtype, q_dtype, is_prefill=False, has_sinks=has_sinks)
545
+ prefill_use_trtllm = use_trtllm_attention(num_qo_heads,
546
+ num_kv_heads,
547
+ num_prefill_tokens,
548
+ max_seq_len,
549
+ cache_dtype,
550
+ q_dtype,
551
+ is_prefill=True,
552
+ has_sinks=has_sinks)
553
+ decode_use_trtllm = use_trtllm_attention(num_qo_heads,
554
+ num_kv_heads,
555
+ num_decode_tokens,
556
+ max_seq_len,
557
+ cache_dtype,
558
+ q_dtype,
559
+ is_prefill=False,
560
+ has_sinks=has_sinks)
561
562
attn_metadata = FlashInferMetadata(
563
num_actual_tokens=num_actual_tokens,
@@ -654,19 +665,18 @@ class FlashInferImpl(AttentionImpl):
665
raise ValueError(
666
"Sinks must have the same number of heads as the number of "
667
f"heads in the layer. Expected {num_heads}, but got "
657
- f"{sinks.shape[0]}."
658
- )
668
+ f"{sinks.shape[0]}.")
669
self.sinks = sinks
670
661
- self.support_trtllm_attn = (supports_trtllm_attention() and
662
- num_heads % num_kv_heads == 0)
671
+ self.support_trtllm_attn = (supports_trtllm_attention()
672
+ and num_heads % num_kv_heads == 0)
673
self.bmm1_scale: Optional[float] = None
674
self.bmm2_scale: Optional[float] = None
675
676
def fused_output_quant_supported(self, dtype: torch.dtype, static: bool,
677
group_shape: GroupShape):
668
- supported_quant_type = (dtype == FP8_DTYPE and static and
669
- group_shape == GroupShape.PER_TENSOR)
678
+ supported_quant_type = (dtype == FP8_DTYPE and static
679
+ and group_shape == GroupShape.PER_TENSOR)
680
return (self.support_trtllm_attn
681
and self.kv_cache_dtype.startswith("fp8")
682
and supported_quant_type)
@@ -731,7 +741,8 @@ class FlashInferImpl(AttentionImpl):
741
# Insert FP8 quant for query
742
num_tokens, num_heads, head_size = query.shape
743
query, _ = ops.scaled_fp8_quant(
734
- query.reshape((num_tokens, num_heads * head_size)).contiguous(),
744
+ query.reshape(
745
+ (num_tokens, num_heads * head_size)).contiguous(),
746
layer._q_scale)
747
query = query.reshape((num_tokens, num_heads, head_size))
748