onnxruntime
c7b64756 - Skip redundant present_key/value copy when aliased to external KV cache (follow-up to #29715) (#31150)

Commit
6 days ago
Skip redundant present_key/value copy when aliased to external KV cache (follow-up to #29715) (#31150) ## Summary Follow-up to #29715 (cuDNN SDPA decode tier for the ONNX standard `Attention` CUDA kernel). On the external-KV-cache path (4-D BNSH, `nonpad_kv_seqlen`), all four CUDA attention backends (Flash, cuDNN SDPA, Memory-Efficient, unfused) unconditionally `cudaMemcpyAsync`'d the entire K/V cache into `present_key`/`present_value`, even when the caller IOBinds `present_key`/`present_value` to the SAME device buffer as the K/V cache input — the documented `TensorScatter` + IOBinding production pattern (mirroring `TensorScatter`'s own `.MayInplace(0, 0)` self-copy skip and `GroupQueryAttention`'s `past_key`==`present_key` aliasing). This contradicted the file's own PERFORMANCE NOTE and was technically undefined behavior (`cudaMemcpyAsync` requires non-overlapping src/dst; a full self-copy is maximal overlap). ## What changed - Added `llm_attention_detail::CopyKVToPresent(src, dst, stream)`, a shared helper that skips the D2D copy via a pointer-equality check when `present_*` aliases the K/V cache buffer, with a greppable VERBOSE log tag (`present_copy_skipped`) for test observability. Applied at all 8 call sites (K+V × 4 backends), only in the 4-D BNSH branches (3-D BSNH always needs a layout-changing transpose and can never alias). - Added a defensive size-equality `ORT_ENFORCE` inside the helper: proven safe today (present shape only equals K/V's shape when `past_sequence_length == 0`), but guards against a future caller reusing this helper outside that precondition. - `TestAttentionPresentKVCopySkip` (8 parameterized tests) covers all 4 backends × aliased/ non-aliased, asserting both the copy-skip fires (or doesn't) via the log tag AND that the intended backend actually dispatched (not a silent MATH fallback), plus output/present_key/ present_value correctness in both cases. ## Notes Also includes (unpushed until now) a NOTE documenting the Phase 3 (cuDNN SDPA prefill chunking) investigation and its abandonment — chunking added complexity without a clear latency win at the shapes profiled, so it was not pursued further. ## Testing - Full `test_onnx_attention` suite: 342/342 passed (334 pre-existing + 8 new). - New test class re-verified individually with `-v`; confirmed via dispatch-log scraping that flash/efficient/cudnn/math parameterizations each hit their intended backend on an A100 (SM80), not a MATH fallback. - `lintrunner` clean. This PR went through an internal multi-agent review pass (readability, correctness, adversarial, spec/invariant, cross-module integration, and QA execution) before being opened; findings from that pass are already incorporated. --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a1011908-d09e-411a-8979-fcf0dc18e20e Copilot-Session: d745b4ab-b999-4523-b05c-7c4f463a0cab
Author
Parents
Loading