Current environment
Python:3.12
afd-plugin commit: 85e826e
vLLM version: 0.26.0
Reproduction
-
Start AFD with cudagraph mode FULL_DECODE_ONLY.
-
Allow startup to capture an FFN graph for a uniform decode batch whose FFN
graph key is K.
-
Send a chunked-prefill request whose Attention execution is eager, but whose
FFN graph key is also K.
For example, with two Attention ranks mapped to one FFN rank:
captured decode: 4 + 4 tokens -> FFN key 8
eager prefill: 7 + 1 tokens -> FFN key 8
The FFN worker receives is_graph_capturing=False and is_warmup=False, finds
the cached graph for K, and replays it.
Expected behavior
The FFN runner must replay an FFN graph only when the paired Attention step is
also executing the corresponding graph-replay path.
An eager Attention prefill step must run the FFN eagerly, even if an FFN graph
with the same shape key already exists.
For NPU, the graph key must not alias distinct Attention-peer token layouts
when the captured CAM communication or MoE execution depends on that layout.
Actual behavior and logs
The FFN runners decide replay from “graph exists” alone.
GPU:
GPUFFNModelRunner.execute_model()
graph_run_mode(... graph_exists=cuda_graph_info is not None)
NPU:
AFDNPUFFNModelRunner.execute_model()
graph_run_mode(... graph_exists=graph_info is not None)
The AFD control payload contains only is_graph_capturing and is_warmup. It does
not carry whether Attention is currently replaying a graph. Therefore, after a
decode graph has been captured, an eager prefill with the same FFN key can take
the FFN replay path.
Observed in NPU consistency testing:
- FULL: 24/24 outputs matched.
- FULL_DECODE_ONLY: 22/24 outputs matched.
- Minimum character 3-gram cosine similarity: 0.857.
The consistency result is evidence of a correctness problem, but does not by
itself prove that every mismatch is caused by this replay-policy bug.
For a plain GPU token-wise FFN, replaying a same-layout graph may still produce
correct values. Incorrect output requires a semantic mismatch in graph-captured
state, such as stale receive buffers, different peer token layout, stale
routing/mask/scale tensors, or communication metadata that was fixed at
capture time.
AFD configuration
Before submitting
Current environment
Reproduction
Start AFD with cudagraph mode FULL_DECODE_ONLY.
Allow startup to capture an FFN graph for a uniform decode batch whose FFN
graph key is K.
Send a chunked-prefill request whose Attention execution is eager, but whose
FFN graph key is also K.
For example, with two Attention ranks mapped to one FFN rank:
captured decode: 4 + 4 tokens -> FFN key 8
eager prefill: 7 + 1 tokens -> FFN key 8
The FFN worker receives is_graph_capturing=False and is_warmup=False, finds
the cached graph for K, and replays it.
Expected behavior
The FFN runner must replay an FFN graph only when the paired Attention step is
also executing the corresponding graph-replay path.
An eager Attention prefill step must run the FFN eagerly, even if an FFN graph
with the same shape key already exists.
For NPU, the graph key must not alias distinct Attention-peer token layouts
when the captured CAM communication or MoE execution depends on that layout.
Actual behavior and logs
AFD configuration
Before submitting