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

[titaiwangms]

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.

[Copilot]

Pull request overview

This PR improves the CUDA implementation of the ONNX-domain Attention op (opset 24 external KV-cache path) by avoiding a redundant—and in the aliased case undefined—device-to-device cudaMemcpyAsync when present_key/present_value are bound to the same device buffer as the K/V cache. It also adds Python coverage that attempts to prove the copy-skip behavior via a stable log tag plus backend-dispatch observability.

Changes:

• Add llm_attention_detail::CopyKVToPresent() helper in attention.cc that skips the present-cache D2D copy when src and dst alias, and logs a greppable present_copy_skipped tag.
• Apply the helper at all 4 backends’ 4-D BNSH present-cache population sites (Flash, cuDNN SDPA, MEA, unfused).
• Add parameterized Python tests to validate both aliased vs non-aliased behavior across the 4 backends and to assert the intended backend dispatch via attention debug info.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
onnxruntime/core/providers/cuda/llm/attention.cc	Introduces CopyKVToPresent and routes 4-D BNSH present population through it across all CUDA backends; updates performance/behavior notes.
onnxruntime/test/python/transformers/test_onnx_attention/test_tensorscatter_attention.py	Adds new parameterized tests that exercise alias vs non-alias present-cache binding and assert copy-skip via logs + backend dispatch via debug output.

[tianleiwu]

Confirming this — and the root cause is a bit different from the ordering described above.

IOBinding::GetOutputs() returns outputs_, which is populated in bind order, not graph-output order (onnxruntime/core/session/IOBinding.cc:90-96):

auto it = mapped_output_names_.emplace(name, output_names_.size());
if (/* new */) {
  output_names_.push_back(name);
  outputs_.push_back(ml_value);

_run_tensorscatter_attention_4d binds in this order: output(0), updated_key_cache(1), updated_value_cache(2), present_key(3), present_value(4). So indices 1/2 are the TensorScatter outputs regardless of the graph declaration order.

Why this matters more than "weakened coverage": updated_key_cache holds the in-place scattered cache, which is exactly ref_present_k. So both assert_allclose calls pass vacuously, and test_copy_still_runs_when_present_not_aliased never validates the D2D copy result — the single thing that test exists to check. A regression that skipped the copy unconditionally would still be green there.

present_k_ort / present_v_ort are already in scope in that branch, so the fix is to read from them directly, matching how the other helpers in this file do it (lines 804-806):

present_k = present_k_ort.numpy()
present_v = present_v_ort.numpy()

[tianleiwu]

Overall

Nice, well-scoped optimization. I verified the past_sequence_length == 0 precondition the helper depends on actually holds at all 8 call sites (Flash/MEA via present_kv_already_populated == false, cuDNN via the explicit ORT_ENFORCE at the top of RunCudnnSdpaAttention, unfused via present_already_populated == false) — so present_* really is shape-identical to K/V there and the skip is a faithful identity, not just a size coincidence. Restricting to the !is_bsnh branches is right. Collapsing 8 duplicated cudaMemcpyAsync blocks into one helper is a readability win on its own.

Also a genuine plus on the test side: asserting the dispatched tier via SdpaKernel=... instead of trusting numeric correctness alone. Every non-MATH parameterization ORs in a MATH fallback bit, so without that assertion three of the four cases could silently degrade and still pass green.

Verdict: COMMENT. One issue that makes half the new coverage vacuous, plus a few suggestions.

Highest-priority item

The non-aliased branch reads io_binding.get_outputs()[1]/[2]. IOBinding::GetOutputs() returns outputs_, which is populated in bind order (onnxruntime/core/session/IOBinding.cc:90-96), not graph-output order — and this helper binds output, updated_key_cache, updated_value_cache, present_key, present_value. So indices 1/2 are the TensorScatter outputs. Since updated_key_cache holds the scattered cache (== ref_present_k), the assertions pass vacuously and test_copy_still_runs_when_present_not_aliased never validates the D2D copy result — the one thing it exists to check. present_k_ort/present_v_ort are already in scope; read from them directly the way the other helpers in this file do (lines 804-806). Replied with detail on the existing thread rather than opening a duplicate.

Inline suggestions

Left five inline comments covering: the LOGS_DEFAULT vs. session-logger choice (which also unblocks a non-flaky test assertion), ORT_ENFORCE in a Status-returning helper, the volume of Phase 2/3 retrospective prose in ComputeInternal, the process-global negative log assertion, and the backend-availability predicates.

Nitpicks (no action required)

• The dst == nullptr early return in CopyKVToPresent is dead code — all 8 call sites already gate on present_key != nullptr / present_value != nullptr.
• Pointer equality only detects exact aliasing; a same-size, partially-overlapping dst would still reach cudaMemcpyAsync with overlapping ranges. Essentially unreachable through IOBinding, and the size check narrows it further — worth one honest clause in the comment rather than a code change.
• ~30 lines of comment preamble on a 15-line helper is disproportionate; the argument compresses to ~5 lines plus a link to this PR.
• The VERBOSE log fires twice per decode step (K and V) on the steady-state path.
• finally restores set_default_logger_severity(_ORT_LOG_SEVERITY_WARNING) — a hardcoded constant rather than the severity in effect on entry. Moot if the global mutation goes away (see the inline comment on the negative assertion).

[tianleiwu]

LOGS_DEFAULT routes to the process-wide default logger rather than the session logger. Two consequences:

1. A user setting session_options.log_severity_level cannot control this message; conversely, raising the default severity floods every session in the process. The new test has to work around this by mutating global logger state.
2. The tag cannot be attributed to a specific session, which makes the test's negative assertion inherently racy (see my comment on assertNotIn).

OpKernelContext::Logger() is available at every call site — all four Run*Attention methods already take context. Threading the logger into the helper and using LOGS(logger, VERBOSE) matches the ORT kernel convention and lets the test anchor on session_options.logid, exactly the way the CUDA-graph capture test in this same file already does.

[tianleiwu]

This function returns Status, but ORT_ENFORCE throws instead of returning one. ORT_RETURN_IF_NOT keeps the failure on the Status path that every caller already propagates via ORT_RETURN_IF_ERROR:

ORT_RETURN_IF_NOT(src->SizeInBytes() == dst->SizeInBytes(),
                  "CopyKVToPresent requires identical src/dst sizes; this only holds when "
                  "past_sequence_length == 0 (see callers' gating).");

Purely defensive here — I confirmed the precondition holds at all 8 call sites — so this is consistency, not correctness.

[tianleiwu]

These two NOTE (Phase N — investigated and abandoned) blocks add ~60 lines of design-retrospective narrative to the body of ComputeInternal. Both already close with "See issue #29714 for the full writeup and benchmark data" — which is the right home for the chunk-size sweep, the A100/cuDNN 9.8 timings, and the SM90 auto-enable caveat.

Two reasons to trim to 2-3 lines each (what was tried, why rejected, link to #29714):

• The numbers are version- and hardware-specific and will rot silently as cuDNN and hardware move, whereas an issue is expected to be a point-in-time record.
• 60 lines of prose between the dispatch gate and the USE_FLASH_ATTENTION block pushes the actual control flow apart for every future reader of this function.

The Phase 3 note is also outside this PR's stated scope, as the description acknowledges.

[tianleiwu]

This negative assertion scans captured fd 2, which is process-global, while the helper has just raised the default logger severity to VERBOSE process-wide. Any concurrently running session that hits CopyKVToPresent with an aliased buffer emits the tag into this capture and fails the assertion — a false failure under parallel test execution.

This file already documents and solves exactly this hazard for the CUDA-graph capture test:

"fd 2 is a process-global descriptor, so a bare substring match could false-pass on a capture line logged by an unrelated concurrently-running session (e.g. under parallel test execution). The logid anchors the match to this test's own session."

The same mitigation is unavailable here only because the message goes through LOGS_DEFAULT. Switching the kernel to context->Logger() (see my comment on that line) unlocks it: set session_options.logid, then regex-match logid + tag on the same line. That also removes the need for set_default_logger_severity() entirely, since session_options.log_severity_level is already VERBOSE — and with it the hardcoded _ORT_LOG_SEVERITY_WARNING restore in the finally.

[tianleiwu]

The gate for the memory-efficient backend is not just SM level — it is the onnxruntime_USE_MEMORY_EFFICIENT_ATTENTION build flag plus the ORT_DISABLE_MEMORY_EFFICIENT_ATTENTION env override (both already used elsewhere in this suite, e.g. test_gqa.py:1427, test_mha.py:1381). On a build without MEA compiled in, assertEqual("EFFICIENT_ATTENTION", sdpa_kernel) hard-fails instead of skipping.

The flash case has the same gap in the other direction: has_flash_attention() is just has_cuda_device(80) and says nothing about USE_FLASH_ATTENTION / ORT_DISABLE_FLASH_ATTENTION.

Since the aliasing logic under test is backend-independent, the cheapest robust fix is to treat an unexpected MATH result as a skipTest for the non-MATH parameterizations, or probe availability once in setUpClass.