[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:
co-authored by
shyeh25
Po-Han Huang
parent
0027af2eac
commit
7b791c9534
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user