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