diff --git a/onnxruntime/core/providers/cpu/llm/attention.cc b/onnxruntime/core/providers/cpu/llm/attention.cc index 1b0f37feab9fb..150fe7b478c1e 100644 --- a/onnxruntime/core/providers/cpu/llm/attention.cc +++ b/onnxruntime/core/providers/cpu/llm/attention.cc @@ -347,7 +347,15 @@ void AttentionBase::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 diff --git a/onnxruntime/test/providers/cpu/llm/attention_op_test.cc b/onnxruntime/test/providers/cpu/llm/attention_op_test.cc index cf76ca0fa00f8..03a47664c632a 100644 --- a/onnxruntime/test/providers/cpu/llm/attention_op_test.cc +++ b/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("is_causal", static_cast(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 q(batch_size * q_num_heads * q_sequence_length * head_size, 0.0f); + std::vector k(batch_size * kv_num_heads * kv_sequence_length * head_size, 0.0f); + std::vector v(batch_size * kv_num_heads * kv_sequence_length * head_size, 0.0f); + std::fill_n(v.begin(), head_size, 1.0f); + + test.AddInput("Q", {batch_size, q_num_heads, q_sequence_length, head_size}, q); + test.AddInput("K", {batch_size, kv_num_heads, kv_sequence_length, head_size}, k); + test.AddInput("V", {batch_size, kv_num_heads, kv_sequence_length, head_size}, v); + test.AddOptionalInputEdge(); // attn_mask + test.AddOptionalInputEdge(); // past_key + test.AddOptionalInputEdge(); // past_value + test.AddInput("nonpad_kv_seqlen", {batch_size}, {kv_sequence_length}); + + std::vector expected_y(batch_size * q_num_heads * q_sequence_length * head_size, + 1.0f / static_cast(kv_sequence_length)); + test.AddOutput("Y", {batch_size, q_num_heads, q_sequence_length, head_size}, expected_y, false, 0, + 1e-4f); + test.AddOptionalOutputEdge(); // present_key + test.AddOptionalOutputEdge(); // present_value + + std::vector> 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 q(batch_size * q_num_heads * q_sequence_length * head_size, 0.0f); + std::vector k(batch_size * kv_num_heads * kv_sequence_length * head_size, 0.0f); + std::vector 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 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(), std::initializer_list(), + std::vector(), std::vector(), + 1, -1, std::numeric_limits::quiet_NaN(), std::numeric_limits::quiet_NaN(), -1, + TensorType::kFloat, + y, std::vector(), std::vector(), std::vector(), + 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. diff --git a/onnxruntime/test/python/transformers/test_onnx_attention/common.py b/onnxruntime/test/python/transformers/test_onnx_attention/common.py index 349f68f0ac5ce..d1a5d833d12a4 100644 --- a/onnxruntime/test/python/transformers/test_onnx_attention/common.py +++ b/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) diff --git a/onnxruntime/test/python/transformers/test_onnx_attention/test_gqa.py b/onnxruntime/test/python/transformers/test_onnx_attention/test_gqa.py index 542094a9ac4ee..eca86e429b597 100644 --- a/onnxruntime/test/python/transformers/test_onnx_attention/test_gqa.py +++ b/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 --- diff --git a/onnxruntime/test/python/transformers/test_onnx_attention/test_mha.py b/onnxruntime/test/python/transformers/test_onnx_attention/test_mha.py index e2d9acbd0c500..16a00085f224f 100644 --- a/onnxruntime/test/python/transformers/test_onnx_attention/test_mha.py +++ b/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(