Fix parameter name in flash_fn call

LoserCheems · web-flow · commit 1ce2be0c7289 · 2025-10-10T14:27:06.000+08:00
diff --git a/flash_dmattn/integrations/modeling_flash_dynamic_mask_attention_utils.py b/flash_dmattn/integrations/modeling_flash_dynamic_mask_attention_utils.py
@@ -129,7 +129,7 @@ def _flash_dynamic_mask_attention_forward(
         attention_mask = attention_mask.scatter(-1, topk_indices, topk_values != min_dtype)
 
     out = flash_fn(
-        query_states, key_states, value_states, attn_mask=attention_mask, attn_bias=attention_bias, scale=softmax_scale, is_causal=is_causal
+        query_states, key_states, value_states, attn_mask=attention_mask, attn_bias=attention_bias, softmax_scale=softmax_scale, is_causal=is_causal
     )
 
     return out[0] if isinstance(out, tuple) else out

Original file line number	Diff line number	Diff line change
`@@ -129,7 +129,7 @@ def _flash_dynamic_mask_attention_forward(`
`129`	`129`	`attention_mask = attention_mask.scatter(-1, topk_indices, topk_values != min_dtype)`
`130`	`130`
`131`	`131`	`out = flash_fn(`
`132`		`- query_states, key_states, value_states, attn_mask=attention_mask, attn_bias=attention_bias, scale=softmax_scale, is_causal=is_causal`
	`132`	`+ query_states, key_states, value_states, attn_mask=attention_mask, attn_bias=attention_bias, softmax_scale=softmax_scale, is_causal=is_causal`
`133`	`133`	`)`
`134`	`134`
`135`	`135`	`return out[0] if isinstance(out, tuple) else out`