Skip to content

Emit maskless is_causal=1 + nonpad_kv_seqlen static-cache graph once upstream lands #345

Description

@titaiwangms

Goal

Once the upstream bottom-right is_causal support lands and is pinnable, emit the maskless static-cache attention graph directly:

Attention(..., is_causal=1, nonpad_kv_seqlen=<seqlens_k>)   # no explicit attn_mask, no past_key

This is the end-state for the static-cache (--static-cache, ONNX Attention + TensorScatter in-place KV-cache) path. It is strictly better than the explicit-mask form on two axes: fewer nodes and Flash-eligible (an explicit attn_mask forces the Memory-Efficient Attention kernel; the maskless form routes to Flash for fp16/bf16).

Background

mobius #340 originally proposed "Option-Y" as a bridge: is_causal=0 + an explicit offset-aware causal mask + nonpad_kv_seqlen, which forces ORT to the MEA kernel — the only currently-runnable static-cache config, because today's opset-24 Attention CUDA kernel rejects is_causal=1 + nonpad_kv_seqlen (no past_key, S_q != total_kv) with NOT_IMPLEMENTED.

Option-Y was dropped, not landed (the bug-fix half of #340 shipped separately as #351). So there is nothing to revert — this issue now tracks emitting the maskless end-state directly.

Dependency chain

The blocker is not "PR merged" but "a released ORT containing the kernel fix, pinned by mobius":

  1. Fix Attention is_causal causal-mask alignment + composed is_causal/attn_mask NaN robustness for external (static) KV cache (#8054) onnx/onnx#8068 — spec/errata: bottom-right is_causal + nonpad_kv_seqlen (external KV cache)
  2. Update CPU/CUDA Attention kernels for bottom-right is_causal with nonpad_kv_seqlen (no past_key), and composed is_causal + attn_mask, per onnx/onnx#8068 microsoft/onnxruntime#28904 / Fix Attention is_causal bottom-right alignment for external KV cache (onnx#8068, #28904) microsoft/onnxruntime#28958 — CPU/CUDA Attention kernel implementing it (maskless is_causal=1 + nonpad_kv_seqlen becomes Flash-eligible)
  3. an ORT release cut that contains the above
  4. mobius bumps its ORT pin past that release (the node-count gate baselines on the pinned ORT)

Action (when the pin can bump)

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions