Fix CPU Attention causal mask alignment

onnxruntime/core/providers/cpu/llm/attention.cc

@@ -347,7 +347,15 @@ void AttentionBase<T>::ComputeAttentionProbs(T* attention_probs,
 
   T* mask_data = nullptr;
   bool delete_mask_data = false;
-  bool causal = parameters.is_causal && parameters.q_sequence_length > 1;
+  // In the nonpad_kv_seqlen path, q_len=1 is external KV-cache decode with
+  // bottom-right alignment. The single query's causal frontier is the valid
+  // length, so nonpad masking alone leaves exactly all valid keys visible.
+  // Keep causal=false for that case to avoid applying the batch-shared
+  // upper-left overlay used by the no-nonpad path.
+  bool causal = parameters.is_causal &&
+                (parameters.has_nonpad_kv_seqlen
+                     ? parameters.q_sequence_length > 1
+                     : !(parameters.q_sequence_length == 1 && parameters.past_sequence_length > 0));
   // When nonpad_kv_seqlen is present the causal frontier is offset-aware
   // (bottom-right) and per-batch, so it cannot be baked into the batch-shared mask
   // buffer here; it is applied per-batch in the main loop below. Skip the top-left

onnxruntime/test/providers/cpu/llm/attention_op_test.cc

@@ -2332,6 +2332,45 @@ TEST(AttentionTest, Attention_Causal_NonPadKVSeqLen_Decode_BottomRight) {
   test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
 }
 
+// q_len=1 with nonpad_kv_seqlen keeps bottom-right decode behavior: the single
+// query attends every valid key. If the no-nonpad upper-left overlay is applied
+// here, this returns 1.0 instead of the expected 1/6.
+TEST(AttentionTest, Attention_Causal_NonPadKVSeqLen_SingleQueryKeepsBottomRight_CPU) {
+  OpTester test("Attention", 24, onnxruntime::kOnnxDomain);
+  test.AddAttribute<int64_t>("is_causal", static_cast<int64_t>(1));
+
+  constexpr int batch_size = 1;
+  constexpr int q_num_heads = 1;
+  constexpr int kv_num_heads = 1;
+  constexpr int q_sequence_length = 1;
+  constexpr int kv_sequence_length = 6;
+  constexpr int head_size = 8;
+
+  std::vector<float> q(batch_size * q_num_heads * q_sequence_length * head_size, 0.0f);
+  std::vector<float> k(batch_size * kv_num_heads * kv_sequence_length * head_size, 0.0f);
+  std::vector<float> v(batch_size * kv_num_heads * kv_sequence_length * head_size, 0.0f);
+  std::fill_n(v.begin(), head_size, 1.0f);
+
+  test.AddInput<float>("Q", {batch_size, q_num_heads, q_sequence_length, head_size}, q);
+  test.AddInput<float>("K", {batch_size, kv_num_heads, kv_sequence_length, head_size}, k);
+  test.AddInput<float>("V", {batch_size, kv_num_heads, kv_sequence_length, head_size}, v);
+  test.AddOptionalInputEdge<bool>();   // attn_mask
+  test.AddOptionalInputEdge<float>();  // past_key
+  test.AddOptionalInputEdge<float>();  // past_value
+  test.AddInput<int64_t>("nonpad_kv_seqlen", {batch_size}, {kv_sequence_length});
+
+  std::vector<float> expected_y(batch_size * q_num_heads * q_sequence_length * head_size,
+                                1.0f / static_cast<float>(kv_sequence_length));
+  test.AddOutput<float>("Y", {batch_size, q_num_heads, q_sequence_length, head_size}, expected_y, false, 0,
+                        1e-4f);
+  test.AddOptionalOutputEdge<float>();  // present_key
+  test.AddOptionalOutputEdge<float>();  // present_value
+
+  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+  execution_providers.push_back(DefaultCpuExecutionProvider());
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
+}
+
 // Continued / chunked prefill (S_q=2) into a partially-filled static cache.
 // nonpad=[4], S_q=2 -> offset = 4 - 2 = 2: query 0 attends keys {0,1,2}, query 1
 // attends {0,1,2,3}. The old top-left alignment would mask everything past the
@@ -3087,9 +3126,41 @@ TEST(AttentionTest, Attention4DSoftCapOutputQkRawLogits) {
 
 // ============================================================================
 // Causal alignment tests: verify upper-left (no past) vs lower-right (with past)
-// These are CUDA-only tests that validate the causal masking fix.
+// These tests validate causal mask alignment across CPU and CUDA.
 // ============================================================================
 
+// Test: Causal + cross-attention (S_q=1, S_kv=6, no past)
+// ONNX spec mandates upper-left alignment: q0 attends only to kv[0].
+// This covers GitHub issue #29020, where CPU skipped causal masking for S_q=1.
+TEST(AttentionTest, Attention4DCausalSingleQueryCrossAttentionUpperLeft) {
+  int batch_size = 1;
+  int q_num_heads = 1;
+  int q_sequence_length = 1;
+  int head_size = 8;
+  int kv_sequence_length = 6;
+  int kv_num_heads = 1;
+  int v_head_size = 8;
+  int past_sequence_length = 0;
+
+  std::vector<float> q(batch_size * q_num_heads * q_sequence_length * head_size, 0.0f);
+  std::vector<float> k(batch_size * kv_num_heads * kv_sequence_length * head_size, 0.0f);
+  std::vector<float> v(batch_size * kv_num_heads * kv_sequence_length * v_head_size, 0.0f);
+  for (int i = 0; i < v_head_size; ++i) {
+    v[i] = 1.0f;
+  }
+  std::vector<float> y(batch_size * q_num_heads * q_sequence_length * v_head_size, 1.0f);
+
+  RunTest4D(batch_size, q_num_heads, q_sequence_length, head_size, kv_sequence_length, kv_num_heads,
+            v_head_size, past_sequence_length,
+            q, k, v, std::vector<float>(), std::initializer_list<bool>(),
+            std::vector<float>(), std::vector<float>(),
+            1, -1, std::numeric_limits<float>::quiet_NaN(), std::numeric_limits<float>::quiet_NaN(), -1,
+            TensorType::kFloat,
+            y, std::vector<float>(), std::vector<float>(), std::vector<float>(),
+            false, false, true  // disable_cpu, disable_cuda, disable_dml
+  );
+}
+
 // Test: Causal + cross-attention (S_q=3, S_kv=5, no past)
 // ONNX spec mandates upper-left alignment: q_i attends to kv[0..i].
 // V is identity-like so output directly reveals which KV positions were attended.

onnxruntime/test/python/transformers/test_onnx_attention/common.py

@@ -666,12 +666,11 @@ def attention_past_func(
 # #################################################################################################
 
 
-def construct_causal_mask(seqlen_q, seqlen_k, device):
-    """Construct a causal mask for attention."""
+def construct_causal_mask(seqlen_q, seqlen_k, device, past_seqlen=0):
+    """Construct a causal mask for ONNX Attention upper-left alignment."""
     row_idx = rearrange(torch.arange(seqlen_q, device=device, dtype=torch.long), "s -> s 1")
     col_idx = torch.arange(seqlen_k, device=device, dtype=torch.long)
-    # Causal: positions can only attend to earlier positions
-    return col_idx > row_idx + seqlen_k - seqlen_q
+    return col_idx > row_idx + past_seqlen
 
 
 def attention_ref(
@@ -682,6 +681,7 @@ def attention_ref(
     attn_bias=None,
     causal=False,
     softcap=0.0,
+    past_seqlen=0,
 ):
     """
     Reference implementation of scaled dot-product attention with GQA support.
@@ -694,6 +694,7 @@ def attention_ref(
         attn_bias: Additive attention bias [broadcastable to batch, num_heads, seq_q, seq_k]
         causal: Whether to apply causal masking
         softcap: Softcap value for attention scores (0.0 = disabled)
+        past_seqlen: Number of past K/V tokens before q[0] for causal masking.
 
     Returns:
         output: Attention output [batch, seq_q, num_heads, head_size]
@@ -724,7 +725,7 @@ def attention_ref(
         scores.masked_fill_(rearrange(~key_padding_mask, "b s -> b 1 1 s"), float("-inf"))
 
     if causal:
-        causal_mask = construct_causal_mask(seqlen_q, seqlen_k, q.device)
+        causal_mask = construct_causal_mask(seqlen_q, seqlen_k, q.device, past_seqlen=past_seqlen)
         scores.masked_fill_(causal_mask, float("-inf"))
 
     attention = torch.softmax(scores, dim=-1)

onnxruntime/test/python/transformers/test_onnx_attention/test_gqa.py

@@ -289,6 +289,7 @@ def parity_check_gqa_past(
         key_padding_mask=key_padding_mask,
         causal=causal,
         softcap=config.softcap,
+        past_seqlen=config.past_kv_sequence_length,
     )
     out_ref_np = out_ref.to(torch.float32).detach().cpu().numpy()
 
@@ -548,6 +549,7 @@ def parity_check_gqa_past_with_padding(
         key_padding_mask=key_padding_mask,
         causal=config.is_causal == 1,
         softcap=config.softcap,
+        past_seqlen=config.past_kv_sequence_length,
     )
 
     # --- ONNX Runtime Path ---

onnxruntime/test/python/transformers/test_onnx_attention/test_mha.py

@@ -298,6 +298,7 @@ def parity_check_mha_past(
         attn_bias=attn_bias_ref,
         causal=causal,
         softcap=config.softcap,
+        past_seqlen=config.past_kv_sequence_length,
     )
     out_ref_np = out_ref.to(torch.float32).detach().cpu().numpy()
 
@@ -2393,7 +2394,13 @@ def test_mha_past_asymmetric_v_head_size(self):
         full_k_bsnh = full_k_bnsh.transpose(1, 2)
         full_v_bsnh = full_v_bnsh.transpose(1, 2)
 
-        out_ref, _ = attention_ref(q=q, k=full_k_bsnh, v=full_v_bsnh, causal=True)
+        out_ref, _ = attention_ref(
+            q=q,
+            k=full_k_bsnh,
+            v=full_v_bsnh,
+            causal=True,
+            past_seqlen=config.past_kv_sequence_length,
+        )
 
         # ORT path — should fall back to unfused (not crash in MEA)
         out_ort, present_k, present_v = attention_past_func(