A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix] Complete one-shot fused all-reduce PDL at end to avoid NaN (#45448)
Ace Eldeib committed
Jun 17, 2026 at 20:54 UTC
b4092176b9bc76839f76b0e91972a52e8b14ea2b
1 file changed
+11
-1
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+11
-1
@@ -238,7 +238,17 @@ if flashinfer_comm is not None:
238
use_oneshot=use_oneshot,
239
fp32_acc=fp32_acc,
240
weight_bias=weight_bias,
241
- trigger_completion_at_end=num_tokens > PDL_ADVANCE_LAUNCH_TOKENS,
241
+ # The one-shot Lamport all-reduce signals PDL completion before its
242
+ # output buffer is committed when trigger_completion_at_end is
243
+ # False, so the next PDL-launched kernel can read the uninitialized
244
+ # Lamport buffer and produce NaN. This only fires for
245
+ # num_tokens <= PDL_ADVANCE_LAUNCH_TOKENS (the batch=1 / spec-decode
246
+ # shapes, where the one-shot path is always selected). Complete at
247
+ # the end for the one-shot path; the two-shot path is synchronized
248
+ # and keeps the early completion. Related one-shot instability in
249
+ # the same kernel: flashinfer-ai/flashinfer#1223.
250
+ trigger_completion_at_end=use_oneshot
251
+ or num_tokens > PDL_ADVANCE_LAUNCH_TOKENS,
252
)
253
254
def call_trtllm_fused_allreduce_norm_fake(