[aciddelgado]

Description

Enable  com.microsoft::PagedAttention  on the WebGPU EP with  is_causal=0  when  local_window_size <= 0 .

This is needed by block-drafting models such as DFlash2, where all query tokens submitted in one speculative block must attend bidirectionally to the complete live context. Previously, the WebGPU kernel rejected every non-causal configuration and then hard-coded causal behavior when constructing  WebgpuAttentionParameters .

Changes

• Store and propagate the  is_causal  attribute through WebGPU PagedAttention into the existing FlashAttention  is_unidirectional  shader specialization.
• Preserve the existing causal and causal local-window behavior.
• Continue to reject  is_causal=0  with  local_window_size > 0 ; query-relative left-window masking for that combination is intentionally out of scope.
• Generalize the C++ reference implementation to compare causal prefix visibility with non-causal full-context visibility.
• Add focused WebGPU parity tests covering:
• direct paged decode with a multi-token speculative block;
• fused paged prefill;
• gather-then-FlashAttention fallback;
• ragged requests with different query and past lengths.
• Replace the Python non-causal rejection test with reference parity coverage and add a dedicated test for the unsupported non-causal local-window combination.

No WGSL redesign is required: the dense and paged FlashAttention shaders already support causal and non-causal specializations. This change only restores the missing host-side attribute plumbing.

Masking behavior

•  is_causal=1 : query row  i  attends through  past_length + i .
•  is_causal=0 : every query row attends through  past_length + query_length , including later tokens in the same submitted block.
•  is_causal=0 && local_window_size>0 : remains explicitly unsupported.

Validation

•  git diff --check 
• Python syntax compilation for  test_paged_attention.py 
• Source-level checks for attribute propagation and focused test coverage

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Copilot review overview

🟡 Changes recommended

The causal flag is discarded before shader specialization, and malformed attribute values are not rejected.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 1

Open (1)

• WebGPU attention ignores non-causal mode · New

What changed in this PR

Enables non-causal WebGPU PagedAttention for full-context blocks and expands parity coverage.

Changes:

• Stores and propagates is_causal.
• Generalizes reference masking.
• Adds non-causal WebGPU tests.

File	Review
onnxruntime/​test/​python/​transformers/​test_paged_attention.py	Adds non-causal parity and invalid local-window coverage.
onnxruntime/​test/​contrib_ops/​paged_attention_op_test.cc	Generalizes reference masking and adds focused WebGPU tests.
onnxruntime/​contrib_ops/​webgpu/​bert/​paged_attention.h	Stores causal-mode state.
onnxruntime/​contrib_ops/​webgpu/​bert/​paged_attention.cc	Critical: parameter conversion still forces causal masking, so non-causal execution remains broken. Moderate: is_causal must reject values other than 0 or 1.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Copilot review overview

🟡 Changes recommended

Invalid is_causal values can silently produce incorrect non-causal behavior.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 1

Open (1)

• Validate causal attribute before converting integer to bool · New

Resolved since last review (1)

• WebGPU attention ignores non-causal mode