[Feat] Support platform-graph capture and replay for KDA chunk

[zhiyuan1i]

Summary

Add use_graph + max_num_seqs to chunk_kda so the forward and backward can be captured in a platform graph (CUDA graph). The varlen chunk list is built on-device at a static shape via prepare_chunk_indices_static, kernels early-return on sentinel rows, and backward intermediates come from persistent buffers via get_static_buffer. The eager path (use_graph=False) is unchanged.

Also supports cp_context (context parallelism) in graph capture/replay: FLACPContext gains device scalars for pre/post num_ranks, the merge kernel handles NUM_RANKS_ON_DEVICE, and the pre_process wrapper always launches.

NPU backends explicitly reject use_graph=True with NotImplementedError.

Changes

• fla/ops/utils/graph.py (new, 27 lines): get_static_buffer hands out persistent intermediates whose addresses stay valid for the graph's lifetime.
• fla/ops/utils/index.py: prepare_chunk_indices_static builds chunk_indices/chunk_offsets with on-device ops at a fixed shape for graph recording.
• fla/ops/kda/chunk.py: use_graph + max_num_seqs entry point; varlen chunk list built on-device; eager path unchanged.
• fla/ops/kda/chunk_fwd.py / chunk_bwd.py: kernels early-return on sentinel rows; backward intermediates from persistent buffers.
• fla/ops/kda/chunk_intra.py / chunk_intra_token_parallel.py: use_graph guard for tokens beyond cu_seqlens[-1].
• fla/ops/common/chunk_delta_h.py: use_graph support in the shared delta-rule state recurrence.
• fla/ops/cp/chunk_delta_h.py / context.py: CP graph capture/replay support.
• fla/ops/utils/cumsum.py: zeros_like under use_graph so dA/dbias are not polluted by dirty padding rows.
• NPU backends: explicit NotImplementedError on use_graph=True; chunk_offsets no longer recomputed when provided; zeros_like parity.

Tests

• tests/ops/test_kda_graph.py (386 lines): 16 tests covering fwd/bwd graph capture, varlen, state, gate, partial tokens.
• tests/context_parallel/test_cp_kda_graph.py (179 lines): CP graph capture/replay with world=2, 64/64 checks pass.

Related

• RFC: [RFC] Graph capture support for prefill and training: GDN/KDA chunk and convolution paths #1155 (tracking issue for CUDA graph support across fla)

Checklist

• I have read CONTRIBUTING.md and follow its conventions (code style, docstrings, commit prefixes).
• I have read AGENTS.md and, where my change matches its scope, the relevant skill under .agents/skills.
• Dependent tests pass locally or in CI, and new behavior is covered by tests where applicable (tick as N/A for changes with no testable code, e.g. docs-only).
• Kernel changes include same-hardware before/after benchmark numbers, dense + varlen where applicable (tick as N/A when no kernel code changed).
• This PR is minor/cosmetic-only (typo, formatting, style-only tweaks) — tick only if it is, and justify below.