[Prathik Rao (prathikr)]

WebGPU GatedDeltaNet Compact State Updates

Changes

This PR adds compact state-update capture to the native WebGPU
com.microsoft.GatedDeltaNet kernel, which previously rejected
state_update_capacity > 0.

• Validates state_update_capacity and the required [batch] capture_count input.
• Allocates the packed float32 state_update output and includes its buffers in WebGPU binding
  accounting.
• Captures the clamped per-row prefix of decay, shared-key, and delta factors in the recurrent
  shader. These factors replay S = decay * S + outer(key, delta).
• Preserves inverse-GQA key sharing and existing value-channel tiling.
• Uses the recurrent path while capture is active because parallel prefill does not emit per-token
  transitions.
• Honors state_update_active=0 by clearing the entire output; uncaptured entries remain
  unspecified when capture is active.
• Updates the registered schema text and generated contrib-operator documentation to advertise
  WebGPU compact-state support.

Test Coverage

Tests now compare only the captured decay/key/delta ranges instead of assuming uncaptured storage is
zero. The inactive test still verifies that the complete output is zero-filled.

Coverage includes ragged capacity-7 cases with zero, partial, and full capture counts, plus an FP16
Qwen-like case (Hq=2, Hv=6, head_size_qk=head_size_v=128). The larger case exercises the
FP16 aliases, value tiling, shared-key head mapping, and binding-packing path used by GenAI.

DFlash2 Prerequisite

This PR adds the ORT kernel prerequisite; it does not enable end-to-end WebGPU DFlash2 by itself.
GenAI also needs:

• builder export with nonzero state_update_capacity (#2453,
  #2585); and
• compact-state replay after partial acceptance (#2454),
  including WebGPU ReplayStateUpdates (#2606).

Before advertising WebGPU DFlash2 support, an end-to-end GenAI run should cover full, partial, and
zero draft acceptance, match non-speculative output/state, and confirm that neither target execution
nor replay silently falls back from WebGPU.

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Copilot review overview

🟡 Changes recommended

The generated contrib-operator documentation must be updated before approval.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 1 · 1

Open (2)

• Regenerate contrib-op documentation for schema change · New
• Avoid clearing full state_update buffer during active capture · New

What changed in this PR

Adds WebGPU compact recurrent-state update support for GatedDeltaNet, matching the CUDA contract for speculative decoding.

Changes:

• Captures decay, key, and delta transitions in WGSL.
• Adds capacity validation, output handling, and WebGPU tests.
• Updates the operator schema source.

File	Review
onnxruntime/​test/​contrib_ops/​gated_delta_net_op_test.cc	Adds compact-update reference logic and WebGPU tests.
onnxruntime/​core/​graph/​contrib_ops/​bert_defs.cc	Critical: Regenerate docs/ContribOperators.md to reflect the schema change and prevent documentation validation failure.
onnxruntime/​contrib_ops/​webgpu/​bert/​gated_delta_net.wgsl.template	Implements transition capture shader logic.
onnxruntime/​contrib_ops/​webgpu/​bert/​gated_delta_net.h	Adds capture program state and clearing declarations.
onnxruntime/​contrib_ops/​webgpu/​bert/​gated_delta_net.cc	Wires capture execution. Nit: Restrict full-buffer clearing to the inactive path.

---

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

[github-actions]

You can commit the suggested changes from lintrunner.

[Baiju Meswani (baijumeswani)]

Thanks for adding the WebGPU compact-state capture path. The kernel logic looks aligned with the CUDA layout and capture semantics. I left a few inline comments about test coverage and the remaining GenAI integration needed before this can be considered end-to-end DFlash2 support. I would also wait for the full CI matrix to finish.

[Baiju Meswani (baijumeswani)]

Thanks for addressing the earlier comments. The production kernel changes look good by inspection, but the updated test has one remaining coverage issue. I would fix this and wait for the WebGPU CI jobs to finish before approval.