[mmshad]

Summary

• A cached forward that appends more than one token to a cache that already holds keys now calls SDPA with a lower-right causal mask (torch.nn.attention.bias.causal_lower_right), so each appended token attends to every cached key and to the new keys up to its own position. With is_causal=True, SDPA aligned the triangle to the first key, and the appended rows were wrong while the cache still advanced (Appending several tokens to a non-empty KV cache gives wrong logits #541).
• _attention_with_weights offsets its triangle by the cached keys the same way.
• Every other SDPA path keeps its previous call: training and prefill into an empty cache, one-token decode steps, and masked forwards. With sdpa_backend = "flash" in bf16 on CUDA, the append raised No available kernel; it now runs on the flash kernel.
• On torch 2.11 the bias constructor is a graph break: a cached multi-token append compiles by default but not with fullgraph=True. No compile call site in the repo passes KV caches.
• docs/architecture/model.md lists the new case.

Testing

• On an H200, a 4-layer model (8 query heads, 2 KV heads, batch 2) appending chunks of 40, 32 and 24 tokens: on main the appended rows differ from a full forward by up to 1.24; with this change they match exactly in bf16 (auto, flash, efficient, cudnn, math) and within 7.2e-7 in fp32 (auto, efficient, math). The full forward, prefill plus one-token decode steps and greedy generate() tokens are bitwise identical to main on auto, efficient and math in bf16 and fp32, and on cudnn in bf16. A compiled model's appended chunks differ from the eager full forward by 1.2e-2 in bf16, the same as a compiled forward without a cache.
• New tests: tests/unit/test_masking.py::TestKVCacheDecode appends chunks (3+2+3, 1+7, 5+1+2) to a filled cache with GQA and batch 2, fused and with captured weights, and checks an appended chunk's captured weights against the full forward's rows; tests/integration/test_kv_cache_append.py runs the append on CUDA under auto, flash and efficient in bf16 and efficient and math in fp32. Without the SDPA change 3 unit tests fail, without the capture change 4, and the integration test fails 5 of 5 on main.
• On 4 H200s: tests/unit 3725 passed, 4 skipped, 5 xfailed; tests/integration (not slow) 128 passed; tests/distributed 129 passed, 2 skipped, 1 xfailed; tests/e2e --e2e 42 passed, 1 skipped, 2 xfailed. ruff clean; pyright 0 errors.

Closes #541

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
kempnerforge/model/attention.py	97.82% <100.00%> (+1.23%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.