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
[NVIDIA] flashinfer TRTLLM attention prefill token limit (#25998)
Signed-off-by: jasonlizhengjian <jason.li@centml.ai> Signed-off-by: jasonlizhengjian <jasonlizhengjian@gmail.com>
Jason Li committed
Oct 5, 2025 at 16:24 UTC
6b6e98775f2452a105fa2b0b4b318ffe0cffb06d
1 file changed
+12
-5
vllm/utils/flashinfer.py
+12
-5
@@ -283,11 +283,18 @@ def use_trtllm_attention(
283
284
if force_use_trtllm is None:
285
# Environment variable not set - use auto-detection
286
- use_trtllm = (
287
- num_tokens <= 256 and max_seq_len <= 131072 and kv_cache_dtype == "auto"
288
- )
289
- if use_trtllm:
290
- logger.warning_once("Using TRTLLM attention (auto-detected).")
286
+ if is_prefill:
287
+ # Prefill auto-detection
288
+ use_trtllm = max_seq_len <= 131072 and kv_cache_dtype == "auto"
289
+ if use_trtllm:
290
+ logger.warning_once("Using TRTLLM prefill attention (auto-detected).")
291
+ else:
292
+ # Decode auto-detection
293
+ use_trtllm = (
294
+ num_tokens <= 256 and max_seq_len <= 131072 and kv_cache_dtype == "auto"
295
+ )
296
+ if use_trtllm:
297
+ logger.warning_once("Using TRTLLM decode attention (auto-detected).")
298
return use_trtllm
299
300
# Environment variable is set to 1 - respect it