[4/5] feat: capture the FFN experts at bucketed padded shapes

[specture724]

Fourth of a five-PR stack. Based on #334 — the diff shown here is only this
PR's own change; review #332, #333 and #334 first.

1. [1/5] feat: NVSHMEM symmetric window for GPU AFD transport #332 — NVSHMEM symmetric window (transport)
2. [2/5] feat: the async GPU connector on the symmetric window #333 — the async GPU connector
3. [3/5] feat: run DeepSeek-V4 on the async GPU connector #334 — both GPU runners + DeepSeek-V4
4. this PR — bucketed padded FFN expert graphs
5. [5/5][WIP] feat: request-aligned DBO ubatch splitting (flag off by default) #336 — request-aligned DBO ubatch splitting (flag off by default)

What this adds

The connector-driven FFN side never learns the next work item's shape ahead of
time — there is no control plane — which is why it ran eagerly. It does not need
to: a grouped GEMM takes its grouping from a device-side count vector, not from
its row count, so one row count can be captured and every smaller item padded up
to it, with the padding charged to the last expert and its output sliced off.

Sizing that capture is the whole problem. The upper bound is every one of a
token's topk partials landing on one rank; with 256 experts over two FFN ranks a
rank really receives about half that, so a single graph at the bound charges
every replay ~2x. The graphs are a ladder of buckets built as multiples of the
expected item size, clustered just above 1.0 and 0.5 where routing scatter
actually puts items, and an item takes the smallest bucket that holds it — or
runs eager when the bucket would pad it past MAX_REPLAY_PADDING_RATIO, since a
replay pays for its whole bucket while eager pays launches.

Two capture-order traps, both faulting on the first replay, both found the hard
way:

1. vLLM's WorkspaceManager grows the shared MoE scratch by freeing the old
   buffer, so any growth after a capture leaves that graph dangling. Capture
   runs largest-bucket-first behind one eager call at the ceiling.
2. Capturing with all rows on one expert reserves too little per-expert block
   padding for a replay whose routing touches every expert. Capture counts are
   spread evenly instead.

Bounded by free memory, not by the token budget

The ladder is sized from max_num_batched_tokens * topk, so its memory grows
with the token budget. At 8192 tokens the default ladder captured 65 GiB per
DeepSeek-V4 FFN rank, and the connector's NVSHMEM heap — set up after
capture — then failed with cuMemCreate failed. Capture now checks free memory
before each MoE layer and stops at the first layer whose graphs would cut into
AFD_FFN_GRAPH_MEM_RESERVE_GIB (default 8). The remaining layers run eagerly
through the existing no-graph fallback.

Validated on DeepSeek-V4-Flash 2A2F, 8192 tokens, default ladder: capture
stopped at 37/43 layers (58.2 GiB), NVSHMEM initialized, and the server
served correct completions. At 2048 tokens every layer still fits, so the
behaviour measured below is unchanged.

Measurements, including where it does not pay

Pure prefill, DeepSeek-V4 2A2F on 4x L20X, 2048-token steps:

FFN mode	time
one graph at the worst-case shape	27.6 s
bucketed graphs	19.0 s
eager	17.9 s

Bucketing removes a 1.55x regression; it does not beat eager, because the
connector issues about one grouped GEMM per item and there are no launches to
save. FFN_EAGER=1 is therefore the V4 default and this path is opt-in.

Separately, on the Attention side in the decode regime, FULL_DECODE_ONLY
graphs are a real win — DeepSeek-V2-Lite 1A1F, 2x L20X: 36.6 s eager vs 28.3 s
with graphs (+22.6%), replay-verified. That is an existing vLLM mechanism this
PR does not change; noting it because it is the reason not to leave graphs off
for decode-heavy serving.

Testing

test_cuda_graph.py (bucket ladder, shared-row scaling, bucket selection),
test_ffn_model_runner.py (replay, padding slice-off, eager fallbacks, capture
ordering), and tests/e2e/async_gpu_ffn_padded_graph.py (1 GPU) which replays
against an eager run of the same routing.

Full unit suite unchanged against main: same 13 pre-existing failures, none
added. pre-commit clean over the PR range.

🤖 Generated with Claude Code