File tree Expand file tree Collapse file tree 1 file changed +6
-7
lines changed Expand file tree Collapse file tree 1 file changed +6
-7
lines changed Original file line number Diff line number Diff line change @@ -285,17 +285,16 @@ def use_trtllm_attention(
285285 # Environment variable not set - use auto-detection
286286 if is_prefill :
287287 # Prefill auto-detection
288- use_trtllm = ( max_seq_len <= 131072 and kv_cache_dtype == "auto" )
288+ use_trtllm = max_seq_len <= 131072 and kv_cache_dtype == "auto"
289289 if use_trtllm :
290- logger .warning_once (
291- "Using TRTLLM prefill attention (auto-detected)." )
290+ logger .warning_once ("Using TRTLLM prefill attention (auto-detected)." )
292291 else :
293292 # Decode auto-detection
294- use_trtllm = (num_tokens <= 256 and max_seq_len <= 131072
295- and kv_cache_dtype == "auto" )
293+ use_trtllm = (
294+ num_tokens <= 256 and max_seq_len <= 131072 and kv_cache_dtype == "auto"
295+ )
296296 if use_trtllm :
297- logger .warning_once (
298- "Using TRTLLM decode attention (auto-detected)." )
297+ logger .warning_once ("Using TRTLLM decode attention (auto-detected)." )
299298 return use_trtllm
300299
301300 # Environment variable is set to 1 - respect it
You can’t perform that action at this time.
0 commit comments