[DeepSeek-V4] Fix idle-rank dummy-extend sparse-prefill crash under DP breakable CUDA graph (#31705)

2026-07-19 14:42:04 -07:00
parent d4801be447
commit 688a6d23f1
1 changed files with 2 additions and 0 deletions
@@ -94,6 +94,8 @@ def _get_logical_forward_mode(forward_batch: ForwardBatch) -> ForwardMode:
    # from a reused/padded ForwardBatch turn an empty rank into TARGET_VERIFY.
    if forward_batch.forward_mode.is_idle():
        return forward_batch.forward_mode
+    if forward_batch.forward_mode == ForwardMode.EXTEND:
+        return forward_batch.forward_mode
    return (
        getattr(forward_batch, "_original_forward_mode", None)
        or forward_batch.forward_mode