Repository navigation
[WebGPU] Batch input uploads for built-in and plugin EPs - #33491
Conversation
Extract the shared uploader and built-in CopyTensors optimization from the QLoRA work so input-transfer performance can be reviewed independently. Group CPU-to-GPU inputs in staging buffers capped at 16 MiB, including on UMA, while retaining single-copy and mixed-transfer fallbacks. Include the existing round-trip test for empty and unaligned tensors, chunk boundaries, oversized inputs, and repeated writes. All five extracted files match the source branch exactly. Scoped lint and whitespace checks pass; standalone build and runtime tests have not been rerun on this branch. Source commits: 8eab658 and 034acaa; final source tree: 2971861.
Extract plugin batching from a49f019. Group same-device CPU-to-GPU inputs on a shared non-null WebGPU session stream using the bounded staging uploader, including on UMA. Preserve mapped-buffer, mixed-transfer, single-copy, and streamless fallbacks and existing command ordering. Include five existing submission-count and data-integrity regression tests covering repeated uploads, staging boundaries, empty inputs, oversized inputs, and single-input fallback. All four extracted files match source commit 2971861. Scoped lint and whitespace checks pass; standalone build and runtime tests have not been rerun on this branch.
There was a problem hiding this comment.
🟡 Changes recommended
Built-in batching accepts mapped destinations and bypasses their existing copy-and-unmap fallback, causing validation failures.
What changed in this PR
Extracts WebGPU input-upload batching into a shared implementation for built-in and plugin execution providers.
Changes:
- Groups eligible uploads into staging buffers capped at 16 MiB.
- Adds plugin batching on a shared owning session stream.
- Adds six regression tests for data integrity, boundary cases, and submission counts.
| File | Description |
|---|---|
onnxruntime/test/providers/webgpu/plugin/test_webgpu_allocators.cc |
Tests session uploads and submission counts. |
onnxruntime/test/providers/webgpu/data_transfer_test.cc |
Tests upload integrity across chunks and repeated writes. |
onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc |
Checks plugin batch eligibility. |
onnxruntime/core/providers/webgpu/ep/sync_stream.h |
Declares stream-based batch uploads. |
onnxruntime/core/providers/webgpu/ep/sync_stream.cc |
Routes uploads through the owning provider. |
onnxruntime/core/providers/webgpu/data_transfer.h |
Declares built-in batch transfer support. |
onnxruntime/core/providers/webgpu/data_transfer.cc |
Routes eligible built-in transfers to batching. |
onnxruntime/core/providers/webgpu/buffer_manager.h |
Defines the shared batch-upload interface. |
onnxruntime/core/providers/webgpu/buffer_manager.cc |
Implements staging chunks and oversized-upload fallback. |
Review effort: Balanced
Check destination map state during the built-in batch preflight and fall back for the entire list before copying. This preserves the existing direct memcpy/unmap behavior for mapped inputs while retaining batching for unmapped inputs. Add coverage for mapped destinations first, last, and in both positions, plus repeated uploads after unmapping. The regression fails before the fix and passes after it. Verified built-in and plugin Vulkan builds, 6 built-in transfer tests, 62 enabled plugin tests, and scoped lint.
|
Thanks for working on this improvement and sharing the performance numbers! The reduction in upload overhead looks very promising. I have a few questions to better understand the tradeoffs and scope:
Thanks again for the optimization and for helping clarify these points! |
|
I compared both approaches locally on Intel Arc 140V with Vulkan. Here are the results and scope clarifications:
|
Xiaofei Han (xiaofeihan1)
left a comment
There was a problem hiding this comment.
Thanks for verifying this! LGTM, approved.
Keep the upstream upload implementation and regression coverage from microsoft#33491, including mapped-buffer fallback, and remove the superseded plugin upload test from the QLoRA diff.
### Description <!-- Describe your changes. --> Automatically use `Queue.WriteBuffer` for eligible CPU-to-GPU uploads up to and including 1 MiB. - Flush deferred readers before queue writes to preserve read-before-overwrite ordering. - Preserve mapped-buffer handling and staging for larger, unaligned, and capture/replay uploads. - Retain existing completion, validation, device-loss, and failure handling. - Add focused upload tests and document ordering, source-lifetime guarantees, and the size threshold. ### Motivation and Context <!-- - Why is this change required? What problem does it solve? - If it fixes an open issue, please link to the issue here. --> Profiling Qwen/DFlash2 on A100 identified substantial host overhead from allocating, mapping, copying, and submitting a fresh staging buffer for each upload. Queue writes let Dawn reuse staging allocations and batch pending transfers with subsequent queue work, without adding CPU completion waits. The conservative 1 MiB cap is the largest consistently beneficial size measured on A100/Vulkan. Larger transfers can regress, so they retain the existing staging path rather than extrapolating the small-upload improvement. **Validation:** On Linux/A100/Vulkan, 80 native tests passed with one unrelated build-mode skip, and 100 GenAI tests passed. Both 512-token baseline and DFlash2 generation runs preserved exact tokens and speculative counters. Formatting and lint checks passed. ### Update: Rebased onto latest main, including merged PR #33491, and reran A100/Vulkan benchmarks to measure the incremental benefit of `Queue.WriteBuffer`. **Results**: - Baseline: **3.3% faster** (short), **6.8% faster** (long) - DFlash2 width 3: **2.5% faster** (short), **3.1% faster** (long) All generated tokens and speculative counters matched. Native and plugin tests passed. Existing batched uploads remain unaffected, and larger transfers retain the staging path. This confirms that the optimization provides additional performance benefits beyond #33491. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
### Description Adds `com.microsoft.MatMulNBitsQLoRA` fusion for the base projection, LoRA down/up projections, and residual addition. Supported raw `MatMulNBits`/`Add` patterns fuse automatically on WebGPU at extended graph optimization or higher. Adapters remain replaceable; base weights are unchanged. Empty packed defaults let one raw model run on CPU and fuse on WebGPU without storing full-size zero weights. Fusion preserves initializer shapes and data: existing empty defaults stay empty, and full-size zero defaults retain their computation. General empty-weight `MatMulNBits` support is split into **#33559** and remains included here until that PR merges. The current WebGPU kernel targets FP16/FP32, INT4 base weights, rank-32 INT8 adapters, and block size 32. The operator definition allows other `MatMulNBits` quantization settings and arbitrary ranks. ### Motivation Reduce GPU dispatches and duplicated work in quantized LoRA inference while preserving runtime adapter switching. ### Optimization method - **Prefill:** share activation quantization between base/down projections on eligible DP4A paths; fuse up/Add into the base output. Use subgroup-matrix acceleration where available. - **Decode:** compute base/down together where supported, followed by one up/Add dispatch; retain fallback paths. - **Empty adapters:** skip LoRA computation and temporaries in the fused WebGPU kernel. Input upload batching was split into **#33491 and has been merged**. It is already included in both sides of the measurements below. ### Performance RTX 5080, built-in WebGPU/D3D12, FP16, accuracy level 4, batch 1, capture enabled, 128 generated tokens with readback. October 23 medians of ten requests (two reversed-order rounds, three warmups each). Original and converted/fused graphs use the same PR runtime (`507021fd21`). TTFT is time to first token. | Prompt tokens | Adapter | TTFT ms: unfused → optimized | Decode tok/s: unfused → optimized | |---:|---|---:|---:| | 128 | Off | 32.87 → 31.22 | 197.24 → 235.14 | | 128 | On | 79.50 → 40.81 | 112.87 → 199.15 | | 1024 | Off | 171.21 → 162.79 | 195.42 → 232.69 | | 1024 | On | 221.92 → 173.21 | 112.44 → 196.35 | Adapter-on TTFT falls **22–49%** and decode improves approximately **75%**; a second adapter shows similar gains. Adapter-off decode improves **19%**. Capture-disabled runs also improve; session creation remains approximately 6.8 seconds. In six sample prompts, adapter-off outputs match the original; with adapters active, 4/6 generated sequences match (logit RMSE 0.076). Broader quality validation remains necessary. ### Follow-ups - Compose with QKV/MLP fusion to close the approximately 10% adapter-off decode gap against a pruned base-only graph. - Investigate the repeatable approximately 1.5 ms / 5% increase in 128-token TTFT on unchanged models versus main; decode is essentially unchanged. - Support more ranks/quantization settings and direct `DequantizeLinear`/`MatMul` adapter-pattern matching. Adapter-on/off states require separate graph-capture identities; see [onnxruntime-genai#2740](microsoft/onnxruntime-genai#2740). Latest local validation: build and lint passed; 37 targeted tests passed, with one subgroup-matrix test skipped for device capability. --------- Co-authored-by: qjia7 <4221210+qjia7@users.noreply.github.com>

Description
Performance
Intel Arc 140V / Vulkan, previous A/B measurement: 44.84 MB of adapter data across 784 tensors: 95.3 → 10.8 ms upload host-call time (~8.8× faster).
Validation