[Feat]: support multistream for CAMP2p - #146
Conversation
yujuancao07
changed the title
|
multistream: |
| # Keep all attention DP ranks aligned before entering model forward. | ||
| # In graph mode this is immediately before graph replay, preventing a | ||
| # late attention rank from surfacing as A2E/Dispatch waits on FFN. | ||
| dist.barrier(group=get_dp_group().cpu_group) |
|
|
||
|
|
||
| def npu_stream_switch_within_graph( | ||
| current_stream: Any, |
There was a problem hiding this comment.
Type do not use Any, Alse please add docstring.
jiangkuaixue123
left a comment
jiangkuaixue123
left a comment
There was a problem hiding this comment.
Could we simplify the nested condition here?
| # therefore treat the per-stage metadata list as a single-batch dict. | ||
| update_standard_full_graph = forward_context.ubatch_slices is None | ||
|
|
||
| if self.enable_enpu: |
There was a problem hiding this comment.
Could we combine these two conditions to reduce the nesting, for example, if self.enable_enpu and update_standard_full_graph:?
jiangkuaixue123
left a comment
jiangkuaixue123
left a comment
There was a problem hiding this comment.
Please consider encapsulating the FFN multistream lifecycle to keep the runner maintainable.
| ) | ||
| num_tokens = _ffn_token_count_for_rank(self.connector, num_tokens_across_dp) | ||
| rank_ffn_output = None | ||
| multistream_enabled = self.ffn_multistream_enabled and num_stages > 1 |
There was a problem hiding this comment.
The multistream lifecycle is currently spread across __init__, this loop, and _send_ffn_output (stream/event allocation, resizing, stage-reuse waits, recorded state, final draining, and connector arguments). Could we encapsulate it in a small helper such as FFNMultistreamManager in npu/multistream.py? It could own the stream and events and expose operations such as begin(stage_ids, num_stages), wait_before_reuse(stage_idx), send_context(layer_idx, stage_idx), and wait_all(). This would keep _ffn_forward focused on model execution and let _send_ffn_output accept one typed context instead of three loosely related arguments, making the synchronization invariants easier to test and maintain.
2 participants
Purpose
Add configurable multi-stream execution to the Ascend CAMP2P connector
so that A2F/F2A communication can overlap with model computation.
Attention-side and FFN-side multi-stream execution can be enabled independently
through
connector_extra_config.Issue
Scope
This feature applies only to
CAMP2pAFDConnectoron NPU.In scope:
the current graph runtime.
Out of scope:
CAMAsyncAFDConnector.P2pNcclAFDConnector.Implementation Notes
The implementation applies only to
CAMP2pAFDConnector.Both switches are located under CAMP2P
connector_extra_config:{ "connector_extra_config": { "is_attn_multistream": true, "is_ffn_multistream": true, "attn_core_num": 8, "ffn_core_num": 8 } }Test Plan
Focused multi-stream tests:
pytest -q tests/unit -k multistream -vvTest Result