[Manoj Kumar Thapa (iammanoj807)]

Description

With is_causal=1 and a past_key cache, the CPU Attention kernel (opsets 23 and 24) skips the causal mask whenever q_sequence_length == 1:

: !(parameters.q_sequence_length == 1 && parameters.past_sequence_length > 0));

That is a no-op only for a one-token decode step (kv_sequence_length == 1), where the query may see every key. When the step appends two or more new KV tokens, the query should see only keys t <= past_sequence_length, but it also sees the newer keys after it.

The spec defines the causal mask with a past cache as j <= i + past_sequence_length. Opset 23 says "query i attends key j iff j <= i + past_sequence_length". Opset 24 says "offset = past_sequence_length when past_key is provided" (onnx/onnx#8068). For q_len > 1 this kernel already follows that rule: its overlay loop masks t >= past_sequence_length + s + 1. ORT's Python test reference (construct_causal_mask in test_onnx_attention/common.py) uses the same rule.

This PR keeps the shortcut only when kv_sequence_length == 1 too, so a normal decode step still skips the mask buffer. Every other case goes through the existing overlay. The nonpad_kv_seqlen branch is unchanged, because nonpad_kv_seqlen cannot be combined with past_key.

Repro (opset 23; opset 24 gives the same):

import numpy as np
import onnxruntime as ort
from onnx import TensorProto, helper
from onnx.reference import ReferenceEvaluator

P, S_kv, D = 2, 2, 4  # 2 cached keys, 2 new keys, 1 query
node = helper.make_node("Attention", ["Q", "K", "V", "", "past_key", "past_value"],
                        ["Y", "present_key", "present_value"], is_causal=1)
shapes = {"Q": [1, 1, 1, D], "K": [1, 1, S_kv, D], "V": [1, 1, S_kv, D], "past_key": [1, 1, P, D], "past_value": [1, 1, P, D]}
inputs = [helper.make_tensor_value_info(n, TensorProto.FLOAT, s) for n, s in shapes.items()]
outputs = [helper.make_tensor_value_info(n, TensorProto.FLOAT, None) for n in ("Y", "present_key", "present_value")]
model = helper.make_model(helper.make_graph([node], "g", inputs, outputs), opset_imports=[helper.make_opsetid("", 23)])
model.ir_version = 10

one_hot = np.eye(P + S_kv, D, dtype=np.float32)  # the value of key t is one-hot at t
feeds = {"Q": np.zeros(shapes["Q"], np.float32), "K": np.zeros(shapes["K"], np.float32),
         "V": one_hot[P:].reshape(shapes["V"]), "past_key": np.zeros(shapes["past_key"], np.float32),
         "past_value": one_hot[:P].reshape(shapes["past_value"])}

sess = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"])
print("onnxruntime", ort.__version__, sess.run(["Y"], feeds)[0].ravel())
print("onnx reference", ReferenceEvaluator(model).run(["Y"], feeds)[0].ravel())

onnxruntime 1.32.0 [0.25 0.25 0.25 0.25]
onnx reference [0.33333334 0.33333334 0.33333334 0.        ]

The query is at position 2, so it should attend keys 0, 1 and 2 equally. ORT 1.32.0 also attends key 3. With this change ORT returns [1/3, 1/3, 1/3, 0].

Adds AttentionTest.Attention4DCausalSingleQueryWithPastMultiTokenKV_CPU, which covers opsets 23 and 24, with and without a boolean attn_mask. It fails on main (0.25 instead of 1/3, and 1/3 instead of 0.5 with the mask) and passes with this change.

Motivation and Context

This is the past-cache half of #29020. #29050 fixed q_len=1, kv_len>1 without a past cache and left this case as it was. The Copilot review on #29050 pointed at this line: "skipping causal here would incorrectly allow the single query to attend to keys beyond its logical position".

The trigger is narrow. A normal decode step (q_len == kv_len == 1) is not affected. It needs a step that appends more KV tokens than it has queries.

Testing (macOS arm64, Release build, CPU EP):

• AttentionTest.*: 78 passed and 12 skipped (CUDA only). On main, 77 pass and the new test fails.
• All *Attention* gtests: the same pass/fail set as main, plus the new test. Some contrib-op tests fail on both because my partial build did not copy their testdata models.
• Old vs new comparison of the CPU kernel against a float64 implementation of the spec, over 50,688 configurations, run with graph optimizations both disabled and at ORT_ENABLE_ALL. They cover:
  • opsets 23 and 24; batch 1–2; MHA and GQA head counts;
  • q_len 1–3, kv_len 1–4, past 0/1/3, is_causal 0/1;
  • no mask; bool 2D/3D/4D masks, including fully masked rows; float masks;
  • 3D and 4D inputs, softcap, fp32 and fp16, every qk_matmul_output_mode, and an explicit scale.
• Results of that comparison:
  • In the 46,464 configurations outside is_causal=1, past > 0, q_len == 1, kv_len >= 2, every output (Y, present_key, present_value, qk_matmul_output) is bit-identical to main.
  • Inside that case, this branch matches the spec in all 4,224 configurations, and main gives different outputs in 3,716 of them. The other 508 have a bool mask that already hides the extra keys.
• The Attention cases in onnx/backend/test/case/node/attention.py (from the onnx submodule; 87 of the 93 run on the CPU EP) give identical results on main and this branch.
• lintrunner: no issues.

CUDA EP: not changed by this PR. On a Tesla T4 with onnxruntime-gpu 1.32.0, the unfused CUDA path follows the spec. The memory-efficient (MEA) path uses the offset total_sequence_length - q_sequence_length when a past cache is present, so it is wrong whenever kv_len != q_len (see the comment below). Flash needs sm80 and was not tested. The new test runs on CPU only.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

🟢 Approval recommended

No active review issues require changes to the CPU Attention causal-masking fix.

No open issues 🎉

What changed in this PR

Fixes CPU ONNX Attention causal masking when a single query uses a past cache and receives multiple new key/value tokens. This completes the past-cache case left unchanged by #29050.

Changes:

• Restricts the causal-mask shortcut to single-token query and key/value decode steps.
• Adds CPU regression coverage for opsets 23 and 24, with and without a boolean attention mask.

No actionable findings.

File	Description
onnxruntime/​test/​providers/​cpu/​llm/​attention_op_test.cc	Adds regression checks for causal attention outputs and present-cache contents.
onnxruntime/​core/​providers/​cpu/​llm/​attention.cc	Applies causal masking when a single query appends multiple key/value tokens to a past cache.

_{Review effort: Balanced}

[Manoj Kumar Thapa (iammanoj807)]

@microsoft-github-policy-service agree

[Manoj Kumar Thapa (iammanoj807)]

I also checked the CUDA EP, which this PR does not change, on a Tesla T4 (sm75) with onnxruntime-gpu 1.32.0. Inputs were is_causal=1 with past_key, past 1 or 3, q_len 1–3, kv_len 1–4, opsets 23 and 24, and CPU fallback disabled. The results were compared against j <= i + past_sequence_length. When kv_len == q_len, every path matches. When kv_len != q_len, the result depends on which kernel runs:

CUDA path on the T4	kv_len != q_len
Unfused (head_size 4; fp32 GQA)	matches in every case
MEA (head_size 8 or 64: fp16, and fp32 MHA)	wrong in every case (216 of 216)

So the same model gives different results for head_size 4 and head_size 8. This also happens for q_len=2, kv_len=1.

MEA sets causal_from_top_left = (past_sequence_length == 0), so with a past cache it uses the offset total_sequence_length - q_sequence_length instead of past_sequence_length. These are equal only when kv_len == q_len. The unfused path passes past_kv_length = past_sequence_length and is correct. Flash is bottom-right only too, so I expect the same there, but I could not test it (it needs sm80).

One possible fix is to send is_causal + past_key + kv_len != q_len to the unfused path. I can't build the CUDA EP locally, so I kept this PR CPU-only. I'm happy to open a follow-up if that approach sounds right.