fix(kvcache): trace-faithful PositionalKVCache.update (#763)

[michalharakal]

Summary

PositionalKVCache.update is not trace-faithful: it copies the incoming K/V tensor data into a heap FloatArray (writeAt) and reads it back via ctx.fromData (sliceView), bypassing ctx.ops. Under tracing the symbolic K/V carry no data, so toComputeGraph(embedConstants=true) bakes an all-zero KV buffer as stablehlo.constant and disconnects the computed k_proj/v_proj — the exported decoder attends over K=V=0.

Eager inference is fine (the buffer holds real data). In export the bug was masked in plain decoders (the unnormalized FFN dominates the residual stream) and exposed by Gemma's sandwichNorms (post_ffw_norm normalizes the FFN, so the lost attention becomes significant → ~1.4×/block logit error). Tracked as SKaiNET#763.

Fix

When ctx.isRecording, wire K/V functionally through ops.concat (the same history AppendKVCache already uses) instead of the raw heap buffer, so the StableHLO export carries the real projections. The eager fast-path (heap buffer) is unchanged. Side effect: the traced graph no longer surfaces dangling KV-cache buffer leaves as graph outputs.

Verification

End-to-end via skainet-iree-conformance gemma-decoder — a real 2-block Gemma3 decoder with qk-norm + sandwichNorms + layer-output-scale: export → iree-compile → iree-run-module matches the SKaiNET-CPU oracle at max_abs_err 3.8e-6 (was ~5.3). Adds KvCacheTraceFidelityTest (asserts a traced sandwich decoder bakes no zero KV-cache constants).

Note: the other PositionalKVCache-derived shared variants (SharedPositionalKVCache, PaddedSharedPositionalKVCache, OwnerReadOnlyKVCache) still use the buffer sliceView on read and would need the same treatment for kv-shared exports — out of scope here (gemma-decoder uses the non-shared cache).

🤖 Generated with Claude Code