diff --git a/afd_plugin/connectors/npu/camp2p.py b/afd_plugin/connectors/npu/camp2p.py index d1831723..a000c889 100644 --- a/afd_plugin/connectors/npu/camp2p.py +++ b/afd_plugin/connectors/npu/camp2p.py @@ -51,6 +51,9 @@ send_control_payload, ) from afd_plugin.distributed import init_afd_process_group, topology_from_config +from afd_plugin.v1.worker.npu.multistream import ( + npu_stream_switch_within_graph, +) if TYPE_CHECKING: from vllm.config import VllmConfig @@ -63,6 +66,8 @@ "core_num", "attn_core_num", "ffn_core_num", + "is_attn_multistream", + "is_ffn_multistream", "compute_gate_on_attention", "quant_mode", }, @@ -77,6 +82,10 @@ class CAMP2PExtraInfo(ConnectorExtraInfo): core_num: Default number of AIV cores used by each AFD role. attn_core_num: Optional Attention-role override for ``core_num``. ffn_core_num: Optional FFN-role override for ``core_num``. + is_attn_multistream: Run Attention-side A2F communication on a + dedicated NPU stream. + is_ffn_multistream: Run FFN-side F2A communication on a dedicated NPU + stream. compute_gate_on_attention: Whether Attention computes MoE gate outputs. quant_mode: CAM quantization mode; the current runtime supports only 0. """ @@ -84,6 +93,8 @@ class CAMP2PExtraInfo(ConnectorExtraInfo): core_num: int = 8 attn_core_num: int | None = None ffn_core_num: int | None = None + is_attn_multistream: bool = False + is_ffn_multistream: bool = False compute_gate_on_attention: bool = False quant_mode: int = 0 @@ -117,6 +128,14 @@ def from_mapping(cls, raw: Mapping[str, Any] | None) -> CAMP2PExtraInfo: raw.get("ffn_core_num"), field_name="ffn_core_num", ), + is_attn_multistream=coerce_extra_bool( + raw.get("is_attn_multistream", False), + field_name="is_attn_multistream", + ), + is_ffn_multistream=coerce_extra_bool( + raw.get("is_ffn_multistream", False), + field_name="is_ffn_multistream", + ), compute_gate_on_attention=coerce_extra_bool( raw.get("compute_gate_on_attention", False), field_name="compute_gate_on_attention", @@ -145,6 +164,8 @@ def validate_supported(self) -> None: def to_mapping(self) -> dict[str, Any]: result: dict[str, Any] = { "core_num": self.core_num, + "is_attn_multistream": self.is_attn_multistream, + "is_ffn_multistream": self.is_ffn_multistream, "compute_gate_on_attention": self.compute_gate_on_attention, "quant_mode": self.quant_mode, } @@ -595,24 +616,35 @@ def send_ffn_output( if states.atten_batch_size is None: raise RuntimeError("CAMP2P FFN side is missing A2E atten_batch_size") ubatch_idx = int(kwargs.get("ubatch_idx", context.metadata.stage_idx)) + multistream_enable = bool(kwargs.get("multistream_enable", False)) + comm_stream = kwargs.get("comm_stream") + comm_event = kwargs.get("comm_event") group_ep = _get_group_ep( ubatch_idx, self.hccl_comm_name, self.hccl_comm_name2, self.hccl_comm_name3, ) - torch.ops.afd_ascend.e2a( - ffn_output, - states.atten_batch_size, - states.batch_size, - states.h, - states.k, - self.ffn_size, - self.attn_size, - self.world_rank, - group_ep, - states.aiv_num, - ) + current_stream = torch.npu.current_stream() + with npu_stream_switch_within_graph( + current_stream, + comm_stream, + multistream_enable, + ): + torch.ops.afd_ascend.e2a( + ffn_output, + states.atten_batch_size, + states.batch_size, + states.h, + states.k, + self.ffn_size, + self.attn_size, + self.world_rank, + group_ep, + states.aiv_num, + ) + if multistream_enable and comm_event is not None: + comm_event.record(comm_stream) return None @@ -845,20 +877,34 @@ def send_attn_output_impl( hccl_comm_name3, ) - outputs = torch.ops.afd_ascend.a2e( - hidden_states, - None, - None, - transfer_state.batch_size, - transfer_state.h, - transfer_state.k, - ffn_size, - attn_size, - world_rank, - group_ep, - transfer_state.aiv_num, - compute_gate, + forward_context = get_forward_context() + multistream_enable = bool( + getattr(forward_context, "afd_multistream_enabled", False) ) + comm_stream = getattr(forward_context, "afd_comm_stream", None) + comm_event = getattr(forward_context, "afd_comm_event", None) + current_stream = torch.npu.current_stream() + with npu_stream_switch_within_graph( + current_stream, + comm_stream, + multistream_enable, + ): + outputs = torch.ops.afd_ascend.a2e( + hidden_states, + None, + None, + transfer_state.batch_size, + transfer_state.h, + transfer_state.k, + ffn_size, + attn_size, + world_rank, + group_ep, + transfer_state.aiv_num, + compute_gate, + ) + if multistream_enable and comm_event is not None: + comm_event.record(comm_stream) transfer_state.atten_batch_size = outputs[3] forward_context = get_forward_context() forward_context.cam_afdtransfer_state = transfer_state @@ -907,6 +953,12 @@ def recv_ffn_output_impl( hccl_comm_name2, hccl_comm_name3, ) + forward_context = get_forward_context() + if bool(getattr(forward_context, "afd_multistream_enabled", False)): + comm_event = getattr(forward_context, "afd_comm_event", None) + if comm_event is None: + raise RuntimeError("CAMP2P Attention multistream requires an event") + comm_event.wait(torch.npu.current_stream()) output = torch.ops.afd_ascend.e2a( ref_tensor, transfer_state.atten_batch_size, diff --git a/afd_plugin/v1/worker/npu/attention_model_runner.py b/afd_plugin/v1/worker/npu/attention_model_runner.py index dcc71974..6de7416c 100644 --- a/afd_plugin/v1/worker/npu/attention_model_runner.py +++ b/afd_plugin/v1/worker/npu/attention_model_runner.py @@ -161,6 +161,12 @@ def _model_forward(self, *args: Any, **kwargs: Any) -> Any: self._install_afd_metadata_on_forward_context(forward_context) self._install_async_moe_ubatch_metadata_on_forward_context(forward_context) + # 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. + if self.afd_config.connector == "CAMP2pAFDConnector": + dist.barrier(group=get_dp_group().cpu_group) + ( num_tokens_padded, input_ids, @@ -180,20 +186,29 @@ def _model_forward(self, *args: Any, **kwargs: Any) -> Any: } run_model = partial(self.model, **model_inputs) + # The Ascend ubatch wrapper captures and replays all stage forwards in + # one NPUGraph. Its stage-local contexts intentionally use + # CUDAGraphMode.NONE, so they do not register entries in Ascend's + # standard full-graph parameter tables. Updating the outer context would + # 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: - self._update_full_graph_params_if_needed( - forward_context, - num_tokens_padded, - positions, - ) + if update_standard_full_graph: + self._update_full_graph_params_if_needed( + forward_context, + num_tokens_padded, + positions, + ) hidden_states = run_model() else: hidden_states = run_model() - self._update_full_graph_params_if_needed( - forward_context, - num_tokens_padded, - positions, - ) + if update_standard_full_graph: + self._update_full_graph_params_if_needed( + forward_context, + num_tokens_padded, + positions, + ) if ( forward_context.flash_comm_v1_enabled diff --git a/afd_plugin/v1/worker/npu/ffn_model_runner.py b/afd_plugin/v1/worker/npu/ffn_model_runner.py index c14f67d5..20dbf00d 100644 --- a/afd_plugin/v1/worker/npu/ffn_model_runner.py +++ b/afd_plugin/v1/worker/npu/ffn_model_runner.py @@ -32,6 +32,7 @@ AFDForwardContextMetadata, AFDTransferContext, ) +from afd_plugin.connectors.npu.camp2p import CAMP2PExtraInfo from afd_plugin.v1.worker.attention_model_runner import ( _resolve_world_ranks, _with_dp_derived_afd_rank, @@ -76,6 +77,20 @@ def __init__(self, vllm_config: VllmConfig, device: object) -> None: self.afd_config, ) self.num_layers = int(self.model_config.hf_config.num_hidden_layers) + connector_extra_info = self.connector.extra_info + self.ffn_multistream_enabled = ( + isinstance(connector_extra_info, CAMP2PExtraInfo) + and connector_extra_info.is_ffn_multistream + ) + configured_ubatches = self.parallel_config.num_ubatches or 1 + self.ffn_comm_stream = ( + torch.npu.Stream(device=device) if self.ffn_multistream_enabled else None + ) + self.ffn_comm_events = ( + [torch.npu.Event() for _ in range(configured_ubatches)] + if self.ffn_multistream_enabled + else [] + ) self.use_aclgraph = _use_npu_aclgraph(vllm_config, self) self._acl_graphs: dict[tuple, dict[str, Any]] = {} self.graph_pool = ( @@ -223,6 +238,11 @@ def _ffn_forward( ) 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 + required_event_count = max(stage_ids) + 1 + while multistream_enabled and len(self.ffn_comm_events) < required_event_count: + self.ffn_comm_events.append(torch.npu.Event()) + event_recorded = [False] * required_event_count # Build DP-level token counts for vLLM's forward context. # num_tokens_across_dp has ffn_size entries (AFD-level, one per @@ -242,7 +262,12 @@ def _ffn_forward( aclgraph_runtime_mode=aclgraph_runtime_mode, ) as forward_context: for layer_idx in _ffn_layer_indices(self): + layer_multistream = multistream_enabled and layer_idx > 0 for stage_idx in stage_ids: + if multistream_enabled and event_recorded[stage_idx]: + self.ffn_comm_events[stage_idx].wait( + torch.npu.current_stream(), + ) payload = self.connector.recv_attn_output( ubatch_idx=stage_idx, layer_idx=layer_idx, @@ -268,7 +293,20 @@ def _ffn_forward( rank_ffn_output, context, stage_idx=stage_idx, + multistream_enable=layer_multistream, + comm_stream=self.ffn_comm_stream, + comm_event=self.ffn_comm_events[stage_idx] + if layer_multistream + else None, ) + if layer_multistream: + event_recorded[stage_idx] = True + + if multistream_enabled: + current_stream = torch.npu.current_stream() + for stage_idx in stage_ids: + if event_recorded[stage_idx]: + self.ffn_comm_events[stage_idx].wait(current_stream) return rank_ffn_output def _ffn_forward_connector_driven(self) -> Any: @@ -435,20 +473,27 @@ def _send_ffn_output( context: AFDTransferContext, *, stage_idx: int, + multistream_enable: bool = False, + comm_stream=None, + comm_event=None, ) -> None: - if not isinstance(ffn_output, AFDF2ATransferPayload): - connector.send_ffn_output( - ffn_output, - context, - ubatch_idx=stage_idx, + kwargs: dict[str, object] = {"ubatch_idx": stage_idx} + if multistream_enable: + kwargs.update( + multistream_enable=True, + comm_stream=comm_stream, + comm_event=comm_event, ) - return - kwargs: dict[str, object] = {"ubatch_idx": stage_idx} - if ffn_output.shared_output is not None: - kwargs["expand_x_shared"] = ffn_output.shared_output + if isinstance(ffn_output, AFDF2ATransferPayload): + output_to_send = ffn_output.routed_output + if ffn_output.shared_output is not None: + kwargs["expand_x_shared"] = ffn_output.shared_output + else: + output_to_send = ffn_output + connector.send_ffn_output( - ffn_output.routed_output, + output_to_send, context, **kwargs, ) diff --git a/afd_plugin/v1/worker/npu/forward_context.py b/afd_plugin/v1/worker/npu/forward_context.py index 522acaea..8b7d533b 100644 --- a/afd_plugin/v1/worker/npu/forward_context.py +++ b/afd_plugin/v1/worker/npu/forward_context.py @@ -27,6 +27,9 @@ def create_ascend_forward_context( cudagraph_runtime_mode: CUDAGraphMode | None = None, batch_descriptor: BatchDescriptor | None = None, skip_compiled: bool = False, + afd_comm_stream=None, + afd_comm_event=None, + afd_multistream_enabled: bool = False, ) -> ForwardContext: if cudagraph_runtime_mode is None: cudagraph_runtime_mode = CUDAGraphMode.NONE @@ -130,6 +133,9 @@ def create_ascend_forward_context( new_forward_context.mc2_mask = mc2_mask new_forward_context.dbo_enabled = True + new_forward_context.afd_comm_stream = afd_comm_stream + new_forward_context.afd_comm_event = afd_comm_event + new_forward_context.afd_multistream_enabled = afd_multistream_enabled return new_forward_context diff --git a/afd_plugin/v1/worker/npu/multistream.py b/afd_plugin/v1/worker/npu/multistream.py new file mode 100644 index 00000000..8abd5cb9 --- /dev/null +++ b/afd_plugin/v1/worker/npu/multistream.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""NPU stream helpers used by CAMP2P communication.""" + +from __future__ import annotations + +from contextlib import AbstractContextManager, nullcontext + +import torch + + +def npu_stream_switch_within_graph( + current_stream: torch.npu.Stream | None, + target_stream: torch.npu.Stream | None, + enabled: bool, +) -> AbstractContextManager[None]: + """Switch to ``target_stream`` after ``current_stream`` when enabled. + + Return a no-op context when multi-stream execution is disabled. Both + streams are required when it is enabled. + """ + if not enabled: + return nullcontext() + if current_stream is None or target_stream is None: + raise RuntimeError( + "CAMP2P multistream requires compute and communication streams", + ) + target_stream.wait_stream(current_stream) + return torch.npu.stream(target_stream) + + +__all__ = ["npu_stream_switch_within_graph"] diff --git a/afd_plugin/v1/worker/npu/npu_ubatch_wrapper.py b/afd_plugin/v1/worker/npu/npu_ubatch_wrapper.py index d0632511..7e0b1323 100644 --- a/afd_plugin/v1/worker/npu/npu_ubatch_wrapper.py +++ b/afd_plugin/v1/worker/npu/npu_ubatch_wrapper.py @@ -28,6 +28,9 @@ from vllm_ascend.compilation.acl_graph import ACLGraphWrapper from vllm_ascend.utils import enable_sp +from afd_plugin.config import parse_afd_config +from afd_plugin.connectors import AFDConnectorFactory +from afd_plugin.connectors.npu.camp2p import CAMP2PExtraInfo from afd_plugin.v1.worker.npu.forward_context import ( create_ascend_forward_context, ) @@ -63,7 +66,25 @@ def __init__( self.runnable = runnable self.vllm_config = vllm_config self.compilation_config = vllm_config.compilation_config - self.comm_stream = torch.npu.Stream(device=device) + afd_config = parse_afd_config(vllm_config) + self.attn_multistream_enabled = False + if afd_config.connector == "CAMP2pAFDConnector": + extra_info = AFDConnectorFactory.parse_connector_extra_info( + afd_config.connector, + vllm_config, + ) + if not isinstance(extra_info, CAMP2PExtraInfo): + raise TypeError("CAMP2P connector returned unexpected extra config") + self.attn_multistream_enabled = extra_info.is_attn_multistream + self.comm_stream = ( + torch.npu.Stream(device=device) if self.attn_multistream_enabled else None + ) + configured_ubatches = vllm_config.parallel_config.num_ubatches or 1 + self.comm_events = ( + [torch.npu.Event() for _ in range(configured_ubatches)] + if self.attn_multistream_enabled + else [] + ) self.ready_barrier = threading.Barrier(3) self.cudagraphs: dict[int, AscendNPUGraphMetaData] = {} self.cudagraph_wrapper = None @@ -193,6 +214,10 @@ def _make_ubatch_metadata( cudagraph_runtime_mode, ) -> list[AscendUbatchMetadata]: cur_forward_context = get_forward_context() + while self.attn_multistream_enabled and len(self.comm_events) < len( + ubatch_slices + ): + self.comm_events.append(torch.npu.Event()) forward_contexts = [] for i, _ubatch_slice in enumerate(ubatch_slices): forward_contexts.append( @@ -208,6 +233,11 @@ def _make_ubatch_metadata( cudagraph_runtime_mode=cudagraph_runtime_mode, ubatch_num=i, skip_compiled=cur_forward_context.skip_compiled, + afd_comm_stream=self.comm_stream, + afd_comm_event=self.comm_events[i] + if self.attn_multistream_enabled + else None, + afd_multistream_enabled=self.attn_multistream_enabled, ) ) diff --git a/tests/unit/connectors/test_camp2p_connector.py b/tests/unit/connectors/test_camp2p_connector.py index 757b0538..592bf402 100644 --- a/tests/unit/connectors/test_camp2p_connector.py +++ b/tests/unit/connectors/test_camp2p_connector.py @@ -171,6 +171,24 @@ def test_camp2p_extra_info_coerces_integer_bool_values(): ) +def test_camp2p_extra_info_parses_independent_multistream_switches(): + defaults = CAMP2PExtraInfo.from_mapping({}) + attn = CAMP2PExtraInfo.from_mapping({"is_attn_multistream": True}) + ffn = CAMP2PExtraInfo.from_mapping({"is_ffn_multistream": True}) + + assert defaults.is_attn_multistream is False + assert defaults.is_ffn_multistream is False + assert attn.is_attn_multistream is True + assert attn.is_ffn_multistream is False + assert ffn.is_attn_multistream is False + assert ffn.is_ffn_multistream is True + + +def test_camp2p_extra_info_rejects_invalid_multistream_switch(): + with pytest.raises(TypeError, match="is_attn_multistream must be a boolean"): + CAMP2PExtraInfo.from_mapping({"is_attn_multistream": "sometimes"}) + + def test_camp2p_connector_uses_role_specific_core_num(monkeypatch): torch = pytest.importorskip("torch") monkeypatch.setattr( diff --git a/tests/unit/v1/worker/test_npu_multistream.py b/tests/unit/v1/worker/test_npu_multistream.py new file mode 100644 index 00000000..546791be --- /dev/null +++ b/tests/unit/v1/worker/test_npu_multistream.py @@ -0,0 +1,53 @@ +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest + +pytest.importorskip("torch") + +from afd_plugin.v1.worker.npu import multistream + + +class _FakeStream: + def __init__(self): + self.waited_for = None + + def wait_stream(self, stream): + self.waited_for = stream + + +def test_stream_switch_waits_for_compute_stream(monkeypatch): + comm_stream = _FakeStream() + monkeypatch.setattr( + multistream.torch, + "npu", + SimpleNamespace(stream=nullcontext), + raising=False, + ) + + with multistream.npu_stream_switch_within_graph( + "compute", + comm_stream, + enabled=True, + ): + pass + + assert comm_stream.waited_for == "compute" + + +def test_stream_switch_disabled_accepts_missing_streams(): + with multistream.npu_stream_switch_within_graph(None, None, enabled=False): + pass + + +@pytest.mark.parametrize( + ("compute_stream", "comm_stream"), + [(None, object()), (object(), None)], +) +def test_stream_switch_requires_both_streams(compute_stream, comm_stream): + with pytest.raises(RuntimeError, match="requires compute and communication"): + multistream.npu_stream_switch_within_graph( + compute_stream, + comm_stream, + enabled=True, + ) diff --git a/tests/unit/v1/worker/test_npu_runtime.py b/tests/unit/v1/worker/test_npu_runtime.py index ad323fd2..44c08a5b 100644 --- a/tests/unit/v1/worker/test_npu_runtime.py +++ b/tests/unit/v1/worker/test_npu_runtime.py @@ -241,9 +241,24 @@ def _new_ffn_runner(): runner = object.__new__(AFDNPUFFNModelRunner) runner.prof = None runner.device = SimpleNamespace(type="npu") + runner.ffn_multistream_enabled = False + runner.ffn_comm_stream = None + runner.ffn_comm_events = [] return runner +def _queue_ffn_inputs(connector, *, layers, stages): + for layer_idx in range(layers): + for stage_idx in range(stages): + metadata = AFDTransferMetadata.create_attention_metadata( + layer_idx=layer_idx, + stage_idx=stage_idx, + seq_len=1, + ) + connector.attn_outputs.append(("hidden", metadata)) + return {stage_idx: _FakeDPMetadata([1]) for stage_idx in range(stages)} + + def _new_ffn_worker(): _require_npu_runtime() from afd_plugin.v1.worker.npu.ffn_worker import AFDNPUFFNWorker @@ -634,6 +649,9 @@ def test_npu_create_ascend_forward_context_marks_current_ubatch(monkeypatch): vllm_config=vllm_config, ubatch_slices=ubatch_slices, ubatch_num=1, + afd_comm_stream="attention-comm-stream", + afd_comm_event="ubatch-1-event", + afd_multistream_enabled=True, ) child_metadata = new_forward_context.additional_kwargs["afd_metadata"] @@ -641,6 +659,9 @@ def test_npu_create_ascend_forward_context_marks_current_ubatch(monkeypatch): assert new_forward_context.num_ubatches == 2 assert new_forward_context.num_tokens == 3 assert child_metadata.stage_idx == 1 + assert new_forward_context.afd_comm_stream == "attention-comm-stream" + assert new_forward_context.afd_comm_event == "ubatch-1-event" + assert new_forward_context.afd_multistream_enabled is True def test_npu_ffn_runner_executes_eager_ffn_step(monkeypatch): @@ -671,6 +692,95 @@ def test_npu_ffn_runner_executes_eager_ffn_step(monkeypatch): ] +@pytest.mark.parametrize( + ("enabled", "stages"), + [(True, 1), (False, 2)], +) +def test_npu_ffn_multistream_requires_switch_and_multiple_stages(enabled, stages): + runner = _new_ffn_runner() + runner.vllm_config = _vllm_config(role="ffn") + runner.afd_config = SimpleNamespace(compute_gate_on_attention=False) + runner.connector = _FakeFFNConnector() + runner.model = _FakeModel() + runner.num_layers = 2 + runner.max_num_tokens = 1 + runner.ffn_multistream_enabled = enabled + dp_metadata = _queue_ffn_inputs(runner.connector, layers=2, stages=stages) + + runner._ffn_forward(dp_metadata_list=dp_metadata) + + assert all( + "multistream_enable" not in kwargs + for _output, _metadata, kwargs in runner.connector.ffn_outputs + ) + + +def test_npu_ffn_multistream_waits_before_stage_reuse(monkeypatch): + _require_npu_runtime() + from afd_plugin.v1.worker.npu import ffn_model_runner + + trace = [] + + class FakeEvent: + def __init__(self, stage_idx): + self.stage_idx = stage_idx + + def record(self, stream): + assert stream == "comm" + trace.append(("record", self.stage_idx)) + + def wait(self, stream): + assert stream == "compute" + trace.append(("wait", self.stage_idx)) + + class RecordingConnector(_FakeFFNConnector): + def send_ffn_output(self, ffn_output, metadata, **kwargs): + if kwargs.get("multistream_enable"): + kwargs["comm_event"].record(kwargs["comm_stream"]) + super().send_ffn_output(ffn_output, metadata, **kwargs) + + @contextmanager + def fake_ascend_forward_context(**_kwargs): + yield SimpleNamespace(additional_kwargs={}, dp_metadata=None) + + monkeypatch.setattr( + ffn_model_runner, + "ascend_forward_context", + fake_ascend_forward_context, + ) + monkeypatch.setattr(ffn_model_runner, "_set_moe_layer_index", lambda *_args: None) + monkeypatch.setattr( + ffn_model_runner.torch.npu, + "current_stream", + lambda: "compute", + ) + + runner = _new_ffn_runner() + runner.vllm_config = _vllm_config(role="ffn") + runner.afd_config = SimpleNamespace(compute_gate_on_attention=False) + runner.connector = RecordingConnector() + runner.model = _FakeModel() + runner.num_layers = 3 + runner.max_num_tokens = 1 + runner.ffn_multistream_enabled = True + runner.ffn_comm_stream = "comm" + runner.ffn_comm_events = [FakeEvent(0), FakeEvent(1)] + dp_metadata = _queue_ffn_inputs(runner.connector, layers=3, stages=2) + + runner._ffn_forward(dp_metadata_list=dp_metadata) + + assert trace == [ + ("record", 0), + ("record", 1), + ("wait", 0), + ("record", 0), + ("wait", 1), + ("record", 1), + ("wait", 0), + ("wait", 1), + ] + + def test_npu_ffn_runner_dp_path_invokes_model_with_hidden_states_and_layer(monkeypatch): from afd_plugin.connectors.npu.async_cam import AFDAsyncTransferState