Optimize ONNX Attention KV cache with ConcatNewToPast and add release-build kernel safety

[titaiwangms]

This pull request introduces a major optimization to the CUDA attention kernel by fusing the key/value (KV) cache concatenation step into a single kernel, reducing memory overhead and kernel launches during autoregressive decoding. The update also clarifies and improves documentation and safety around attention mask handling, especially regarding right-padding validity and mask conversion. Support for native __nv_bfloat16 types is expanded for better hardware utilization.

Performance and Kernel Fusion Improvements:

• The attention kernel now uses LaunchConcatNewToPastKV to fuse the past KV copy and new token append into a single CUDA kernel, eliminating the previous memset and strided copy overhead. This reduces memory traffic and kernel launch overhead during decoding, matching the efficiency of the contrib GQA path. [1] [2]
• Support for native __nv_bfloat16 types is added to the fused concat kernel, enabling direct hardware arithmetic on supported GPUs and aligning with GQA's early type conversion pattern. [1] [2]

Attention Mask Handling and Documentation:

• The logic and documentation for converting boolean attention masks to sequence lengths are clarified. In debug builds, non-contiguous (non-right-padded) masks trigger a CUDA kernel assertion; in release builds, the output is memory-safe but may be semantically incorrect, as only leading True values are counted. [1] [2] [3] [4]
• Comments and docstrings are updated to explain the handling of offsets in mask-to-seqlen conversion, and the broadcasting behavior in bias addition kernels is clarified. [1] [2]

Code Quality and Minor Fixes:

• Includes minor code style improvements, such as using using instead of typedef for CUDA type aliases, and adds missing header includes for type conversion. [1] [2]
• Adds missing includes and updates comments for clarity in several files. [1] [2] [3]
• Adds preprocessor guards to prevent unused variable warnings in certain build configurations.

[Copilot]

Pull request overview

This PR optimizes CUDA LLM attention decode by replacing the previous multi-step KV cache preparation (memset + strided copies + Flash internal append) with a fused “concat new KV to past” kernel, while also hardening several CUDA kernels to avoid unsafe behavior in release builds (where device asserts may be compiled out).

Changes:

• Use LaunchConcatNewToPastKV to pre-populate present_key/value in ONNX Attention’s Flash decode path and call mha_fwd_kvcache with k_new/v_new=nullptr.
• Add release-safety clamping/break behavior in CUDA kernels (TensorScatter write indices; boolean mask contiguity handling) to prevent OOB access/hard faults without host-side synchronization.
• Remove an extra debug assert in the rotary embedding CUDA implementation.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
onnxruntime/core/providers/cuda/llm/attention.cc	Switch Flash decode-with-past to fused KV concat kernel and skip Flash internal KV append.
onnxruntime/core/providers/cuda/llm/attention_mask_impl.cu	Make non-contiguous bool masks stop scanning after the first invalid transition (plus clamp seqlens to non-negative).
onnxruntime/core/providers/cuda/llm/attention_mask_impl.h	Update comments documenting debug-only asserts and release-build behavior for invalid masks.
onnxruntime/core/providers/cuda/llm/tensorscatter_impl.cu	Clamp negative / OOB write indices in-kernel for release safety.
onnxruntime/core/providers/cuda/llm/tensorscatter.cc	Update documentation/comments about in-kernel validation and release behavior.
onnxruntime/core/providers/cuda/llm/rotary_embedding_impl.cu	Remove redundant assert (retains runtime enforcement).

Comments suppressed due to low confidence (1)

onnxruntime/core/providers/cuda/llm/tensorscatter_impl.cu:43

• In the release-build path (when CUDA_KERNEL_ASSERT is compiled out), clamping out-of-bounds write_indices causes all OOB writes to be redirected to max_seq_len - 1. That can silently corrupt the last cache position (potentially overwriting valid data multiple times) instead of preserving the original output for invalid indices. Consider making invalid indices a no-op (skip the write) and/or clamping wi to the largest valid start index (max_seq_len - seq_len) so that the entire update range stays in-bounds without collapsing onto the last element.

  int64_t wi = (write_indices != nullptr) ? write_indices[batch_idx] : 0;
  CUDA_KERNEL_ASSERT(wi >= 0);
  wi = max(wi, static_cast<int64_t>(0));  // Clamp for release safety (CUDA_KERNEL_ASSERT is debug-only)
  int64_t cache_pos;
  if (circular) {
    cache_pos = (wi + seq_idx) % max_seq_len;
  } else {
    cache_pos = wi + seq_idx;
    CUDA_KERNEL_ASSERT(cache_pos < max_seq_len);
    cache_pos = min(cache_pos, max_seq_len - 1);  // Clamp for release safety (CUDA_KERNEL_ASSERT is debug-only)
  }

  int64_t out_offset = prefix_idx * (max_seq_len * suffix_count) + cache_pos * suffix_count + suffix_idx;
  output_data[out_offset] = update_data[id];

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[tianleiwu]

Please consider using LaunchUnpackRoPEAppend in onnxruntime/contrib_ops/cuda/bert/group_query_attention_qkv.cuh if possible. That's unified preprocessing for GQA/MHA.

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 4 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 3 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 3 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 5 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

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

Comments suppressed due to low confidence (1)

onnxruntime/core/providers/cuda/llm/attention_mask_impl.cu:86

• ConvertMaskToSeqlensKernel treats any row with mask_row[0]==false as an all-false (seq_len=0) mask and skips scanning the rest of the row. That means an invalid/non-contiguous pattern like [False, True, ...] will not trigger CUDA_KERNEL_ASSERT in debug builds and will be silently misinterpreted as length 0 in release builds. If right-padding convention validation is intended, consider checking the remaining elements are all False when mask_row[0] is False (at least under CUDA_KERNEL_ASSERT/debug).

  if (!mask_row[0]) {
    // Entire row is padding (all-false mask)
    seq_len = 0;
  } else {

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[titaiwangms]

Addressing @tianleiwu's Review Feedback

Hi @tianleiwu — here's how each of your feedback items has been addressed:

1. Consider using LaunchUnpackRoPEAppend for unified preprocessing

• Status: ✅ Acknowledged
• Resolution: Added TODO at attention.cc:315-318 documenting this as a future optimization (fuse RoPE + mask + concat into ~2 kernel launches). Out of scope for this PR.
• Commit: d093340

2. Avoid zero-filling present KV tail — I/O impacts performance

• Status: ✅ Resolved
• Resolution: Reverted tail-zeroing. Flash respects seqlens_k bounds so trailing positions are never read. Saves ~134MB I/O per decode step.
• Commit: d093340

3. Use OrtToCudaType instead of ToCudaType for native bf16

• Status: ✅ Addressed
• Resolution: All Flash path casts now use OrtToCudaType<T>::type (NativeCudaT → __nv_bfloat16), consistent with GQA pattern. Unfused path correctly retains ToCudaType for contrib_ops interface compatibility.
• Commit: 79e5ca9

All 3 items addressed. Please let me know if anything needs further changes!