[yzhangcs]

Restores #1294 after the October 15 repository-history incident. Original contribution by taking-lying-flat; earlier review and discussion remain in #1294. This recovery draft uses the existing source branch at ccea6748 and preserves its commits. Validation details below are retained from the original PR and have not been rerun for this recovery. Current-main integration and readiness need confirmation before review resumes.

Summary

With use_gate_in_kernel=True and no A_log, unbounded recurrent and fused-gate KDA paths used exp(1) as the decay scale, while chunk KDA rejected the call. Use a unit scale consistently across chunk, recurrent, fused-gate, and Triton-Ascend paths, equivalent to passing an explicit zero A_log.

The PyTorch reference and existing parameterized tests cover this optional-argument contract. Explicit A_log, bounded-gate formulas, parameter shapes, and checkpoint layouts retain their existing behavior. This is a correctness fix for callers that omit A_log; it does not claim a training-throughput improvement.

Test plan

107 passed, 8 deselected on an RTX A1000 Laptop GPU (4 GiB), PyTorch 2.12.1+cu130, and Triton 3.7.1. Backend dispatch is disabled to exercise the changed Triton kernels directly. The original tolerances and NaN-memory-poisoning fixture are retained. All pre-commit hooks and git diff --check passed.

FLA_CI_ENV=0 FLA_DISABLE_BACKEND_DISPATCH=1 python -m pytest \
  tests/ops/test_kda.py::test_chunk \
  tests/ops/test_kda.py::test_fused_recurrent_gate_in_kernel \
  tests/ops/test_kda.py::test_gate \
  -k 'not (test_chunk and (D100 or D128 or T2048 or HV8))' \
  -q --tb=short -o log_cli=false
pre-commit run --all-files
git diff --check

The selection includes the omitted-A_log regressions, explicit-parameter controls, bounded and unbounded gates, FP32/FP16/BF16, output and final-state comparisons, and supported gradients. Large chunk cases are excluded to fit local GPU memory. Two explicit-A_log BF16 chunk cases emitted dA warnings under the existing warning=True checks (relative RMS errors 0.006700 and 0.003616); no tolerances were changed. Full-model, context-parallel, and local NPU tests were not run.

The prior Ascend CI run at f85ebbbc stopped on three activation stride failures (sigmoid, swiglu, and powglu), before completing KDA coverage. This branch now includes the upstream fix from #1374. The refreshed Ascend run remains the NPU validation gate.

Benchmark / NCU (kernel changes only)

Previously recorded same-machine measurements compare CUDA fix 039ba807 with baseline 0259f80c on an RTX A1000 Laptop GPU, PyTorch 2.13.0+cu130, and Triton 3.7.1. The CUDA formulas are unchanged by the subsequent test consolidation and mainline merges.

Workload: B=1, T=73, H=1, HV=2, K=V=64, FP32, explicit zero A_log, nonzero initial state, final-state output enabled, and varlen boundaries [0, 3, 73]. Median latency uses triton.testing.do_bench(fn, warmup=100, rep=300, quantiles=[0.5, 0.2, 0.8]).

Path	Before (µs)	After (µs)
Chunk, dense	52.224	52.224
Chunk, varlen	53.248	53.248
Recurrent, dense	76.800	77.152
Recurrent, varlen	73.728	74.752
Gate forward + backward	9.216	8.192

Performance is neutral within the measurement spread. Nsight Compute --set full could not collect counters because the driver returned ERR_NVGPUCTRPERM. The table records the original CUDA-fix validation; current-branch correctness is checked separately above.

Breaking changes

No parameter shapes or checkpoint layouts change. Previously rejected chunk calls now work. Unbounded recurrent and fused-gate calls that omit A_log change from the unintended exp(1) scale to unit scale. Pass an explicit ones tensor as A_log to reproduce the former recurrent scaling.

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.