[yuhengcai1]

ChunkCausalAttention updates its cumulative KV and K-sum tensors in place after earlier chunks have used them in matrix products. In multi-chunk FP32/FP64 backward, this either raises a tensor-version error or, with non-reentrant checkpointing, returns incorrect gradients while the forward output remains correct.

Replace the two in-place updates with out-of-place additions so each chunk retains the prefix values needed by autograd. The forward formula and accumulation order are unchanged.

Add 12 regression cases comparing the actual attention module with an independent dense chunk-causal reference. They cover one/two/three chunks, direct/non-reentrant checkpoint execution, FP64 and FP32 with Q/K normalization, and nontrivial RoPE. Each checks forward values plus input and every learned parameter gradient using the same random output probe.

Validation with repository-pinned PyTorch 2.9.1, torchvision 0.24.1, and timm 0.6.13 on CPU:

• Before: 4 passed / 8 failed. All four single-chunk controls pass; four direct multi-chunk cases raise version errors; four checkpoint cases fail the gradient comparison.
• After: 12 passed.
• Reintroducing either in-place update separately makes all eight multi-chunk cases fail again.
• All applicable repository-pinned Python pre-commit hooks and git diff --check pass.

The local runner imports the complete, unmodified Sana source modules with namespace parents to bypass unrelated package initializers. It uses the real timm base constructor, normal committed pytest functions, PyTorch autograd, and checkpoint implementation; no attention arithmetic or classes are mocked. Source hashes and JUnit results were captured. Validation is limited to CPU FP32/FP64: no CUDA/mixed-precision or end-to-end model training run is claimed.

[lawrence-cj]

Thanks @yuhengcai1, LGTM. The forward is unchanged, and the out-of-place adds keep the chunk prefixes autograd needs.