fix(cuda-graph): keep replay metadata dynamic across tensor-parallel ranks

[junjiewang253-ctrl]

Summary

Keep static-KV decode metadata device-dynamic during CUDA Graph replay and ensure tensor-parallel row reductions complete before downstream work consumes their output.

Root cause

CUDA Graph capture parsed KV offsets and attention shapes on the host, freezing capture-time decode decisions. Replay could therefore reuse stale KV and attention metadata. On the current upstream tensor-parallel path, RowParallel also returned the result of an asynchronous all-reduce without waiting for stream completion, which could let ranks advance out of phase during graph-enabled decode.

Fix

• Keep replay sequence metadata device-resident.
• Use graph-safe static KV-cache update and attention operations.
• Preserve eager fallback for unsupported layouts, devices, dtypes, and KV quantization.
• Synchronize the current stream after TP RowParallel compute_linear_allreduce() completes and before returning its output.

CUDA Graph remains opt-in; this change does not modify the default enable policy.

Correctness

Fresh current-upstream builds were validated with the same 12 exact prompts in each cell:

• TP1 Graph OFF: 12/12 token-exact
• TP1 Graph ON: 12/12 token-exact
• TP2 Graph OFF: 12/12 token-exact
• TP2 Graph ON: 12/12 token-exact
• TP4 Graph OFF: 12/12 token-exact
• TP4 Graph ON: 12/12 token-exact

Overall: 6/6 valid cells, 72/72 token-exact, with zero request errors, timeouts, OOMs, output collapse, loops, or abnormal lengths.

A targeted ablation reproduced the TP2 Graph ON hang without the RowParallel synchronization, while the synchronization overlay completed 12/12 exact requests. The committed candidate then passed the full six-cell matrix.

Build and tests

• Fresh InfiniCore and InfiniLM compile/link
• Python extension import and zero-device smoke
• Focused static replay-metadata tests: 3/3
• ELF dependency, build-id, diff, ancestry, allowed-file, and public-content checks

Limitations

• CUDA Graph is not enabled by default.
• This PR makes no production performance claim.
• This PR does not claim a universal comparison with vLLM.
• The synchronization is scoped to TP RowParallel all-reduce completion.

[Copilot]

Pull request overview

This PR addresses correctness issues when using CUDA Graphs for static-KV decode with tensor-parallelism by ensuring decode metadata remains device-dynamic at replay time and by preventing TP row-parallel outputs from being consumed before their asynchronous all-reduce completes.

Changes:

• Add a CUDA-Graph recording path in static attention that uses graph-safe KV update + paged-attention ops with device-resident (dynamic) metadata.
• Gate static batching graph compilation on supported device/layout/dtype/KV-quant settings, and initialize replay metadata tensors as I32 device tensors (including block tables).
• Synchronize the current stream after TP RowParallel compute_linear_allreduce() before returning the output.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
test/static/test_static_graph_dynamic_metadata.py	Adds focused unit tests that assert graph recording forwards device-dynamic metadata into graph-safe ops and that replay inputs are copied into device-resident tensors.
csrc/layers/linear/linear.cpp	Ensures TP RowParallel forward waits for async all-reduce completion via stream synchronization.
csrc/layers/attention/backends/static_attn.hpp	Declares a private forward_graph_ helper for the graph-recording attention path.
csrc/layers/attention/backends/static_attn.cpp	Implements the graph-recording attention path using kv_caching_ and paged_attention_ with device-resident metadata.
csrc/engine/compiler/static_batching_compiler.cpp	Adds support checks for static graph attention and initializes/updates replay metadata tensors to remain device-dynamic.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[junjiewang253-ctrl]

Hi maintainers @Vincent777 , gentle ping on #540.

This PR is still scoped to two correctness issues:

CUDA Graph replay reusing stale host-resolved KV/attention metadata;
TP RowParallel returning before the asynchronous all-reduce has completed.

The branch has a focused correctness validation matrix (TP1/TP2/TP4 × Graph OFF/ON, 72/72 token-exact) and a targeted ablation for the TP2 Graph ON hang.

If this fix is still in scope, I’d appreciate a review when convenient. Since CI does not run automatically for fork PRs, could a maintainer also trigger /test or /retest?

I’m happy to split the CUDA Graph metadata fix and the TP synchronization fix into separate PRs if that would make review easier. Thanks!

[github-actions]

⛔ Only repository members can run test.