onnxruntime
e64a8ae8 - [WebGPU] Fuse rank-32 quantized LoRA with MatMulNBits (#33370)

Commit
4 days ago
[WebGPU] Fuse rank-32 quantized LoRA with MatMulNBits (#33370) ### 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](https://github.com/microsoft/onnxruntime-genai/pull/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>
Author
Parents
Loading