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
[Kernel] Simplify `get_kv_cache_layout` and cache `use_trtllm_attention` env-dependent bit (#22735)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Aug 16, 2025 at 02:14 UTC
070da660c1bf9e7a7be8b9efeff4b06f91c7342f
2 files changed
+42
-22
vllm/utils/flashinfer.py
+31
-15
@@ -148,6 +148,31 @@ def has_nvidia_artifactory() -> bool:
148
return False
149
150
151
+@functools.cache
152
+def supports_trtllm_attention() -> tuple[bool, Optional[str]]:
153
+ """Cache result which only depends on the environment"""
154
+ # This is a lambda, call it once
155
+ env_value = envs.VLLM_USE_TRTLLM_ATTENTION
156
+
157
+ # Requires SM100 and NVIDIA artifactory to be accessible to download cubins
158
+ if not (current_platform.is_device_capability(100)
159
+ and has_nvidia_artifactory()):
160
+ return False, env_value
161
+
162
+ if env_value is not None:
163
+ logger.info_once("VLLM_USE_TRTLLM_ATTENTION is set to %s", env_value)
164
+ # Environment variable is set - respect it
165
+ # Making the conditional check for zero because
166
+ # the path is automatically enabled if the batch size condition
167
+ # is satisfied.
168
+ use_trtllm = (env_value == "1")
169
+ if use_trtllm:
170
+ logger.info_once("Using TRTLLM attention.")
171
+ return use_trtllm, env_value
172
+
173
+ return True, None
174
+
175
+
176
def use_trtllm_attention(
177
num_tokens: int,
178
max_seq_len: int,
@@ -157,9 +182,8 @@ def use_trtllm_attention(
182
attn_head_size: Optional[int],
183
has_sinks: bool = False,
184
) -> bool:
160
- # Requires SM100 and NVIDIA artifactory to be accessible to download cubins
161
- if not (current_platform.is_device_capability(100)
162
- and has_nvidia_artifactory()):
185
+ use_trtllm, env_value = supports_trtllm_attention()
186
+ if not use_trtllm:
187
return False
188
189
# Check if the dimensions are supported by TRTLLM decode attention
@@ -174,18 +198,7 @@ def use_trtllm_attention(
198
"Using TRTLLM attention (required for attention sinks).")
199
return True
200
177
- env_value = envs.VLLM_USE_TRTLLM_ATTENTION
178
- if env_value is not None:
179
- logger.info_once("VLLM_USE_TRTLLM_ATTENTION is set to %s", env_value)
180
- # Environment variable is set - respect it
181
- # Making the conditional check for zero because
182
- # the path is automatically enabled if the batch size condition
183
- # is satisfied.
184
- use_trtllm = (env_value == "1")
185
- if use_trtllm:
186
- logger.info_once("Using TRTLLM attention.")
187
- return use_trtllm
188
- else:
201
+ if env_value is None:
202
# Environment variable not set - use auto-detection
203
use_trtllm = (num_tokens <= 256 and max_seq_len < 131072
204
and kv_cache_dtype == "auto")
@@ -193,6 +206,9 @@ def use_trtllm_attention(
206
logger.warning_once("Using TRTLLM attention (auto-detected).")
207
return use_trtllm
208
209
+ # Environment variable is set to 1 - respect it
210
+ return True
211
+
212
213
if has_flashinfer():
214
vllm/v1/attention/backends/utils.py
+11
-7
@@ -248,19 +248,23 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):
248
249
@functools.lru_cache
250
def get_kv_cache_layout():
251
+ # Format specified by the code.
252
global _KV_CACHE_LAYOUT_OVERRIDE
252
- # Override with format specified by the user.
253
+
254
+ if _KV_CACHE_LAYOUT_OVERRIDE is not None:
255
+ cache_layout = _KV_CACHE_LAYOUT_OVERRIDE
256
+ logger.info_once("`_KV_CACHE_LAYOUT_OVERRIDE` variable detected. " \
257
+ "Setting KV cache layout to %s.", cache_layout)
258
+ return cache_layout
259
+
260
+ # Format specified by the user.
261
cache_layout = envs.VLLM_KV_CACHE_LAYOUT
262
+ # When neither the user nor the override specified a layout, get default
263
if cache_layout is None:
255
- if envs.VLLM_USE_TRTLLM_ATTENTION:
256
- cache_layout = "HND"
257
- else:
258
- cache_layout = get_kv_connector_cache_layout()
264
+ cache_layout = get_kv_connector_cache_layout()
265
else:
266
logger.info_once("`VLLM_KV_CACHE_LAYOUT` environment variable " \
267
"detected. Setting KV cache layout to %s.", cache_layout)
262
- if _KV_CACHE_LAYOUT_OVERRIDE is not None:
263
- cache_layout = _KV_CACHE_LAYOUT_OVERRIDE
268
return cache_layout
269
270