[Feature Request] WebGPU LinearAttention state_window for speculative/MTP rollback

[FU-max-boop]

Describe the feature request

I would like to implement the existing state_window contract for WebGPU LinearAttention.

This is not a new operator or schema proposal. #31157 added the contract and its CUDA implementation for speculative/MTP decoding; CPU and WebGPU currently retain explicit state_window = 0 checks. The current WebGPU constructor therefore rejects state_window > 0.

The contract I propose to preserve exactly is:

• state_window = 0: legacy state shape [B, H_kv, d_k, d_v].
• state_window = W, 1 <= W <= 8: state shape [W, B, H_kv, d_k, d_v].
• The window is right-aligned: slot j is the recurrent state after prefix length T - W + j + 1; slot W - 1 is the final state.
• Only past_state[W - 1] initializes the next call. Earlier past slots are not read.
• When T < W, leading slots without a corresponding token are zero.

For a first implementation, I propose a deliberately narrow boundary:

• WebGPU LinearAttention only.
• Reuse the shared linear_attention_helper parsing and shape validation.
• Extend the existing WGSL token loop, which already keeps the recurrent state private while processing tokens sequentially, to write only the trailing checkpoints.
• No CausalConvWithState, GenAI scheduler, graph schema, or model-export changes in the first PR.

This scope does not claim end-to-end MTP support for every hybrid model: models that also need windowed causal-convolution state would still require that separate EP implementation and scheduler integration.

Describe scenario use case

Speculative/MTP decoding verifies several draft tokens in one forward pass. If only a prefix is accepted, a KV cache can discard its rejected tail, but a recurrent linear-attention state cannot be truncated after the fact: after verifying all tokens, it represents the full sequence. Without intermediate checkpoints, the accepted prefix must be rerun to reconstruct the correct carry state.

state_window makes that rollback explicit and bounded for the trailing prefixes represented by the window. The op only produces recurrent-state checkpoints; the caller must select a represented slot and place it in the next call's past_state[W - 1] by rebinding, gathering, or copying. A zero-token acceptance still requires retaining the original past_state. The feature therefore enables caller-managed checkpoint/select/discard semantics rather than providing a scheduler or commit operation itself.

Correctness plan

I would generalize the existing CUDA prefix oracle and cover:

• W = 1, 4, 8, including T < W.
• FP32/FP16, batch 1/2, standard and inverse GQA, and all four update rules.
• Output parity with state_window = 0.
• Every slot against a separately evaluated prefix.
• Poisoned earlier past slots to prove that only W - 1 is read.
• Leading-zero behavior when T < W.
• Select-a-slot-and-resume parity for every represented checkpoint prefix, plus preservation of the original past state for a zero-token acceptance.
• Aliased and non-aliased recurrent-state bindings.

Performance and memory gates

The feature should remain evidence-gated because checkpoint writes and state memory scale with W. For FP16 with B=1, H_kv=32, and d_k=d_v=128, one slot is about 1 MiB per layer; W=4 across 30 recurrent layers is about 120 MiB of state.

My local benchmark would use 20 warmups, 200 measured iterations, and 3 independent repetitions, reporting W=0/1/4/8 p50/p95, checkpoint-select-and-resume versus accepted-prefix rerun, incremental live GPU bytes, and copy/allocation counts where observable without changing the WebGPU allocator. Initial GO gates would be correctness first, no more than 5% W=1 steady-state regression, no more than 15% W=4 verification overhead, and no unbounded live-tensor or scratch growth beyond the declared window and bounded kernel scratch. Allocator-reserved memory may remain cached.

For rollback value, I would report the measured break-even rejection probability

p* = (verify_W - verify_0) / E[prefix_rerun - checkpoint_select | rejection]

and use p* <= 20% as the initial local GO gate.

Local feasibility baseline and overlap

On macOS arm64 / Apple M5 Pro at 4d308dacbbb385fcba9911cd9c07f5603d65cbd6, a native ORT WebGPU build completed with static WGSL templates enabled:

python tools/ci_build/build.py --build_dir build-webgpu-baseline --config Release \
  --enable_onnx_tests --use_webgpu --wgsl_template static --cmake_generator Ninja \
  --cmake_extra_defines CMAKE_OSX_ARCHITECTURES=arm64 --update --skip_submodule_sync \
  --skip_pip_install --build --parallel 10 --target onnxruntime_test_all \
  --compile_no_warning_as_error
cmake --build build-webgpu-baseline/Release --target onnxruntime_provider_test --parallel 10

The non-CUDA build instantiated 40 filtered tests: 39 EP-executing cases plus one schema-bound rejection case. All 40 passed in 4.3 seconds with no skips:

onnxruntime_provider_test --gtest_filter='ContribOpLinearAttentionTest.*' --gtest_color=no

The 39 operational cases use a helper that selects DefaultWebGpuExecutionProvider() before the CPU fallback and pass it explicitly to OpTester, so this exercises the WebGPU EP rather than only proving compilation.

As of 2026-08-24, I also searched open issues/PRs for state_window and WebGPU LinearAttention. I did not find an active change touching the proposed WebGPU LinearAttention files; the active state-window work I found is in CausalConvWithState, which this proposal excludes.

Would this WebGPU-first boundary and the existing CUDA contract be acceptable? If so, I can start with the minimal kernel/test PR and keep scheduler/model integration out of scope.