[RFC] Add Graph Capture Support for FLA Prefill and Training Paths on Ascend NPUs #1233

Proposal

Enable graph capture and replay for the GDN, KDA, and short-convolution paths in Flash Linear Attention (FLA) on Ascend NPUs.

The initial scope covers:

• GDN/KDA chunk attention forward;
• GDN/KDA chunk attention backward;
• Short convolution forward and backward;
• Complete GatedDeltaNet and KDA layer paths;
• Single-device Ascend NPU execution.

Context Parallel graph support will be implemented and validated in a subsequent phase.

This RFC is based on [Issue #1155](#1155). The issue proposes a backend-neutral graph contract based on static grids, static tensor addresses, device-side metadata, and capture-safe kernels. This RFC applies that contract to the Ascend NPU backend.

The first Ascend implementation will use TorchNPU's NPUGraph API. FLA operators will provide graph-safe fixed-capacity execution paths, while graph creation, buffer ownership, capture, replay, and bucket selection remain the responsibility of the caller.

Motivation

The current FLA varlen paths contain several dynamic operations:

1. Kernel grids depend on the actual number of chunks or tokens.
2. chunk_indices are generated dynamically from cu_seqlens.
3. Metadata shapes depend on the current input.
4. Input, output, and intermediate tensor addresses may change between calls.
5. A single layer may require many host-side kernel launches.

In eager mode, the CPU repeatedly performs metadata construction, grid calculation, kernel submission, memory management, and synchronization. For short prefill requests and short-sequence training, this host-side overhead can become a significant portion of end-to-end latency.

Graph capture allows the complete fixed-capacity execution sequence to be captured once and replayed many times. Subsequent executions only update static input buffers and submit the captured graph, reducing host dispatch and synchronization overhead.

Principles

• The initial scope is GDN/KDA chunk attention and short-convolution forward/backward paths.
• Decode-specific paths, including fused_recurrent and convolution cache-update paths, are out of scope.
• use_graph=False must preserve the existing eager behavior.
• Graph capture must not change the mathematical computation, state semantics, reduction semantics, or dtype behavior.
• Graph mode changes scheduling metadata, memory lifetime, and runtime submission, not the underlying algorithm.
• Inputs exceeding the configured capacity may be routed to the existing eager path by a higher-level caller.
• An explicit use_graph=True request must fail clearly when the backend, dtype, shape, or execution path is unsupported.
• FLA operators must not silently fall back to eager execution after graph capture has started.
• Graph buckets and graph lifetimes are owned by the caller.
• The initial implementation targets single-device Ascend NPU execution.
• Standard Attention changes are not part of this proposal.

Core design

Fixed-capacity execution

For a graph bucket, define:

• T_max: maximum number of tokens;
• N_max: maximum number of sequences;
• BT: chunk size;
• NT_max: maximum number of chunks.

The maximum chunk count is bounded by:

NT_max = ceil(T_max / BT) + N_max - 1

Graph mode uses:

• cu_seqlens with a fixed shape of [N_max + 1];
• chunk_indices with a fixed shape of [NT_max, 2];
• chunk_offsets with a fixed shape of [N_max + 1];
• fixed input, output, gradient, and intermediate buffers;
• a fixed kernel grid based on NT_max.

Unused chunk slots are filled with a sentinel value:

[-1, 0]

The kernel checks the sentinel before calculating any input, state, or output pointer:

i_n, i_t = load(chunk_indices[program_id])

if i_n < 0:
    return

# Calculate valid pointers and execute the original computation

This allows the graph to use a fixed grid while supporting different actual sequence layouts during replay.

Device-side metadata

The current host-side dynamic metadata path is not suitable for graph capture. Graph mode will generate fixed-shape metadata on the device:

static cu_seqlens
    -> device-side metadata generation
    -> chunk_indices [NT_max, 2]
    -> chunk_offsets [N_max + 1]

Each replay may update the static cu_seqlens buffer using copy_(). The current chunk layout is then regenerated on the device without returning data to the host.

The graph path must avoid:

• .item();
• .tolist();
• dynamic nonzero() results used for allocation;
• dynamic output shapes;
• using the actual chunk count to determine the launch grid;
• host-side loops over individual sequences;
• allocation during graph capture.

Static buffer contract

Graph replay records tensor addresses. Therefore all graph-related tensors must be allocated before capture:

static inputs
static outputs
static gradients
static chunk metadata
static intermediate states

Replay updates data in place:

static_inputs.copy_(new_inputs)
static_cu_seqlens.copy_(new_cu_seqlens)
graph.replay()

Tensor objects must not be replaced between capture and replay.

Graph-only buffers that are only partially written must be explicitly cleared or masked to prevent data from a previous replay from affecting the current result.

TorchNPU graph runtime

The first implementation will use the TorchNPU graph API:

graph = torch.npu.NPUGraph()

# Warmup: complete JIT, autotuning, and workspace initialization
run_step(static_inputs, use_graph=True)

with torch.npu.graph(graph):
    static_output = run_step(static_inputs, use_graph=True)

static_inputs.copy_(new_inputs)
graph.replay()

FLA operators will not create or own the full graph object. The caller is responsible for:

• selecting a graph bucket;
• allocating static buffers;
• running warmup;
• creating and holding the graph;
• updating static inputs;
• replaying the graph;
• managing graph and buffer lifetime.

The graph-safe contract should remain usable by future ACL Graph integrations and higher-level serving frameworks.

API proposal

The existing use_graph API should be reused and extended to GDN and short convolution:

chunk_kda(
    ...,
    use_graph=True,
    max_num_seqs=N_max,
    chunk_indices=static_chunk_indices,
    chunk_offsets=static_chunk_offsets,
)

chunk_gated_delta_rule(
    ...,
    use_graph=True,
    max_num_seqs=N_max,
    chunk_indices=static_chunk_indices,
    chunk_offsets=static_chunk_offsets,
)

causal_conv1d(
    ...,
    use_graph=True,
    max_num_seqs=N_max,
    chunk_indices=static_chunk_indices,
    chunk_offsets=static_chunk_offsets,
)

The API contract is:

• use_graph=False uses the existing eager path;
• use_graph=True selects the fixed-capacity graph-safe path;
• max_num_seqs must match cu_seqlens.shape[0] - 1;
• externally supplied metadata must match the expected shape, dtype, device, and layout;
• graph mode must not regenerate dynamic-shape metadata;
• low-level operators must not silently fall back;
• high-level wrappers may select eager execution for inputs exceeding the configured capacity.

Implementation plan

Step 0: Ascend graph capability validation

Add runtime capability tests covering:

• NPUGraph capture and replay;
• static address reuse;
• copy_() input updates;
• graph pool and stream behavior;
• searchsorted, where, cumsum, and related metadata operations;
• repeated replay of static chunk metadata;
• fixed-grid Triton-Ascend kernels;
• custom autograd forward and backward;
• JIT, autotuning, and workspace initialization before capture.

This phase must produce a capability matrix with explicit supported, failed, and unsupported cases.

Step 1: Shared metadata and buffer contract

Improve the common graph infrastructure:

• static chunk-index generation;
• static chunk-offset generation;
• zero-length tail sequence handling;
• int32 and int64 cu_seqlens coverage;
• graph buffer namespaces;
• fixed intermediate buffers;
• graph replay buffer clearing;
• protection against stale metadata caused by identity caching.

Exit criteria:

• the same static buffers can replay multiple sequence layouts;
• metadata matches the eager implementation;
• no host-side dynamic shape or synchronization is required;
• eager behavior remains unchanged.

Step 2: KDA single-device graph support

KDA will be the first complete operator path.

The implementation will update:

• the KDA public operator;
• KDA forward and backward wrappers;
• intra, WY, gate, and state paths;
• the Triton-Ascend backend verifier;
• shared chunk_delta_h, chunk_o, KKT, and triangular-solve paths.

The implementation must ensure:

• all kernels use a fixed NT_max;
• sentinel checks precede all pointer calculations;
• partial-token replays do not contaminate outputs or gradients;
• states and intermediates use fixed N_max and NT_max capacities;
• existing Ascend tl.dot lhs reload/pristine-copy protections are preserved;
• the mathematical computation and dtype semantics remain unchanged.

Exit criteria:

• KDA forward graph capture and replay pass;
• KDA forward/backward graph capture and replay pass;
• gate, GVA, state, partial-token, and sequence-layout cases pass;
• graph results match eager results on the same NPU.

Step 3: GDN and short-convolution graph support

GDN must support the complete forward pipeline:

gate / chunk cumsum
    -> forward intra
    -> recompute_w_u
    -> forward h
    -> forward o

The backward path must cover the corresponding recomputation, state, WY, gate, and gradient kernels.

Short convolution must cover:

• dense and varlen inputs;
• forward and backward;
• bias, activation, and residual paths;
• initial and final states;
• partial-token replay;
• fixed chunk metadata;
• fixed state buffers.

Decode update and cache-update kernels remain unchanged.

Step 4: Layer integration

Apply the graph contract to the complete layer path:

projection
    -> short convolution
    -> GDN/KDA chunk attention
    -> output projection

The caller should own a layer-level graph rather than creating one graph for each sub-operator. This ensures consistent:

• stream ordering;
• buffer lifetime;
• capture boundaries;
• replay behavior;
• intermediate tensor addresses.

Exit criteria:

• GatedDeltaNet layer capture/replay passes;
• KDA layer capture/replay passes;
• outputs, final states, and gradients match the eager path.

Step 5: Context Parallel graph support

Context Parallel support will be implemented after single-device graph support is stable.

This phase will validate:

• HCCL collective capture behavior;
• lockstep capture and replay across ranks;
• device-resident rank metadata;
• stream, port, and device mapping stability;
• zigzag layout compatibility.

If the current TorchNPU/CANN combination does not provide reliable collective capture, Context Parallel will retain eager execution and will not be advertised as graph-supported.

Tests

The following tests should be added:

tests/ops/test_npu_graph_runtime.py
tests/ops/utils/test_graph_index_npu.py
tests/ops/test_kda_graph_npu.py
tests/ops/test_gdn_graph_npu.py
tests/modules/test_conv_graph_npu.py
tests/layers/test_gated_deltanet_graph_npu.py
tests/layers/test_kda_layer_graph_npu.py
tests/ops/test_graph_backend_dispatch_npu.py
tests/context_parallel/test_cp_kda_graph_npu.py

Each operator should cover:

• dense batches;
• varlen sequences;
• zero-length tail sequences;
• actual_t < T_max;
• different sequence splits;
• BT=32 and BT=64;
• fp16;
• bf16 where supported by the Ascend backend;
• initial and final states;
• gate branches;
• forward-only execution;
• forward plus backward;
• replay after changing inputs, states, gates, beta, gradients, and cu_seqlens;
• non-zero padding data;
• multiple replays to detect stale metadata.

Correctness comparisons must use the same Ascend device and software stack for both eager and graph execution.

Graph tests should follow this sequence:

1. Allocate fixed-shape static inputs.
2. Allocate fixed-shape static cu_seqlens.
3. Run eager warmup to initialize JIT, autotuning, and workspaces.
4. Capture the graph once.
5. Prepare another valid input and sequence layout.
6. Update static buffers using copy_().
7. Replay and synchronize.
8. Run an eager reference using the same input content.
9. Compare outputs, states, gradients, and padding regions.
10. Repeat with multiple sequence layouts.

Failures must be categorized as:

• unsupported graph operator;
• graph address or lifetime failure;
• stale or dynamically generated metadata;
• sentinel, offset, or padding correctness failure;
• pre-existing eager backend failure;
• distributed rank, port, or collective failure.

Graph failures must not be hidden by unconditional skips.

Performance evaluation

At least four execution modes should be compared:

1. eager-live: actual input sizes and actual chunk counts;
2. eager-fixed: fixed-capacity execution without graph replay;
3. graph-replay: graph replay cost only;
4. graph-update-replay: input, metadata, and gradient updates plus replay.

The benchmark should report:

• p50 and p95 latency;
• warmup, capture, replay, and update time;
• host launch count;
• NPU kernel count;
• token and chunk utilization;
• peak and resident memory;
• forward, forward/backward, and layer-level performance;
• graph/eager numerical differences.

The primary performance target is to reduce host dispatch and scheduling overhead. Graph capture does not change the algorithmic complexity and may not improve long-sequence workloads where computation dominates.

The benchmark must also quantify:

• fixed-capacity padding cost;
• graph memory overhead;
• eager-path regression;
• graph update overhead;
• break-even sequence length.

CI plan

The Ascend CI should add separate graph jobs.

Single-device graph job

• Use one fixed NPU device;
• run graph runtime, metadata, KDA, GDN, convolution, and layer tests serially;
• disable xdist;
• use a fixed ASCEND_RT_VISIBLE_DEVICES value.

Context Parallel graph job

• Use at least two NPU devices;
• use torchrun or multiprocessing;
• manage rank, port, and device mapping independently;
• do not combine with xdist-based single-device tests.

Eager regression job

Continue running the existing eager test suite to ensure that:

use_graph=False

preserves NPU, CUDA, and CPU behavior.

The NPU workflow change-detection rules should cover:

fla/ops/**
fla/modules/**
fla/layers/**
tests/ops/**
tests/modules/**
tests/layers/**
tests/context_parallel/**
pyproject.toml
tests/conftest.py
.github/workflows/ascend-a2-ci.yml

The initial validation environment should be recorded in CI artifacts:

CANN 9.1.0
Python 3.12
torch 2.9.0+cpu
torch_npu 2.9.0.post6
triton-ascend 3.2.2
Ascend 910 series

Limitations

• Graph execution is limited to the configured T_max and N_max.
• Inputs exceeding the configured capacity must use eager execution.
• Each graph bucket requires static buffers and graph memory.
• Different (T_max, N_max) configurations require separate captures.
• JIT, autotuning, and workspace initialization must complete before capture.
• Fixed-capacity execution may process padding tasks.
• Unsupported NPU operators cannot silently fall back during capture.
• Decode cache-update paths are not included in the first phase.
• Context Parallel graph support is not guaranteed in the first phase.
• The first implementation should not be combined with torch.compile graph capture.
• Performance benefits must be demonstrated on real Ascend hardware and cannot be inferred from CUDA results.

Rationale

This proposal does not introduce a new Attention algorithm or modify the mathematical behavior of FLA. Its purpose is to transform dynamic varlen chunk execution into a fixed-capacity, fixed-address, fixed-grid execution graph.

The expected benefits come from reducing:

• host-side kernel launch overhead;
• dynamic metadata construction;
• runtime scheduling;
• repeated synchronization;
• repeated allocation and setup work.

The expected benefit is greatest for short prefill and short-sequence training workloads with many relatively small kernels. For long sequences, computation dominates and the benefit may be smaller. For low-utilization buckets, padding and memory overhead may offset the replay benefit.

The implementation will therefore be considered successful only when it demonstrates:

• graph/eager numerical equivalence;
• no regression in the default eager path;
• acceptable fixed-capacity memory cost;
• stable graph replay behavior;
• measurable host-dispatch or end-to-end improvements on Ascend hardware;
• clear behavior for unsupported configurations.

References

• [FLA Issue #1155: Graph capture support for prefill and training]([RFC] Graph capture support for prefill and training: GDN/KDA chunk and convolution paths #1155)
• [TorchNPU NPUGraph developer notes](https://github.com/Ascend/pytorch/blob/master/docs/zh/developer_notes/npugraph.md)
• [Ascend PyTorch adapter](https://github.com/Ascend/pytorch)