fix: gate GPU FFN graph replay on Attention run mode

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)),
     )
 

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)

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,

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

afd_plugin/v1/worker/cuda_graph.py

@@ -137,14 +137,15 @@ def graph_run_mode(
     *,
     is_warmup: bool,
     is_graph_capturing: bool,
+    is_graph_replaying: bool,
     graph_enabled: bool,
     graph_exists: bool,
 ) -> AFDGraphRunMode:
     if is_warmup:
         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
 

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,
         )

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":

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,

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.
 

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,
         )

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()
 

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