onnxruntime
ce9bb081 - [WebGPU EP] implements chunked gated delta net optimization (#33143)

Commit
31 days ago
[WebGPU EP] implements chunked gated delta net optimization (#33143) # ๐Ÿš€ WebGPU Chunked Gated Delta Net Prefill Optimization ## โšก Motivation Prefill in WebGPU `GatedDeltaNet` previously relied on purely sequential token-by-token recurrence. While well-suited for single-token decode, this serial dependency chain underutilized GPU compute during prompt processing, creating a noticeable Time-To-First-Token (TTFT) bottleneck on models like Qwen 3.8-27B. ## ๐Ÿ”ง Optimization We overhauled prefill by breaking the sequential token chain into parallel blocks and maximizing GPU on-chip memory reuse: - **Chunk Factorization (Parallel Recurrence)**: Grouped tokens into chunks of 16. Within each chunk, intra-token interactions form a lower-triangular system $(I + M)$ that we solve in parallel via forward substitution. This turns token-by-token recurrence into parallel matrix ops, leaving only a fast state carry between chunks. - **Collaborative Shared Memory Tiling**: Cooperatively cached $Q, K, W,$ and $Q_g$ tiles into fast on-chip workgroup shared memory (acting like an L1 cache) instead of letting 128 threads redundantly fetch from slow VRAM. This cut inner-loop global DRAM reads by **128x**. - **4-Wide Vectorization**: Unrolled inner dot products and accumulator loops 4 elements at a time, allowing the GPU compiler to issue 4-wide fused multiply-accumulate (FMA) instructions to saturate Apple Silicon SIMD execution units. - **Bounded Workspace Allocation**: Fixed sequence budgeting so scratch memory scales linearly with sequence length (`O(seq_len)`) per sequence rather than quadratically (`O(batch^2)`), staying safely under the 64 MiB hardware buffer cap. ## ๐ŸŽฏ Scope This PR adds an experimental chunkwise prefill route to contrib `GatedDeltaNet` on WebGPU: - Activates on Apple adapters with Qwen gating, sigmoid beta, L2-normalized Q/K, and gated-delta updates when `chunk_size=16` is explicitly configured. - Decode steps, unnormalized inputs, unsupported head dims, state-capture modes, and the default `chunk_size=64` remain on the existing recurrent path. ## โœ… Validation 1. **๐Ÿงช ONNX Runtime Unit Tests** - Added unit test cases in `GatedDeltaNetWebGpuTest` covering partial chunks, ragged/empty sequences, packed QKV, and optional final state retention. - All 20 WebGPU Gated Delta Net operator and planning tests pass (`20/20 passed`). - C++ formatting verified clean via `lintrunner --take CLANGFORMAT`. 2. **๐Ÿ“Š GenAI Model Benchmark (Qwen 3.8-27B)** - Measured on Apple M4 Pro (16-core GPU, 64 GB unified memory) with 1077 prompt tokens, batch size 1, greedy decoding across 2 warmups and 5 fresh-process measured runs. - **TTFT reduced by 14.05%** (from `14.920 s` down to `12.823 s`, saving `2.10 s` per request). - Decode throughput (`8.48 tok/s`) and peak memory (`16.30 GiB`) remained steady, with 100% token output match across benchmark runs. --------- Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: Prathik Rao <prathikrao@microsoft.com>
Author
Parents
Loading