Fix: explicit not none check for tensors in flash attention (huggingface#39639)

jeffrey-dot-li · zaristei · commit 9896f0df3481 · 2025-09-09T15:37:40.000-07:00
fix: explicit not none check for tensors
diff --git a/src/transformers/modeling_flash_attention_utils.py b/src/transformers/modeling_flash_attention_utils.py
@@ -398,7 +398,7 @@ def _flash_attention_forward(
     query_states, key_states, value_states = fa_peft_integration_check(
         query_states, key_states, value_states, target_dtype
     )
-    use_mask = position_ids is not None or all([cu_seq_lens_q, cu_seq_lens_k, max_length_q, max_length_k])
+    use_mask = position_ids is not None or all(k is not None for k in [cu_seq_lens_q, cu_seq_lens_k, max_length_q, max_length_k])
     if attention_mask is not None:
         q, k, v, idx, (cu_q, cu_k), (mq, mk) = _upad_input(
             query_states, key_states, value_states, attention_mask, query_length, unpad_fn

Original file line number	Diff line number	Diff line change
`@@ -398,7 +398,7 @@ def _flash_attention_forward(`
`398`	`398`	`query_states, key_states, value_states = fa_peft_integration_check(`
`399`	`399`	`query_states, key_states, value_states, target_dtype`
`400`	`400`	`)`
`401`		`- use_mask = position_ids is not None or all([cu_seq_lens_q, cu_seq_lens_k, max_length_q, max_length_k])`
	`401`	`+ use_mask = position_ids is not None or all(k is not None for k in [cu_seq_lens_q, cu_seq_lens_k, max_length_q, max_length_k])`
`402`	`402`	`if attention_mask is not None:`
`403`	`403`	`q, k, v, idx, (cu_q, cu_k), (mq, mk) = _upad_input(`
`404`	`404`	`query_states, key_states, value_states, attention_mask, query_length, unpad_fn`