fix fa3 scale bug (#10961)

phlrain · web-flow · commit 842684706b85 · 2025-08-16T10:30:56.000+08:00
diff --git a/paddlenlp/transformers/deepseek_v2/modeling.py b/paddlenlp/transformers/deepseek_v2/modeling.py
@@ -1306,9 +1306,9 @@ def forward(
         )
 
         q_head_dim = query_states.shape[-1]
-        softmax_scale = softmax_scale * (q_head_dim**0.5)
 
         if FA_VERSION == 2:
+            softmax_scale = softmax_scale * (q_head_dim**0.5)
             query_states = query_states * softmax_scale
             kv_seq_len = value_states.shape[1]
             v_num_heads = value_states.shape[2]

Original file line number	Diff line number	Diff line change
`@@ -1306,9 +1306,9 @@ def forward(`
`1306`	`1306`	`)`
`1307`	`1307`
`1308`	`1308`	`q_head_dim = query_states.shape[-1]`
`1309`		`- softmax_scale = softmax_scale * (q_head_dim**0.5)`
`1310`	`1309`
`1311`	`1310`	`if FA_VERSION == 2:`
	`1311`	`+ softmax_scale = softmax_scale * (q_head_dim**0.5)`
`1312`	`1312`	`query_states = query_states * softmax_scale`
`1313`	`1313`	`kv_seq_len = value_states.shape[1]`
`1314`	`1314`	`v_num_heads = value_states.shape[2]`