Fix missing idle-batch handling in prepare_mlp_sync_batch_raw (#25404)

2026-05-24 23:53:54 -07:00
parent 87e69d57c4
commit ca029e816b
1 changed files with 5 additions and 1 deletions
@@ -147,7 +147,11 @@ def prepare_mlp_sync_batch_raw(
    offload_tags: set[str],
 ):
    # Check if other DP workers have running batches
-    if local_batch is None or local_batch.forward_mode.is_prebuilt():
+    if (
+        local_batch is None
+        or local_batch.forward_mode.is_prebuilt()
+        or local_batch.forward_mode.is_idle()
+    ):
        num_tokens = 0
        num_tokens_for_logprob = 0
    elif local_batch.forward_mode.is_decode():