fix: gate GPU FFN graph replay on Attention run mode

[Betelgeu]

Purpose

FFN used to replay whenever a matching graph key existed, even if Attention ran eager. Attention now sends is_graph_replaying; FFN only replays when that flag is set and a local graph exists.

Issue

• Related issue(s): [Bug]: FFN graph cache replays during eager prefill in FULL_DECODE_ONLY #261

Scope

• In scope: GPU/NPU Attention publish and FFN consume of is_graph_replaying.
• Out of scope: graph-key layout changes; hardware e2e.

Test Result

Focused CPU unit tests passed. NPU runtime tests skipped without Ascend.

[jiangkuaixue123]

Blocking: this adds only a placeholder mode to NPU-created payloads, while AFDNPUFFNModelRunner.execute_model() still calls the shared graph_run_mode() with the removed is_warmup and is_graph_capturing keyword arguments (around line 164). With this PR, every NPU step reaching that call raises TypeError: graph_run_mode() got an unexpected keyword argument is_warmup before selecting eager/capture/replay. Please either preserve backward compatibility in the helper or update the NPU consumer and call site to pass attention_run_mode; deferring the NPU behavior change must not leave the existing path unusable.

[Betelgeu]

Blocking: this adds only a placeholder mode to NPU-created payloads, while AFDNPUFFNModelRunner.execute_model() still calls the shared graph_run_mode() with the removed is_warmup and is_graph_capturing keyword arguments (around line 164). With this PR, every NPU step reaching that call raises TypeError: graph_run_mode() got an unexpected keyword argument is_warmup before selecting eager/capture/replay. Please either preserve backward compatibility in the helper or update the NPU consumer and call site to pass attention_run_mode; deferring the NPU behavior change must not leave the existing path unusable.

Fixed in the rewrite: NPU FFN now passes is_graph_replaying into graph_run_mode() together with is_warmup / is_graph_capturing, and NPU Attention publishes the flag instead of a placeholder.

[jiangkuaixue123]

The GPU-side direction is sound, and the five modified GPU/control-plane unit-test modules pass locally (35 tests). However, this PR currently breaks the existing NPU FFN execution path: the shared graph_run_mode() API now requires attention_run_mode, while AFDNPUFFNModelRunner.execute_model() still passes the removed lifecycle keywords. I reproduced the resulting TypeError and left the exact finding inline. Even if full NPU run-mode synchronization remains out of scope, this shared API change must keep NPU executable. Please update the NPU call path or retain a compatible helper contract, and add a CPU-safe regression test covering that invocation. Separately, the PR DCO check is currently ACTION_REQUIRED.

[jiangkuaixue123]

cc @lirx-pd

[lirx-pd]

Several parts of this PR appear redundant or unrelated to that requirement:
1.AFDGraphRunMode duplicates existing lifecycle state. WARMUP and CAPTURE are already represented by is_warmup and is_graph_capturing
2.The patch size is disproportionate to the bug.

Can a smaller solution, that adds an is_graph_replaying field to AFDControlPayload, propagates it through the workers, and uses is_graph_replaying and graph_enabled and graph_exists as the FFN replay condition, solve the problem?

[Betelgeu]

Several parts of this PR appear redundant or unrelated to that requirement: 1.AFDGraphRunMode duplicates existing lifecycle state. WARMUP and CAPTURE are already represented by is_warmup and is_graph_capturing 2.The patch size is disproportionate to the bug.

Can a smaller solution, that adds an is_graph_replaying field to AFDControlPayload, propagates it through the workers, and uses is_graph_replaying and graph_enabled and graph_exists as the FFN replay condition, solve the problem?

Thanks for the review! Reworked to the smaller approach:

• add is_graph_replaying on AFDControlPayload
• FFN replay only when is_graph_replaying and graph_enabled and graph_exists

[Betelgeu]

The GPU-side direction is sound, and the five modified GPU/control-plane unit-test modules pass locally (35 tests). However, this PR currently breaks the existing NPU FFN execution path: the shared graph_run_mode() API now requires attention_run_mode, while AFDNPUFFNModelRunner.execute_model() still passes the removed lifecycle keywords. I reproduced the resulting TypeError and left the exact finding inline. Even if full NPU run-mode synchronization remains out of scope, this shared API change must keep NPU executable. Please update the NPU call path or retain a compatible helper contract, and add a CPU-safe regression test covering that invocation. Separately, the PR DCO check is currently ACTION_REQUIRED.

Thanks for the review! The rewrite updates the NPU Attention/FFN path to publish and consume is_graph_replaying, and graph_run_mode() keeps the lifecycle kwargs so the NPU call site no longer breaks.

Also added Signed-off-by for DCO.

[jiangkuaixue123]

Thanks for the PR! Have you tested this change in an actual GPU or NPU environment?

In our current usage, this situation generally should not occur. FULL_DECODE_ONLY is mainly used for graph execution during the decode stage. With max_num_batched_tokens and the graph capture sizes configured appropriately, Attention and FFN should normally enter graph mode together, rather than only one side doing so.

If this PR passes the relevant E2E tests in a real hardware environment, I’ll be happy to merge it.

[Betelgeu]

Thanks for the PR! Have you tested this change in an actual GPU or NPU environment?感谢您的 PR！您是否在实际的 GPU 或 NPU 环境中测试过这项更改？

In our current usage, this situation generally should not occur. FULL_DECODE_ONLY is mainly used for graph execution during the decode stage. With max_num_batched_tokens and the graph capture sizes configured appropriately, Attention and FFN should normally enter graph mode together, rather than only one side doing so.在我们目前的用法中，这种情况通常不会发生。FULL_DECODE_ONLY 主要用于解码阶段的图执行。 max_num_batched_tokens FULL_DECODE_ONLY 图捕获大小配置得当，Attention 和 FFN 通常应该同时进入图模式，而不是只有一方进入。

If this PR passes the relevant E2E tests in a real hardware environment, I’ll be happy to merge it.如果这个 PR 在真实的硬件环境中通过了相关的 E2E 测试，我将很乐意合并它。

Ran GPU E2E on H20 for this change:

• DeepSeek-V2-Lite afd-graph (FULL_DECODE_ONLY, 2A1F)
• GSM8K-7 passed (exact_match 0.2857)

I don’t have Ascend access for NPU E2E right now. Happy to run more GPU cases if needed.