Skip to content

[NVIDIA] flashinfer TRTLLM attention prefill token limit #25998

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
+13 −4
Open

[NVIDIA] flashinfer TRTLLM attention prefill token limit #25998

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
978f575
Fix: Separate prefill and decode logic for TRTLLM attention auto-dete…
jasonlizhengjian Sep 24, 2025
6397563
Merge branch 'vllm-project:main' into fix/flashinfer-prefill-token-limit
jasonlizhengjian Sep 30, 2025
f808ba9
Fix TRTLLM attention logging to show both prefill and decode messages
jasonlizhengjian Oct 1, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 13 additions & 4 deletions vllm/utils/flashinfer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -274,10 +274,19 @@ def use_trtllm_attention(

if force_use_trtllm is None:
# Environment variable not set - use auto-detection
use_trtllm = (num_tokens <= 256 and max_seq_len <= 131072
and kv_cache_dtype == "auto")
if use_trtllm:
logger.warning_once("Using TRTLLM attention (auto-detected).")
if is_prefill:
# Prefill auto-detection
use_trtllm = (max_seq_len <= 131072 and kv_cache_dtype == "auto")
if use_trtllm:
logger.warning_once(
"Using TRTLLM prefill attention (auto-detected).")
else:
# Decode auto-detection
use_trtllm = (num_tokens <= 256 and max_seq_len <= 131072
and kv_cache_dtype == "auto")
if use_trtllm:
logger.warning_once(
"Using TRTLLM decode attention (auto-detected).")
return use_trtllm

# Environment variable is set to 1 - respect it
Expand Down