fix(critical): v2 kernel softmax normalization was commented out — outputs were wrong

qwen3_6_scripts/prefix_prefill.py line 435:
  BEFORE: # acc /= l_i[:, None]  (commented out = BUG)
  AFTER:  acc = acc / l_i[:, None]  (restored)

Impact: _fwd_kernel_flash_attn_v2 was producing unnormalized attention
output — every prefill with context length > BLOCK_M would have had
incorrect softmax weights, causing wrong generation quality. This
directly affects the effect test (偏差 ≤ ±4% benchmark).

Root cause: the v1 kernel (_fwd_kernel) does online normalization
(p_scale = beta/l_i_new), but v2 uses acc_scale = alpha only and
defers normalization to the end. Someone commented out the final
division, breaking v2.

NOTE: The file that actually gets deployed is qwen3_6_scripts/,
NOT vllm/. Previous commits edited vllm/ which has no effect
on the built Docker image.
This commit is contained in:
muh-bot
2026-08-07 02:46:46 +00:00
parent 09e7751d27
commit 0caabf285b

View File

@@ -432,7 +432,12 @@ if triton.__version__ >= "2.1.0":
l_i = l_i_new
m_i = m_i_new
# acc /= l_i[:, None]
# BUG FIX: v2 kernel accumulates unnormalized softmax weights.
# Without this final division, output = sum(softmax_unnorm * V)
# instead of the correct sum(softmax_normalized * V).
# v1 kernel does online normalization inside the loop (p_scale/acc_scale).
# v2 defers normalization — it MUST happen here.
acc = acc / l_i[:, None]
# initialize pointers to output
off_o = (
(cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs +