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

[Ti-Tai Wang (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.

[Tianlei Wu (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).

[Ti-Tai Wang (titaiwangms)]

Pushed a follow-up commit (f5ba8ea) addressing Tianlei Wu (@tianleiwu)'s review plus findings from an additional internal review pass:

Highest-priority fix (confirmed real bug)

• The non-aliased-branch test assertion read io_binding.get_outputs()[1]/[2], which — due to IOBinding::GetOutputs()'s bind-order semantics — actually read TensorScatter's updated_key_cache/updated_value_cache, not present_key/present_value. This made test_copy_still_runs_when_present_not_aliased pass vacuously (never validating the D2D copy path it exists to check). Now reads from present_k_ort/present_v_ort directly.

Other findings addressed

• ORT_ENFORCE → ORT_RETURN_IF_NOT in CopyKVToPresent (a Status-returning helper).
• LOGS_DEFAULT(VERBOSE) (process-global) → session-scoped logging via a new KernelSessionLogger helper, since OpKernelContext::Logger() isn't reachable through the CUDA EP's shared-provider bridge. Test log assertions now anchor on session_options.logid + the log line's delimited logger_id field (verified against ostream_sink.cc's actual format), replacing the racy process-global substring match and the global set_default_logger_severity(...) mutation.
• TestAttentionPresentKVCopySkip's dispatch-tier assertion now distinguishes "not compiled into this build" from a hard failure via skipTest, with ORT_DISABLE_FLASH_ATTENTION/ORT_DISABLE_MEMORY_EFFICIENT_ATTENTION forced off beforehand so an observed MATH fallback can't be attributed to an ambient env-var disable.
• Trimmed the Phase 2/3 investigation notes in ComputeInternal to a few lines each, pointing to [Perf] Add a cuDNN SDPA decode tier to the ONNX standard Attention CUDA kernel #29714 for the full writeup.
• Documented a known limitation of KernelSessionLogger: it assumes the EP instance is exclusively owned by one session's lifetime (true for the standard CUDA EP registration path; would not hold for a manually shared IExecutionProvider, as e.g. XNNPACK's EP supports).

Validation: incremental rebuild; TestAttentionPresentKVCopySkip 8/8 passed, whole test_tensorscatter_attention.py 121/121 passed on an A100. clang-format/ruff clean.

[Tianlei Wu (tianleiwu)]

Overall

The exact-alias copy skip is correct at all eight 4-D BNSH call sites, and the latest revision addresses the earlier separate-buffer, error-path, session-attribution, and backend-routing issues. One supported configuration still needs correction: in the CUDA Plugin EP path, the new accessor does not reach the session's logger, so the tag falls back to the plugin-wide logger and the session-scoped observability contract does not hold. I also left a focused test-coverage suggestion for independently proving both the K and V skips.

Verdict: REQUEST CHANGES because the logger path should work consistently in both legacy/shared-provider and CUDA Plugin EP builds.

[Tianlei Wu (tianleiwu)]

This does not retrieve the session logger in a CUDA Plugin EP build. In that path, the adapter's OpKernelInfo::GetExecutionProvider() returns KernelInfoCache::ep_impl_ (include/onnxruntime/ep/adapter/op_kernel_info.h:51,114), which is the internal CUDA shim created by CreateCudaPluginProvider. The session OrtLogger is instead stored on CudaEp::logger_ (cuda_ep.cc:120-127), and no SetLogger call copies it onto the shim. Consequently GetLogger() is null here and the fallback is the plugin-wide default logger: SessionOptions.log_severity_level = VERBOSE will not enable this tag, and the session logid will not identify it, so the new aliased tests fail if run against the plugin EP and users cannot control the message per session.

Please use the plugin's KernelInfo_GetLogger/KernelContext_GetLogger path (the adapter exposes the former through Info().GetKernelInfo().GetLogger()), while retaining an appropriate legacy/shared-provider path, or add one common accessor that returns the actual session logger in both build modes.

[Ti-Tai Wang (titaiwangms)]

Fixed in d296155. CUDA Plugin EP builds now retrieve the session logger through KernelInfo_GetLogger rather than the internal CUDA shim. Retrieval failures suppress only this diagnostic, and emission uses ORT_CXX_LOGF_NOEXCEPT; legacy/shared-provider builds retain the EP session-logger path.

[Tianlei Wu (tianleiwu)]

Each run calls CopyKVToPresent twice (K and V), but assertRegex passes after only one matching tag. If a later change restores the unconditional self-copy at just the value call site, the key call still emits this tag and every numeric assertion remains green because a self-copy is value-preserving. That leaves half of the eight claimed backend/copy call sites unproven. Please count the session-anchored matches and require exactly two, or label the helper calls as key/value and assert that both labels appear.

[Ti-Tai Wang (titaiwangms)]

Fixed in d296155. The aliased test now counts session-logid-anchored matches and requires exactly two records (K and V); the non-aliased test explicitly requires zero.

[Copilot]

🟡 Changes recommended

The hot path adds uncached logging API overhead, while backend tests can skip genuine routing regressions.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

onnxruntime/core/providers/cuda/llm/attention.cc:1766

• This NOTE narrates the abandoned prototype and profiling history, which repository guidance assigns to the issue or PR (AGENTS.md:57-61). Retain the durable decode-only rationale without embedding the experiment record in the dispatch implementation.

  // NOTE (Phase 3 — cuDNN SDPA prefill via fixed-size query-row chunking — investigated and
  // abandoned): A chunked cuDNN prefill dispatch (fixed-size, left-padded query-row chunks, needed
  // to keep cuDNN's graph-plan cache key stable) was prototyped and found correct, but measured only
  // ~0.94x-1.15x versus the existing MATH fallback on A100, and the auto-enable gate only fires on
  // SM>=90 hardware unavailable for this investigation — not worth the complexity. See issue #29714

• Files reviewed: 2/2 changed files
• Comments generated: 3
• Review effort level: Balanced

[Copilot]

🟢 Approval recommended

The optimization is correctly guarded and thoroughly covered across CUDA attention backends.

Review details

• Files reviewed: 3/3 changed files
• Comments generated: 0 new
• Review effort level: Balanced

[Ti-Tai Wang (titaiwangms)]

cc Tianlei Wu (@tianleiwu)