[WebGPU EP] Add native WebGPU GatedDeltaNet kernel

[Copilot]

Description

Adds a correctness-first native WebGPU implementation of com.microsoft::GatedDeltaNet opset 1.

• Kernel
  • Recurrent FP32 state update with float/float16 I/O.
  • Supports rank-3 packed and rank-4 dense layouts, device-resident ragged cu_seqlens, inverse GQA, and aliased input/output state.
  • Supports all update rules, default/explicit scale, Qwen decay, sigmoid beta, and fused Q/K L2 normalization.

S = exp(g) * S
delta = beta * (v - Sᵀk)
S = S + outer(k, delta)
o = scale * Sᵀq

• Integration

  • Registers GatedDeltaNet with the WebGPU contrib kernel registry.
  • Documents WebGPU float/float16 support and scalar-decay restriction.

• Compatibility and tests

  • Adds WebGPU parity coverage for dense, packed, ragged, inverse-GQA, fused Qwen options, FP32, and FP16.
  • Supports omitted final_state.
  • Rejects state_update_capacity > 0 explicitly; compact state capture remains CUDA-only.

Motivation and Context

GatedDeltaNet was CUDA-only, preventing native WebGPU execution of fused Qwen gated-delta attention graphs. This adds a native recurrent WebGPU path without changing the existing operator schema or CUDA ABI.

[azure-pipelines]

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

[Copilot]

🟡 Changes recommended

Head-index overflow, inconsistent shape validation, missing generated documentation, and coverage gaps remain.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds a native WebGPU kernel for com.microsoft::GatedDeltaNet.

Changes:

• Implements FP32/FP16 dense, packed, ragged, inverse-GQA, and recurrent-state execution.
• Registers and documents WebGPU support.
• Adds WebGPU parity and validation tests.

File summaries

File	Description
onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc	Adds WebGPU coverage.
onnxruntime/core/graph/contrib_ops/bert_defs.cc	Documents provider support and restrictions.
onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc	Registers the kernel.
onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.wgsl.template	Implements the recurrent WGSL shader.
onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.h	Defines kernel and program classes.
onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc	Implements validation, binding, and dispatch.

Review details

• Files reviewed: 6/6 changed files
• Comments generated: 5
• Review effort level: Balanced

---

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

[Copilot]

🟡 Changes recommended

Padded dispatches can access buffers out of bounds, and some documented input combinations exceed guaranteed WebGPU binding limits.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc:120

• state_update_active is schema-valid as an optional input when capacity is zero (only capture_count is explicitly forbidden), and the CUDA kernel accepts and validates it. This condition adds an undocumented WebGPU-only rejection even though positive capacity is the stated restriction. Keep rejecting capture_count, but accept/validate a [1] active input, mark input 10 as a CPU input, and ignore its value because capture is unavailable.

  ORT_RETURN_IF_NOT(capture_count == nullptr && state_update_active == nullptr,
                    "capture_count and state_update_active require state_update_capacity > 0");

• Files reviewed: 7/7 changed files
• Comments generated: 4
• Review effort level: Balanced

[Jiajia Qin (qjia7)]

Review frame

• Problem/feature validity: Validated. Before this PR, com.microsoft::GatedDeltaNet had only CUDA registrations, so fused Qwen gated-delta graphs could not remain native on WebGPU.
• Risk/scope: Deep. This adds a recurrent GPU kernel with generated WGSL, optional and aliased state, ragged device offsets, inverse GQA, adapter binding limits, and FP16/FP32 numerical behavior.
• Direction gate: Pass. A correctness-first FP32 recurrence is an appropriate first WebGPU implementation. The PR keeps validation and dispatch in the EP kernel, uses a preprocessing pass only when the recurrence would exceed WebGPU's guaranteed storage-binding limit, and reuses the writable final-state binding for aliased state. The blocker below is an implementation gap in that direction.

Confirmed findings

C1: Route all tensor accesses through the segmented-buffer helpers

onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.wgsl.template:71

The new shaders directly index every storage variable (query[...], key[...], value[...], parameters[...], output[...], and the state buffers); the parameter shader does the same at gated_delta_net_params.wgsl.template:32-46. ORT's ProgramManager::CalculateSegmentsForInputsAndOutputs splits any tensor larger than maxStorageBufferBindingSize into multiple bindings, and ShaderHelper makes those bindings reachable only through ShaderVariableHelper::GetByOffset/SetByOffset. Raw indexing continues to address only the first binding.

For a concrete schema-valid case on an adapter with the WebGPU-guaranteed 128 MiB storage-binding limit, rank-4 linear FP16 inputs with total_tokens=262145, Hq=Hv=1, Dk=256, Dv=1, no initial state, and omitted final state make each query/key buffer 134,218,240 bytes. The host validation accepts those dimensions and the six resulting physical bindings remain below the guaranteed limit, but accesses for the final token index past the first query/key segment instead of selecting query1/key1. Depending on robustness mode, this produces invalid results or a device/validation failure.

Please pass the ShaderVariableHelper objects into both WGSL templates (WGSL_TEMPLATE_VARIABLE/#use .getByOffset .setByOffset) and replace every raw storage-buffer access with getByOffset or setByOffset. Optional bindings can follow the existing placeholder-helper convention used by turbo_quant_hadamard.cc: pass a required variable when the optional binding is absent, while guarding every use with the matching template parameter. Add a regression that forces segmentation with ep.webgpuexecutionprovider.maxStorageBufferBindingSize (or uses a real >128 MiB buffer) and compares output across the segment boundary.

Non-blocking follow-ups

F1: Add common packing support for binding-constrained programs

onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc:82

The supported ragged Qwen path with initial and final state binds exactly eight logical storage tensors in the recurrence pass: query, key, value, cu_seqlens, packed parameters, initial state, output, and final state. Physical bindings, however, are counted after large tensors are segmented. ShaderHelper::AddVariableImpl adds each segment to numbers_storage_buffers_ and rejects the program once it exceeds the adapter's maxStorageBuffersPerShaderStage.

For example, on a compliant adapter exposing the minimum eight storage bindings and 128 MiB per storage binding, a valid FP16 ragged Qwen call with total_tokens=262145, Hq=Hv=1, Dk=256, and Dv=1 segments both query and key. The recurrence therefore needs ten physical bindings and fails shader generation with Too many storage buffers in shader, even though this exact ragged-Qwen-with-state combination is advertised as supported and the adapter can hold each tensor as a multi-segment WebGPU buffer.

This is a common limitation rather than something GatedDeltaNet should solve privately. A standalone WebGPU infrastructure change should support packing several small logical tensors into one physical storage buffer, expose aligned buffer views through the same GetByOffset/SetByOffset abstraction, and account for read/write aliasing, temporary lifetime, and copy cost. Until that exists, keep the current parameter packing and avoid documenting an unconditional guarantee for large segmented tensors on eight-binding adapters.

Optimization options

These are non-blocking follow-ups; they need representative timestamp-query measurements before changing dispatch policy.

P1: Tile multiple V dimensions in each workgroup

The current dispatch assigns one workgroup to each (batch, value_head, value_index). Every value_index for the same value head therefore reloads the same Q/K vector and repeats Q/K normalization, dot(K,Q), reductions, and gate exponentiation. With head_size_v=128, that head-level work is repeated 128 times.

The highest-value optimization is to let one workgroup own a small V tile, such as 4 or 8 columns. Q/K and dot(K,Q) can then be shared while each lane keeps multiple state values. This preserves the recurrence and FP32 state contract while reducing storage reads, reductions, exponentials, and dispatch count.

P2: Select workgroup size from head_size_qk

The fixed 256-thread workgroup leaves half or three quarters of its lanes idle for common head_size_qk values of 128 or 64. Consider selecting the next power of two at or above head_size_qk, capped at 256, and include the resulting workgroup size in the program variant. Benchmark 64/128/256 variants across Metal, D3D12, and Vulkan before choosing thresholds.

P3: Keep an adapter-aware direct-parameter decode path

The parameter-preparation pass is the correct portable fallback when the recurrence would exceed the adapter's binding limit, but its additional dispatch and temporary tensor can dominate single-token decode. When the actual adapter has enough storage bindings for the active optional inputs and their segments, consider binding decay, beta, and Qwen parameters directly. Retain packed parameters for binding-constrained adapters and include the choice in the cache identity.

P4: Treat parallel prefill as a separate optimization

The serial token loop is appropriate for a correctness-first and decode-oriented implementation, but long prefill remains sequential within each state row. A later PR could add a chunked or scan-based prefill path while keeping this recurrence as the decode and fallback implementation. That change should be justified with representative Qwen prefill measurements because it would substantially increase algorithm and workspace complexity.

Clarifications

None.

Test coverage

The macOS arm64 Release WebGPU job at the requested head executed and passed all eight added GatedDeltaNetWebGpuTest cases, including all update rules, FP16/FP32, dense/ragged layouts, explicit scale, omitted final state, and aliased state. The WGSL template checks and the rest of CI are also green. Those cases keep every tensor below one storage-binding segment, so they do not exercise C1 or the binding-capacity follow-up. Linux WebGPU CI is build-only and adds no runtime evidence for these limits.

Verdict

The feature is valid and the overall design direction is appropriate. C1 blocks merge because schema-valid segmented tensors access only the first physical binding. No clarification requests remain. Common multi-tensor binding packing and the performance options above should be handled as measured follow-up work rather than expanding this correctness-first PR.

[kunal-vaishnavi]

Follow-up to review #5165372626: C1 is fixed in 49e561e; all GatedDeltaNet tensor accesses now use segmented-buffer helpers. The inline reply is at #32510 (comment).

I’m treating F1 and P1–P4 as the requested next scope: common multi-tensor packing infrastructure, V-dimension tiling, adaptive workgroup sizing, adapter-aware direct parameter bindings, and a measured prefill strategy.

[kunal-vaishnavi]

Follow-up to review #5165372626: a2537fa implements P2 and P3: the recurrent kernel now selects the next power-of-two workgroup size for head_size_qk, and it keeps direct decay/beta/Qwen bindings when their physical segment count fits the adapter's actual storage-binding limit, using parameter packing only when necessary.

F1 requires shared WebGPU tensor-packing infrastructure (logical typed views, alignment, segmentation, aliasing, and lifetime management), and P1/P4 require independently benchmarked algorithm changes. I have not represented those as completed in this PR because the review explicitly identifies them as separate, non-blocking measured follow-up work.

[kunal-vaishnavi]

Update to review #5165372626: b36bd2279 adds the common packed-buffer-view infrastructure and P1 V-channel tiling; a2537fa adds P2 adaptive workgroup sizing and P3 adapter-aware direct parameter bindings; 2f6a9d... (current tip) uses the packed views to copy Q/K/V into a shared backing buffer when segmentation would exceed the adapter binding limit. C1 remains fixed by segmented accessor helpers.

P4 cannot be safely represented by another dispatch of the existing recurrence: delta/gated-delta prefill requires an associative affine-state scan with dense Dk×Dk transform composition and bounded workspace, while WGSL has no cross-workgroup synchronization. I need confirmation to expand this PR with that new algorithm and its required workspace/precision contract rather than introduce a racy or semantically different path.

[Tianlei Wu (tianleiwu)]

Review summary (head f111a98)

The dispatch-count defect from the previous round is fixed, and all 89 checks are green. Since c13bfc0 the PR also grew a general WebGPU framework capability (ProgramInput/Output::BufferView), a parallel chunked prefill scan for the linear rule, and two unrelated plugin-EP build fixes.

What checks out

• Barrier placement in gated_delta_net.wgsl.template is correct. Lane 0's read of retrieved/pre_output/key_query[0] is followed by a workgroupBarrier() before delta[] is consumed, and the end-of-token barrier separates that read from the next iteration's writes. All loop trip counts (first_token/last_token, chunk_start/chunk_end) are workgroup-uniform, so every barrier sits in uniform control flow.
• Output semantics match the reference on both paths: pre_output + delta * (k·q) is algebraically S_new^T q, and the prefill output pass applies the state update before the reduction.
• The parallel-prefill decomposition is sound. chunk_state[c] is written before accumulating chunk_contribution[c], giving the exclusive prefix state, and the carry ping-pong parity (chunk_base / chunks_per_pass) & 1 alternates the read/write buffers correctly across passes. The guard list for use_parallel_prefill conservatively excludes every rule whose state transition is not additive.
• Segment accounting for the new view feature is consistent end to end: non-owner views get 0 segments in CalculateSegmentsForInputsAndOutputs, AddVariableImpl only charges numbers_storage_buffers_ for owners, GenerateSourceCode declares bindings only for owners, and WebGpuContext::Run skips zero-segment entries so binding order still matches the generated WGSL. Views inherit the owner's segments_, and the storage_offset is added before the chunk division, so a view straddling a segment boundary works. Adding View=V{owner}@{offset} to the cache key is necessary and correct.
• Host validation mirrors the CUDA kernel (rank-1 a_log/dt_bias, initial_state required for rank-3 uniform packing, hv % hq == 0), so provider selection does not change which models are accepted.

Findings — one latent out-of-bounds read plus four suggestions, left inline.

#	Severity	Issue
1	High	Aliased-state read of the final_state output bypasses the segmented accessor; out-of-bounds when state exceeds maxStorageBufferBindingSize.
2	Suggestion	The new BufferView path has no deterministic test coverage on CI adapters.
3	Suggestion	Baking the view offset into the shader text fragments the pipeline cache, and the unconditional + 0 rewrites every existing WebGPU shader.
4	Suggestion	Hard-coded 64 MiB prefill workspace cap, not derived from device limits.
5	Suggestion	The segmented test costs ~128 MiB x 2 of device memory and ~1 GB of host RAM.

Scope: compute_context.h and webgpu_execution_provider.h/.cc carry plugin-EP build fixes unrelated to GatedDeltaNet, and BufferView is a general framework feature. AGENTS.md asks for <= 10 files per PR. Splitting the framework change into its own PR with its own tests would make both easier to review and to revert. Not blocking.

see latest review