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
[Bugfix][Flashinfer] fix VLLM_USE_TRTLLM_ATTENTION issue for models with diff hyperparameters (#25924)
Signed-off-by: elvischenv <219235043+elvischenv@users.noreply.github.com>
elvischenv committed
Oct 9, 2025 at 03:54 UTC
b82f4307c9b9910b77cdb5043fc0a81ce4c459a1
2 files changed
+40
-36
vllm/v1/attention/backends/flashinfer.py
+27
-18
@@ -37,7 +37,6 @@ from vllm.utils import cdiv, is_pin_memory_available
37
from vllm.utils.flashinfer import (
38
can_use_trtllm_attention,
39
flashinfer_disable_q_quantization,
40
- supports_trtllm_attention,
40
use_trtllm_attention,
41
)
42
from vllm.v1.attention.backends.utils import (
@@ -323,15 +322,13 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
322
# VLLM_FLASHINFER_DISABLE_Q_QUANTIZATION is set to 1. Otherwise, try to
323
# use fp8 q if kv cache is fp8, and will fall back to model dtype
324
# if TRTLLM attention kernel is not used when building attn metadata
326
- if supports_trtllm_attention() and not flashinfer_disable_q_quantization():
325
+ can_use_trtllm = can_use_trtllm_attention(self.num_qo_heads, self.num_kv_heads)
326
+ if can_use_trtllm and not flashinfer_disable_q_quantization():
327
self.q_data_type = self.kv_cache_dtype
328
else:
329
self.q_data_type = self.model_config.dtype
330
331
- supports_spec_as_decode = can_use_trtllm_attention(
332
- self.num_qo_heads, self.num_kv_heads
333
- )
334
- self._init_reorder_batch_threshold(1, supports_spec_as_decode)
331
+ self._init_reorder_batch_threshold(1, supports_spec_as_decode=can_use_trtllm)
332
333
self._cascade_wrapper = None # Wrapper for cascade attention
334
@@ -344,7 +341,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
341
self.window_left = self.global_hyperparameters.window_left
342
self.logits_soft_cap = self.global_hyperparameters.logits_soft_cap
343
self.has_sinks = self.global_hyperparameters.has_sinks
347
- if self.has_sinks and not supports_trtllm_attention():
344
+ if self.has_sinks and not can_use_trtllm:
345
raise NotImplementedError(
346
"FlashInfer backend currently does not support attention "
347
"sinks, please use trtllm on blackwell or flash attention on "
@@ -548,16 +545,30 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
545
has_sinks=self.has_sinks,
546
has_spec=uses_spec_reorder,
547
)
551
- if self.has_sinks and not (prefill_use_trtllm and decode_use_trtllm):
552
- raise NotImplementedError(
553
- "FlashInfer backend currently does not support attention "
554
- "sinks, please use trtllm on blackwell or flash attention on "
555
- "earlier GPUs."
556
- )
548
558
- # If TRTLLM attention is not used, the q quantization is not supported.
559
- # Fall back to use model dtype.
549
if not (prefill_use_trtllm and decode_use_trtllm):
550
+ if self.has_sinks:
551
+ raise NotImplementedError(
552
+ "FlashInfer backend currently does not support attention "
553
+ "sinks, please use trtllm on blackwell or flash attention "
554
+ "on earlier GPUs."
555
+ )
556
+
557
+ if not self.global_hyperparameters.has_same_window_lefts:
558
+ raise ValueError(
559
+ "Window left is not the same for all layers. "
560
+ "One potential fix is to set disable_sliding_window=True"
561
+ )
562
+
563
+ assert self.global_hyperparameters.has_same_all_params, (
564
+ "FlashInfer backend currently only supports models in which "
565
+ "all layers share the same values for the following "
566
+ "hyperparameters: `window_left`, `logits_soft_cap`, "
567
+ "`sm_scale`."
568
+ )
569
+
570
+ # The q quantization is not supported for non-trtllm attention,
571
+ # fall back to model dtype.
572
self.q_data_type = self.model_config.dtype
573
574
attn_metadata = FlashInferMetadata(
@@ -772,9 +783,7 @@ class FlashInferImpl(AttentionImpl):
783
)
784
self.sinks = sinks
785
775
- self.support_trtllm_attn = (
776
- supports_trtllm_attention() and num_heads % num_kv_heads == 0
777
- )
786
+ self.support_trtllm_attn = can_use_trtllm_attention(num_heads, num_kv_heads)
787
self.bmm1_scale: float | None = None
788
self.bmm2_scale: float | None = None
789
self.o_sf_scale: float | None = None
vllm/v1/attention/backends/utils.py
+13
-18
@@ -4,7 +4,7 @@ import abc
4
import enum
5
import functools
6
from abc import abstractmethod
7
-from dataclasses import dataclass, fields, make_dataclass
7
+from dataclasses import dataclass, field, fields, make_dataclass
8
from typing import (
9
TYPE_CHECKING,
10
Any,
@@ -233,7 +233,7 @@ class AttentionCGSupport(enum.Enum):
233
"""Cudagraph always supported; supports mixed-prefill-decode"""
234
UNIFORM_BATCH = 2
235
"""Cudagraph supported for batches the only contain query lengths that are
236
- the same, this can be used for spec-decode
236
+ the same, this can be used for spec-decode
237
i.e. "decodes" are 1 + num_speculative_tokens"""
238
UNIFORM_SINGLE_TOKEN_DECODE = 1
239
"""Cudagraph supported for batches the only contain query_len==1 decodes"""
@@ -395,6 +395,9 @@ class PerLayerParameters:
395
logits_soft_cap: Optional[float]
396
sm_scale: float
397
has_sinks: bool = False
398
+ # has same params for all layers
399
+ has_same_window_lefts: Optional[bool] = field(default=None, compare=False)
400
+ has_same_all_params: Optional[bool] = field(default=None, compare=False)
401
402
403
def get_per_layer_parameters(
@@ -446,20 +449,12 @@ def infer_global_hyperparameters(
449
param_sets = list(per_layer_params.values())
450
global_params = param_sets[0]
451
449
- # trtllm attention doesn't need global hyper params so disable the check
450
- if not envs.VLLM_USE_TRTLLM_ATTENTION:
451
- for params in param_sets:
452
- if params.window_left != global_params.window_left:
453
- raise ValueError(
454
- "Window left is not the same for all layers. "
455
- "One potential fix is to set disable_sliding_window=True"
456
- )
457
- assert params == global_params, (
458
- "FlashInfer backend currently only supports models in which all"
459
- "layers share the same values "
460
- "for the following hyperparameters:"
461
- "`window_left`, `logits_soft_cap`, `sm_scale`."
462
- )
452
+ global_params.has_same_window_lefts = all(
453
+ params.window_left == global_params.window_left for params in param_sets
454
+ )
455
+ global_params.has_same_all_params = all(
456
+ params == global_params for params in param_sets
457
+ )
458
459
return global_params
460
@@ -925,8 +920,8 @@ def create_fast_prefill_custom_backend(
920
):
921
def __init__(self, metadata, common_attn_metadata):
922
# Shallow copy all fields in metadata cls
928
- for field in fields(metadata.__class__):
929
- setattr(self, field.name, getattr(metadata, field.name))
923
+ for _field in fields(metadata.__class__):
924
+ setattr(self, _field.name, getattr(metadata, _field.name))
925
926
# Set additional fields that will be used in model code
927
assert (