You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
This introduces important optimizations for audio models, which tends to rely heavily on nodes such as DFT, STFT, ConvTranspose and for some LSTM & GRU.
Motivation and Context
WebGPU didn't have those node implemented (such as STFT) or they were not really optimized (DFT only optimized for sizes <=512, slow ConvTranspose, LSTM & GRU).
The speed gain was tested across 4 popular audio model architectures: HTDemucs, UVR, SCNet, MP-SENet on 2 different hardware accelerators: NVIDIA RTX 3090, and Apple M3.
The measured speed gain before and after these changes:
NVIDIA RTX 3090 on Windows 11 (CPU: Intel 12700K):
Rewrite the WebGPU DFT as a Stockham mixed-radix FFT computed in workgroup
memory, inspired by the MLX execution provider:
- real-input transforms use N/2-point complex FFTs with packing/unpacking
(forward onesided and inverse);
- lengths above the shared-memory limit use a four-step FFT;
- non-smooth lengths fall back to Bluestein (chirp-z);
- the first and last stages read and write global memory directly.
Add an STFT kernel (opset 17) that fuses framing and windowing into the FFT
loads.
On an RTX 3090, htdemucs6 (1 STFT) goes from 1132 ms to 1057 ms end to end
with this change alone; scnetlargestarrytong runs its 3 STFTs on the GPU.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
For group-1, 1D/2D ConvTranspose with at least 16 input channels, compute
X[M, C_in] x W[C_in, kH*kW*C_out] with the tiled MatMul, then gather the
valid taps with a col2im kernel. Constant weights are prepacked to the GEMM
layout (K and N padded to multiples of 4 for vec4 loads); the input is
K-padded and processed in row chunks to respect buffer limits.
This is what models use for iSTFT (overlap-add). On an RTX 3090, htdemucs6
goes from 1057 ms to 265 ms (iSTFT node 812 ms -> 19 ms), identical output.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Compute the input projection of all timesteps and directions with one GEMM,
then for each step a batched h * R^T GEMM followed by a fused gate kernel
(bias, peephole, clip, input_forget, sequence_lens, layout 0/1). W^T and
R^T are prepacked; the recurrent state stays in f32. The previous per-step
kernel is kept as a fallback when the projection buffer exceeds the device
binding limits.
On an RTX 3090: scnetlargestarrytong 10.56 s -> 0.525 s (CPU EP: 7.68 s),
17_HP-Wind_Inst-UVR 457 ms -> 342 ms (LSTM 207 ms -> 18 ms).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
Critical float16 precision loss and overflow in LSTM and ConvTranspose remain unresolved, alongside missing multi-chunk accumulation coverage and excessive CPU reference-test workload.
This PR improves audio-model inference in ONNX Runtime’s WebGPU provider by expanding FFT support, adding STFT, and introducing matrix-multiplication-based paths for ConvTranspose and LSTM.
Changes:
Replaces the DFT fallback with a shared FFT engine supporting larger and irregular transform lengths, also used by STFT.
Adds packed-weight and matrix-multiplication paths for wide-channel ConvTranspose and LSTM.
Adds reference-based operator tests and documents WebGPU STFT support.
Address review feedback on the WebGPU GEMM paths:
- ConvTranspose: for float16, the packed weights, the K-padded input, the
column buffer and the accumulation across row chunks are f32, so per-tap
partial sums can neither overflow nor be rounded before later taps cancel
them. Multi-chunk float16 runs accumulate in an f32 buffer and the last
chunk writes the output; the bias is bound on the first chunk only.
- LSTM: the input projection X * W^T is computed and stored in f32 (W^T is
prepacked in f32, float16 X is converted first), matching the per-step
kernel, so the bias can cancel it without float16 rounding.
- Tests: ConvTranspose with a reduced storage binding size, forcing four
row chunks with contributions across every boundary (f32 and float16);
float16 regressions for cancelling ConvTranspose taps and for an LSTM
projection cancelled by the bias.
- Signal tests: the double-precision reference uses a radix-2 FFT (with
Bluestein for other lengths) above 512 points instead of the O(N^2) sum.
float32 results and timings are unchanged (htdemucs6 260 ms,
17_HP-Wind_Inst-UVR 339 ms on an RTX 3090, outputs bit-identical).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Hi Jiajia Qin (@qjia7)adrastogiYang Gu (@gyagp)Joshua Lochner (@xenova) , this PR speeds up the WebGPU EP on audio models (DFT/STFT, ConvTranspose via GEMM, GEMM-based LSTM & GRU; e.g. htdemucs 1132 → 265 ms, SCNet 12.4 s → 0.53 s on an RTX 3090). As a first-time contributor I can’t start the CI: could you approve the workflows and run /azp run, and review when you have time? The Copilot comments are addressed in 60979d3. Thanks!
Robin Lobel (divideconcept)
changed the title
WebGPU optimizations for audio models (DFT, STFT, ConvTranspose, LSTM)
[WebGPU] optimizations for audio models (DFT, STFT, ConvTranspose, LSTM)
Nov 8, 2026
The WebGPU GRU ran two kernels per timestep and direction, each thread
looping over the whole input and hidden dimension with uncoalesced reads.
Like the LSTM rewrite, the input projection X * W^T for every timestep and
direction is now one GEMM (in f32). Per timestep:
- linear_before_reset = 1 (the PyTorch export default): a single fused
kernel computes H_prev * R^T in workgroup tiles and applies the gates,
for every direction at once.
- linear_before_reset = 0: a batched GEMM for the z / r rows, a kernel for
the gates and r (.) H_prev, a GEMM with Rh^T and a kernel for the new
state.
W^T and R^T are prepacked in f32 when constant. The per-step kernels remain
as a fallback when the input projection exceeds the storage binding limit.
MP-SENet (8 bidirectional GRUs, hidden size 128) on an RTX 3090: 213 ms ->
42 ms per inference (GRU GPU time 179 ms -> 16 ms), outputs within 119 dB
SNR of the CPU EP.
New tests check larger random cases against a double-precision reference:
both reset modes, bidirectional and reverse, sequence_lens, initial_h,
batch-first layout, clip and float16.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Robin Lobel (divideconcept)
changed the title
[WebGPU] optimizations for audio models (DFT, STFT, ConvTranspose, LSTM)
[WebGPU] optimizations for audio models (DFT, STFT, ConvTranspose, LSTM, GRU)
Nov 8, 2026
- ConvTranspose PrePack: use if constexpr on the channels-last template
parameter (MSVC C4127 is an error in CI).
- ConvTranspose test: convert the double reference to float explicitly
(C4244).
- GRU / LSTM reference tests: skip the WebGPU-only cases before building
the OpTester; a debug-build OpTester that never runs traps in its
destructor.
- GRU reference tests: exclude the CUDA EP, whose GRU disagrees with the
reference when linear_before_reset = 1 and biases are given.
- Wide ConvTranspose tests: allow 1e-2 absolute error for CUDA TF32.
- Remove a duplicate blank line flagged by clang-format.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Thanks adrastogi for running the CI! All the failures came from this PR. They are fixed in 3093429:
Windows builds: two MSVC warnings treated as errors.
C4127 in conv_transpose.cc, now an if constexpr on the template parameter.
C4244 in the new ConvTranspose test, now an explicit double → float conversion.
Debug builds (Linux, macOS, Android): the new GRU and LSTM reference tests skipped the WebGPU-only cases after building the OpTester. A debug OpTester that never runs traps in its destructor ("Test was not executed"). The skip now happens before the tester is built.
CUDA plugin EP: the wide ConvTranspose tests drift by up to ~4e-3, which looks like TF32. The absolute tolerance is now 1e-2.
Lint: removed a duplicate blank line.
One finding outside this PR: the existing CUDA GRU kernel disagrees with a double-precision reference when linear_before_reset = 1 and biases are given. The output is already wrong at the first timestep, while the CPU and WebGPU EPs match the reference. The existing CUDA GRU test has no bias, so it doesn't catch this. CUDA EP was excluded from the new GRU reference tests (GRUTest.LargerForwardLinearBeforeReset, GRUTest.LargerBidirectionalAllOptions) rather than change the CUDA kernel here.
Could you approve the workflows again for the new commit? Thanks!
- GRU reference tests: also exclude kCudaNHWCExecutionProvider, which
runs the same CUDA GRU kernel, and the DML EP, which does not zero the
outputs past sequence_lens.
- LSTM reference tests: exclude the DML EP for the same reason.
- Wide ConvTranspose float16 tests: exclude XNNPACK, whose float16
deconvolution misses the bias on these shapes.
- Wide ConvTranspose tests with pads, dilations or output_padding:
exclude the CUDA EP when it runs as the dynamic plugin EP; the CUDA
plugin EP on Windows gets them wrong while the built-in CUDA EP passes.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Pushed 05cc3b6, which excludes the EPs that disagree with the new reference tests:
GRU: CUDA NHWC and DML. The NHWC provider runs the same CUDA GRU kernel, and DML does not zero the outputs past sequence_lens.
LSTM: DML, for the same reason.
ConvTranspose float16: XNNPACK, which misses the bias on these shapes.
ConvTranspose with pads, dilations or output_padding: the CUDA EP when it runs as the plugin EP. The plugin gets these wrong on Windows; the built-in CUDA EP passes.
The remaining Linux CI (GCC 14 graph.cc) and Android minimal binary-size failures are unrelated to this PR and also fail on main. The graph.cc one is fixed on main by #34068.
Hi Robin Lobel (@divideconcept) — once the CI workflow runs appear for the current commit, you can comment exactly /run bots on this PR to trigger them. This permission applies only to #33978, not to your other PRs. Please post a new command after a new commit if CI needs approval again.
Please allow up to 10 minutes for the command to take effect. If nothing happens after 10 minutes, post /run bots again. If it still does not work, please let me know.
Before requesting human review, please make sure Copilot has reviewed the latest commit and reports either “Approval recommended. No open issues” or “Needs a closer look. No open issues”. If it reports open issues, please address them and obtain an updated Copilot review first. Thanks.
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
Empty-spatial ConvTranspose inputs can leave outputs unwritten, while DFT scratch allocations can exceed device buffer limits even when inputs and outputs fit.
… input
- FFT: Bluestein and four-step scratch buffers hold every transform row,
so a large batch could ask for a buffer beyond maxBufferSize even when
the input and output fit. Run the rows in chunks whose scratch fits
both maxBufferSize and maxStorageBufferBindingSize; a row_offset
uniform maps each chunk's rows to the operator input and output.
- ConvTranspose GEMM path: an empty spatial input (m == 0) left the
non-empty output unwritten. Run one col2im, which reads no columns and
writes the bias (or zeros).
- Tests: Bluestein DFT and STFT with a reduced binding size that splits
the rows into chunks (strided rows and chunks crossing STFT batches
included), and an empty-input ConvTranspose with and without bias.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
The new DFT/STFT tests fail in WebGPU plugin builds by selecting a null execution provider, while chunked ConvTranspose introduces quadratic col2im traversal overhead.
Each row chunk dispatches col2im over the entire output and traverses the tap loops before rejecting other chunks' contributions, making the gather quadratic in input length at a fixed binding limit. For float32 X=[1,16,262144], W=[16,1,4096], stride 1 and a 128 MiB binding, 32 chunks perform about 34 billion row-range checks for roughly 1 billion contributions. Restrict each chunk's output/tap traversal to its reachable region, initializing bias separately, or retain the direct path when chunking would cause this amplification.
- ConvTranspose GEMM path: with several column chunks, each chunk's
col2im ran over the whole output and rejected the other chunks' taps,
which made long inputs quadratic. A first pass now writes the bias (or
zeros) to the whole output, each chunk gathers only into the output
box its input rows reach, and a last pass converts the f32 sums to a
narrower output type. One chunk keeps the single full-output pass.
- Tests: 2D chunked ConvTranspose (float32 and float16) with chunks
inside one input row, across rows and across batches.
- DFT/STFT chunked-row tests: create the WebGPU EP before the test and
skip when it is unavailable; plugin builds cannot cap the binding size.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
Human assessment is needed because float32 packing and accumulation can exceed WebGPU buffer limits for valid float16 workloads, with risks spanning initialization, runtime path selection, and fallback behavior.
All tests now pass, and all Copilot remarks have been addressed.
On the remaining Copilot "closer look" point: for float16 inputs, ConvTranspose, LSTM and GRU keep their intermediates in float32 (packed weights, input projections, partial sums). This is intentional and fixes the float16 precision and overflow issues raised earlier by Copilot. It means these buffers are twice the size of the float16 tensors. The float32 path already allocates buffers of the same size, and the large ConvTranspose column buffer is already chunked to the binding limit. A float16 model would only hit the device buffer limits through these buffers if its tensors were already close to them.
Could you take a look and merge when you have time? Thanks!
- DFT/STFT: chunked rows reuse one set of scratch buffers instead of
allocating new ones per chunk, which the buffer manager would keep
alive until the queue is submitted.
- ConvTranspose: the GEMM path requires the packed (padded, f32 for
float16) filter to fit one buffer and one binding, both when
prepacking and at run time; otherwise the direct kernel runs. When a
float16 output needs column chunks and its f32 accumulator would not
fit, the output is gathered in one pass from X and the packed weights.
- LSTM/GRU: W and R are prepacked to f32 only when the f32 copy fits,
and the GEMM path also requires it; otherwise the stepwise kernel
runs on the original weights.
- Tests: float16 ConvTranspose with the fused gather (1D and 2D) and
with a packed filter over the binding limit.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Robin Lobel (divideconcept)
changed the title
Description
This introduces important optimizations for audio models, which tends to rely heavily on nodes such as DFT, STFT, ConvTranspose and for some LSTM & GRU.
Motivation and Context
WebGPU didn't have those node implemented (such as STFT) or they were not really optimized (DFT only optimized for sizes <=512, slow ConvTranspose, LSTM & GRU).
The speed gain was tested across 4 popular audio model architectures: HTDemucs, UVR, SCNet, MP-SENet on 2 different hardware accelerators: NVIDIA RTX 3090, and Apple M3.
The measured speed gain before and after these changes:
NVIDIA RTX 3090 on Windows 11 (CPU: Intel 12700K):
Apple M3 on macOS 14.7:
Op usage: