[Bugfix] Keep a shared MAX_LEN prefill CUDA graph bucket when the graph captures a DP gather (MegaMoE sparse-DP hang) (#37933)

Co-authored-by: shyeh25 <206795756+shyeh25@users.noreply.github.com>
Co-authored-by: Po-Han Huang (NVIDIA) <53919306+nvpohanh@users.noreply.github.com>
This commit is contained in:
YAMY
2026-09-09 10:52:16 -07:00
committed by GitHub
co-authored by shyeh25 Po-Han Huang
parent 0027af2eac
commit 7b791c9534
5 changed files with 106 additions and 11 deletions
+4
View File
@@ -565,6 +565,10 @@ class DpFlags(_FlagGroupBase):
# Hybrid-SSM models materialize idle ranks via the MAX_LEN fabricated-row
# conversion (set when hf_config has hybrid_override_pattern).
max_len_with_idle: bool = False
# Set while the prefill CUDA graph runner captures; latched by the DP
# gather/scatter helpers, whose captured geometry needs one shared bucket.
capturing_prefill_graph: bool = False
prefill_graph_has_dp_gather: bool = False
# DP gathered-buffer allocation metadata (model hidden size / dtype /
# device), set by initialize_dp_attention alongside the flags above.
buffer_hidden_size: Any = None