Repository navigation
[WebGPU] Fuse rank-32 quantized LoRA with MatMulNBits - #33370
Draft
Jiajia Qin (qjia7) wants to merge 42 commits into
Draft
Jiajia Qin (qjia7) wants to merge 42 commits into
Jiajia Qin (qjia7) wants to merge 42 commits into
Conversation
Skip uploader teardown before acquiring telemetry locks when Windows is detaching DLLs for process termination. WinHTTP callbacks and worker threads are already gone at that point, so draining them can wait forever. Keep normal object destruction and explicit DLL unload cleanup. Validated against captured native hang stacks and identical machine code at all resolved frames. A pending-uploader fault injection hangs the unchanged runtime and exits in 20 of 20 fixed trials. Both WebGPU modes completed 1024-prompt, 32-output-token native profiles with LoRA on and off; explicit unload also passed.
Add an opt-in MatMulNBitsLoRA operator and graph fusion for INT4/block-32 base weights with replaceable FP16 rank-32 adapters. Empty adapters execute only the base projection. Match original and offline-transposed layouts, absorb private adapter Transpose nodes, and preserve overridable inputs. Use two fusion implementations on the same transposed model: - fuse_down: base + T, followed by base + T * Bt^T. - fuse_up: T, followed by base with an integrated T * Bt^T epilogue. Default session.webgpu.lora_fusion_mode to auto: select fuse_up for M > 1 prefill and fuse_down for M == 1 decode on each invocation. Use X's per-matrix row dimension, independent of batch size. Retain explicit fuse_down/fuse_up overrides for comparisons. A one-token prompt uses the decode path. Mode selection applies to transposed adapters only. Both implementations read transposed B and T as vec4 FP16 components. Share the decode GEMV reduction and prefill down projection, preserving FP16 rounding, 128-thread decode producers and 64-thread standalone output. Keep DP4A activation quantization for the base, with original FP16 X for T. Name internal flags has_lora_down/has_lora_up by the fused projection. Retain compatible original-layout and base-kernel fallback paths. Validation on Windows/D3D12, RTX 5080, built-in and dynamic-plugin WebGPU: - 26 focused native tests pass with no skips; cover omitted mode, auto, both forced paths, decode, prefill, batches and empty/partial adapters. - 192 auto/default full-model cases and 4024 full-vocabulary logit vectors match the saved original unfused FP16 reference bitwise; rewind matches. - Cover adapter on/off and graph capture on/off for both providers. - Traces/shader dumps confirm 196 fused projections switching from prefill fuse_up to decode fuse_down, vec4 bindings, no adapter transpose, and base-only work with LoRA off. - All 72 native 1024-input/300-output benchmark sequences match exactly.
Add opt-in graph fusion and WebGPU kernels for INT4 base projections with block-32 INT8 adapters. Preserve replaceable packed weights, original INT8 arithmetic, partial overrides, and base-only execution for empty adapters. Use fuse_up for prefill and fuse_down for one-row NVIDIA decode, with packed vec4 loads. Add strict numerical and graph-matching tests and document native GenAI usage.
Contributor
|
Hi Jiajia Qin (@qjia7), we tested this branch on an NVIDIA GeForce RTX 4070 (Windows D3D12, built-in WebGPU) and ran into two numerical tolerance issues in the unit test suites:
|
Retain FP16/FP32 MatMulNBitsQLoRA and remove floating-only schema, transformer, kernels and shader branches. Preserve QLoRA dispatch policy and packed adapter arithmetic. Validation: 38 native tests, 32 exact full-model comparisons, WASM syntax checks and lint. Isolated built-in WebGPU before/after timings show no repeatable material regression; record sanitized results.
Jiajia Qin (qjia7)
changed the title
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Both joint down-projection paths need bounds checks to prevent padded dispatches from accessing GPU buffers out of bounds.
1 open issue
What changed in this PR
Adds opt-in rank-32 quantized LoRA fusion to ONNX Runtime’s WebGPU provider, keeping adapters replaceable while reducing projection and transfer overhead.
Changes:
- Adds the QLoRA operator, graph matcher, optimized kernels, and packaging documentation.
- Batches CPU-to-GPU uploads and expands numerical, adapter-switching, and capture tests.
- Includes Windows telemetry shutdown protection and lifecycle tests.
| File | Description |
|---|---|
| onnxruntime/test/providers/webgpu/data_transfer_test.cc | Tests batched uploads and reuse. |
| onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test.cc | Drives isolated lifecycle checks. |
| onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.def | Exports the test entry point. |
| onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.cc | Verifies telemetry teardown behavior. |
| onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc | Tests fusion, overrides, serialization, and capture. |
| onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc | Tests QLoRA kernels and numerical accuracy. |
| onnxruntime/test/contrib_ops/matmul_8bits_test.cc | Adds narrow-output INT8 regressions. |
| onnxruntime/test/contrib_ops/matmul_4bits_test.cc | Adds INT4 regressions and updates a comment. |
| onnxruntime/test/contrib_ops/matmul_2bits_test.cc | Runs FP16 cases and adds INT2 regressions. |
| onnxruntime/core/providers/webgpu/data_transfer.h | Declares batched transfers. |
| onnxruntime/core/providers/webgpu/data_transfer.cc | Routes eligible transfers through batching. |
| onnxruntime/core/providers/webgpu/buffer_manager.h | Declares upload batches. |
| onnxruntime/core/providers/webgpu/buffer_manager.cc | Implements bounded staging batches. |
| onnxruntime/core/platform/windows/telemetry_1ds.cc | Detects Windows process shutdown. |
| onnxruntime/core/platform/telemetry_1ds.h | Enables lifecycle test access. |
| onnxruntime/core/platform/telemetry_1ds.cc | Skips teardown during process termination. |
| onnxruntime/core/platform/telemetry_1ds_platform.h | Declares shutdown detection. |
| onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.h | Declares the QLoRA transformer. |
| onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc | Matches projections and preserves adapter defaults. |
| onnxruntime/core/optimizer/graph_transformer_utils.cc | Registers opt-in fusion. |
| onnxruntime/core/graph/contrib_ops/contrib_defs.cc | Defines the operator schema. |
| onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc | Registers the WebGPU kernel. |
| onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.wgsl.template | Adds fused up-projection output. |
| onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h | Extends projection interfaces. |
| onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc | Integrates QLoRA dispatch paths. |
| onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template | Adds wide-tile adapter output. |
| onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h | Declares QLoRA helpers. |
| onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc | Implements QLoRA execution and shaders. |
| onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h | Adds shared quantization workspace and interfaces. |
| onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc | Reuses quantized inputs and selects narrow tiles. |
| onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_k_parallel.wgsl.template | Implements joint quantized down projection. |
| docs/WebGPU_MatMulNBitsQLoRA_Prototype.md | Documents packaging, dispatch, and capture requirements. |
| docs/WebGPU_MatMulNBitsQLoRA_Performance.md | Records performance and numerical limitations. |
| docs/ContribOperators.md | Documents the new operator. |
| cmake/onnxruntime_unittests.cmake | Registers telemetry lifecycle tests. |
Review effort: Balanced
Compute the adapter down projection with a narrow subgroup-matrix tile, reuse the packed input, and fuse the up projection and addition into the subgroup base kernel. Preserve the DP4A and tiled paths and cover subgroup dispatch selection in tests. Verified on Intel Vulkan: built onnxruntime and onnxruntime_provider_test; all 17 MatMulNBitsQLoRA and WebGpuDataTransfer tests passed. Changed-file lintrunner checks passed.
Use the existing bounded staging upload path for multi-tensor copies on UMA too. UMA capability alone does not mean destination buffers are mapped; serial uploads unnecessarily submit each tensor separately. Existing mapped-buffer and single-copy fallbacks remain unchanged. Verified on Intel Vulkan: all 5 WebGpuDataTransfer tests passed, including chunked uploads and buffer reuse; all 12 MatMulNBitsQLoRA tests also passed. Build and changed-file lint passed.
Batch same-device CPU-to-GPU inputs on a shared non-null WebGPU session stream using the existing bounded staging uploader, including on UMA. Preserve mixed-transfer, mapped-buffer, single-copy, and streamless fallbacks, and retain allocation clears and deferred command ordering. Add real submission-count and data-integrity tests for repeated uploads, staging chunk boundaries, empty inputs, oversized tensors, and single-input fallback. The batching tests fail on the control plugin with four submissions instead of one or two. Verified Intel Vulkan: production and provider-test builds, lint, 67 plugin tests passed with 6 intentional capability-specific skips; all 17 built-in QLoRA/data-transfer tests passed. Fixed 1024-token adapter-on prefill improved from 848.1 to 716.7 ms across paired run orders (15.5%); decode remained about 35 tokens/s. All work remains local.
Merge PR microsoft#33370 at be2c49c, including its latest main integration, Dawn upgrade, subgroup-ID opt-in, and padded-dispatch guards. Preserve fused subgroup prefill with shared narrow-down input, UMA upload batching, and owning-stream plugin upload batching. Combine upstream wave64 tiling with the shared tiling helper; bound the fused 256-row tile K depth to keep shared memory within 32 KiB. Update dispatch-test expectations for the expanded cache key. Verified: built-in Vulkan runtime and provider tests build; all 30 focused Vulkan tests pass with no skips; all 89 WGSL tooling tests pass; lint passes for the local changes. No local build-system changes.
Merge origin/main at a7a1172. Retain upstream's extracted wide-tile helper while preserving fused QLoRA inputs, cache keys, and dispatch guards. Make serialized adapter-state tests use an exact nonzero base, correct half-open random bounds for fixed packed weights, and verify independent expected outputs without relaxing tolerances. Verified both built-in and plugin Vulkan builds. All 36 added tests pass across configurations. Built-in suites: 33 focused, 50 broader matmul (7 existing skips), and 12 fusion tests; fusion repeated twice. Plugin: 70 passed, 6 expected internal-API skips covered by built-in tests. WGSL tooling: 89 passed. Scoped lint and PR-relative whitespace checks pass.
kunal-vaishnavi
pushed a commit
to microsoft/onnxruntime-genai
that referenced
this pull request
Oct 22, 2026
Activating a LoRA adapter after a generator has captured decode currently reuses the base path's `gpu_graph_id`. Graph replay bypasses input processing and kernel selection, so decode can silently ignore the newly activated adapter. This change includes the active-adapter revision in capture identity, alongside input length and execution variant. It reuses the process-wide `GraphAnnotationIds` allocator to prevent collisions between generators sharing a session, retains older captures for cleanup, and uses the actual capture ID for recurrent-state capture fixup. Adapter activation remains additive; it does not recompute earlier KV-cache entries. The new Python regression is included in the regular test runner. It captures the base path, activates A and then A+B on the same generator after rewind, compares full logits with fresh active generators, and verifies a fresh base generator sharing the session. It also checks distinct capture IDs and adapter unloading. Validation on Windows / NVIDIA RTX 5080: - Release native library and Python wheel built successfully. Formatting and whitespace checks pass. - WebGPU regression: 2 passed (capture off/on); the two CUDA parameterizations were not run locally. - Full-model reproduction before the fix: capture-off comparisons were exact, but captured decode differed from a fresh active generator by up to 2.82764 in logits. After the fix, all six forced-prefix logit rows match exactly with built-in WebGPU, and with the plugin under capture. - Isolated native performance checks cover FP16 and INT8 adapters, on/off, at 128 and 1024 prompt tokens. Two initially slower cases were checked in a three-configuration comparison separating the ORT update from GenAI. Final 1024-token confirmation changes versus the previous build were +0.14% decode / -0.41% TTFT for FP16 adapter-off and -0.25% decode / +0.15% TTFT for INT8 adapter-on. No repeatable regression was observed on this GPU. Related ORT operator and capture coverage: microsoft/onnxruntime#33370. This GenAI fix addresses capture identity independently of that operator implementation.
Accept zero-sized packed defaults while retaining logical output shapes, and skip unused parameter packing and computation. Allow QLoRA fusion from empty adapter defaults with dynamic axes, and cover correctness, overrides, and graph replay.
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.
Jiajia Qin (qjia7)
added a commit
that referenced
this pull request
Oct 28, 2026
### Description Support empty packed weights in CPU and WebGPU `MatMulNBits`. With `weight_prepacked=0`, one or more dimensions of B's packed shape may be zero; the other dimensions must retain their normal values. The output keeps the logical shape determined by A, K and N, and contains zeros or the broadcast bias. Unused quantization parameters may also be empty. CPU skips unused parameter packing and matrix multiplication; WebGPU records a fill dispatch so graph replay rewrites the output correctly. QKV/MLP fusion remains enabled: when a projection has empty weights, the fused operator uses the shared zero/bias fallback while computing the nonempty projections. Includes operator documentation and focused tests for FP16/FP32, 2/4/8-bit weights, constant/runtime inputs, bias, malformed shapes, empty/nonempty overrides, and graph capture. ### Motivation and Context Optional projections, such as disabled LoRA adapters, should not require full-size zero weight tensors. Zero-sized defaults let the same raw `MatMulNBits` graph run on CPU and WebGPU, with nonempty weights supplied when needed. Split from #33370. That PR retains these changes until this PR merges; QLoRA-specific operators, fusion and shaders remain there. ### Validation Built the standalone branch against main with built-in WebGPU on Windows x64. All 16 empty-weight tests, 28 existing QKV/MLP fusion regressions, and 8 existing nonempty-weight regressions passed; 4 Arm64-specific tests were skipped. Coverage includes fused empty/nonempty overrides, MLP bias, residual outputs, and capture replay. FP16 WebGPU cases require confirmed ShaderF16 support. Lint and `git diff --check` passed.
Restore upstream tiled DP4A configuration while preserving the fused LoRA-up output stage and K-parallel narrow-N path. Update documentation and test comments to match. Validated with D3D12 provider tests, Muffin performance comparisons, and full-model token/logit checks with adapters off/on and graph capture off/on.
This branch has not been deployed
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Description
Adds
com.microsoft.MatMulNBitsQLoRAfusion for the base projection, LoRA down/up projections, and residual addition. Supported rawMatMulNBits/Addpatterns 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
MatMulNBitssupport 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
MatMulNBitsquantization settings and arbitrary ranks.Motivation
Reduce GPU dispatches and duplicated work in quantized LoRA inference while preserving runtime adapter switching.
Optimization method
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.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
DequantizeLinear/MatMuladapter-pattern matching.Adapter-on/off states require separate graph-capture identities; see onnxruntime-genai#2740.
Latest local validation: build and lint passed; 37 targeted tests passed, with one subgroup-matrix test skipped for device capability.