From 43bf26b5ee32c77e285c7c2146e2e4b57002827a Mon Sep 17 00:00:00 2001 From: wuxiaojia Date: Wed, 26 Aug 2026 21:32:52 +0800 Subject: [PATCH] fix: gate FFN graph replay on Attention is_graph_replaying FFN used to replay whenever a cached graph key existed. Attention now publishes is_graph_replaying on the control payload so GPU and NPU FFN only replay when Attention is actually replaying. Signed-off-by: wuxiaojia Co-authored-by: Cursor --- afd_plugin/connectors/metadata.py | 14 +++++--- afd_plugin/v1/worker/attention_metadata.py | 2 ++ .../v1/worker/attention_model_runner.py | 13 ++++++- .../v1/worker/attention_model_runner_v2.py | 35 +++++++++++++------ afd_plugin/v1/worker/cuda_graph.py | 3 +- afd_plugin/v1/worker/ffn_model_runner.py | 2 ++ afd_plugin/v1/worker/ffn_worker.py | 2 ++ .../v1/worker/npu/attention_model_runner.py | 14 +++++++- .../worker/npu/attention_model_runner_v2.py | 35 +++++++++++++------ afd_plugin/v1/worker/npu/ffn_model_runner.py | 4 +++ afd_plugin/v1/worker/npu/ffn_worker.py | 2 ++ tests/unit/connectors/test_p2p_connector.py | 2 ++ tests/unit/v1/worker/test_cuda_graph.py | 27 ++++++++++++++ tests/unit/v1/worker/test_ffn_model_runner.py | 24 ++++++++++++- tests/unit/v1/worker/test_npu_runtime.py | 32 ++++++++++++++++- 15 files changed, 180 insertions(+), 31 deletions(-) diff --git a/afd_plugin/connectors/metadata.py b/afd_plugin/connectors/metadata.py index 67447beb..3b839f0d 100644 --- a/afd_plugin/connectors/metadata.py +++ b/afd_plugin/connectors/metadata.py @@ -123,6 +123,9 @@ class AFDControlPayload: is_warmup: Whether the payload belongs to a warmup step. This is separate from graph capture because warmup may prepare state without representing a real serving step. + is_graph_replaying: Whether the Attention side is currently replaying a + captured graph for this step. FFN may replay only when this is True + and a matching local graph exists. is_profile: Whether the payload belongs to an initial memory-profile forward. NPU FFN workers use this to preserve Ascend's balanced dummy MoE routing in the split-process AFD execution model. @@ -131,6 +134,7 @@ class AFDControlPayload: dp_metadata_list: dict[int, AFDDPMetadata] is_graph_capturing: bool is_warmup: bool + is_graph_replaying: bool = False is_profile: bool = False def __post_init__(self) -> None: @@ -316,10 +320,10 @@ def encode_control_payload(payload: AFDControlPayload) -> bytes: """Serialize an ``AFDControlPayload`` to a compact JSON byte string. Only ``num_tokens_across_dp_cpu`` / ``max_tokens_across_dp_cpu`` per stage - and the graph-capturing, warmup, and profile flags are carried. These are the - fields the FFN-side connectors read back after decode. A plugin-owned - minimal schema keeps the wire format decoupled from vLLM-internal DP - metadata objects. + and the graph-capturing, warmup, replay, and profile flags are carried. + These are the fields the FFN-side connectors read back after decode. A + plugin-owned minimal schema keeps the wire format decoupled from + vLLM-internal DP metadata objects. """ metadata_payload: dict[str, dict[str, int | list[int]]] = {} for stage_idx, dp_metadata in payload.dp_metadata_list.items(): @@ -332,6 +336,7 @@ def encode_control_payload(payload: AFDControlPayload) -> bytes: "dp_metadata_list": metadata_payload, "is_graph_capturing": bool(payload.is_graph_capturing), "is_warmup": bool(payload.is_warmup), + "is_graph_replaying": bool(payload.is_graph_replaying), "is_profile": bool(payload.is_profile), } return json.dumps(wire_payload, separators=(",", ":"), sort_keys=True).encode( @@ -361,6 +366,7 @@ def decode_control_payload(payload_bytes: bytes) -> AFDControlPayload: dp_metadata_list=dp_metadata_list, is_graph_capturing=bool(payload.get("is_graph_capturing", False)), is_warmup=bool(payload.get("is_warmup", False)), + is_graph_replaying=bool(payload.get("is_graph_replaying", False)), is_profile=bool(payload.get("is_profile", False)), ) diff --git a/afd_plugin/v1/worker/attention_metadata.py b/afd_plugin/v1/worker/attention_metadata.py index 88994a57..f633dc59 100644 --- a/afd_plugin/v1/worker/attention_metadata.py +++ b/afd_plugin/v1/worker/attention_metadata.py @@ -99,10 +99,12 @@ def send_dp_metadata( # Keep the V1 object.__new__ test seam and older graph lifecycle # callers compatible with runners created before this mixin existed. is_graph_capturing = getattr(self, "_afd_is_graph_capturing", False) + is_graph_replaying = getattr(self, "_afd_is_graph_replaying", False) payload = AFDControlPayload( dp_metadata_list=dp_metadata_list, is_graph_capturing=is_graph_capturing, is_warmup=is_warmup, + is_graph_replaying=is_graph_replaying, is_profile=self._afd_is_profile, ) self.connector.control_plane.update_state_from_dp_metadata(payload) diff --git a/afd_plugin/v1/worker/attention_model_runner.py b/afd_plugin/v1/worker/attention_model_runner.py index 94c4aefd..b559be68 100644 --- a/afd_plugin/v1/worker/attention_model_runner.py +++ b/afd_plugin/v1/worker/attention_model_runner.py @@ -82,6 +82,7 @@ def __init__( ) self._is_warmup = False self._afd_is_graph_capturing = False + self._afd_is_graph_replaying = False self._afd_pending_metadata: AFDForwardContextMetadata | None = None self._afd_suppress_metadata_send = False self._afd_transaction_counter = 0 @@ -233,6 +234,11 @@ def _determine_batch_execution_and_padding( force_num_active_loras, num_encoder_reqs, ) + self._afd_is_graph_replaying = ( + not bool(getattr(self, "_is_warmup", False)) + and not bool(getattr(self, "_afd_is_graph_capturing", False)) + and cudagraph_mode == CUDAGraphMode.FULL + ) args = ( num_tokens, @@ -346,7 +352,12 @@ def execute_model( intermediate_tensors: IntermediateTensors | None = None, ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors | None: step_afd_gpu_profiler(self.prof) - return super().execute_model(scheduler_output, intermediate_tensors) + previous_is_graph_replaying = getattr(self, "_afd_is_graph_replaying", False) + self._afd_is_graph_replaying = False + try: + return super().execute_model(scheduler_output, intermediate_tensors) + finally: + self._afd_is_graph_replaying = previous_is_graph_replaying def _dummy_run( self, diff --git a/afd_plugin/v1/worker/attention_model_runner_v2.py b/afd_plugin/v1/worker/attention_model_runner_v2.py index 2c01ae0d..75e74ac4 100644 --- a/afd_plugin/v1/worker/attention_model_runner_v2.py +++ b/afd_plugin/v1/worker/attention_model_runner_v2.py @@ -181,18 +181,27 @@ def run_fullgraph( desc: v2_cudagraph_utils.BatchExecutionDescriptor, ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]] | IntermediateTensors: # ### PATCH START: publish one AFD pre-replay payload. - padded_tokens = int(desc.num_tokens) - metadata = runner.build_afd_metadata(None, real_tokens) - metadata.tokens_lens = [padded_tokens] - runner._afd_pending_metadata = metadata - runner._afd_suppress_metadata_send = True - runner._is_warmup = False - runner._afd_is_graph_capturing = False - runner.send_dp_metadata( - runner.build_capture_dp_metadata(padded_tokens), - None, + previous_is_graph_replaying = getattr( + runner, + "_afd_is_graph_replaying", + False, ) - result = original_run_fullgraph(desc) + try: + padded_tokens = int(desc.num_tokens) + metadata = runner.build_afd_metadata(None, real_tokens) + metadata.tokens_lens = [padded_tokens] + runner._afd_pending_metadata = metadata + runner._afd_suppress_metadata_send = True + runner._is_warmup = False + runner._afd_is_graph_capturing = False + runner._afd_is_graph_replaying = True + runner.send_dp_metadata( + runner.build_capture_dp_metadata(padded_tokens), + None, + ) + result = original_run_fullgraph(desc) + finally: + runner._afd_is_graph_replaying = previous_is_graph_replaying # ### PATCH END: publish one AFD pre-replay payload. return result @@ -223,6 +232,8 @@ def _use_afd_execution_context( previous_suppress_send = runner._afd_suppress_metadata_send previous_is_warmup = runner._is_warmup previous_is_graph_capturing = runner._afd_is_graph_capturing + previous_is_graph_replaying = getattr(runner, "_afd_is_graph_replaying", False) + runner._afd_is_graph_replaying = False replay_scope = ( _use_afd_fullgraph_replay_hook(runner, real_tokens) @@ -242,6 +253,7 @@ def _use_afd_execution_context( runner._afd_suppress_metadata_send = previous_suppress_send runner._is_warmup = previous_is_warmup runner._afd_is_graph_capturing = previous_is_graph_capturing + runner._afd_is_graph_replaying = previous_is_graph_replaying class AFDAttentionModelRunnerV2(AFDMetadataProviderMixin, GPUModelRunnerV2): @@ -286,6 +298,7 @@ def __init__( ) self._is_warmup = False self._afd_is_graph_capturing = False + self._afd_is_graph_replaying = False self._afd_pending_metadata: AFDForwardContextMetadata | None = None self._afd_suppress_metadata_send = False self._afd_transaction_counter = 0 diff --git a/afd_plugin/v1/worker/cuda_graph.py b/afd_plugin/v1/worker/cuda_graph.py index 157ffabb..71323dd7 100644 --- a/afd_plugin/v1/worker/cuda_graph.py +++ b/afd_plugin/v1/worker/cuda_graph.py @@ -137,6 +137,7 @@ def graph_run_mode( *, is_warmup: bool, is_graph_capturing: bool, + is_graph_replaying: bool, graph_enabled: bool, graph_exists: bool, ) -> AFDGraphRunMode: @@ -144,7 +145,7 @@ def graph_run_mode( return AFDGraphRunMode.WARMUP if is_graph_capturing: return AFDGraphRunMode.CAPTURE - if graph_enabled and graph_exists: + if is_graph_replaying and graph_enabled and graph_exists: return AFDGraphRunMode.REPLAY return AFDGraphRunMode.EAGER diff --git a/afd_plugin/v1/worker/ffn_model_runner.py b/afd_plugin/v1/worker/ffn_model_runner.py index 1349da77..4582d229 100644 --- a/afd_plugin/v1/worker/ffn_model_runner.py +++ b/afd_plugin/v1/worker/ffn_model_runner.py @@ -139,6 +139,7 @@ def execute_model( dp_metadata_list: dict[int, DPMetadata | AFDDPMetadata] | None = None, is_graph_capturing: bool = False, is_warmup: bool = False, + is_graph_replaying: bool = False, ) -> None: step_afd_gpu_profiler(self.prof) if dp_metadata_list is None: @@ -148,6 +149,7 @@ def execute_model( run_mode = graph_run_mode( is_warmup=is_warmup, is_graph_capturing=is_graph_capturing, + is_graph_replaying=is_graph_replaying, graph_enabled=bool(self.use_cuda_graph), graph_exists=cuda_graph_info is not None, ) diff --git a/afd_plugin/v1/worker/ffn_worker.py b/afd_plugin/v1/worker/ffn_worker.py index 5803fff7..4800eab9 100644 --- a/afd_plugin/v1/worker/ffn_worker.py +++ b/afd_plugin/v1/worker/ffn_worker.py @@ -190,6 +190,7 @@ def _run_ffn_server_loop(self) -> None: dp_metadata_list = payload.dp_metadata_list is_attn_graph_capturing = payload.is_graph_capturing is_warmup = payload.is_warmup + is_graph_replaying = payload.is_graph_replaying if self.model_runner.use_cuda_graph and ( is_warmup or is_attn_graph_capturing @@ -204,6 +205,7 @@ def _run_ffn_server_loop(self) -> None: dp_metadata_list=dp_metadata_list, is_graph_capturing=is_attn_graph_capturing, is_warmup=is_warmup, + is_graph_replaying=is_graph_replaying, ) if self.device.type == "cuda": diff --git a/afd_plugin/v1/worker/npu/attention_model_runner.py b/afd_plugin/v1/worker/npu/attention_model_runner.py index cc3f524a..d412663b 100644 --- a/afd_plugin/v1/worker/npu/attention_model_runner.py +++ b/afd_plugin/v1/worker/npu/attention_model_runner.py @@ -155,6 +155,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): # so Attention and FFN weight loading overlap; see that method. self._is_warmup = False self._afd_is_graph_capturing = False + self._afd_is_graph_replaying = False self._afd_pending_metadata: AFDForwardContextMetadata | None = None self._afd_suppress_metadata_send = False self._afd_transaction_counter = 0 @@ -179,11 +180,14 @@ def execute_model( ) -> ModelRunnerOutput | IntermediateTensors | None: step_afd_npu_profiler(self.prof) # ### PATCH START: AFD live execution scope + previous_is_graph_replaying = getattr(self, "_afd_is_graph_replaying", False) + self._afd_is_graph_replaying = False self._afd_live_execution = True try: result = super().execute_model(scheduler_output, intermediate_tensors) finally: self._afd_live_execution = False + self._afd_is_graph_replaying = previous_is_graph_replaying # ### PATCH END: AFD live execution scope return result @@ -1541,19 +1545,22 @@ def _send_dp_metadata( dp_metadata_list = {0: dp_metadata} is_warmup = bool(self._is_warmup) is_graph_capturing = bool(self._afd_is_graph_capturing) + is_graph_replaying = bool(getattr(self, "_afd_is_graph_replaying", False)) payload = AFDControlPayload( dp_metadata_list=dp_metadata_list, is_graph_capturing=is_graph_capturing, is_warmup=is_warmup, + is_graph_replaying=is_graph_replaying, ) self.connector.control_plane.update_state_from_dp_metadata(payload) logger.warning( "AFD NPU Attention send_dp_metadata decision; world_rank=%d " - "key=%s is_graph_capturing=%s is_warmup=%s", + "key=%s is_graph_capturing=%s is_warmup=%s is_graph_replaying=%s", self.connector.world_rank, _dp_metadata_debug_key(dp_metadata_list), is_graph_capturing, is_warmup, + is_graph_replaying, ) self.connector.control_plane.send_dp_metadata_list(payload) @@ -1873,6 +1880,11 @@ def dispatch_cudagraph( num_paddings=batch_descriptor.num_tokens - num_tokens, runtime_mode=str(cudagraph_mode), ) + self._afd_is_graph_replaying = ( + not bool(getattr(self, "_is_warmup", False)) + and not bool(getattr(self, "_afd_is_graph_capturing", False)) + and cudagraph_mode == CUDAGraphMode.FULL + ) return ( cudagraph_mode, batch_descriptor, diff --git a/afd_plugin/v1/worker/npu/attention_model_runner_v2.py b/afd_plugin/v1/worker/npu/attention_model_runner_v2.py index 21d33ae5..656f4f12 100644 --- a/afd_plugin/v1/worker/npu/attention_model_runner_v2.py +++ b/afd_plugin/v1/worker/npu/attention_model_runner_v2.py @@ -82,18 +82,27 @@ def run_fullgraph( desc: v2_cudagraph_utils.BatchExecutionDescriptor, ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]] | IntermediateTensors: # ### PATCH START: publish one AFD pre-replay payload. - padded_tokens = int(desc.num_tokens) - metadata = runner.build_afd_metadata(None, real_tokens) - metadata.tokens_lens = [padded_tokens] - runner._afd_pending_metadata = metadata - runner._afd_suppress_metadata_send = True - runner._is_warmup = False - runner._afd_is_graph_capturing = False - runner.send_dp_metadata( - runner.build_capture_dp_metadata(padded_tokens), - None, + previous_is_graph_replaying = getattr( + runner, + "_afd_is_graph_replaying", + False, ) - result = original_run_fullgraph(desc) + try: + padded_tokens = int(desc.num_tokens) + metadata = runner.build_afd_metadata(None, real_tokens) + metadata.tokens_lens = [padded_tokens] + runner._afd_pending_metadata = metadata + runner._afd_suppress_metadata_send = True + runner._is_warmup = False + runner._afd_is_graph_capturing = False + runner._afd_is_graph_replaying = True + runner.send_dp_metadata( + runner.build_capture_dp_metadata(padded_tokens), + None, + ) + result = original_run_fullgraph(desc) + finally: + runner._afd_is_graph_replaying = previous_is_graph_replaying # ### PATCH END: publish one AFD pre-replay payload. return result @@ -156,6 +165,7 @@ def __init__( self._is_warmup = False self._afd_is_graph_capturing = False self._afd_is_profile = False + self._afd_is_graph_replaying = False self._afd_pending_metadata: AFDForwardContextMetadata | None = None self._afd_suppress_metadata_send = False self._afd_transaction_counter = 0 @@ -351,8 +361,10 @@ def execute_model( previous_suppress_send = self._afd_suppress_metadata_send previous_is_warmup = self._is_warmup previous_is_graph_capturing = self._afd_is_graph_capturing + previous_is_graph_replaying = getattr(self, "_afd_is_graph_replaying", False) previous_is_profile = self._afd_is_profile self._afd_is_profile = bool(is_profile) + self._afd_is_graph_replaying = False replay_scope = ( _use_afd_fullgraph_replay_hook( @@ -381,6 +393,7 @@ def execute_model( self._afd_suppress_metadata_send = previous_suppress_send self._is_warmup = previous_is_warmup self._afd_is_graph_capturing = previous_is_graph_capturing + self._afd_is_graph_replaying = previous_is_graph_replaying self._afd_is_profile = previous_is_profile # ### PATCH END: scope AFD metadata provider/replay and profiler step. diff --git a/afd_plugin/v1/worker/npu/ffn_model_runner.py b/afd_plugin/v1/worker/npu/ffn_model_runner.py index 84c6223c..2b6e9281 100644 --- a/afd_plugin/v1/worker/npu/ffn_model_runner.py +++ b/afd_plugin/v1/worker/npu/ffn_model_runner.py @@ -111,6 +111,7 @@ def execute_ffn_step( is_graph_capturing: bool = False, is_warmup: bool = False, is_profile: bool = False, + is_graph_replaying: bool = False, ) -> None: if dp_metadata_list is None: raise RuntimeError("AFD NPU FFN requires dp_metadata_list") @@ -131,6 +132,7 @@ def execute_ffn_step( self.execute_model( dp_metadata_list=dp_metadata_list, is_profile=is_profile, + is_graph_replaying=is_graph_replaying, ) return None @@ -153,6 +155,7 @@ def execute_model( is_graph_capturing: bool = False, is_warmup: bool = False, is_profile: bool = False, + is_graph_replaying: bool = False, ) -> None: step_afd_npu_profiler(self.prof) if dp_metadata_list is None: @@ -164,6 +167,7 @@ def execute_model( run_mode = graph_run_mode( is_warmup=is_warmup and graph_enabled, is_graph_capturing=is_graph_capturing and graph_enabled, + is_graph_replaying=is_graph_replaying, graph_enabled=graph_enabled, graph_exists=graph_info is not None, ) diff --git a/afd_plugin/v1/worker/npu/ffn_worker.py b/afd_plugin/v1/worker/npu/ffn_worker.py index c01173f7..40f5456f 100644 --- a/afd_plugin/v1/worker/npu/ffn_worker.py +++ b/afd_plugin/v1/worker/npu/ffn_worker.py @@ -153,12 +153,14 @@ def _run_ffn_server_loop(self) -> None: is_attn_graph_capturing = payload.is_graph_capturing is_warmup = payload.is_warmup is_profile = payload.is_profile + is_graph_replaying = payload.is_graph_replaying self.model_runner.execute_ffn_step( dp_metadata_list=dp_metadata_list, is_graph_capturing=is_attn_graph_capturing, is_warmup=is_warmup, is_profile=is_profile, + is_graph_replaying=is_graph_replaying, ) torch.npu.synchronize() diff --git a/tests/unit/connectors/test_p2p_connector.py b/tests/unit/connectors/test_p2p_connector.py index 84744584..13250390 100644 --- a/tests/unit/connectors/test_p2p_connector.py +++ b/tests/unit/connectors/test_p2p_connector.py @@ -460,6 +460,7 @@ def test_p2p_dp_metadata_serialization_uses_json_payload(): dp_metadata_list={7: metadata}, is_graph_capturing=True, is_warmup=False, + is_graph_replaying=True, is_profile=True, ), ) @@ -475,6 +476,7 @@ def test_p2p_dp_metadata_serialization_uses_json_payload(): assert _tolist(decoded[7].cu_tokens_across_sp(1)) == [3, 8] assert decoded_payload.is_graph_capturing is True assert decoded_payload.is_warmup is False + assert decoded_payload.is_graph_replaying is True assert decoded_payload.is_profile is True diff --git a/tests/unit/v1/worker/test_cuda_graph.py b/tests/unit/v1/worker/test_cuda_graph.py index a9400072..36ac7195 100644 --- a/tests/unit/v1/worker/test_cuda_graph.py +++ b/tests/unit/v1/worker/test_cuda_graph.py @@ -6,7 +6,9 @@ from afd_plugin.v1.worker.cuda_graph import ( FULL_DECODE_ONLY, + AFDGraphRunMode, cudagraph_mode_name, + graph_run_mode, make_ffn_graph_key, validate_cuda_graph_mode, ) @@ -199,3 +201,28 @@ def test_make_ffn_graph_key_dp1_tp1_unchanged(): fallback=32, ) assert key == ((0, (8,)),) + + +@pytest.mark.parametrize( + ("is_graph_replaying", "graph_exists", "expected"), + [ + (False, True, AFDGraphRunMode.EAGER), + (True, True, AFDGraphRunMode.REPLAY), + (True, False, AFDGraphRunMode.EAGER), + ], +) +def test_graph_run_mode_requires_attention_replaying( + is_graph_replaying, + graph_exists, + expected, +): + assert ( + graph_run_mode( + is_warmup=False, + is_graph_capturing=False, + is_graph_replaying=is_graph_replaying, + graph_enabled=True, + graph_exists=graph_exists, + ) + is expected + ) diff --git a/tests/unit/v1/worker/test_ffn_model_runner.py b/tests/unit/v1/worker/test_ffn_model_runner.py index 3f45bc46..c0d198d0 100644 --- a/tests/unit/v1/worker/test_ffn_model_runner.py +++ b/tests/unit/v1/worker/test_ffn_model_runner.py @@ -611,12 +611,34 @@ def test_ffn_runner_replays_cuda_graph_when_key_exists(): make_ffn_graph_key(dp_metadata): {"graph": graph}, } - runner.execute_model(dp_metadata_list=dp_metadata) + runner.execute_model( + dp_metadata_list=dp_metadata, + is_graph_replaying=True, + ) assert graph.replay_count == 1 assert runner.connector.ffn_outputs == [] +def test_ffn_runner_skips_replay_when_attention_is_eager(): + runner = _runner_with_connector_and_model(_FakeModel()) + runner.use_cuda_graph = True + graph = _FakeGraph() + dp_metadata = {0: _FakeDPMetadata([1])} + runner._cuda_graphs = { + make_ffn_graph_key(dp_metadata): {"graph": graph}, + } + metadata = _metadata() + runner.connector.attn_outputs.append(_payload("hidden", metadata)) + + runner.execute_model(dp_metadata_list=dp_metadata) + + assert graph.replay_count == 0 + assert runner.connector.ffn_outputs == [ + ("ffn(hidden, layer=0)", metadata), + ] + + def test_ffn_runner_cuda_graph_miss_falls_back_to_eager(): runner = _runner_with_connector_and_model(_FakeModel()) runner.use_cuda_graph = True diff --git a/tests/unit/v1/worker/test_npu_runtime.py b/tests/unit/v1/worker/test_npu_runtime.py index df7fe6b7..5f800ca7 100644 --- a/tests/unit/v1/worker/test_npu_runtime.py +++ b/tests/unit/v1/worker/test_npu_runtime.py @@ -1601,12 +1601,42 @@ def test_npu_ffn_runner_replays_acl_graph_when_key_exists(): graph = _FakeGraph() runner._acl_graphs = {runner._make_graph_key(dp_metadata): {"graph": graph}} - runner.execute_model(dp_metadata_list=dp_metadata) + runner.execute_model( + dp_metadata_list=dp_metadata, + is_graph_replaying=True, + ) assert graph.replay_count == 1 assert runner.connector.ffn_outputs == [] +def test_npu_ffn_runner_skips_replay_when_attention_is_eager(monkeypatch): + _patch_ffn_forward_context(monkeypatch) + runner = _new_ffn_runner() + runner.vllm_config = _vllm_config(role="ffn") + runner.connector = _FakeFFNConnector() + runner.model = _FakeModel() + runner.num_layers = 1 + runner.max_num_tokens = 1 + runner.use_aclgraph = True + dp_metadata = {0: _FakeDPMetadata([1])} + graph = _FakeGraph() + runner._acl_graphs = {runner._make_graph_key(dp_metadata): {"graph": graph}} + metadata = AFDTransferMetadata.create_attention_metadata( + layer_idx=0, + stage_idx=0, + seq_len=1, + ) + runner.connector.attn_outputs.append(("hidden", metadata)) + + runner.execute_model(dp_metadata_list=dp_metadata) + + assert graph.replay_count == 0 + assert runner.connector.ffn_outputs == [ + ("npu-ffn(hidden, layer=0)", metadata, {"ubatch_idx": 0}), + ] + + def test_npu_ffn_runner_graph_key_uses_ffn_aggregated_token_counts(): runner = _new_ffn_runner() runner.connector = _FakeFFNConnector(attn_size=8, ffn_size=4)