[FightingZhen]

Summary

Reduce repeated compilation in ordinary KDA, KDA graph, and recurrent KDA benchmarks on Ascend. Ordinary KDA took 643 s in the example baseline CI run.

• Reuse a dedicated Triton cache keyed by the pinned software stack and checkout revision (github.sha, the merge commit for PR runs), with a stack-specific restore prefix.
• Prewarm on selected main-branch pushes. Check the main cache every Monday, Wednesday, and Friday at 00:00 Asia/Shanghai, and prewarm on a cache miss.
• Add stage timings and benchmark result checks; save the cache only after all three KDA benchmark result checks pass.

Benchmark coverage, warmup/measurement settings, operator scheduling, and prior-step failure behavior are preserved. Ordinary KDA checks result coverage and finite positive timings; graph KDA checks eager/graph numerical consistency; recurrent KDA compares outputs and final states against the naive reference.

Test plan

• Local environment: Ascend910, CANN 9.2.0, Python 3.11.17, PyTorch 2.10.0+cpu, torch_npu 2.10.0.post7, Triton-Ascend 3.6.0.

• FLA_CI_ENV=1 python -m pytest tests/benchmarks/test_graph_benchmark.py -q: 15 passed.

• Workflow syntax, cache isolation, failure propagation, result-validation checks, and Python lint/format checks passed locally.

• Cold/warm ordinary KDA runs produced all 12 results; graph numerical validation passed.

• Recurrent KDA completed all 16 cases, both multibuffer configurations, and 7 measurement rounds. Numerical checks passed in all three local runs.

• Additional local numerical checks on Ascend910_9382 / CANN 9.2.0 / Python 3.12.15: existing dense BF16 and varlen FP16 KDA cases passed (2 passed, no skips) with FLA_CI_ENV=0 and confirmed triton_ascend forward/backward routing. Existing tolerances are unchanged; the dense gate dA comparison retains its existing warning-only behavior (ratio 0.003180 versus 0.003).

• Latest revision: 053afe97. A2 runs: correctness, benchmarks. Same-runner cold/warm CI comparison remains pending.

Benchmark / NCU

Local command wall times, comparing an empty cache with reuse of the same cache:

Workload	Empty cache	Reused cache
Ordinary KDA: 6 shapes × fwd/fwdbwd, BF16	167.75 s	25.08 s
KDA graph: B1/T128/H2/D64/N2, fwdbwd, FP16	81.56 s	27.45 s
Recurrent KDA: dense/packed, FP32, both multibuffer settings	63.08 s	39.33 / 40.48 s

Commands used with MPLBACKEND=Agg, FLA_BENCH_WARMUP_MS=25, FLA_BENCH_REP_MS=100, and FLA_BENCH_OP_WARMUP_ITERS=3:

FLA_CI_ENV=1 python -m benchmarks.ops.run \
  --op chunk_kda --base '' --json kda.json

FLA_CI_ENV=1 python -m benchmarks.ops.run --graph \
  --op chunk_kda --no-base --modes fwdbwd \
  --custom-shapes '{"smoke":{"B":1,"T":128,"H":2,"D":64,"N":2}}' \
  --warmup 3 --iterations 10 --json graph.json

FLA_CI_ENV=0 python -m benchmarks.ops.benchmark_kda_recurrent \
  --json recurrent.json

Each workload used a separate empty TRITON_CACHE_DIR, then reused that directory. Ordinary/graph measurements were collected on a59bc5df plus instrumentation; recurrent measurements were collected on 08040acc.

These are local cache comparisons, not kernel speedups or confirmed end-to-end CI savings. Same-runner CI cold/warm validation is still needed. NCU: N/A; no kernel code changed.

Breaking changes

None. Public APIs, numerics, and result schemas remain unchanged.

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.
• I understand this PR exceeds 500 changed lines and have explained its single purpose and why it should stay together below.

[FightingZhen]

Comparing the example baseline CI with the current PR CI, both running on Ascend910B3 with the same CANN 9.1.0 / PyTorch 2.10.0 / torch_npu 2.10.0.post7 / Triton-Ascend 3.6.0 stack:

KDA workload	Baseline command wall time	Current PR
Ordinary KDA: 6 BF16 shapes × fwd/fwdbwd	643 s	52 s
KDA graph smoke: FP16, fwdbwd	Not separately timed	39 s
Recurrent KDA: dense/packed, FP32, both multibuffer settings	Not included	153 s

The ordinary KDA command completed 591 s sooner, a 91.9% reduction, with the same 12 shape/mode results. Its steady-state benchmark medians changed by −1.8% to +2.6%, consistent with the main benefit coming from reduced compilation/preparation time. The graph command lacks a separately timed baseline, so no precise before/after speedup is claimed for it. Recurrent KDA adds coverage absent from the example baseline and has no corresponding comparison.

This PR run restored an existing cache through the software-stack prefix and saved the updated archive under the checkout SHA. Restore/save steps took approximately 7 s / 10 s, and all three KDA result checks passed.

These observations concern individual KDA benchmark commands. The runs use different source revisions and runner instances, so a controlled same-revision cold/warm comparison is still needed to isolate the cache benefit precisely.

[zheliuyu]

@FightingZhen ascend-a2-ci.yml can use the same restore-then-save pattern. Please give it a separate prefix: that job is the 8-card pytest matrix, a different kernel set from this KDA benchmark cache. Save only after both pytest steps pass.