A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 17:00 UTC
[Bugfix][Kernel] fix merge attn states when both prefix and suffix are empty (#28181)
Signed-off-by: courage17340 <courage17340@163.com>
courage17340 committed
Nov 6, 2025 at 17:52 UTC
981cadb35c199c3c018474fd5a3d05bfa387610f
1 file changed
+26
csrc/attention/merge_attn_states.cu
+26
@@ -46,6 +46,32 @@ __global__ void merge_attn_states_kernel(
46
s_lse = std::isinf(s_lse) ? -std::numeric_limits<float>::infinity() : s_lse;
47
48
const float max_lse = fmaxf(p_lse, s_lse);
49
+
50
+ /* In certain edge cases, MLA can produce p_lse = s_lse = -inf;
51
+ continuing the pipeline then yields NaN. Root cause: with chunked prefill
52
+ a batch may be split into two chunks; if a request in that batch has no
53
+ prefix hit, every LSE entry for that request’s position is -inf, and at
54
+ this moment we merge cross-attention at first. For now we simply emit
55
+ prefix_output (expected to be all zeros) and prefix_lse (-inf) to fix
56
+ this problem.
57
+ */
58
+ if (std::isinf(max_lse)) {
59
+ if (pack_offset < head_size) {
60
+ // Pack 128b load
61
+ pack_128b_t p_out_pack = reinterpret_cast<const pack_128b_t*>(
62
+ prefix_head_ptr)[pack_offset / pack_size];
63
+
64
+ // Pack 128b storage
65
+ reinterpret_cast<pack_128b_t*>(output_head_ptr)[pack_offset / pack_size] =
66
+ p_out_pack;
67
+ }
68
+ // We only need to write to output_lse once per head.
69
+ if (output_lse != nullptr && pack_idx == 0) {
70
+ output_lse[head_idx * num_tokens + token_idx] = max_lse;
71
+ }
72
+ return;
73
+ }
74
+
75
p_lse = p_lse - max_lse;
76
s_lse = s_lse - max_lse;
77
const float p_se = expf(p_lse);