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, 22:35 UTC
[Bugfix] Check NVIDIA artifactory is accessible before using flashinfer cubin kernels (#21893)
Michael Goin committed
Aug 1, 2025 at 08:28 UTC
f81c1bb05504672ddd66905161c6ada549fd4b85
4 files changed
+93
-99
vllm/attention/backends/flashinfer.py
+2
-44
@@ -44,9 +44,9 @@ from vllm.attention.layer import Attention
44
from vllm.attention.ops.paged_attn import PagedAttention
45
from vllm.config import VllmConfig, get_layers_from_vllm_config
46
from vllm.logger import init_logger
47
-from vllm.platforms import current_platform
47
from vllm.utils import (async_tensor_h2d, get_kv_cache_torch_dtype,
48
make_tensor_with_pad)
49
+from vllm.utils.flashinfer import use_trtllm_decode_attention
50
51
logger = init_logger(__name__)
52
@@ -56,7 +56,6 @@ if TYPE_CHECKING:
56
57
58
class FlashInferBackend(AttentionBackend):
59
- cached_sm100a_supported: Optional[bool] = None
59
60
@staticmethod
61
def get_name() -> str:
@@ -123,47 +122,6 @@ class FlashInferBackend(AttentionBackend):
122
else:
123
raise ValueError(f"Unrecognized FP8 dtype: {kv_cache_dtype}")
124
126
- @staticmethod
127
- def use_trtllm_decode_attention(
128
- batch_size: int,
129
- max_seq_len: int,
130
- kv_cache_dtype: str,
131
- num_qo_heads: Optional[int],
132
- num_kv_heads: Optional[int],
133
- attn_head_size: Optional[int],
134
- ) -> bool:
135
- if FlashInferBackend.cached_sm100a_supported is None:
136
- FlashInferBackend.cached_sm100a_supported = (
137
- current_platform.has_device_capability(100))
138
- if not FlashInferBackend.cached_sm100a_supported:
139
- return False
140
- # Check if the dimensions are supported by TRTLLM decode attention
141
- if (attn_head_size is None or num_qo_heads is None
142
- or num_kv_heads is None or num_qo_heads // num_kv_heads > 8
143
- or num_qo_heads % num_kv_heads != 0 or attn_head_size != 128):
144
- return False
145
- env_value = envs.VLLM_USE_TRTLLM_DECODE_ATTENTION
146
- if env_value is not None:
147
- logger.info_once("VLLM_USE_TRTLLM_DECODE_ATTENTION is set to %s",
148
- env_value)
149
- # Environment variable is set - respect it
150
- # Making the conditional check for zero because
151
- # the path is automatically enabled if the batch size condition
152
- # is satisfied.
153
- no_use_trtllm = (env_value == "0")
154
- if not no_use_trtllm:
155
- logger.info_once("Using TRTLLM decode attention.")
156
- return not no_use_trtllm
157
- else:
158
- # Environment variable not set - use auto-detection
159
- use_trtllm = (FlashInferBackend.cached_sm100a_supported
160
- and batch_size <= 256 and max_seq_len < 131072
161
- and kv_cache_dtype == "auto")
162
- if use_trtllm:
163
- logger.warning_once(
164
- "Using TRTLLM decode attention (auto-detected).")
165
- return use_trtllm
166
-
125
126
@dataclass
127
class PerLayerParameters:
@@ -1156,7 +1114,7 @@ class FlashInferImpl(AttentionImpl):
1114
assert decode_meta.decode_wrapper._sm_scale == softmax_scale
1115
# TODO: @pavanimajety Remove this once the switch happens
1116
# inside flashinfer.
1159
- if not FlashInferBackend.use_trtllm_decode_attention(
1117
+ if not use_trtllm_decode_attention(
1118
num_decode_tokens, attn_metadata.max_decode_seq_len,
1119
kv_cache_dtype, attn_metadata.num_qo_heads,
1120
attn_metadata.num_kv_heads, attn_metadata.head_dim):
vllm/utils/flashinfer.py
+80
-1
@@ -10,12 +10,25 @@ import contextlib
10
import functools
11
import importlib
12
import importlib.util
13
-from typing import Any, Callable, NoReturn
13
+import os
14
+from typing import Any, Callable, NoReturn, Optional
15
16
+import requests
17
+
18
+import vllm.envs as envs
19
from vllm.logger import init_logger
20
+from vllm.platforms import current_platform
21
22
logger = init_logger(__name__)
23
24
+# This is the storage path for the cubins, it can be replaced
25
+# with a local path for testing.
26
+# Referenced from https://github.com/flashinfer-ai/flashinfer/blob/0c9a92c3d9a7e043ab6f3f7b2273269caf6ab044/flashinfer/jit/cubin_loader.py#L35 # noqa: E501
27
+FLASHINFER_CUBINS_REPOSITORY = os.environ.get(
28
+ "FLASHINFER_CUBINS_REPOSITORY",
29
+ "https://edge.urm.nvidia.com/artifactory/sw-kernelinferencelibrary-public-generic-local/", # noqa: E501
30
+)
31
+
32
33
@functools.cache
34
def has_flashinfer() -> bool:
@@ -108,6 +121,70 @@ def has_flashinfer_cutlass_fused_moe() -> bool:
121
return True
122
123
124
+@functools.cache
125
+def has_nvidia_artifactory() -> bool:
126
+ """Return ``True`` if NVIDIA's artifactory is accessible.
127
+
128
+ This checks connectivity to the kernel inference library artifactory
129
+ which is required for downloading certain cubin kernels like TRTLLM FHMA.
130
+ """
131
+ try:
132
+ # Use a short timeout to avoid blocking for too long
133
+ response = requests.get(FLASHINFER_CUBINS_REPOSITORY, timeout=5)
134
+ accessible = response.status_code == 200
135
+ if accessible:
136
+ logger.debug_once("NVIDIA artifactory is accessible")
137
+ else:
138
+ logger.warning_once(
139
+ "NVIDIA artifactory returned failed status code: %d",
140
+ response.status_code)
141
+ return accessible
142
+ except Exception as e:
143
+ logger.warning_once("Failed to connect to NVIDIA artifactory: %s", e)
144
+ return False
145
+
146
+
147
+def use_trtllm_decode_attention(
148
+ num_tokens: int,
149
+ max_seq_len: int,
150
+ kv_cache_dtype: str,
151
+ num_qo_heads: Optional[int],
152
+ num_kv_heads: Optional[int],
153
+ attn_head_size: Optional[int],
154
+) -> bool:
155
+ # Requires SM100 and NVIDIA artifactory to be accessible to download cubins
156
+ if not (current_platform.is_device_capability(100)
157
+ and has_nvidia_artifactory()):
158
+ return False
159
+
160
+ # Check if the dimensions are supported by TRTLLM decode attention
161
+ if (attn_head_size is None or num_qo_heads is None or num_kv_heads is None
162
+ or num_qo_heads // num_kv_heads > 8
163
+ or num_qo_heads % num_kv_heads != 0 or attn_head_size != 128):
164
+ return False
165
+
166
+ env_value = envs.VLLM_USE_TRTLLM_DECODE_ATTENTION
167
+ if env_value is not None:
168
+ logger.info_once("VLLM_USE_TRTLLM_DECODE_ATTENTION is set to %s",
169
+ env_value)
170
+ # Environment variable is set - respect it
171
+ # Making the conditional check for zero because
172
+ # the path is automatically enabled if the batch size condition
173
+ # is satisfied.
174
+ no_use_trtllm = (env_value == "0")
175
+ if not no_use_trtllm:
176
+ logger.info_once("Using TRTLLM decode attention.")
177
+ return not no_use_trtllm
178
+ else:
179
+ # Environment variable not set - use auto-detection
180
+ use_trtllm = (num_tokens <= 256 and max_seq_len < 131072
181
+ and kv_cache_dtype == "auto")
182
+ if use_trtllm:
183
+ logger.warning_once(
184
+ "Using TRTLLM decode attention (auto-detected).")
185
+ return use_trtllm
186
+
187
+
188
__all__ = [
189
"has_flashinfer",
190
"flashinfer_trtllm_fp8_block_scale_moe",
@@ -117,4 +194,6 @@ __all__ = [
194
"autotune",
195
"has_flashinfer_moe",
196
"has_flashinfer_cutlass_fused_moe",
197
+ "has_nvidia_artifactory",
198
+ "use_trtllm_decode_attention",
199
]
vllm/v1/attention/backends/flashinfer.py
+3
-46
@@ -17,8 +17,8 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
17
AttentionType)
18
from vllm.config import VllmConfig
19
from vllm.logger import init_logger
20
-from vllm.platforms import current_platform
20
from vllm.utils import cdiv
21
+from vllm.utils.flashinfer import use_trtllm_decode_attention
22
from vllm.v1.attention.backends.flash_attn import use_cascade_attention
23
from vllm.v1.attention.backends.utils import (
24
AttentionMetadataBuilder, CommonAttentionMetadata, get_kv_cache_layout,
@@ -38,7 +38,6 @@ logger = init_logger(__name__)
38
class FlashInferBackend(AttentionBackend):
39
40
accept_output_buffer: bool = True
41
- cached_sm100a_supported: Optional[bool] = None
41
42
@classmethod
43
def get_supported_dtypes(cls) -> list[torch.dtype]:
@@ -98,48 +97,6 @@ class FlashInferBackend(AttentionBackend):
97
raise ValueError(f"Unknown cache layout format {cache_layout}.")
98
return stride_order
99
101
- @staticmethod
102
- def use_trtllm_decode_attention(
103
- batch_size: int,
104
- max_seq_len: int,
105
- kv_cache_dtype: str,
106
- num_qo_heads: int,
107
- num_kv_heads: int,
108
- attn_head_size: int,
109
- ) -> bool:
110
- if FlashInferBackend.cached_sm100a_supported is None:
111
- FlashInferBackend.cached_sm100a_supported = (
112
- current_platform.has_device_capability(100))
113
- if not FlashInferBackend.cached_sm100a_supported:
114
- return False
115
- if (num_qo_heads // num_kv_heads > 8
116
- or num_qo_heads % num_kv_heads != 0 or attn_head_size != 128):
117
- return False
118
- env_value = envs.VLLM_USE_TRTLLM_DECODE_ATTENTION
119
- if env_value is not None:
120
- logger.info_once("VLLM_USE_TRTLLM_DECODE_ATTENTION is set to %s",
121
- env_value)
122
- # Environment variable is set - respect it
123
- # Making the conditional check for zero because
124
- # the path is automatically enabled if the batch size condition
125
- # is satisfied.
126
- no_use_trtllm = env_value == "0"
127
- if not no_use_trtllm:
128
- logger.info_once(
129
- "VLLM_USE_TRTLLM_DECODE_ATTENTION is set to 1, "
130
- "using TRTLLM decode attention.")
131
- return not no_use_trtllm
132
- else:
133
- # Environment variable not set - use auto-detection
134
- # Only supports attention head size of 128
135
- use_trtllm = (FlashInferBackend.cached_sm100a_supported
136
- and batch_size <= 256 and max_seq_len < 131072
137
- and kv_cache_dtype == "auto")
138
- if use_trtllm:
139
- logger.warning_once(
140
- "Using TRTLLM decode attention (auto-detected).")
141
- return use_trtllm
142
-
100
@staticmethod
101
def get_fp8_dtype_for_flashinfer(kv_cache_dtype: str) -> torch.dtype:
102
if kv_cache_dtype in ("fp8", "fp8_e4m3"):
@@ -352,7 +309,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
309
310
if num_decodes > 0:
311
attn_metadata.decode_wrapper = self._get_decode_wrapper()
355
- if not FlashInferBackend.use_trtllm_decode_attention(
312
+ if not use_trtllm_decode_attention(
313
num_decodes, attn_metadata.max_seq_len,
314
self.cache_config.cache_dtype,
315
attn_metadata.num_qo_heads, attn_metadata.num_kv_heads,
@@ -636,7 +593,7 @@ class FlashInferImpl(AttentionImpl):
593
decode_query = query[:num_decode_tokens]
594
assert decode_query.shape[0] == num_decode_tokens
595
assert decode_wrapper is not None
639
- if not FlashInferBackend.use_trtllm_decode_attention(
596
+ if not use_trtllm_decode_attention(
597
attn_metadata.num_decodes, attn_metadata.max_seq_len,
598
self.kv_cache_dtype, attn_metadata.num_qo_heads,
599
attn_metadata.num_kv_heads, attn_metadata.head_dim):
vllm/v1/attention/backends/mla/common.py
+8
-8
@@ -209,6 +209,7 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
209
UnquantizedLinearMethod)
210
from vllm.platforms import current_platform
211
from vllm.utils import cdiv, round_down
212
+from vllm.utils.flashinfer import has_nvidia_artifactory
213
from vllm.v1.attention.backends.utils import (
214
AttentionMetadataBuilder, CommonAttentionMetadata,
215
get_per_layer_parameters, infer_global_hyperparameters,
@@ -379,17 +380,16 @@ M = TypeVar("M", bound=MLACommonMetadata)
380
381
382
def use_flashinfer_prefill() -> bool:
382
- if flashinfer_available and not envs.VLLM_USE_CUDNN_PREFILL:
383
- # For blackwell default to flashinfer prefill if its available since
384
- # its faster than FA2.
385
- return current_platform.has_device_capability(100)
386
- return False
383
+ # For blackwell default to flashinfer prefill if its available since
384
+ # it is faster than FA2.
385
+ return (flashinfer_available and not envs.VLLM_USE_CUDNN_PREFILL
386
+ and current_platform.is_device_capability(100))
387
388
389
def use_cudnn_prefill() -> bool:
390
- if flashinfer_available and envs.VLLM_USE_CUDNN_PREFILL:
391
- return current_platform.has_device_capability(100)
392
- return False
390
+ return (flashinfer_available and envs.VLLM_USE_CUDNN_PREFILL
391
+ and current_platform.is_device_capability(100)
392
+ and has_nvidia_artifactory())
393
394
395
# Currently 394MB, this can be tuned based on GEMM sizes used.