[Gajesh2007]

Two-pass vector SDPA casts unnormalized partial sums to the input dtype before combining them. BF16 can lose cancellation residuals and FP16 can overflow even when the final attention output is representable. This patch keeps partial outputs in FP32 through the second pass and casts only the final result. Kernel entry names change so a stale shader cannot silently satisfy the new buffer contract.

The host intermediate allocation, ordinary/GQA first pass, second pass and Metal declarations change together. Python regression cases cover cancellation and overflow across D64/D128/D256, 32/128 partitions and masked/unmasked paths. Inputs and public output dtypes are unchanged. Low-precision intermediate storage grows to four bytes per element.

Before

flowchart LR
  QKV[Input Q K V] --> First[Two-pass vector SDPA first pass]
  First --> Narrow[Store unnormalized partials in input dtype]
  Narrow --> Lost[Cancellation residual lost or FP16 overflow]
  Lost --> Reduce[Second pass combines rounded partials]

Loading

After

flowchart LR
  QKV[Input Q K V] --> First[FP32 first-pass accumulation]
  First --> Partial[FP32 intermediate buffer]
  Partial --> Reduce[FP32 second-pass combination]
  Reduce --> Output[Cast final result to public output dtype]
  Names[New fp32partials entry names] --> Contract[Host and shader contract checked together]

Loading

Validation

• Reviewed candidate source manifest e7b769e2713e48c53fb62fbe74e449bb409c8fc1b7e5f036b14e808a7c3da3f5: three Swift analytic tests passed all 50 internal subcases with zero skips. The new Python tests are included but are not claimed as executed.
• An isolated canonical native SwiftPM build against the same exact core/Swift/native source and reviewed fresh metallib passed all 13 existing Qwen D256 cases, with unchanged assertions and tolerances. Coverage includes BF16/FP32 histories near 4096 and 5523/5585, original-query FP32 reference bounds, planted boundary/GQA faults, and exact fixed/segmented storage/output equality.
• BF16 decode-case relative L2 was 0.0016175–0.0016959 for both corrected contiguous and paged arms, below the existing 0.01 bound. These are B1, one-query synthetic precision controls. Remaining real-model trajectory/MTP checks retain their own strict results; this is not whole-release certification.
• Full source/dependency/compiled-graph/resource identities and owned child retirement were verified. No model weights were loaded by the 13-case run.

This is a new focused follow-up to merged #15. The companion Swift PR #21 regenerates embedded shader sources and pins this core commit; native PR #141 consumes that Swift revision. Parent PR #851 carries serving policy, regression integration and release evidence. No signing workflow or release publication is part of this PR.

Evidence: candidate D256 precision report, published with the companion parent update.

[Gajesh2007]

@codex review this