diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index d45736fe9e6f1..b96f499ead85b 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -73,6 +73,7 @@ Do not modify directly.* * com.microsoft.MatMulIntegerToFloat * com.microsoft.MatMulNBits * com.microsoft.MatMulNBitsMlp + * com.microsoft.MatMulNBitsQLoRA * com.microsoft.MatMulNBitsQkv * com.microsoft.MaxpoolWithMask * com.microsoft.MoE @@ -4503,6 +4504,69 @@ This version of the operator has been available since version 1 of the 'com.micr +### **com.microsoft.MatMulNBitsQLoRA** + + Computes MatMulNBits(X,W,scales) + MatMulNBits(MatMulNBits(X,lora_A,lora_A_scales),lora_B,lora_B_scales). Each projection follows MatMulNBits block quantization and bit packing, with implicit zero points 2^(bits-1) for the base and 2^(lora_bits-1) for the adapters. No bias or explicit zero points are applied. The adapter rank R may be any positive integer and need not be a multiple of lora_block_size. Partial final blocks are padded; only the logical K or R elements participate in each projection. An empty rank dimension in either packed adapter disables the adapter, while scale defaults are retained. Intermediate projections and block scales use the activation type before the final addition. Adapters remain replaceable. Execution providers may support a subset of the defined types, quantization settings, and ranks. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
K : int (required)
+
Positive input feature dimension of the base and down-projection weight matrices.
+
N : int (required)
+
Positive output feature dimension of the base and up-projection weight matrices.
+
accuracy_level : int
+
Accuracy policy following MatMulNBits: 0 (unset), 1 (fp32), 2 (fp16), 3 (bf16), or 4 (int8). Applies to the base and eligible adapter down projections; the up projection uses level 0.
+
bits : int
+
Base weight bit width, following MatMulNBits (2, 4, or 8).
+
block_size : int
+
Base quantization block size along K, following MatMulNBits: a power of two, at least 16.
+
lora_bits : int
+
Adapter weight bit width, following MatMulNBits (2, 4, or 8).
+
lora_block_size : int
+
Adapter quantization block size along K for lora_A and R for lora_B: a power of two, at least 16.
+
+ +#### Inputs + +
+
X : T
+
Input of rank at least one, with last dimension K, following MatMulNBits.
+
W : TQ
+
Packed base weights [N,ceil(K/block_size),block_size*bits/8].
+
scales : T
+
Base block scales [N,ceil(K/block_size)]; may be flattened.
+
lora_A : TQ
+
Packed down weights [R_A,ceil(K/lora_block_size),lora_block_size*lora_bits/8], R_A=R or 0.
+
lora_A_scales : T
+
Down block scales [R,ceil(K/lora_block_size)]; retained when the adapter is empty.
+
lora_B : TQ
+
Packed up weights [N,R_blocks,lora_block_size*lora_bits/8], R_blocks=ceil(R/lora_block_size) or 0.
+
lora_B_scales : T
+
Up block scales [N,ceil(R/lora_block_size)]; retained when the adapter is empty.
+
+ +#### Outputs + +
+
Y : T
+
Input shape with the last dimension replaced by N.
+
+ +#### Type Constraints + +
+
T : tensor(float16), tensor(float)
+
Activations and block scales share the same floating-point type (FP16 or FP32).
+
TQ : tensor(uint8)
+
Packed quantized weights.
+
+ + ### **com.microsoft.MatMulNBitsQkv** MatMulNBitsQkv fuses either SimplifiedLayerNormalization (RMSNorm) diff --git a/docs/WebGPU_MatMulNBitsQLoRA.md b/docs/WebGPU_MatMulNBitsQLoRA.md new file mode 100644 index 0000000000000..895c588986e61 --- /dev/null +++ b/docs/WebGPU_MatMulNBitsQLoRA.md @@ -0,0 +1,236 @@ +# WebGPU MatMulNBitsQLoRA + +`com.microsoft.MatMulNBitsQLoRA` combines a quantized base projection with a +quantized LoRA adapter during inference. Adapter weights remain separate and +replaceable; the operator does not merge them into the base weights or train them. +It reduces dispatches and intermediate memory traffic by sharing input preparation +and computing the adapter contribution inside existing MatMulNBits shaders. + +## Operator contract and current support + +The [operator definition](ContribOperators.md#com.microsoft.matmulnbitsqlora) +follows MatMulNBits quantization and packing rules. It allows arbitrary positive +adapter ranks, including ranks not divisible by eight or the block size. Partial +blocks are padded, and only logical elements participate in the projections. + +The current WebGPU kernel and graph matcher implement this subset: + +- FP16 or FP32 activations, scales, intermediate projections, and output, all + using the same floating type. +- INT4 base weights and INT8 adapter weights, all with block size 32. +- Active adapter rank 32, or empty packed adapter weights as described below. +- Input `X` of rank at least two, positive `K` divisible by 32, and positive `N`. + `N` need not be divisible by four; scalar output tails are supported. +- Implicit zero points: 8 for the base and 128 for the adapters. No bias, + explicit zero points, or group indices. + +These are implementation limits, not additional schema restrictions. Explicit +nodes reject unsupported configurations; the transformer leaves unsupported +patterns unfused. The fused kernel is available only on WebGPU. Both built-in and +plugin WebGPU can use it; plugin deployments need an ORT host with the schema and +transformer as well as a provider library with the kernel. + +For the current implementation, the active inputs are: + +| Input | Storage type | Shape | +| --- | --- | --- | +| `X` | FP16 or FP32 | `[..., M, K]` | +| `W` | Packed INT4 in UINT8 | `[N, K/32, 16]` | +| `scales` | Same as `X` | `[N, K/32]` or flattened | +| `lora_A` | Packed INT8 in UINT8 | `[32, K/32, 32]` | +| `lora_A_scales` | Same as `X` | `[32, K/32]` | +| `lora_B` | Packed INT8 in UINT8 | `[N, 1, 32]` | +| `lora_B_scales` | Same as `X` | `[N, 1]` | + +Let `A_down` and `B_up` be the logical dequantized matrices of shapes `[K, 32]` +and `[32, N]`. The adapter bytes represent signed values offset by 128, multiplied +by their block scales. The computation is: + +```text +base = MatMulNBits(X, W, scales) +T = X * A_down +delta = T * B_up +Y = base + delta +``` + +"Down" projects the input features to the adapter rank; "up" projects that rank +to the output features. `delta` is the adapter's contribution to the output. +`T` has shape `[..., M, 32]`, and `Y` has shape `[..., M, N]`. Dequantized weights +and intermediate projection results use the activation type before addition. +Any LoRA scaling factor must already be represented in the weights/scales; the +operator has no separate adapter multiplier. + +## Automatic graph fusion + +[MatMulNBitsQLoRAFusion](../onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc) +runs at `ORT_ENABLE_EXTENDED` or higher when WebGPU provides the fused kernel. +It matches this graph, accepting either order of the final Add inputs: + +```mermaid +flowchart LR + X --> Base["MatMulNBits: INT4 base"] + X --> Down["MatMulNBits: INT8 down"] + Down --> Up["MatMulNBits: INT8 up"] + Base --> Add + Up --> Add + Add --> Y +``` + +All four nodes must be assigned to WebGPU. The down projection must consume the +same `X` as the base projection. The base, down, and up outputs must each have +one consumer and must not be graph outputs. Packed shapes must match the layout +above, with logical dimensions established by the MatMulNBits `K`/`N` attributes. +Both adapter nodes must use `accuracy_level=0`; the fused node inherits the base +node's accuracy setting. Bias, explicit zero points, group indices, or prepacked +weights (`weight_prepacked != 0`) prevent matching. + +Fusion reuses the input NodeArgs and preserves initializer shapes, contents, and +overridability. It does not quantize adapters, compact zero defaults, or match +arbitrary DequantizeLinear/Cast/MatMul graphs. Such graphs need an offline +conversion to the packed MatMulNBits representation first. + +No QLoRA-specific enable option is needed. To compare with the raw graph, set +the session configuration entry `optimization.disable_specified_optimizers` to +`MatMulNBitsQLoRAFusion` (append it with a semicolon if other names are present). +Disabling the transformer does not expand explicitly exported fused nodes. + +## Execution and optimization + +[MatMulNBitsQLoRA::ComputeInternal](../onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc) +selects a path from the input shape, accuracy policy, and device capabilities. +Here, `M` is the second-to-last dimension of `X`; a one-token prompt uses the +same selection as decode. Graph fusion does not imply a single GPU dispatch. + +| Case | Execution | +| --- | --- | +| Either packed adapter is empty | Run ordinary base MatMulNBits, before allocating or dispatching adapter work. | +| Prefill (`M > 1`) | Compute `T`, then fuse the up projection and addition into the base shader's output stage. | +| Eligible decode (`M = 1`) | Compute base and down projections in separate workgroups of one dispatch, then run an up-plus-add output shader. | +| Other decode shapes | Compute down and base projections separately, then run the up-plus-add output shader. | + +Packed weights are dequantized where consumed, avoiding full floating-point +adapter-weight tensors. Prefill retains the rank-sized `T` intermediate while +avoiding separate base and adapter-delta output tensors. Decode retains separate +base and `T` outputs; the following dispatch synchronizes them before addition. + +### DP4A + +The [shared DP4A implementation](../onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc) +can quantize `X` once and reuse its packed values and scales for the base and +down projections within an invocation. Its normal eligibility checks, including +`accuracy_level=4`, still apply. Eligible prefill therefore uses quantization, +down projection, and base-plus-up dispatches. FP32 short prefill can also share +this preparation when the base selects DP4A. + +The rank-32 down projection uses the existing K-parallel kernel selection: +`M < 4`, or narrow output (`N <= 32`) with `K >= 2048`. Narrow reductions use +FP32 accumulation. Other shapes use the existing tiled DP4A kernel; the LoRA +output stage extends its computation without changing its matrix tiling. + +For FP32 decode with `2048 <= K <= 8192` and a DP4A base, the joint kernel adds +down-projection workgroups to the base dispatch. Both projections read the same +quantized `X`; up-plus-add is a separate dispatch. Prepared inputs are scoped to +the invocation and are not reused as cached values across different inputs. + +### Subgroup matrix + +The [subgroup-matrix implementation](../onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc) +adds the up projection using matrix instructions in the base shader's output +stage. On a compatible FP16 `8x16x16` configuration, a specialized path packs `X` +once for both projections, with row padding sufficient for both tile sizes. +The narrow down projection uses FP32 matrix accumulation and stores `T` in the +activation type; the base projection then consumes `T` for the fused up stage. + +Selection requires supported component types, matrix dimensions, subgroup size, +and device limits. Choosing a Vulkan backend alone does not guarantee support. +If the shared-packing specialization is unavailable, normal MatMulNBits selection +computes the down projection and selects an eligible base-plus-up shader. + +### Floating-point decode and shader helpers + +For one total input row, `2048 <= K <= 8192`, and a base that does not select +DP4A, the generic base shader appends four down-projection workgroups. A compatible +32-lane subgroup path reuses input values in registers and reduces partial sums; +other devices use workgroup memory. A variable subgroup range requires explicit +subgroup-size control before requesting 32 lanes. + +[GenerateQLoRADown and GenerateQLoRAUp](../onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h) +append reusable WGSL fragments to a caller's shader; they do not dispatch GPU work. +The [down template](../onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora_down.wgsl.template) +supplies the generic joint-decode body. The +[up template](../onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora_up.wgsl.template) +supplies dequantization and adapter-delta helpers used by fused prefill and the +separate decode output shader. Subgroup-matrix shaders reuse the dequantization +helper and implement the multiplication with their own matrix instructions. + +## Optional adapters and portable models + +Use the raw three-MatMulNBits-plus-Add pattern for a model shared by CPU and +WebGPU. CPU executes the original operators; WebGPU fuses eligible patterns. +Saving the optimized graph serializes explicit MatMulNBitsQLoRA nodes, which +require its WebGPU kernel. + +To export a compact adapter-off default, use overridable initializers with these +shapes and also list them as graph inputs: + +| Initializer | Default shape | Contents | +| --- | --- | --- | +| `lora_A` | `[0, K/32, 32]` | Empty packed weights | +| `lora_B` | `[N, 0, 32]` | Empty packed weights | +| `lora_A_scales` | `[32, K/32]` | Retained scale defaults | +| `lora_B_scales` | `[N, 1]` | Retained scale defaults | + +Declare the packed graph inputs' adapter axes as dynamic: axis 0 for `lora_A` +and axis 1 for `lora_B`. This allows active rank-32 weights to override the empty +defaults. Either empty packed weight disables the adapter. Scale-only or +single-weight overrides therefore retain the base result; both packed weights +activate the adapter using supplied scales or their defaults. Scale defaults may +also be empty, but activation then requires supplying their active shapes and data. + +This convention relies on empty-weight MatMulNBits support in the raw graph. +Full-size zero adapter weights still execute arithmetic, including its NaN/Inf +behavior. The transformer deliberately preserves them instead of converting them +to empty tensors. + +## Numerical behavior, capture, and limitations + +`accuracy_level` follows MatMulNBits policy for the base and eligible down +projections. Level 4 permits DP4A activation quantization, including in an +optimized down projection whose raw node uses level 0. Quantization, reduction +order, and accumulation precision can change results relative to the raw graph; +fusion does not promise bitwise equivalence. `enableMatmulFp32Accumulation` +controls accumulation in the shaders that honor that option; it does not disable +activation quantization. Evaluate fused-versus-raw differences separately from +the change caused by quantizing an originally floating-point adapter. + +Graph replay does not re-evaluate the empty-adapter branch. Native callers must +use distinct `gpu_graph_id` values for adapter-off and adapter-on captures, and +for other changes to shapes or execution paths. Each capture needs compatible, +stable GPU bindings and buffer lifetimes. Updating adapter values in the same +buffers can reuse a capture when shapes, layout, and dispatches remain compatible; +replacing a buffer or changing the path requires recapture. WebGPU does not derive +the capture ID from adapter state. + +The current QKV/MLP fusions do not compose with MatMulNBitsQLoRA. An empty adapter +uses standalone base MatMulNBits, so it may have different performance from a +separate base-only graph that benefits from those fusions. Supporting arbitrary +active ranks, additional quantization settings, and composed base fusions requires +further kernel/transformer work; the broader operator schema already allows rank +and quantization extensions. + +## Validation + +From a native WebGPU build's test output directory, run: + +```text +onnxruntime_provider_test.exe --gtest_filter="*QLoRA*:*MatMulNBitsInputPreparation*" +onnxruntime_test_all.exe --gtest_filter="*MatMulNBitsQLoRA*" +``` + +[Provider tests](../onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc) and +[fusion tests](../onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc) +cover FP16/FP32 numerics, prefill/decode, batching and tails, empty and partial +adapters, serialization, capture replay, and patterns that must remain unfused. +Hardware-specific tests skip when the required features are absent; those skips +do not validate the corresponding shader path. Full-model evaluation should also +check logits and generated tokens with each intended adapter and capture mode. diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template index 501a8d1d9dc09..5459a918e44b0 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template @@ -5,6 +5,7 @@ #param block_size #param n_bits #param has_bias +#param has_qlora_up #param has_zero_points #param is_qualcomm #param qc_wave128 @@ -13,6 +14,17 @@ #use .getByOffset .setByOffset +fn write_qlora_output(offset: u32, value: vec4) { +#if has_qlora_up + let row = (offset * 4u) / uniforms.N; + let col = (offset * 4u) % uniforms.N; + let delta = qlora_delta4(row, col); + output.setByOffset(offset, value + delta); +#else + output.setByOffset(offset, value); +#endif +} + #include "quantization/dp4a_matmul_common.wgsl.template" // This shader implements co-operative matrix multiply. The key idea here is to @@ -478,15 +490,15 @@ $MAIN { acc_element_t(bias[b_global + 14 + b_bias_offset]), acc_element_t(bias[b_global + 15 + b_bias_offset]) ); - output.setByOffset(output_idx, vec4(vec4(lane_outputs[0], lane_outputs[1], lane_outputs[2], lane_outputs[3]) + bias_vec1)); - output.setByOffset(output_idx+1, vec4(vec4(lane_outputs[4], lane_outputs[5], lane_outputs[6], lane_outputs[7]) + bias_vec2)); - output.setByOffset(output_idx+2, vec4(vec4(lane_outputs[8], lane_outputs[9], lane_outputs[10], lane_outputs[11]) + bias_vec3)); - output.setByOffset(output_idx+3, vec4(vec4(lane_outputs[12], lane_outputs[13], lane_outputs[14], lane_outputs[15]) + bias_vec4)); + write_qlora_output(output_idx, vec4(vec4(lane_outputs[0], lane_outputs[1], lane_outputs[2], lane_outputs[3]) + bias_vec1)); + write_qlora_output(output_idx+1, vec4(vec4(lane_outputs[4], lane_outputs[5], lane_outputs[6], lane_outputs[7]) + bias_vec2)); + write_qlora_output(output_idx+2, vec4(vec4(lane_outputs[8], lane_outputs[9], lane_outputs[10], lane_outputs[11]) + bias_vec3)); + write_qlora_output(output_idx+3, vec4(vec4(lane_outputs[12], lane_outputs[13], lane_outputs[14], lane_outputs[15]) + bias_vec4)); #else - output.setByOffset(output_idx, vec4(vec4(lane_outputs[0], lane_outputs[1], lane_outputs[2], lane_outputs[3]))); - output.setByOffset(output_idx+1, vec4(vec4(lane_outputs[4], lane_outputs[5], lane_outputs[6], lane_outputs[7]))); - output.setByOffset(output_idx+2, vec4(vec4(lane_outputs[8], lane_outputs[9], lane_outputs[10], lane_outputs[11]))); - output.setByOffset(output_idx+3, vec4(vec4(lane_outputs[12], lane_outputs[13], lane_outputs[14], lane_outputs[15]))); + write_qlora_output(output_idx, vec4(vec4(lane_outputs[0], lane_outputs[1], lane_outputs[2], lane_outputs[3]))); + write_qlora_output(output_idx+1, vec4(vec4(lane_outputs[4], lane_outputs[5], lane_outputs[6], lane_outputs[7]))); + write_qlora_output(output_idx+2, vec4(vec4(lane_outputs[8], lane_outputs[9], lane_outputs[10], lane_outputs[11]))); + write_qlora_output(output_idx+3, vec4(vec4(lane_outputs[12], lane_outputs[13], lane_outputs[14], lane_outputs[15]))); #endif #else #if has_bias @@ -517,15 +529,15 @@ $MAIN { acc_element_t(bias[b_global + 14 + b_bias_offset]), acc_element_t(bias[b_global + 15 + b_bias_offset]) ); - output.setByOffset(output_idx, vec4(lane_output1 + bias_vec1)); - output.setByOffset(output_idx+1, vec4(lane_output2 + bias_vec2)); - output.setByOffset(output_idx+2, vec4(lane_output3 + bias_vec3)); - output.setByOffset(output_idx+3, vec4(lane_output4 + bias_vec4)); + write_qlora_output(output_idx, vec4(lane_output1 + bias_vec1)); + write_qlora_output(output_idx+1, vec4(lane_output2 + bias_vec2)); + write_qlora_output(output_idx+2, vec4(lane_output3 + bias_vec3)); + write_qlora_output(output_idx+3, vec4(lane_output4 + bias_vec4)); #else - output.setByOffset(output_idx, vec4(lane_output1)); - output.setByOffset(output_idx+1, vec4(lane_output2)); - output.setByOffset(output_idx+2, vec4(lane_output3)); - output.setByOffset(output_idx+3, vec4(lane_output4)); + write_qlora_output(output_idx, vec4(lane_output1)); + write_qlora_output(output_idx+1, vec4(lane_output2)); + write_qlora_output(output_idx+2, vec4(lane_output3)); + write_qlora_output(output_idx+3, vec4(lane_output4)); #endif #endif } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_k_parallel.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_k_parallel.wgsl.template index 2bed945685d16..4f622ffba9276 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_k_parallel.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_k_parallel.wgsl.template @@ -18,12 +18,27 @@ #param n_bits #param broadcast_a_row #param has_zero_points +#param has_qlora_up #param has_bias #param has_weight_idx #param has_weight_idx_indirect +#param has_qlora_down #use .getByOffset .setByOffset +fn write_qlora_output(offset: u32, value: output_element_t) { +#if has_qlora_up + let row = (offset * 1u) / uniforms.N; + let col = (offset * 1u) % uniforms.N; + let delta = qlora_delta(row, col); + output.setByOffset(offset, value + delta); +#else + output.setByOffset(offset, value); +#endif +} + + + #include "quantization/dp4a_matmul_common.wgsl.template" const double_tile_size_k_vec = 2 * tile_size_k_vec; @@ -58,7 +73,51 @@ fn loadSHMA(batch: u32, a_global: u32, kidx_v: u32, col: u32) } } +#if has_qlora_down +// Eight extra workgroups compute rank 32 from the same quantized activation. +// The next dispatch consumes both outputs, so no cross-workgroup barrier is needed. +fn qlora_down(group: u32, local_idx: u32) { + let batch = group / 8u; + let rank_base = (group % 8u) * 4u; + let lane = local_idx % 32u; + let rank = local_idx / 32u; + var partial = f32(0); + for (var start = 0u; start < uniforms.K32; start += 32u) { + if (local_idx < 64u) { loadSHMA(batch, 0u, start * 2u, local_idx); } + workgroupBarrier(); + let block = start + lane; + if (block < uniforms.K32) { + let offset = (rank_base + rank) * uniforms.K32 + block; + // XOR maps unsigned INT8 with zero point 128 to signed packed INT8. + let w0 = lora_a[offset * 2u] ^ vec4(0x80808080u); + let w1 = lora_a[offset * 2u + 1u] ^ vec4(0x80808080u); + partial += SDP8AI(tile_A[lane * 2u], w0, tile_A[lane * 2u + 1u], w1, + scale_A[lane / 4u] * lora_a_scales[offset]); + } + workgroupBarrier(); + } + inter_results[rank][lane] = partial; + workgroupBarrier(); + if (local_idx < 4u) { + var sum = f32(0); + for (var i = 0u; i < 32u; i++) { sum += inter_results[local_idx][i]; } + lora_down[batch * 32u + rank_base + local_idx] = sum; + } +} +#endif + $MAIN { +#if has_qlora_down + let base_groups = uniforms.batch_count * uniforms.dispatch_M * uniforms.num_N_tile; + // Normalization may pad the dispatch beyond the eight adapter groups per batch. + if (workgroup_idx >= base_groups + uniforms.batch_count * 8u) { + return; + } + if (workgroup_idx >= base_groups) { + qlora_down(workgroup_idx - base_groups, local_idx); + return; + } +#endif let batch = workgroup_idx / (uniforms.dispatch_M * uniforms.num_N_tile); if (batch >= uniforms.batch_count) { return; @@ -207,9 +266,9 @@ $MAIN { #if has_bias let bias_value = acc_element_t(bias[b_global + b_bias_offset]); // Downcast to the output element type only at the final store. - output.setByOffset(output_idx, output_element_t(output_value + bias_value)); + write_qlora_output(output_idx, output_element_t(output_value + bias_value)); #else - output.setByOffset(output_idx, output_element_t(output_value)); + write_qlora_output(output_idx, output_element_t(output_value)); #endif } } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc index 2a218752b9b45..bfe80bb044b71 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc @@ -3,6 +3,7 @@ #include "contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h" #include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" +#include "contrib_ops/webgpu/quantization/matmul_nbits_qlora.h" #include "core/providers/webgpu/shader_helper.h" namespace onnxruntime { @@ -19,7 +20,8 @@ Status DP4AMatMulQuantizeProgram::GenerateShaderCode(ShaderHelper& shader) const } Status DP4AMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { - const auto& a = shader.AddInput("input_a", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | ShaderUsage::UseValueTypeAlias); + const auto& a = shader.AddInput( + "input_a", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | ShaderUsage::UseValueTypeAlias); const auto& scales_a = shader.AddInput("scales_a", ShaderUsage::UseUniform); const auto& b = shader.AddInput("input_b", ShaderUsage::UseUniform); const auto& scales_b = shader.AddInput("scales_b", ShaderUsage::UseUniform); @@ -32,23 +34,23 @@ Status DP4AMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { if (has_weight_idx_indirect_) { shader.AddInput("weight_index_indirect", ShaderUsage::UseUniform); } + if (has_qlora_up_) { + shader.AddInput("lora_t"); + shader.AddInput("lora_b"); + shader.AddInput("lora_b_scales"); + ORT_RETURN_IF_ERROR(GenerateQLoRAUp(shader, acc_f32_, true)); + } const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseElementTypeAlias); - return WGSL_TEMPLATE_APPLY(shader, "quantization/dp4a_matmul.wgsl.template", - WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), - WGSL_TEMPLATE_PARAMETER(block_size, block_size_), - WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), - WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), - WGSL_TEMPLATE_PARAMETER(is_qualcomm, is_qualcomm_), - WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), - WGSL_TEMPLATE_PARAMETER(output_type_i32, true), - WGSL_TEMPLATE_PARAMETER(qc_wave128, SubgroupSize() == 128), - WGSL_TEMPLATE_VARIABLE(a, a), - WGSL_TEMPLATE_VARIABLE(b, b), - WGSL_TEMPLATE_VARIABLE(output, output), - WGSL_TEMPLATE_VARIABLE(scales_a, scales_a), - WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); + return WGSL_TEMPLATE_APPLY( + shader, "quantization/dp4a_matmul.wgsl.template", WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), + WGSL_TEMPLATE_PARAMETER(block_size, block_size_), WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), + WGSL_TEMPLATE_PARAMETER(has_qlora_up, has_qlora_up_), WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), + WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), + WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), WGSL_TEMPLATE_PARAMETER(is_qualcomm, is_qualcomm_), + WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), WGSL_TEMPLATE_PARAMETER(output_type_i32, true), + WGSL_TEMPLATE_PARAMETER(qc_wave128, SubgroupSize() == 128), + WGSL_TEMPLATE_VARIABLE(a, a), WGSL_TEMPLATE_VARIABLE(b, b), WGSL_TEMPLATE_VARIABLE(output, output), + WGSL_TEMPLATE_VARIABLE(scales_a, scales_a), WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); } // scale_A components = 1, b components = 4, output components = 1 @@ -66,47 +68,54 @@ Status DP4AMatMulNBitsKParallelProgram::GenerateShaderCode(ShaderHelper& shader) if (has_weight_idx_indirect_) { shader.AddInput("weight_index_indirect", ShaderUsage::UseUniform); } + if (has_qlora_up_) { + shader.AddInput("lora_t"); + shader.AddInput("lora_b"); + shader.AddInput("lora_b_scales"); + ORT_RETURN_IF_ERROR(GenerateQLoRAUp(shader, acc_f32_)); + } const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseElementTypeAlias); - ORT_ENFORCE(WorkgroupSizeX() % tile_size_k_vec_ == 0 && tile_size_k_vec_ % 4 == 0, "tile_size_k_vec_ must evenly divide workgroup size X and be divisible by 4"); + ORT_ENFORCE(WorkgroupSizeX() % tile_size_k_vec_ == 0 && tile_size_k_vec_ % 4 == 0, + "tile_size_k_vec_ must evenly divide workgroup size X and be divisible by 4"); const uint32_t sub_tile_count = WorkgroupSizeX() / tile_size_k_vec_; ORT_ENFORCE(tile_size_ % sub_tile_count == 0, "tile_size_ must be divisible by sub_tile_count"); - return WGSL_TEMPLATE_APPLY(shader, "quantization/dp4a_matmul_k_parallel.wgsl.template", - WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), - WGSL_TEMPLATE_PARAMETER(broadcast_a_row, broadcast_a_row_), - WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), - WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), - WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), - WGSL_TEMPLATE_PARAMETER(output_type_i32, true), - WGSL_TEMPLATE_PARAMETER(single_scale_weights, single_scale_weights_), - WGSL_TEMPLATE_PARAMETER(sub_tile_count, sub_tile_count), - WGSL_TEMPLATE_PARAMETER(tile_size, tile_size_), - WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec_), - WGSL_TEMPLATE_VARIABLE(a, a), - WGSL_TEMPLATE_VARIABLE(b, b), - WGSL_TEMPLATE_VARIABLE(output, output), - WGSL_TEMPLATE_VARIABLE(scales_a, scales_a), - WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); + if (has_qlora_down_) { + shader.AddInput("lora_a"); + shader.AddInput("lora_a_scales"); + shader.AddOutput("lora_down"); + } + + return WGSL_TEMPLATE_APPLY( + shader, "quantization/dp4a_matmul_k_parallel.wgsl.template", WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), + WGSL_TEMPLATE_PARAMETER(broadcast_a_row, broadcast_a_row_), WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), + WGSL_TEMPLATE_PARAMETER(has_qlora_down, has_qlora_down_), WGSL_TEMPLATE_PARAMETER(has_qlora_up, has_qlora_up_), + WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), + WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), + WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), + WGSL_TEMPLATE_PARAMETER(output_type_i32, true), + WGSL_TEMPLATE_PARAMETER(single_scale_weights, single_scale_weights_), + WGSL_TEMPLATE_PARAMETER(sub_tile_count, sub_tile_count), WGSL_TEMPLATE_PARAMETER(tile_size, tile_size_), + WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec_), WGSL_TEMPLATE_VARIABLE(a, a), + WGSL_TEMPLATE_VARIABLE(b, b), WGSL_TEMPLATE_VARIABLE(output, output), WGSL_TEMPLATE_VARIABLE(scales_a, scales_a), + WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); } -Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, - const Tensor* zero_points, const Tensor* bias, - uint32_t batch_count, - uint32_t M, - uint32_t dispatch_M, - uint32_t N, - uint32_t K, - uint32_t block_size, - uint32_t zero_blocks_per_col, - uint32_t min_M_for_tile_optimization, - uint32_t nbits, - onnxruntime::webgpu::ComputeContext& context, - Tensor* y, - const uint32_t weight_index, - const Tensor* weight_index_indirect, DP4AQuantizedInput* quantized_input) { +Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, const Tensor* zero_points, + const Tensor* bias, uint32_t batch_count, uint32_t M, uint32_t dispatch_M, uint32_t N, + uint32_t K, uint32_t block_size, uint32_t zero_blocks_per_col, + uint32_t min_M_for_tile_optimization, uint32_t nbits, + onnxruntime::webgpu::ComputeContext& context, Tensor* y, const uint32_t weight_index, + const Tensor* weight_index_indirect, DP4AQuantizedInput* quantized_input, + const Tensor* lora_t, const Tensor* lora_b, + const Tensor* lora_at, Tensor* lora_down, const Tensor* lora_b_scales, + const Tensor* lora_a_scales) { + const bool has_qlora_down = lora_at != nullptr; + ORT_RETURN_IF_NOT(!has_qlora_down || (lora_a_scales && lora_down && a->IsDataType() && nbits == 4 && + block_size == 32 && M == 1 && dispatch_M == 1 && + !weight_index && !weight_index_indirect && !zero_points && !bias && !lora_t), + "Unsupported joint DP4A QLoRA decode projection."); constexpr uint32_t kVec4Components = 4; constexpr uint32_t kVec2Components = 2; constexpr uint32_t kU32Components = 4; @@ -134,6 +143,7 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor const auto& a_quant = *input.values; const auto& a_scale = *input.scales; + const bool has_qlora_up = lora_t != nullptr; const bool has_zero_points = zero_points != nullptr; const bool has_bias = bias != nullptr; const bool has_weight_idx_indirect = weight_index_indirect != nullptr; @@ -157,10 +167,10 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor const bool acc_f32 = small_n || context.EnableMatmulFp32Accumulation(); DP4AMatMulNBitsKParallelProgram mul_program{ tile_size_k_vec, tile_size_n, nbits, has_zero_points, has_bias, has_weight_idx, - has_weight_idx_indirect, single_scale_weights, broadcast_a, acc_f32}; + has_weight_idx_indirect, single_scale_weights, broadcast_a, acc_f32, has_qlora_up, has_qlora_down}; uint32_t num_N_tile = (N + tile_size_n - 1) / tile_size_n; mul_program.SetWorkgroupSize(128); - mul_program.SetDispatchGroupSize(batch_count * dispatch_M * num_N_tile); + mul_program.SetDispatchGroupSize(batch_count * dispatch_M * num_N_tile + (has_qlora_down ? batch_count * 8 : 0)); mul_program .AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}, {&a_scale, ProgramTensorMetadataDependency::TypeAndRank, 1}, @@ -170,7 +180,7 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor weight_index, dispatch_M}) .AddOutput({y, ProgramTensorMetadataDependency::TypeAndRank, 1}) .CacheHint(nbits, tile_size_k_vec, tile_size_n, has_zero_points, single_scale_weights, has_bias, has_weight_idx, - has_weight_idx_indirect, broadcast_a, acc_f32); + has_weight_idx_indirect, broadcast_a, acc_f32, has_qlora_up, has_qlora_down); if (has_zero_points) { mul_program.AddInput( {zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); @@ -181,6 +191,16 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor if (has_weight_idx_indirect) { mul_program.AddInput({weight_index_indirect, ProgramTensorMetadataDependency::None}); } + if (has_qlora_up) { + mul_program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, 4}); + mul_program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, 16}); + mul_program.AddInput({lora_b_scales, ProgramTensorMetadataDependency::Type}); + } + if (has_qlora_down) { + mul_program.AddInput({lora_at, ProgramTensorMetadataDependency::Type, 16}); + mul_program.AddInput({lora_a_scales, ProgramTensorMetadataDependency::Type}); + mul_program.AddOutput({lora_down, ProgramTensorMetadataDependency::Type}); + } return context.RunProgram(mul_program); } @@ -197,16 +217,19 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor uint32_t num_N_tile = (N + tile_n - 1) / tile_n; bool is_qualcomm = context.AdapterInfo().vendor == std::string_view{"qualcomm"}; const bool acc_f32 = context.EnableMatmulFp32Accumulation(); - DP4AMatMulNBitsProgram mul_program{block_size, nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, is_qualcomm, acc_f32}; + DP4AMatMulNBitsProgram mul_program{ + block_size, nbits, has_zero_points, + has_bias, has_weight_idx, has_weight_idx_indirect, is_qualcomm, acc_f32, has_qlora_up}; mul_program.SetWorkgroupSize(256); if (qc_wave128) { mul_program.SetSubgroupSize(128); } mul_program.SetDispatchGroupSize(batch_count * num_M_tile * num_N_tile); - mul_program.AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}, - {&a_scale, ProgramTensorMetadataDependency::TypeAndRank, 1}, - {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast((nbits / 2) * kU32Components)}, - {scales, ProgramTensorMetadataDependency::TypeAndRank, 1}}) + mul_program + .AddInputs({{&a_quant, ProgramTensorMetadataDependency::TypeAndRank, static_cast(kVec4Components)}, + {&a_scale, ProgramTensorMetadataDependency::TypeAndRank, 1}, + {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast((nbits / 2) * kU32Components)}, + {scales, ProgramTensorMetadataDependency::TypeAndRank, 1}}) .AddUniformVariables({{static_cast(batch_count)}, {static_cast(M)}, {static_cast(N)}, @@ -218,9 +241,12 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor {zero_blocks_per_col}, {weight_index}}) .AddOutput({y, ProgramTensorMetadataDependency::TypeAndRank, reshaped_y_shape, static_cast(kVec4Components)}) - .CacheHint("Block" + std::to_string(block_size), nbits, has_zero_points, is_qualcomm, has_bias, has_weight_idx, has_weight_idx_indirect, acc_f32); + .CacheHint("Block" + std::to_string(block_size), nbits, has_zero_points, is_qualcomm, + has_bias, has_weight_idx, has_weight_idx_indirect, acc_f32, + has_qlora_up); if (has_zero_points) { - mul_program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); + mul_program.AddInput( + {zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } if (has_bias) { mul_program.AddInput({bias, ProgramTensorMetadataDependency::None}); @@ -228,26 +254,23 @@ Status ApplyDP4AMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Tensor if (has_weight_idx_indirect) { mul_program.AddInput({weight_index_indirect, ProgramTensorMetadataDependency::None}); } + if (has_qlora_up) { + mul_program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, 4}); + mul_program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, 16}); + mul_program.AddInput({lora_b_scales, ProgramTensorMetadataDependency::Type}); + } return context.RunProgram(mul_program); } -bool CanApplyDP4AMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& context, - uint64_t accuracy_level, - uint32_t block_size, - uint32_t N, - uint32_t K, - uint32_t components_k, - uint32_t M, - bool has_weight_idx_indirect, - const Tensor* y) { +bool CanApplyDP4AMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& context, uint64_t accuracy_level, + uint32_t block_size, uint32_t N, uint32_t K, uint32_t components_k, uint32_t M, + bool has_weight_idx_indirect, const Tensor* y) { // macOS - Avoid using dp4a on Metal, as it does not appear to have native dp4a support. // https://github.com/gpuweb/gpuweb/issues/2677#issuecomment-1713292226 // Use 'vendor' to check for metal; 'backend' is always WEBGPU when running under wasm. - bool use_dp4a = context.HasFeature(wgpu::FeatureName::Subgroups) && - context.AdapterInfo().vendor != std::string_view{"apple"}; - if (!(accuracy_level == 4 && block_size % 32 == 0 && - components_k == 4 && K % 128 == 0 && N % 16 == 0 && - use_dp4a)) { + bool use_dp4a = + context.HasFeature(wgpu::FeatureName::Subgroups) && context.AdapterInfo().vendor != std::string_view{"apple"}; + if (!(accuracy_level == 4 && block_size % 32 == 0 && components_k == 4 && K % 128 == 0 && N % 16 == 0 && use_dp4a)) { return false; } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h index e627fc7d60dcd..03d06cdf1b61a 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h @@ -30,28 +30,24 @@ class DP4AMatMulQuantizeProgram final : public Program { public: - DP4AMatMulNBitsProgram(uint32_t block_size, uint32_t nbits, - bool has_zero_points, bool has_bias, - bool has_weight_idx, bool has_weight_idx_indirect, bool is_qualcomm, - bool acc_f32) : Program{"DP4AMatMulNBits"}, - block_size_(block_size), - nbits_(nbits), - has_bias_(has_bias), - has_zero_points_(has_zero_points), - has_weight_idx_(has_weight_idx), - has_weight_idx_indirect_(has_weight_idx_indirect), - is_qualcomm_(is_qualcomm), - acc_f32_(acc_f32) {} + DP4AMatMulNBitsProgram(uint32_t block_size, uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, + bool has_weight_idx_indirect, bool is_qualcomm, bool acc_f32, bool has_qlora_up = false) + : Program{has_qlora_up ? "DP4AMatMulNBitsQLoRAUp" : "DP4AMatMulNBits"}, + block_size_(block_size), + nbits_(nbits), + has_bias_(has_bias), + has_zero_points_(has_zero_points), + has_weight_idx_(has_weight_idx), + has_weight_idx_indirect_(has_weight_idx_indirect), + is_qualcomm_(is_qualcomm), + acc_f32_(acc_f32), + has_qlora_up_(has_qlora_up) {} Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( - {"batch_count", ProgramUniformVariableDataType::Uint32}, - {"M", ProgramUniformVariableDataType::Uint32}, - {"N", ProgramUniformVariableDataType::Uint32}, - {"K", ProgramUniformVariableDataType::Uint32}, - {"K8", ProgramUniformVariableDataType::Uint32}, - {"K16", ProgramUniformVariableDataType::Uint32}, - {"num_M_tile", ProgramUniformVariableDataType::Uint32}, - {"num_N_tile", ProgramUniformVariableDataType::Uint32}, + {"batch_count", ProgramUniformVariableDataType::Uint32}, {"M", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}, {"K", ProgramUniformVariableDataType::Uint32}, + {"K8", ProgramUniformVariableDataType::Uint32}, {"K16", ProgramUniformVariableDataType::Uint32}, + {"num_M_tile", ProgramUniformVariableDataType::Uint32}, {"num_N_tile", ProgramUniformVariableDataType::Uint32}, {"zero_blocks_per_col", ProgramUniformVariableDataType::Uint32}, {"weight_idx", ProgramUniformVariableDataType::Uint32}); @@ -64,38 +60,38 @@ class DP4AMatMulNBitsProgram final : public Program { bool has_weight_idx_indirect_; bool is_qualcomm_; bool acc_f32_; + bool has_qlora_up_; }; class DP4AMatMulNBitsKParallelProgram final : public Program { public: - DP4AMatMulNBitsKParallelProgram(uint32_t tile_size_k_vec, uint32_t tile_size, uint32_t nbits, - bool has_zero_points, bool has_bias, - bool has_weight_idx, bool has_weight_idx_indirect, bool single_scale_weights, - bool broadcast_a_row = false, - bool acc_f32 = false) : Program{"DP4AMatMulNBitsKParallelProgram"}, - tile_size_k_vec_(tile_size_k_vec), - tile_size_(tile_size), - nbits_(nbits), - has_bias_(has_bias), - has_zero_points_(has_zero_points), - has_weight_idx_(has_weight_idx), - has_weight_idx_indirect_(has_weight_idx_indirect), - single_scale_weights_(single_scale_weights), - broadcast_a_row_(broadcast_a_row), - acc_f32_(acc_f32) {} + DP4AMatMulNBitsKParallelProgram(uint32_t tile_size_k_vec, uint32_t tile_size, uint32_t nbits, bool has_zero_points, + bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, + bool single_scale_weights, bool broadcast_a_row = false, bool acc_f32 = false, + bool has_qlora_up = false, bool has_qlora_down = false) + : Program{has_qlora_down ? "DP4AMatMulNBitsQLoRADownKParallel" + : has_qlora_up ? "DP4AMatMulNBitsQLoRAUpKParallel" + : "DP4AMatMulNBitsKParallelProgram"}, + tile_size_k_vec_(tile_size_k_vec), + tile_size_(tile_size), + nbits_(nbits), + has_bias_(has_bias), + has_zero_points_(has_zero_points), + has_weight_idx_(has_weight_idx), + has_weight_idx_indirect_(has_weight_idx_indirect), + single_scale_weights_(single_scale_weights), + broadcast_a_row_(broadcast_a_row), + acc_f32_(acc_f32), + has_qlora_up_(has_qlora_up), + has_qlora_down_(has_qlora_down) {} Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( - {"batch_count", ProgramUniformVariableDataType::Uint32}, - {"M", ProgramUniformVariableDataType::Uint32}, - {"N", ProgramUniformVariableDataType::Uint32}, - {"K", ProgramUniformVariableDataType::Uint32}, - {"K16", ProgramUniformVariableDataType::Uint32}, - {"K32", ProgramUniformVariableDataType::Uint32}, - {"block_size", ProgramUniformVariableDataType::Uint32}, - {"num_N_tile", ProgramUniformVariableDataType::Uint32}, + {"batch_count", ProgramUniformVariableDataType::Uint32}, {"M", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}, {"K", ProgramUniformVariableDataType::Uint32}, + {"K16", ProgramUniformVariableDataType::Uint32}, {"K32", ProgramUniformVariableDataType::Uint32}, + {"block_size", ProgramUniformVariableDataType::Uint32}, {"num_N_tile", ProgramUniformVariableDataType::Uint32}, {"zero_blocks_per_col", ProgramUniformVariableDataType::Uint32}, - {"weight_idx", ProgramUniformVariableDataType::Uint32}, - {"dispatch_M", ProgramUniformVariableDataType::Uint32}); + {"weight_idx", ProgramUniformVariableDataType::Uint32}, {"dispatch_M", ProgramUniformVariableDataType::Uint32}); private: uint32_t tile_size_k_vec_; @@ -108,24 +104,20 @@ class DP4AMatMulNBitsKParallelProgram final : public Program()) - .TypeConstraint("T3", DataTypeImpl::GetTensorType()) - .TypeConstraint("T4", DataTypeImpl::GetTensorType()), - MatMulNBits); +ONNX_OPERATOR_KERNEL_EX(MatMulNBits, kMSDomain, 1, kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T1", WebGpuSupportedFloatTypes()) + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) + .TypeConstraint("T3", DataTypeImpl::GetTensorType()) + .TypeConstraint("T4", DataTypeImpl::GetTensorType()), + MatMulNBits); Status MatMulNBitsWideTileProgram::GenerateShaderCode(ShaderHelper& shader) const { const auto& a = shader.AddInput("input_a", ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); @@ -45,6 +42,12 @@ Status MatMulNBitsWideTileProgram::GenerateShaderCode(ShaderHelper& shader) cons if (has_weight_idx_indirect_) { shader.AddInput("weight_index_indirect", ShaderUsage::UseUniform); } + if (has_qlora_up_) { + shader.AddInput("lora_t", ShaderUsage::UseElementTypeAlias); + shader.AddInput("lora_b", ShaderUsage::UseElementTypeAlias); + shader.AddInput("lora_b_scales"); + ORT_RETURN_IF_ERROR(GenerateQLoRAUp(shader, acc_f32_)); + } const auto& output = shader.AddOutput("output", ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); const uint32_t workgroup_size = WorkgroupSizeX() * WorkgroupSizeY() * WorkgroupSizeZ(); @@ -52,20 +55,15 @@ Status MatMulNBitsWideTileProgram::GenerateShaderCode(ShaderHelper& shader) cons ORT_ENFORCE(tile_m_ % (workgroup_size / 8) == 0, "tile_m must be a multiple of workgroup_size / 8."); ORT_ENFORCE(nbits_ == 4 || nbits_ == 8, "Only 4/8 bits are supported for webgpu matmulnbits."); - return WGSL_TEMPLATE_APPLY(shader, "quantization/matmul_nbits_wide_tile.wgsl.template", - WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), - WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), - WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), - WGSL_TEMPLATE_PARAMETER(nbits, nbits_), - WGSL_TEMPLATE_PARAMETER(subgroup_min_size, subgroup_min_size_), - WGSL_TEMPLATE_PARAMETER(tile_m, tile_m_), - WGSL_TEMPLATE_PARAMETER(tile_n, tile_n_), - WGSL_TEMPLATE_VARIABLE(a, a), - WGSL_TEMPLATE_VARIABLE(b, b), - WGSL_TEMPLATE_VARIABLE(output, output), - WGSL_TEMPLATE_VARIABLE(scales, scales)); + return WGSL_TEMPLATE_APPLY( + shader, "quantization/matmul_nbits_wide_tile.wgsl.template", WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), + WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), WGSL_TEMPLATE_PARAMETER(has_qlora_up, has_qlora_up_), + WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), + WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), + WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), WGSL_TEMPLATE_PARAMETER(nbits, nbits_), + WGSL_TEMPLATE_PARAMETER(subgroup_min_size, subgroup_min_size_), WGSL_TEMPLATE_PARAMETER(tile_m, tile_m_), + WGSL_TEMPLATE_PARAMETER(tile_n, tile_n_), WGSL_TEMPLATE_VARIABLE(a, a), WGSL_TEMPLATE_VARIABLE(b, b), + WGSL_TEMPLATE_VARIABLE(output, output), WGSL_TEMPLATE_VARIABLE(scales, scales)); } // Uses the same K-parallel reduction strategy as DP4AMatMulNBitsKParallelProgram. @@ -82,7 +80,21 @@ Status MatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { if (has_weight_idx_indirect_) { shader.AddInput("weight_index_indirect", ShaderUsage::UseUniform); } + if (has_qlora_up_) { + shader.AddInput("lora_t", ShaderUsage::UseElementTypeAlias); + shader.AddInput("lora_b", ShaderUsage::UseElementTypeAlias); + shader.AddInput("lora_b_scales"); + ORT_RETURN_IF_ERROR(GenerateQLoRAUp(shader, acc_f32_)); + } + if (has_qlora_down_) { + shader.AddInput("lora_at"); + shader.AddInput("lora_a_scales"); + } const auto& output = shader.AddOutput("output", ShaderUsage::UseElementTypeAlias); + if (has_qlora_down_) { + shader.AddOutput("lora_down", ShaderUsage::UseElementTypeAlias); + ORT_RETURN_IF_ERROR(GenerateQLoRADown(shader, a, acc_f32_, qlora_down_subgroups_)); + } const uint32_t components_a = a.NumComponents(); const uint32_t components_b = b.NumComponents() / 4; // b is stored as uint32 which includes 4 uint8. @@ -92,28 +104,20 @@ Status MatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader) const { const uint32_t a_length_per_tile = tile_size_k / components_a; uint32_t sub_tile_count = WorkgroupSizeX() / tile_size_k_vec; - return WGSL_TEMPLATE_APPLY(shader, "quantization/matmul_nbits.wgsl.template", - WGSL_TEMPLATE_PARAMETER(a_length_per_tile, a_length_per_tile), - WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), - WGSL_TEMPLATE_PARAMETER(broadcast_a_row, broadcast_a_row_), - WGSL_TEMPLATE_PARAMETER(component_a, components_a), - WGSL_TEMPLATE_PARAMETER(component_b, components_b), - WGSL_TEMPLATE_PARAMETER(elements_in_value_b, elements_in_value_b), - WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), - WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), - WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), - WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), - WGSL_TEMPLATE_PARAMETER(output_type_i32, false), - WGSL_TEMPLATE_PARAMETER(single_scale_weights, single_scale_weights_), - WGSL_TEMPLATE_PARAMETER(sub_tile_count, sub_tile_count), - WGSL_TEMPLATE_PARAMETER(tile_size, tile_size_), - WGSL_TEMPLATE_PARAMETER(tile_size_k, tile_size_k), - WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec), - WGSL_TEMPLATE_VARIABLE(a, a), - WGSL_TEMPLATE_VARIABLE(b, b), - WGSL_TEMPLATE_VARIABLE(output, output), - WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); + return WGSL_TEMPLATE_APPLY( + shader, "quantization/matmul_nbits.wgsl.template", WGSL_TEMPLATE_PARAMETER(a_length_per_tile, a_length_per_tile), + WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_), WGSL_TEMPLATE_PARAMETER(broadcast_a_row, broadcast_a_row_), + WGSL_TEMPLATE_PARAMETER(component_a, components_a), WGSL_TEMPLATE_PARAMETER(component_b, components_b), + WGSL_TEMPLATE_PARAMETER(elements_in_value_b, elements_in_value_b), WGSL_TEMPLATE_PARAMETER(has_bias, has_bias_), + WGSL_TEMPLATE_PARAMETER(has_qlora_up, has_qlora_up_), WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx_), + WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect_), + WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points_), WGSL_TEMPLATE_PARAMETER(n_bits, nbits_), + WGSL_TEMPLATE_PARAMETER(output_type_i32, false), + WGSL_TEMPLATE_PARAMETER(single_scale_weights, single_scale_weights_), + WGSL_TEMPLATE_PARAMETER(sub_tile_count, sub_tile_count), WGSL_TEMPLATE_PARAMETER(tile_size, tile_size_), + WGSL_TEMPLATE_PARAMETER(tile_size_k, tile_size_k), WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec), + WGSL_TEMPLATE_VARIABLE(a, a), WGSL_TEMPLATE_VARIABLE(b, b), WGSL_TEMPLATE_VARIABLE(output, output), + WGSL_TEMPLATE_VARIABLE(scales_b, scales_b)); } bool CanApplyWideTileMatMulNBits(uint32_t M, @@ -147,11 +151,13 @@ Status ApplyWideTileMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* uint32_t batch_count, uint32_t dispatch_M, uint32_t N, uint32_t K, uint32_t block_size, uint32_t zero_blocks_per_col, uint32_t nbits, onnxruntime::webgpu::ComputeContext& context, Tensor* y, - uint32_t weight_index, const Tensor* weight_index_indirect) { + uint32_t weight_index, const Tensor* weight_index_indirect, + const Tensor* lora_t, const Tensor* lora_b, const Tensor* lora_b_scales) { const bool has_zero_points = zero_points != nullptr; const bool has_bias = bias != nullptr; const bool has_weight_idx_indirect = weight_index_indirect != nullptr; const bool has_weight_idx = weight_index > 0 || has_weight_idx_indirect; + const bool has_qlora_up = lora_t != nullptr; const bool acc_f32 = context.EnableMatmulFp32Accumulation(); const uint32_t block_size_per_col = block_size == K * N ? K : block_size; const uint32_t n_blocks_per_col = CeilDiv(K, block_size_per_col); @@ -177,7 +183,7 @@ Status ApplyWideTileMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* : 0u; MatMulNBitsWideTileProgram program{has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, - tile_m, tile_n, nbits, subgroup_min_size, acc_f32}; + tile_m, tile_n, nbits, subgroup_min_size, acc_f32, has_qlora_up}; program.SetWorkgroupSize(workgroup_size); program.SetDispatchGroupSize(num_N_tile, num_M_tile, batch_count); @@ -216,7 +222,14 @@ Status ApplyWideTileMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* {num_N_tile}, {num_M_tile}, {weight_index}}); - program.CacheHint(nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, subgroup_min_size, acc_f32); + program.CacheHint(nbits, has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, subgroup_min_size, + acc_f32, has_qlora_up); + + if (has_qlora_up) { + program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, 4}); + program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, 16}); + program.AddInput({lora_b_scales, ProgramTensorMetadataDependency::Type}); + } return context.RunProgram(program); } @@ -253,6 +266,14 @@ Status MatMulNBits::ComputeInternal(onnxruntime::webgpu::ComputeContext& context return Status::OK(); } + if (empty_weights) { + // Record a fill dispatch so graph replay also overwrites the output. The + // MatMul shader uses [batch, row, column] coordinates, even for a vector A. + auto program = CreateMatMulFillBiasOrZeroBeforeSplitKProgram( + bias, y, false, 1.0f, 1, TensorShape({1, y->Shape().Size() / N_, N_})); + return context.RunProgram(program); + } + return ApplyMatMulNBits(a, b, scales, zero_points, bias, K_, N_, block_size_, accuracy_level_, bits_, context, y, 0); } @@ -273,7 +294,8 @@ Status MatMulNBits::ComputeInternal(onnxruntime::webgpu::ComputeContext& context * @param zero_points Pointer to the zero-point tensor for quantization; must be of type uint8 if provided. * weight_index > 0 is only supported when zero_points is nullptr. * @param bias Pointer to the bias tensor; optional. - * @param K_op The K dimension of the operation (number of columns in 'a' and rows in 'b' before quantization). + * @param K_op The K dimension of the operation (number of columns in 'a' and rows in 'b' before + * quantization). * @param N_op The N dimension of the operation (number of columns in 'b'). * @param block_size_op The block size used for quantization partitioning. * @param accuracy_level Accuracy level influencing the choice of optimized kernel. @@ -290,17 +312,13 @@ Status MatMulNBits::ComputeInternal(onnxruntime::webgpu::ComputeContext& context * - A wide tile program is used when block size, component count, and other criteria are met. * - Otherwise, a default matmul program is used. */ -Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, const Tensor* zero_points, const Tensor* bias, - int64_t K_op, - int64_t N_op, - int64_t block_size_op, - int64_t accuracy_level, - int64_t nbits, - onnxruntime::webgpu::ComputeContext& context, - Tensor* y, - const uint32_t weight_index, - const Tensor* weight_index_indirect, - uint32_t override_M, DP4AQuantizedInput* quantized_input) { +Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, const Tensor* zero_points, + const Tensor* bias, int64_t K_op, int64_t N_op, int64_t block_size_op, int64_t accuracy_level, + int64_t nbits, onnxruntime::webgpu::ComputeContext& context, Tensor* y, + const uint32_t weight_index, const Tensor* weight_index_indirect, uint32_t override_M, + DP4AQuantizedInput* quantized_input, + const Tensor* lora_t, const Tensor* lora_b, const Tensor* lora_a_transposed, + Tensor* lora_a_output, const Tensor* lora_b_scales, const Tensor* lora_a_scales) { TensorShape b_shape({N_op, K_op}); MatMulComputeHelper helper; ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b_shape, false, true)); @@ -316,12 +334,18 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, return context.RunProgram(program); } + const bool has_qlora_up = lora_t != nullptr; + const bool has_qlora_down = lora_a_transposed != nullptr; + ORT_RETURN_IF_NOT(!has_qlora_down || (lora_a_output != nullptr && lora_a_scales != nullptr && !has_qlora_up && + a->Shape().Size() == K_op && nbits == 4 && block_size_op == 32), + "Unsupported joint QLoRA projection."); const bool has_bias = bias != nullptr; const bool has_weight_idx_indirect = weight_index_indirect != nullptr; const bool has_weight_idx = weight_index > 0 || has_weight_idx_indirect; const bool has_zero_points = zero_points != nullptr; if (has_zero_points) { - ORT_ENFORCE(zero_points->DataType() == DataTypeImpl::GetType(), "Currently, only uint8 is supported for zero points, but got ", zero_points->DataType()); + ORT_ENFORCE(zero_points->DataType() == DataTypeImpl::GetType(), + "Currently, only uint8 is supported for zero points, but got ", zero_points->DataType()); } const uint32_t batch_count = onnxruntime::narrow(helper.OutputOffsets().size()); @@ -346,58 +370,82 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, // n_blocks_per_col up to the next multiple of (8/nbits) — the number of zero point // values per byte — so that the linear stride correctly skips byte-boundary padding. const uint32_t zp_elements_per_byte = 8 / static_cast(nbits); - uint32_t zero_blocks_per_col = (n_blocks_per_col + zp_elements_per_byte - 1) / zp_elements_per_byte * zp_elements_per_byte; + uint32_t zero_blocks_per_col = + (n_blocks_per_col + zp_elements_per_byte - 1) / zp_elements_per_byte * zp_elements_per_byte; // apple|intel - Experimental dawn support for subgroup matrix matmul. std::optional subgroup_matrix_config; - if (CanApplySubgroupMatrixMatMulNBits(context, - accuracy_level, - block_size, - batch_count, - N, - K, + if (!has_qlora_down && + CanApplySubgroupMatrixMatMulNBits(context, accuracy_level, block_size, batch_count, N, K, static_cast(nbits), y->DataType() == DataTypeImpl::GetType(), - subgroup_matrix_config, - M, - has_weight_idx_indirect)) { - return ApplySubgroupMatrixMatMulNBits(a, b, scales, zero_points, bias, M, N, K, - static_cast(nbits), zero_blocks_per_col, - *subgroup_matrix_config, context, y, weight_index, - weight_index_indirect); + subgroup_matrix_config, M, has_weight_idx_indirect)) { + return ApplySubgroupMatrixMatMulNBits(a, b, scales, zero_points, bias, M, N, K, static_cast(nbits), + zero_blocks_per_col, *subgroup_matrix_config, context, y, weight_index, + weight_index_indirect, nullptr, lora_t, lora_b, lora_b_scales); } - // On FP32 only GPUs and Qualcomm GPUs, integer math is faster than FP32 therefore always use DP4A independent of length of M. - // DP4A Q2 path now supports custom zero points via a 1024-entry LUT (4 zero-point sections × 256 byte values). - if (CanApplyDP4AMatrixMatMulNBits(context, accuracy_level, block_size, N, K, components_a, - M, has_weight_idx_indirect, y)) { - return ApplyDP4AMatrixMatMulNBits(a, b, scales, zero_points, bias, batch_count, M, dispatch_M, N, K, block_size, zero_blocks_per_col, kMinMForTileOptimization, static_cast(nbits), context, y, weight_index, weight_index_indirect, quantized_input); + // On FP32 only GPUs and Qualcomm GPUs, integer math is faster than FP32 therefore always use DP4A independent of + // length of M. DP4A Q2 path now supports custom zero points via a 1024-entry LUT (4 zero-point sections × 256 byte + // values). + if (!has_qlora_down && + CanApplyDP4AMatrixMatMulNBits(context, accuracy_level, block_size, N, K, components_a, M, + has_weight_idx_indirect, y)) { + return ApplyDP4AMatrixMatMulNBits(a, b, scales, zero_points, bias, batch_count, M, dispatch_M, N, K, block_size, + zero_blocks_per_col, kMinMForTileOptimization, static_cast(nbits), + context, y, weight_index, weight_index_indirect, quantized_input, lora_t, lora_b, + nullptr, nullptr, lora_b_scales); } // This program is optimized for FP32/FP16 Block32 prefill. - if (CanApplyWideTileMatMulNBits(M, K, block_size, nbits, has_weight_idx_indirect, - components_a, components_b)) { + if (!has_qlora_down && CanApplyWideTileMatMulNBits(M, K, block_size, nbits, has_weight_idx_indirect, + components_a, components_b)) { return ApplyWideTileMatMulNBits(a, b, scales, zero_points, bias, batch_count, dispatch_M, N, K, block_size, zero_blocks_per_col, static_cast(nbits), - context, y, weight_index, weight_index_indirect); + context, y, weight_index, weight_index_indirect, lora_t, lora_b, lora_b_scales); } const bool acc_f32 = context.EnableMatmulFp32Accumulation(); // Use tile_size_k_vec=32 by default for better K-dimension parallelism. // Intel devices use 16 as they have different subgroup/cache characteristics. - const uint32_t tile_size_k_vec = - (context.AdapterInfo().vendor == std::string_view{"intel"}) ? 16u : 32u; + const uint32_t tile_size_k_vec = (context.AdapterInfo().vendor == std::string_view{"intel"}) ? 16u : 32u; constexpr uint32_t workgroup_size = 128; constexpr uint32_t tile_size = 8; constexpr uint32_t kU32Components = 4; uint32_t components_b_with_u32 = components_b * kU32Components; uint32_t K_of_b = (n_blocks_per_col * blob_size) / components_b_with_u32; - MatMulNBitsProgram program{tile_size, static_cast(nbits), has_zero_points, has_bias, has_weight_idx, has_weight_idx_indirect, single_scale_weights, tile_size_k_vec, broadcast_a, acc_f32}; + // The decode reduction assigns two ranks to each 32-lane subgroup. Devices + // with a variable subgroup range must explicitly request that width. + const auto& adapter = context.AdapterInfo(); + const bool fixed_subgroup_32 = adapter.subgroupMinSize == 32 && adapter.subgroupMaxSize == 32; + const bool request_subgroup_32 = has_qlora_down && !fixed_subgroup_32 && + context.HasFeature(wgpu::FeatureName::SubgroupSizeControl) && + adapter.subgroupMinSize <= 32 && adapter.subgroupMaxSize >= 32; + const bool qlora_down_subgroups = has_qlora_down && context.HasFeature(wgpu::FeatureName::Subgroups) && + context.HasWGSLLanguageFeature(wgpu::WGSLLanguageFeatureName::SubgroupId) && + (fixed_subgroup_32 || request_subgroup_32); + MatMulNBitsProgram program{tile_size, + static_cast(nbits), + has_zero_points, + has_bias, + has_weight_idx, + has_weight_idx_indirect, + single_scale_weights, + tile_size_k_vec, + broadcast_a, + acc_f32, + has_qlora_up, + has_qlora_down, + qlora_down_subgroups}; + program.SetUseSubgroupId(qlora_down_subgroups); program.SetWorkgroupSize(workgroup_size); + if (qlora_down_subgroups && request_subgroup_32) { + program.SetSubgroupSize(32); + } uint32_t num_N_tile = (N + tile_size - 1) / tile_size; - program.SetDispatchGroupSize(num_N_tile, dispatch_M, batch_count); + program.SetDispatchGroupSize(num_N_tile + (has_qlora_down ? 4 : 0), dispatch_M, batch_count); program .AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, static_cast(components_a)}, {b, ProgramTensorMetadataDependency::TypeAndRank, static_cast(components_b_with_u32)}, @@ -415,7 +463,8 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, {batch_count}, {weight_index}, {dispatch_M}}) - .CacheHint(nbits, has_zero_points, single_scale_weights, has_bias, has_weight_idx, has_weight_idx_indirect, tile_size_k_vec, broadcast_a, acc_f32); + .CacheHint(nbits, has_zero_points, single_scale_weights, has_bias, has_weight_idx, has_weight_idx_indirect, + tile_size_k_vec, broadcast_a, acc_f32, has_qlora_up, has_qlora_down, qlora_down_subgroups); if (has_zero_points) { program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } @@ -425,6 +474,17 @@ Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, if (has_weight_idx_indirect) { program.AddInput({weight_index_indirect, ProgramTensorMetadataDependency::None}); } + if (has_qlora_up) { + program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, 4}); + program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, 16}); + program.AddInput({lora_b_scales, ProgramTensorMetadataDependency::Type}); + } + if (has_qlora_down) { + ORT_RETURN_IF_NOT(components_a == 4, "Joint QLoRA projection requires vec4 X."); + program.AddInput({lora_a_transposed, ProgramTensorMetadataDependency::Type, 16}); + program.AddInput({lora_a_scales, ProgramTensorMetadataDependency::Type}); + program.AddOutput({lora_a_output, ProgramTensorMetadataDependency::Type}); + } return context.RunProgram(program); } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h index e2205f37920dd..30b75f2b40ad6 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h @@ -16,20 +16,29 @@ struct DP4AQuantizedInput; class MatMulNBitsWideTileProgram final : public Program { public: - MatMulNBitsWideTileProgram(bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, uint32_t tile_m, uint32_t tile_n, uint32_t nbits, uint32_t subgroup_min_size, bool acc_f32) - : Program{"MatMulNBitsWideTile"}, has_zero_points_{has_zero_points}, has_bias_{has_bias}, has_weight_idx_{has_weight_idx}, has_weight_idx_indirect_{has_weight_idx_indirect}, tile_m_(tile_m), tile_n_(tile_n), nbits_(nbits), subgroup_min_size_(subgroup_min_size), acc_f32_(acc_f32) {} + MatMulNBitsWideTileProgram(bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, + uint32_t tile_m, uint32_t tile_n, uint32_t nbits, uint32_t subgroup_min_size, bool acc_f32, + bool has_qlora_up = false) + : Program{has_qlora_up ? "MatMulNBitsQLoRAUpWideTile" : "MatMulNBitsWideTile"}, + has_zero_points_{has_zero_points}, + has_bias_{has_bias}, + has_weight_idx_{has_weight_idx}, + has_weight_idx_indirect_{has_weight_idx_indirect}, + tile_m_(tile_m), + tile_n_(tile_n), + nbits_(nbits), + subgroup_min_size_(subgroup_min_size), + acc_f32_(acc_f32), + has_qlora_up_(has_qlora_up) {} Status GenerateShaderCode(ShaderHelper& sh) const override; - WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"Batch", ProgramUniformVariableDataType::Uint32}, - {"M", ProgramUniformVariableDataType::Uint32}, - {"N", ProgramUniformVariableDataType::Uint32}, - {"K_of_a", ProgramUniformVariableDataType::Uint32}, - {"K_of_b", ProgramUniformVariableDataType::Uint32}, - {"n_blocks_per_col", ProgramUniformVariableDataType::Uint32}, - {"zero_blocks_per_col", ProgramUniformVariableDataType::Uint32}, - {"num_N_tile", ProgramUniformVariableDataType::Uint32}, - {"num_M_tile", ProgramUniformVariableDataType::Uint32}, - {"weight_idx", ProgramUniformVariableDataType::Uint32}); + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( + {"Batch", ProgramUniformVariableDataType::Uint32}, {"M", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}, {"K_of_a", ProgramUniformVariableDataType::Uint32}, + {"K_of_b", ProgramUniformVariableDataType::Uint32}, {"n_blocks_per_col", ProgramUniformVariableDataType::Uint32}, + {"zero_blocks_per_col", ProgramUniformVariableDataType::Uint32}, + {"num_N_tile", ProgramUniformVariableDataType::Uint32}, {"num_M_tile", ProgramUniformVariableDataType::Uint32}, + {"weight_idx", ProgramUniformVariableDataType::Uint32}); private: bool has_zero_points_; @@ -41,26 +50,40 @@ class MatMulNBitsWideTileProgram final : public Program { public: - MatMulNBitsProgram(uint32_t tile_size, uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, bool single_scale_weights, uint32_t tile_size_k_vec = 16, bool broadcast_a_row = false, bool acc_f32 = false) - : Program{"MatMulNBits"}, tile_size_(tile_size), nbits_(nbits), has_zero_points_(has_zero_points), has_bias_(has_bias), has_weight_idx_{has_weight_idx}, has_weight_idx_indirect_{has_weight_idx_indirect}, single_scale_weights_(single_scale_weights), tile_size_k_vec_(tile_size_k_vec), broadcast_a_row_(broadcast_a_row), acc_f32_(acc_f32) {} + MatMulNBitsProgram(uint32_t tile_size, uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, + bool has_weight_idx_indirect, bool single_scale_weights, uint32_t tile_size_k_vec = 16, + bool broadcast_a_row = false, bool acc_f32 = false, bool has_qlora_up = false, + bool has_qlora_down = false, bool qlora_down_subgroups = false) + : Program{has_qlora_down ? (qlora_down_subgroups ? "MatMulNBitsQLoRADownSubgroup" : "MatMulNBitsQLoRADown") + : has_qlora_up ? "MatMulNBitsQLoRAUp" + : "MatMulNBits"}, + tile_size_(tile_size), + nbits_(nbits), + has_zero_points_(has_zero_points), + has_bias_(has_bias), + has_weight_idx_{has_weight_idx}, + has_weight_idx_indirect_{has_weight_idx_indirect}, + single_scale_weights_(single_scale_weights), + tile_size_k_vec_(tile_size_k_vec), + broadcast_a_row_(broadcast_a_row), + acc_f32_(acc_f32), + has_qlora_up_(has_qlora_up), + has_qlora_down_(has_qlora_down), + qlora_down_subgroups_(qlora_down_subgroups) {} Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( - {"M", ProgramUniformVariableDataType::Uint32}, - {"N", ProgramUniformVariableDataType::Uint32}, - {"K", ProgramUniformVariableDataType::Uint32}, - {"K_of_a", ProgramUniformVariableDataType::Uint32}, - {"K_of_b", ProgramUniformVariableDataType::Uint32}, - {"block_size", ProgramUniformVariableDataType::Uint32}, + {"M", ProgramUniformVariableDataType::Uint32}, {"N", ProgramUniformVariableDataType::Uint32}, + {"K", ProgramUniformVariableDataType::Uint32}, {"K_of_a", ProgramUniformVariableDataType::Uint32}, + {"K_of_b", ProgramUniformVariableDataType::Uint32}, {"block_size", ProgramUniformVariableDataType::Uint32}, {"blocks_per_col", ProgramUniformVariableDataType::Uint32}, {"zero_blocks_per_col", ProgramUniformVariableDataType::Uint32}, - {"num_N_tile", ProgramUniformVariableDataType::Uint32}, - {"batch_count", ProgramUniformVariableDataType::Uint32}, - {"weight_idx", ProgramUniformVariableDataType::Uint32}, - {"dispatch_M", ProgramUniformVariableDataType::Uint32}); + {"num_N_tile", ProgramUniformVariableDataType::Uint32}, {"batch_count", ProgramUniformVariableDataType::Uint32}, + {"weight_idx", ProgramUniformVariableDataType::Uint32}, {"dispatch_M", ProgramUniformVariableDataType::Uint32}); private: uint32_t tile_size_; @@ -73,6 +96,9 @@ class MatMulNBitsProgram final : public Program { uint32_t tile_size_k_vec_; bool broadcast_a_row_; bool acc_f32_; + bool has_qlora_up_; + bool has_qlora_down_; + bool qlora_down_subgroups_; }; class MatMulNBits final : public WebGpuKernel { @@ -105,12 +131,14 @@ bool CanApplyWideTileMatMulNBits(uint32_t M, uint32_t components_a = 0, uint32_t components_b = 0); -Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, const Tensor* zero_points, const Tensor* bias, - int64_t K_op, int64_t N_op, int64_t block_size_op, int64_t accuracy_level, int64_t bits_op, - onnxruntime::webgpu::ComputeContext& context, Tensor* y, const uint32_t weight_index = 0, - const Tensor* weight_index_indirect = nullptr, - uint32_t override_M = 0, - DP4AQuantizedInput* quantized_input = nullptr); +Status ApplyMatMulNBits(const Tensor* a, const Tensor* b, const Tensor* scales, const Tensor* zero_points, + const Tensor* bias, int64_t K_op, int64_t N_op, int64_t block_size_op, int64_t accuracy_level, + int64_t bits_op, onnxruntime::webgpu::ComputeContext& context, Tensor* y, + const uint32_t weight_index = 0, const Tensor* weight_index_indirect = nullptr, + uint32_t override_M = 0, DP4AQuantizedInput* quantized_input = nullptr, + const Tensor* lora_t = nullptr, const Tensor* lora_b = nullptr, + const Tensor* lora_a_transposed = nullptr, Tensor* lora_a_output = nullptr, + const Tensor* lora_b_scales = nullptr, const Tensor* lora_a_scales = nullptr); } // namespace webgpu } // namespace contrib diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.wgsl.template index b75bbbebcddc2..69e320d1acde5 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.wgsl.template @@ -13,6 +13,7 @@ #param tile_size_k_vec #param tile_size_k #param tile_size +#param has_qlora_up #param has_bias #param has_weight_idx #param has_weight_idx_indirect @@ -264,7 +265,12 @@ $MAIN { output_value += acc_element_t(bias[b_global + b_bias_offset]); #endif // Downcast to the output type only at the final write. +#if has_qlora_up + let delta = qlora_delta(batch * uniforms.M + a_global, b_global); + output.setByOffset(output_idx, output_element_t(output_value) + delta); +#else output.setByOffset(output_idx, output_element_t(output_value)); +#endif } } } // MAIN diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc new file mode 100644 index 0000000000000..acc6f6342e381 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc @@ -0,0 +1,206 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/webgpu/quantization/matmul_nbits_qlora.h" +#include "contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h" +#include "contrib_ops/webgpu/quantization/matmul_nbits.h" +#include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" +#include "contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h" +#include "contrib_ops/webgpu/webgpu_contrib_kernels.h" +#include "core/common/safeint.h" +#include "core/providers/cpu/math/matmul_helper.h" +#include "core/providers/webgpu/shader_helper.h" +#include "core/providers/webgpu/webgpu_supported_types.h" + +namespace onnxruntime::contrib::webgpu { + +Status GenerateQLoRADown(ShaderHelper& shader, const ShaderVariableHelper& x, bool acc_f32, + bool use_subgroups) { + return WGSL_TEMPLATE_APPLY(shader, "quantization/matmul_nbits_qlora_down.wgsl.template", + WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32), + WGSL_TEMPLATE_PARAMETER(use_subgroups, use_subgroups), + WGSL_TEMPLATE_VARIABLE(x, x)); +} + +Status GenerateQLoRAUp(ShaderHelper& shader, bool acc_f32, bool vector_output) { + return WGSL_TEMPLATE_APPLY(shader, "quantization/matmul_nbits_qlora_up.wgsl.template", + WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32), + WGSL_TEMPLATE_PARAMETER(vector_output, vector_output)); +} + +namespace { + +// Single-row execution stores the base projection separately. This shader computes +// the rank-32 up projection and adds it to that base output in one dispatch. +class QLoRAOutputProgram final : public Program { + public: + explicit QLoRAOutputProgram(bool acc_f32) : Program{"QLoRAOutputVec4"}, acc_f32_(acc_f32) {} + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(QLoRAOutputProgram); + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"size", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}, + {"M", ProgramUniformVariableDataType::Uint32}); + Status GenerateShaderCode(ShaderHelper& shader) const override { + shader.AddInput("base"); + shader.AddInput("lora_t"); + shader.AddInput("lora_b"); + shader.AddInput("lora_b_scales"); + shader.AddOutput("output", ShaderUsage::UseElementTypeAlias); + ORT_RETURN_IF_ERROR(GenerateQLoRAUp(shader, acc_f32_)); + shader.MainFunctionBody() + << "if (global_idx >= uniforms.size) { return; }\n" + << "output[global_idx] = base[global_idx] + qlora_delta(global_idx / uniforms.N, global_idx % uniforms.N);\n"; + return Status::OK(); + } + + private: + bool acc_f32_; +}; + +} // namespace + +Status ApplyQLoRAOutput(onnxruntime::webgpu::ComputeContext& context, const Tensor* base, const Tensor* t, + const Tensor* b, const Tensor* scales, Tensor* y) { + const auto size = narrow(y->Shape().Size()); + QLoRAOutputProgram program(context.EnableMatmulFp32Accumulation()); + program + .AddInputs({{base, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten}, + {t, ProgramTensorMetadataDependency::Type, 4}, + {b, ProgramTensorMetadataDependency::Type, 16}, + {scales, ProgramTensorMetadataDependency::Type}}) + .AddOutput({y, ProgramTensorMetadataDependency::Type, ProgramOutput::Flatten}) + .AddUniformVariables( + {{size}, {narrow(b->Shape()[0])}, {narrow(y->Shape()[y->Shape().NumDimensions() - 2])}}) + .SetWorkgroupSize(64) + .SetDispatchGroupSize(CeilDiv(size, 64u)) + .CacheHint(context.EnableMatmulFp32Accumulation()); + return context.RunProgram(program); +} + +class MatMulNBitsQLoRA final : public WebGpuKernel { + public: + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(MatMulNBitsQLoRA); + explicit MatMulNBitsQLoRA(const OpKernelInfo& info) + : WebGpuKernel(info), + k_(info.GetAttr("K")), + n_(info.GetAttr("N")), + accuracy_level_(info.GetAttrOrDefault("accuracy_level", 0)) { + ORT_ENFORCE(info.GetAttr("bits") == 4 && + info.GetAttr("block_size") == 32 && info.GetAttr("lora_bits") == 8 && + info.GetAttr("lora_block_size") == 32, + "WebGPU MatMulNBitsQLoRA supports only bits=4, block_size=32, lora_bits=8, and lora_block_size=32."); + ORT_ENFORCE(k_ > 0 && k_ % 32 == 0 && n_ > 0, + "WebGPU MatMulNBitsQLoRA requires positive K divisible by 32 and positive N."); + } + + Status ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const override { + const auto* x = context.Input(0); + const auto* w = context.Input(1); + const auto* s = context.Input(2); + const auto* aq = context.Input(3); + const auto* as = context.Input(4); + const auto* bq = context.Input(5); + const auto* bs = context.Input(6); + ORT_RETURN_IF_NOT(x->Shape().NumDimensions() >= 2, "WebGPU MatMulNBitsQLoRA requires X rank at least two."); + ORT_RETURN_IF_NOT(aq->Shape().NumDimensions() == 3 && bq->Shape().NumDimensions() == 3, + "Invalid quantized LoRA input ranks."); + ORT_RETURN_IF_NOT( + w->Shape() == TensorShape({n_, k_ / 32, 16}) && s->Shape().Size() == SafeInt(n_) * (k_ / 32), + "Invalid base weight/scales shape."); + // TODO: Support arbitrary adapter ranks with a rank uniform, per-block up scales, + // and vector-tail handling while preserving rank-32 performance. + const auto rank_a = aq->Shape()[0]; + const auto blocks_b = bq->Shape()[1]; + ORT_RETURN_IF_NOT((rank_a == 0 || rank_a == 32) && (blocks_b == 0 || blocks_b == 1), + "WebGPU MatMulNBitsQLoRA supports only adapter rank 32 (or empty adapters)."); + ORT_RETURN_IF_NOT(aq->Shape()[1] == k_ / 32 && + aq->Shape()[2] == 32 && bq->Shape()[0] == n_ && bq->Shape()[2] == 32, + "Invalid packed adapter dimensions."); + MatMulComputeHelper helper; + ORT_RETURN_IF_ERROR(helper.Compute(x->Shape(), TensorShape({n_, k_}), false, true)); + auto* y = context.Output(0, helper.OutputShape()); + if (y->Shape().Size() == 0) return Status::OK(); + if (rank_a == 0 || blocks_b == 0) { + return ApplyMatMulNBits(x, w, s, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, y); + } + ORT_RETURN_IF_NOT(as->Shape() == TensorShape({32, k_ / 32}) && bs->Shape() == TensorShape({n_, 1}), + "Invalid adapter scales shape."); + // With logical dequantized weights, Y = X * W + T * B_up, where T = X * A_down. + // T keeps X's leading dimensions and replaces K with the adapter rank (32). + auto t_shape = x->Shape().AsShapeVector(); + t_shape.back() = 32; + auto t = context.CreateGPUTensor(x->DataType(), TensorShape(t_shape)); + const auto m = narrow(x->Shape()[x->Shape().NumDimensions() - 2]); + // Prefill fuses T * B_up into the base shader's output stage. Single-row + // decode uses ApplyQLoRAOutput after computing the base projection and T. + const bool fuse_up = m > 1; + // FP32 may select DP4A even for M=1. The joint floating-point down + // shader must not replace that base projection's activation quantization. + const bool base_dp4a = CanApplyDP4AMatrixMatMulNBits(context, accuracy_level_, 32, narrow(n_), + narrow(k_), 4, m, false, y); + std::optional base_subgroup_matrix_config; + const auto subgroup_m = narrow(helper.M()); + const bool base_subgroup_matrix = + fuse_up && CanApplySubgroupMatrixMatMulNBits( + context, accuracy_level_, 32, narrow(helper.OutputOffsets().size()), + narrow(n_), narrow(k_), 4, x->IsDataType(), + base_subgroup_matrix_config, subgroup_m, false); + if (base_subgroup_matrix && base_subgroup_matrix_config->Is(8, 16, 16)) { + const auto down_config = SelectSubgroupMatrixConfig( + context, {{wgpu::SubgroupMatrixComponentType::F16, wgpu::SubgroupMatrixComponentType::F32, + 8, 16, 16, 32, true}}); + if (down_config.has_value()) { + return ApplySubgroupMatrixMatMulNBitsQLoRA( + x, w, s, aq, as, bq, bs, subgroup_m, narrow(n_), narrow(k_), + *base_subgroup_matrix_config, *down_config, context, &t, y); + } + } + if (!fuse_up && m == 1 && x->IsDataType() && base_dp4a && k_ >= 2048 && k_ <= 8192) { + auto base = context.CreateGPUTensor(y->DataType(), y->Shape()); + ORT_RETURN_IF_ERROR(ApplyDP4AMatrixMatMulNBits( + x, w, s, nullptr, nullptr, narrow(x->Shape().Size() / k_), 1, 1, narrow(n_), + narrow(k_), 32, narrow(k_ / 32), kMinMForTileOptimization, 4, context, &base, 0, nullptr, + nullptr, nullptr, nullptr, aq, &t, nullptr, as)); + return ApplyQLoRAOutput(context, &base, &t, bq, bs, y); + } + // For one input vector, append four rank-reduction workgroups to the generic + // base dispatch. GenerateQLoRADown supplies their shader body; DP4A has its own. + const bool joint_down = !fuse_up && x->Shape().Size() == k_ && k_ >= 2048 && k_ <= 8192 && !base_dp4a; + if (joint_down) { + auto base = context.CreateGPUTensor(y->DataType(), y->Shape()); + ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, w, s, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, &base, 0, + nullptr, 0, nullptr, nullptr, nullptr, aq, &t, nullptr, as)); + return ApplyQLoRAOutput(context, &base, &t, bq, bs, y); + } + // FP32 small-M prefill can use the existing DP4A down kernel and share + // quantized X with the base projection when accuracy_level permits it. + const auto down_accuracy = + fuse_up && + (m >= kMinMForTileOptimization || (x->IsDataType() && base_dp4a)) + ? accuracy_level_ + : 0; + DP4AQuantizedInput quantized_x; + ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, aq, as, nullptr, nullptr, k_, 32, 32, down_accuracy, 8, context, &t, 0, + nullptr, 0, &quantized_x)); + if (fuse_up) { + return ApplyMatMulNBits(x, w, s, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, y, 0, nullptr, 0, + &quantized_x, &t, bq, nullptr, nullptr, bs); + } + auto base = context.CreateGPUTensor(y->DataType(), y->Shape()); + ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, w, s, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, &base, 0, + nullptr, 0, &quantized_x)); + return ApplyQLoRAOutput(context, &base, &t, bq, bs, y); + } + + private: + int64_t k_; + int64_t n_; + int64_t accuracy_level_; +}; + +ONNX_OPERATOR_KERNEL_EX(MatMulNBitsQLoRA, kMSDomain, 1, kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", WebGpuSupportedFloatTypes()) + .TypeConstraint("TQ", DataTypeImpl::GetTensorType()), + MatMulNBitsQLoRA); + +} // namespace onnxruntime::contrib::webgpu diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h new file mode 100644 index 0000000000000..667b37c3b52b1 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h @@ -0,0 +1,34 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/webgpu/program.h" + +namespace onnxruntime::webgpu { +class ShaderVariableHelper; +} + +namespace onnxruntime::contrib::webgpu { + +// These C++ generation helpers append reusable WGSL fragments to the caller's +// shader. They do not create standalone GPU programs or dispatch GPU work. +// +// Append the rank-32 down-projection body to the generic single-vector decode +// shader. The base projection and T = X * A_down run in the same dispatch using +// separate workgroups. X uses vec4 loads; uniforms.num_N_tile marks the four +// appended adapter workgroups. +Status GenerateQLoRADown(onnxruntime::webgpu::ShaderHelper& shader, const onnxruntime::webgpu::ShaderVariableHelper& x, + bool acc_f32, bool use_subgroups = false); + +// Append WGSL helper functions for up-weight dequantization and delta = T * B_up. +// Generic and DP4A fused prefill shaders and the separate decode output shader +// reuse these functions. Subgroup-matrix shaders reuse only the dequantization +// helper and implement the multiplication with their own matrix instructions. +Status GenerateQLoRAUp(onnxruntime::webgpu::ShaderHelper& shader, bool acc_f32, bool vector_output = false); + +// Dispatch the decode output shader to compute Y = base + T * B_up. +Status ApplyQLoRAOutput(onnxruntime::webgpu::ComputeContext& context, const Tensor* base, const Tensor* t, + const Tensor* b, const Tensor* scales, Tensor* y); + +} // namespace onnxruntime::contrib::webgpu diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora_down.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora_down.wgsl.template new file mode 100644 index 0000000000000..f1854a3e2495f --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora_down.wgsl.template @@ -0,0 +1,119 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#param acc_f32 +#param use_subgroups + +#use .getByOffset + +// Joint single-vector decode: the base projection and T = X * A_down share +// one dispatch. lora_at stores rank-major packed INT8 weights; lora_down is +// the rank-32 intermediate T consumed by the later up projection. +#if acc_f32 +alias qlora_acc_t = f32; +#else +alias qlora_acc_t = output_element_t; +#endif + +#if !use_subgroups +var qlora_x_tile: array, 128>; +var qlora_partials: array, 8>; +#endif + +$MAIN { + // Four appended workgroups compute eight ranks each. Normalizing the dispatch + // to multiple dimensions may add padding workgroups, which must not write T. + let qlora_base_groups = uniforms.num_N_tile; + if (workgroup_idx >= qlora_base_groups + 4u) { return; } + if (workgroup_idx >= qlora_base_groups) { + let rank_base = (workgroup_idx - qlora_base_groups) * 8u; +#if use_subgroups + // Each of the four 32-lane subgroups computes two ranks. Subgroup membership + // need not match consecutive local invocation indices, so use subgroup IDs. + let lane = sg_id; + let rank_local = subgroup_index; + var partial0 = qlora_acc_t(0); + var partial1 = qlora_acc_t(0); + for (var k = lane * 16u; k < uniforms.K; k += 512u) { + var k4 = k / 4u; + // Unrolling keeps the four X vectors in registers for reuse by both ranks + // and fixes the dot-product grouping that determines FP16 rounding. + let x0 = vec4(x.getByOffset(k4)); + k4 += 1u; + let x1 = vec4(x.getByOffset(k4)); + k4 += 1u; + let x2 = vec4(x.getByOffset(k4)); + k4 += 1u; + let x3 = vec4(x.getByOffset(k4)); + let rank0 = rank_base + rank_local + 0u; + let packed0 = lora_at[rank0 * (uniforms.K / 16u) + k / 16u]; + let scale0 = lora_a_scales[rank0 * (uniforms.K / 32u) + k / 32u]; + var sum0 = qlora_acc_t(0); + sum0 += dot(x0, vec4( + (vec4(unpack4xU8(packed0[0])) - vec4(128)) * scale0)); + sum0 += dot(x1, vec4( + (vec4(unpack4xU8(packed0[1])) - vec4(128)) * scale0)); + sum0 += dot(x2, vec4( + (vec4(unpack4xU8(packed0[2])) - vec4(128)) * scale0)); + sum0 += dot(x3, vec4( + (vec4(unpack4xU8(packed0[3])) - vec4(128)) * scale0)); + partial0 += sum0; + let rank1 = rank_base + rank_local + 4u; + let packed1 = lora_at[rank1 * (uniforms.K / 16u) + k / 16u]; + let scale1 = lora_a_scales[rank1 * (uniforms.K / 32u) + k / 32u]; + var sum1 = qlora_acc_t(0); + sum1 += dot(x0, vec4( + (vec4(unpack4xU8(packed1[0])) - vec4(128)) * scale1)); + sum1 += dot(x1, vec4( + (vec4(unpack4xU8(packed1[1])) - vec4(128)) * scale1)); + sum1 += dot(x2, vec4( + (vec4(unpack4xU8(packed1[2])) - vec4(128)) * scale1)); + sum1 += dot(x3, vec4( + (vec4(unpack4xU8(packed1[3])) - vec4(128)) * scale1)); + partial1 += sum1; + } + // FP32 avoids half-precision rounding at each step of the subgroup tree. + let result0 = subgroupAdd(f32(partial0)); + let result1 = subgroupAdd(f32(partial1)); + if (lane == 0u) { + lora_down[rank_base + rank_local] = output_element_t(result0); + lora_down[rank_base + rank_local + 4u] = output_element_t(result1); + } +#else + let lane = local_idx % 32u; + let rank_local = local_idx / 32u; + // Match the generic INT8 MatMulNBits tile, dot grouping and reduction order + // so the FP16 rounding behavior is preserved in this shared-memory fallback. + for (var start = 0u; start < uniforms.K; start += 512u) { + let k4 = start / 4u + local_idx; + var xv = vec4(0); + if (k4 < uniforms.K / 4u) { xv = x.getByOffset(k4); } + qlora_x_tile[local_idx] = xv; + workgroupBarrier(); + for (var rank_offset = 0u; rank_offset < 8u; rank_offset += 4u) { + let rank = rank_base + rank_offset + rank_local; + let k = start + lane * 16u; + if (k < uniforms.K) { + let packed = lora_at[rank * (uniforms.K / 16u) + k / 16u]; + let scale = lora_a_scales[rank * (uniforms.K / 32u) + k / 32u]; + var sum = qlora_acc_t(0); + for (var i = 0u; i < 4u; i++) { + let bv = (vec4(unpack4xU8(packed[i])) - vec4(128)) * scale; + sum += dot(vec4(qlora_x_tile[lane * 4u + i]), vec4(bv)); + } + qlora_partials[rank_offset + rank_local][lane] += sum; + } + } + workgroupBarrier(); + } + if (local_idx < 8u) { + var value = qlora_acc_t(0); + for (var i = 0u; i < 32u; i++) { value += qlora_partials[local_idx][i]; } + lora_down[rank_base + local_idx] = output_element_t(value); + } +#endif + // Adapter workgroups stop here; the remaining workgroups execute the base + // MatMulNBits body appended after this template. + return; + } +} diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora_up.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora_up.wgsl.template new file mode 100644 index 0000000000000..e44cd7957490f --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora_up.wgsl.template @@ -0,0 +1,72 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#param acc_f32 +#param vector_output + +// T = X * A_down is stored in lora_t as eight vec4 values per row (rank 32). +// These helpers compute the adapter contribution delta = T * B_up. The caller +// adds delta to the base projection, either in its output stage or a separate +// decode dispatch. Subgroup-matrix templates reuse only qlora_weight and perform +// the up multiplication with their own matrix instructions. +#if acc_f32 +alias qlora_acc_t = f32; +#else +alias qlora_acc_t = output_element_t; +#endif + +fn qlora_weight(col: u32, r4: u32) -> vec4 { + // Each output column has 32 packed INT8 weights in two vec4 values. + // r4 selects four consecutive ranks; 128 is the implicit zero point. + let packed = lora_b[col * 2u + r4 / 4u][r4 % 4u]; + return (vec4(unpack4xU8(packed)) - vec4(128)) * lora_b_scales[col]; +} + +fn qlora_delta(row: u32, col: u32) -> output_element_t { + var delta = qlora_acc_t(0); + // Match the unfused INT8 up projection's reduction order: wide tiles at M>=4, + // generic tiles otherwise. Regrouping these sums can change FP16 rounding. + if (uniforms.M >= 4u) { + for (var r4 = 0u; r4 < 8u; r4 += 2u) { + let a0 = lora_t[row * 8u + r4]; + let a1 = lora_t[row * 8u + r4 + 1u]; + delta += dot(vec4(a0), vec4(qlora_weight(col, r4))) + + dot(vec4(a1), vec4(qlora_weight(col, r4 + 1u))); + } + } else { + for (var block = 0u; block < 2u; block++) { + var sum = qlora_acc_t(0); + for (var i = 0u; i < 4u; i++) { + let r4 = block * 4u + i; + sum += dot(vec4(lora_t[row * 8u + r4]), vec4(qlora_weight(col, r4))); + } + delta += sum; + } + // The remaining generic-kernel partials are zero, including their sign. + delta += qlora_acc_t(0); + } + return output_element_t(delta); +} + +#if vector_output +// Tiled DP4A writes four adjacent output columns at once. Reuse each T vector +// across those columns while keeping the scalar wide-tile dot grouping. +fn qlora_delta4(row: u32, col: u32) -> vec4 { + var delta = vec4(0); + for (var r4 = 0u; r4 < 8u; r4 += 2u) { + let a0 = vec4(lora_t[row * 8u + r4]); + let a1 = vec4(lora_t[row * 8u + r4 + 1u]); + delta += vec4( + dot(a0, vec4(qlora_weight(col, r4))), + dot(a0, vec4(qlora_weight(col + 1u, r4))), + dot(a0, vec4(qlora_weight(col + 2u, r4))), + dot(a0, vec4(qlora_weight(col + 3u, r4)))) + + vec4( + dot(a1, vec4(qlora_weight(col, r4 + 1u))), + dot(a1, vec4(qlora_weight(col + 1u, r4 + 1u))), + dot(a1, vec4(qlora_weight(col + 2u, r4 + 1u))), + dot(a1, vec4(qlora_weight(col + 3u, r4 + 1u)))); + } + return vec4(delta); +} +#endif diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template index a433f8064f5e5..bc24ccccdac33 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template @@ -2,6 +2,7 @@ // Licensed under the MIT License. #param has_zero_points +#param has_qlora_up #param has_bias #param has_weight_idx #param has_weight_idx_indirect @@ -113,7 +114,12 @@ fn load_scale(row : u32, block_idx : u32) -> output_element_t { fn write_output(batch : u32, row : u32, col : u32, value : output_element_t) { if (batch < uniforms.Batch && row < uniforms.M && col < uniforms.N) { let offset = batch * uniforms.M * uniforms.N + row * uniforms.N + col; +#if has_qlora_up + let delta = qlora_delta(batch * uniforms.M + row, col); + output.setByOffset(offset, value + delta); +#else output.setByOffset(offset, value); +#endif } } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc index 2396dbefd7fa8..7c0f867221459 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc @@ -2,7 +2,10 @@ // Licensed under the MIT License. #include "contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h" + +#include #include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" +#include "contrib_ops/webgpu/quantization/matmul_nbits_qlora.h" #include "core/providers/webgpu/math/subgroup_matrix_config.h" #include "core/providers/webgpu/webgpu_utils.h" @@ -63,10 +66,11 @@ Status GenerateShaderCode16x16x16(ShaderHelper& shader, const ShaderVariableHelper& output, uint32_t nbits, const SubgroupMatrixConfig& config, uint32_t tile_size_m, uint32_t tile_size_n, uint32_t tile_size_k, bool has_zero_points, bool has_bias, - bool has_weight_idx, bool has_weight_idx_indirect) { + bool has_weight_idx, bool has_weight_idx_indirect, bool has_qlora_up) { // Use the configurable tiled shader for the 16x16x16 configs. return WGSL_TEMPLATE_APPLY(shader, "quantization/subgroup_matrix_matmul_nbits_16x16x16_128.wgsl.template", WGSL_TEMPLATE_PARAMETER(has_bias, has_bias), + WGSL_TEMPLATE_PARAMETER(has_qlora_up, has_qlora_up), WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx), WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect), WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points), @@ -87,10 +91,13 @@ Status GenerateShaderCode8x16x16(ShaderHelper& shader, const ShaderVariableHelper& b, const ShaderVariableHelper& scales_b, const ShaderVariableHelper& output, - uint32_t nbits, const SubgroupMatrixConfig& config, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, + uint32_t nbits, const SubgroupMatrixConfig& config, bool has_zero_points, bool has_bias, + bool has_weight_idx, bool has_weight_idx_indirect, bool has_qlora_up, uint32_t tile_m, uint32_t tile_n, bool has_tail_buffer) { return WGSL_TEMPLATE_APPLY(shader, "quantization/subgroup_matrix_matmul_nbits_8x16x16.wgsl.template", + WGSL_TEMPLATE_PARAMETER(acc_f32, config.resultComponentType == wgpu::SubgroupMatrixComponentType::F32), WGSL_TEMPLATE_PARAMETER(has_bias, has_bias), + WGSL_TEMPLATE_PARAMETER(has_qlora_up, has_qlora_up), WGSL_TEMPLATE_PARAMETER(has_tail_buffer, has_tail_buffer), WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx), WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect), @@ -109,9 +116,11 @@ Status GenerateShaderCode8x16x16(ShaderHelper& shader, Status GenerateShaderCode8x8x8(ShaderHelper& shader, const ShaderVariableHelper& a, const ShaderVariableHelper& b, const ShaderVariableHelper& scales_b, - const ShaderVariableHelper& output, uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect) { + const ShaderVariableHelper& output, uint32_t nbits, bool has_zero_points, bool has_bias, + bool has_weight_idx, bool has_weight_idx_indirect, bool has_qlora_up) { return WGSL_TEMPLATE_APPLY(shader, "quantization/subgroup_matrix_matmul_nbits_8x8x8.wgsl.template", WGSL_TEMPLATE_PARAMETER(has_bias, has_bias), + WGSL_TEMPLATE_PARAMETER(has_qlora_up, has_qlora_up), WGSL_TEMPLATE_PARAMETER(has_weight_idx, has_weight_idx), WGSL_TEMPLATE_PARAMETER(has_weight_idx_indirect, has_weight_idx_indirect), WGSL_TEMPLATE_PARAMETER(has_zero_points, has_zero_points), @@ -136,20 +145,29 @@ Status SubgroupMatrixMatMulNBitsProgram::GenerateShaderCode(ShaderHelper& shader if (has_weight_idx_indirect_) { shader.AddInput("weight_index_indirect", ShaderUsage::UseUniform); } + if (has_qlora_up_) { + shader.AddInput("lora_t", ShaderUsage::UseElementTypeAlias); + shader.AddInput("lora_b", ShaderUsage::UseElementTypeAlias); + shader.AddInput("lora_b_scales"); + ORT_RETURN_IF_ERROR(GenerateQLoRAUp(shader, false)); + } const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseElementTypeAlias); if (has_tail_buffer_) { shader.AddOutput("tail_output", ShaderUsage::None); } if (config_.Is(8, 8, 8)) { - return GenerateShaderCode8x8x8(shader, a, b, scales_b, output, nbits_, has_zero_points_, has_bias_, has_weight_idx_, has_weight_idx_indirect_); + return GenerateShaderCode8x8x8(shader, a, b, scales_b, output, nbits_, has_zero_points_, has_bias_, + has_weight_idx_, has_weight_idx_indirect_, has_qlora_up_); } else if (config_.Is(8, 16, 16)) { - return GenerateShaderCode8x16x16(shader, b, scales_b, output, nbits_, config_, has_zero_points_, has_bias_, has_weight_idx_, has_weight_idx_indirect_, tile_m_, tile_n_, has_tail_buffer_); + return GenerateShaderCode8x16x16(shader, b, scales_b, output, nbits_, config_, has_zero_points_, has_bias_, + has_weight_idx_, has_weight_idx_indirect_, has_qlora_up_, tile_m_, tile_n_, + has_tail_buffer_); } else if (config_.Is(16, 16, 16)) { return GenerateShaderCode16x16x16(shader, b, scales_b, output, nbits_, config_, tile_m_, tile_n_, tile_k_, has_zero_points_, has_bias_, has_weight_idx_, - has_weight_idx_indirect_); + has_weight_idx_indirect_, has_qlora_up_); } else { return Status(onnxruntime::common::ONNXRUNTIME, onnxruntime::common::NOT_IMPLEMENTED, "Unsupported subgroup matrix config dimensions."); @@ -266,13 +284,35 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te Tensor* y, const uint32_t weight_index, const Tensor* weight_index_indirect, - const Tensor* prepacked_a) { + const Tensor* prepacked_a, + const Tensor* lora_t, + const Tensor* lora_b, + const Tensor* lora_b_scales, + uint32_t narrow_tile_m) { const bool has_zero_points = zero_points != nullptr; const bool has_bias = bias != nullptr; const bool has_weight_idx_indirect = weight_index_indirect != nullptr; const bool has_weight_idx = weight_index > 0 || has_weight_idx_indirect; + const bool has_qlora_up = lora_t != nullptr; + ORT_RETURN_IF_NOT((lora_b != nullptr) == has_qlora_up && (lora_b_scales != nullptr) == has_qlora_up, + "QLoRA up inputs must be provided together."); auto [tile_m, tile_n, tile_size_k, workgroup_size] = SelectMatMulNBitsTiling(config, M, N, K, has_bias); + if (has_qlora_up && config.Is(16, 16, 16) && tile_m == 256) { + // The fused rank-32 tile needs another 16 KiB of shared memory. Cap BK so the + // wave64 base and up projections together stay within 32 KiB. + tile_size_k = std::min(tile_size_k, 64u); + } + if (narrow_tile_m != 0) { + ORT_ENFORCE(config.Is(8, 16, 16) && nbits == 8 && N == 32 && + narrow_tile_m % config.M == 0 && !has_bias && !has_qlora_up); + tile_m = narrow_tile_m; + tile_n = 32; + workgroup_size = tile_m / config.M * config.subgroupSize; + } + if (config.Is(8, 16, 16) && config.resultComponentType == wgpu::SubgroupMatrixComponentType::F32) { + ORT_ENFORCE(narrow_tile_m != 0, "FP32 subgroup accumulation is implemented for the narrow QLoRA down tile."); + } ORT_ENFORCE(workgroup_size <= context.DeviceLimits().maxComputeWorkgroupSizeX && workgroup_size <= context.DeviceLimits().maxComputeInvocationsPerWorkgroup, "Subgroup matrix MatMulNBits workgroup size exceeds device limits: ", workgroup_size); @@ -297,10 +337,12 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te // padded buffer instead of `y` (which is too short by this point); the crop-copy // dispatch below moves the valid rows into `y`. This keeps the no-bias write-out // free of any bounds-checked workgroup-scratch store. - const bool has_tail_buffer = !has_bias && config.Is(8, 16, 16) && (M % tile_m != 0); + const bool has_tail_buffer = !has_bias && config.Is(8, 16, 16) && + config.resultComponentType != wgpu::SubgroupMatrixComponentType::F32 && + (M % tile_m != 0); SubgroupMatrixMatMulNBitsProgram matmul_program{nbits, config, tile_m, tile_n, tile_size_k, has_zero_points, has_bias, has_weight_idx, - has_weight_idx_indirect, has_tail_buffer}; + has_weight_idx_indirect, has_qlora_up, has_tail_buffer}; matmul_program.SetWorkgroupSize(workgroup_size); // Pin kernels running on variable-size adapters to the subgroup size they were written for. @@ -323,7 +365,7 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te static_cast(config.resultComponentType), config.M, config.N, config.K, config.subgroupSize, tile_m, tile_n, tile_size_k, has_zero_points, has_bias, has_weight_idx, - has_weight_idx_indirect, has_tail_buffer); + has_weight_idx_indirect, has_qlora_up, has_tail_buffer); if (has_zero_points) { matmul_program.AddInput({zero_points, ProgramTensorMetadataDependency::None, {(zero_points->Shape().Size() + 3) / 4}, 4}); } @@ -333,6 +375,11 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te if (has_weight_idx_indirect) { matmul_program.AddInput({weight_index_indirect, ProgramTensorMetadataDependency::None}); } + if (has_qlora_up) { + matmul_program.AddInput({lora_t, ProgramTensorMetadataDependency::Type, 4}); + matmul_program.AddInput({lora_b, ProgramTensorMetadataDependency::Type, 16}); + matmul_program.AddInput({lora_b_scales, ProgramTensorMetadataDependency::Type}); + } Tensor tail_buffer; if (has_tail_buffer) { @@ -366,6 +413,34 @@ Status ApplySubgroupMatrixMatMulNBits(const Tensor* a, const Tensor* b, const Te return Status::OK(); } +Status ApplySubgroupMatrixMatMulNBitsQLoRA( + const Tensor* x, const Tensor* w, const Tensor* scales, + const Tensor* lora_a, const Tensor* lora_a_scales, + const Tensor* lora_b, const Tensor* lora_b_scales, + uint32_t M, uint32_t N, uint32_t K, const SubgroupMatrixConfig& config, + const SubgroupMatrixConfig& down_config, + onnxruntime::webgpu::ComputeContext& context, Tensor* t, Tensor* y) { + ORT_ENFORCE(config.Is(8, 16, 16) && config.needsPrepack); + ORT_ENFORCE(down_config.Is(8, 16, 16) && down_config.needsPrepack && + down_config.componentType == config.componentType && + down_config.resultComponentType == wgpu::SubgroupMatrixComponentType::F32 && + down_config.subgroupSize == config.subgroupSize); + constexpr uint32_t kDownTileM = 32; + const auto base_tiling = SelectMatMulNBitsTiling(config, M, N, K, false); + const auto padded_M = std::max(CeilDiv(M, base_tiling.m) * base_tiling.m, + CeilDiv(M, kDownTileM) * kDownTileM); + // The layout depends on the subgroup matrix dimensions, not the workgroup tile. + // Both projections can read one packed X, padded for the larger workgroup tile. + Tensor packed_x; + ORT_RETURN_IF_ERROR(PrepackMatMulNBitsInput(x, M, K, padded_M, config, context, packed_x)); + ORT_RETURN_IF_ERROR(ApplySubgroupMatrixMatMulNBits( + x, lora_a, lora_a_scales, nullptr, nullptr, M, 32, K, 8, K / 32, down_config, context, t, + 0, nullptr, &packed_x, nullptr, nullptr, nullptr, kDownTileM)); + return ApplySubgroupMatrixMatMulNBits( + x, w, scales, nullptr, nullptr, M, N, K, 4, CeilDiv(K / 32, 2u) * 2u, config, context, y, + 0, nullptr, &packed_x, t, lora_b, lora_b_scales); +} + bool CanApplySubgroupMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& context, uint64_t accuracy_level, uint32_t block_size, @@ -402,12 +477,9 @@ bool CanApplySubgroupMatrixMatMulNBits(onnxruntime::webgpu::ComputeContext& cont return false; } - // Every fp16 preference below has resultComponentType == F16, and the kernels declare - // subgroup_matrix_result to match, so the accumulation inside - // subgroupMatrixMultiplyAccumulate is f16 and there is no variant of this kernel that can - // honour an f32 accumulator request. Decline the path instead of ignoring the option: the - // caller then falls through to a kernel that does honour it. Only fp16 outputs are affected; - // the fp32 config accumulates in f32 already. + // The FP16 configurations below accumulate the base projection in FP16. The separately + // selected FP32-accumulating QLoRA down tile does not support a general base projection. + // Fall back when FP32 accumulation is requested instead of ignoring the option. if (is_fp16 && context.EnableMatmulFp32Accumulation()) { return false; } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h index a0e037152e4a2..c7f6a29b1ea4d 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h @@ -23,8 +23,10 @@ class SubgroupMatrixMatMulNBitsProgram final : public Program buf_b: array; // Small staging buffer for edge-tile output (one 16x16 tile per subgroup) var coopmat_stage: array; +#if has_qlora_up +const kQLoRARank: u32 = 32; +var lora_a: array; +#endif + #if n_bits == 4 // Dequantize one 16-element K chunk from one B row into buf_b. fn dequant_b_to_buf(n_base: u32, k_idx: u32, n_idx: u32, k_chunk: u32) { @@ -255,6 +261,84 @@ $MAIN { workgroupBarrier(); } +#if has_qlora_up + for (var load_unit = local_idx; load_unit < BM * (kQLoRARank / 4u); + load_unit += workgroup_size_x) { + let tile_m = load_unit / (kQLoRARank / 4u); + let r4 = load_unit % (kQLoRARank / 4u); + let global_m = global_base_m + tile_m; + let offset = tile_m * kQLoRARank + r4 * 4u; + if (global_m < uniforms.M) { + let values = lora_t[global_m * (kQLoRARank / 4u) + r4]; + for (var i = 0u; i < 4u; i++) { + lora_a[offset + i] = f16(values[i]); + } + } else { + for (var i = 0u; i < 4u; i++) { + lora_a[offset + i] = f16(0.0); + } + } + } + for (var load_unit = local_idx; load_unit < BN * (kQLoRARank / 4u); + load_unit += workgroup_size_x) { + let tile_n = load_unit / (kQLoRARank / 4u); + let r4 = load_unit % (kQLoRARank / 4u); + let global_n = global_base_n + tile_n; + let offset = tile_n * SHMEM_STRIDE + r4 * 4u; + if (global_n < uniforms.N) { + let values = qlora_weight(global_n, r4); + for (var i = 0u; i < 4u; i++) { + buf_b[offset + i] = f16(values[i]); + } + } else { + for (var i = 0u; i < 4u; i++) { + buf_b[offset + i] = f16(0.0); + } + } + } + workgroupBarrier(); + + for (var k = 0u; k < kQLoRARank; k += kSgMatK) { + let b_base = (warp_c * WN) * SHMEM_STRIDE + k; + var b0: subgroup_matrix_right = subgroupMatrixLoad, col_major>( + &buf_b, b_base, SHMEM_STRIDE); + var b1: subgroup_matrix_right = subgroupMatrixLoad, col_major>( + &buf_b, b_base + 16u * SHMEM_STRIDE, SHMEM_STRIDE); + var b2: subgroup_matrix_right = subgroupMatrixLoad, col_major>( + &buf_b, b_base + 32u * SHMEM_STRIDE, SHMEM_STRIDE); + var b3: subgroup_matrix_right = subgroupMatrixLoad, col_major>( + &buf_b, b_base + 48u * SHMEM_STRIDE, SHMEM_STRIDE); + + let a_base = (warp_r * WM) * kQLoRARank + k; + var a0: subgroup_matrix_left = subgroupMatrixLoad, row_major>( + &lora_a, a_base, kQLoRARank); + var a1: subgroup_matrix_left = subgroupMatrixLoad, row_major>( + &lora_a, a_base + 16u * kQLoRARank, kQLoRARank); + var a2: subgroup_matrix_left = subgroupMatrixLoad, row_major>( + &lora_a, a_base + 32u * kQLoRARank, kQLoRARank); + var a3: subgroup_matrix_left = subgroupMatrixLoad, row_major>( + &lora_a, a_base + 48u * kQLoRARank, kQLoRARank); + + c00 = subgroupMatrixMultiplyAccumulate(a0, b0, c00); + c01 = subgroupMatrixMultiplyAccumulate(a0, b1, c01); + c02 = subgroupMatrixMultiplyAccumulate(a0, b2, c02); + c03 = subgroupMatrixMultiplyAccumulate(a0, b3, c03); + c10 = subgroupMatrixMultiplyAccumulate(a1, b0, c10); + c11 = subgroupMatrixMultiplyAccumulate(a1, b1, c11); + c12 = subgroupMatrixMultiplyAccumulate(a1, b2, c12); + c13 = subgroupMatrixMultiplyAccumulate(a1, b3, c13); + c20 = subgroupMatrixMultiplyAccumulate(a2, b0, c20); + c21 = subgroupMatrixMultiplyAccumulate(a2, b1, c21); + c22 = subgroupMatrixMultiplyAccumulate(a2, b2, c22); + c23 = subgroupMatrixMultiplyAccumulate(a2, b3, c23); + c30 = subgroupMatrixMultiplyAccumulate(a3, b0, c30); + c31 = subgroupMatrixMultiplyAccumulate(a3, b1, c31); + c32 = subgroupMatrixMultiplyAccumulate(a3, b2, c32); + c33 = subgroupMatrixMultiplyAccumulate(a3, b3, c33); + } + workgroupBarrier(); +#endif + // Write out results let dr = global_base_m + warp_r * WM; let dc = global_base_n + warp_c * WN; diff --git a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x16x16.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x16x16.wgsl.template index 9f0a2c4e3b7db..2cf337add7fd4 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x16x16.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x16x16.wgsl.template @@ -13,8 +13,10 @@ // K dimension iterated in kChunkSizeK(32) blocks #param n_bits +#param acc_f32 #param has_zero_points #param has_bias +#param has_qlora_up #param has_tail_buffer #param has_weight_idx #param has_weight_idx_indirect @@ -55,17 +57,27 @@ const kSgMatCountK : u32 = kChunkSizeK / kSgMatDimK; // Number of subgroup_mat alias SgMatLeft = subgroup_matrix_left; alias SgMatRight = subgroup_matrix_right; +#if acc_f32 +alias SgMatResult = subgroup_matrix_result; +// The narrow down projection accumulates over long K in FP32, then rounds T to FP16. +var result_f32 : array; +#else alias SgMatResult = subgroup_matrix_result; +#endif // Shared local memory for dequantized B weights [kTileN, kChunkSizeK], column-major var tile_b : array; #if has_bias -// Scratch space for bounds-checked write-out with bias addition, [kTileM, kTileN], -// row-major. The no-bias path never needs this - see `has_tail_buffer` below. +// Scratch space for bounds-checked write-out with bias addition, [kTileM, kTileN], row-major. var scratch : array; #endif +#if has_qlora_up +const kQLoRARank : u32 = 32; +var lora_a : array; +#endif + // --- Data Loading Functions --- // scale_offset is resolved before the K loop, not recomputed for each load. @@ -176,8 +188,74 @@ $MAIN { workgroupBarrier(); } +#if has_qlora_up + // Reuse the subgroup-matrix accumulators for the rank-32 adapter-up product. + // Loading the row-major adapter activation into workgroup memory avoids the + // standalone subgroup layout and adapter-up dispatches. + for (var load_unit = local_idx; load_unit < kTileM * (kQLoRARank / 4u); + load_unit += workgroup_size_x) { + let row_in_tile = load_unit / (kQLoRARank / 4u); + let r4 = load_unit % (kQLoRARank / 4u); + let global_m = global_base_m + row_in_tile; + if (global_m < uniforms.M) { + let values = lora_t[global_m * (kQLoRARank / 4u) + r4]; + let offset = row_in_tile * kQLoRARank + r4 * 4u; + for (var i = 0u; i < 4u; i++) { + lora_a[offset + i] = f16(values[i]); + } + } + } + for (var load_unit = local_idx; load_unit < kTileN * (kQLoRARank / 4u); + load_unit += workgroup_size_x) { + let col_in_tile = load_unit / (kQLoRARank / 4u); + let r4 = load_unit % (kQLoRARank / 4u); + let global_n = global_base_n + col_in_tile; + if (global_n < uniforms.N) { + let values = qlora_weight(global_n, r4); + let offset = col_in_tile * kQLoRARank + r4 * 4u; + for (var i = 0u; i < 4u; i++) { + tile_b[offset + i] = f16(values[i]); + } + } + } + workgroupBarrier(); + + let lora_a_base = sg_idx * kSgMatDimM * kQLoRARank; + for (var k_idx = 0u; k_idx < kQLoRARank; k_idx += kSgMatDimK) { + var sg_mat_a: SgMatLeft = subgroupMatrixLoad(&lora_a, + lora_a_base + k_idx, + kQLoRARank); + var sg_mat_offset_b = k_idx; + for (var sg_mat_n_idx = 0u; sg_mat_n_idx < kSgMatCountN; sg_mat_n_idx++) { + var sg_mat_b: SgMatRight = subgroupMatrixLoad(&tile_b, + sg_mat_offset_b, + kQLoRARank); + sg_mat_offset_b += kSgMatStrideN; + sg_mat_acc[sg_mat_n_idx] = + subgroupMatrixMultiplyAccumulate(sg_mat_a, sg_mat_b, sg_mat_acc[sg_mat_n_idx]); + } + } + workgroupBarrier(); +#endif + // Write out -#if has_bias +#if acc_f32 + let scratch_store_base = sg_idx * kSgMatDimM * kTileN; + for (var sg_mat_n_idx = 0u; sg_mat_n_idx < kSgMatCountN; sg_mat_n_idx++) { + subgroupMatrixStore(&result_f32, + scratch_store_base + sg_mat_n_idx * kSgMatDimN, + sg_mat_acc[sg_mat_n_idx], + kTileN); + } + workgroupBarrier(); + for (var i = local_idx; i < kTileM * kTileN; i += workgroup_size_x) { + let row = global_base_m + i / kTileN; + let col = global_base_n + i % kTileN; + if (row < uniforms.M && col < uniforms.N) { + output.setByOffset(row * uniforms.N + col, output_element_t(result_f32[i])); + } + } +#elif has_bias let scratch_store_base = sg_idx * kSgMatDimM * kTileN; for (var sg_mat_n_idx = 0u; sg_mat_n_idx < kSgMatCountN; sg_mat_n_idx++) { subgroupMatrixStore(&scratch, @@ -205,7 +283,7 @@ $MAIN { #endif for (var i : u32 = 0; i < ElemsPerThreadRow; i++) { if (out_global_n_base + i < uniforms.N) { - let val = f16(scratch[scratch_base + i]) + bias[bias_offset + out_global_n_base + i]; + let val = output_element_t(scratch[scratch_base + i]) + bias[bias_offset + out_global_n_base + i]; output.setByOffset(out_base + i, val); } } diff --git a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x8x8.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x8x8.wgsl.template index f28360af79c45..6f8927903124e 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x8x8.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x8x8.wgsl.template @@ -7,6 +7,7 @@ #param n_bits #param has_zero_points #param has_bias +#param has_qlora_up #param has_weight_idx #param has_weight_idx_indirect @@ -211,6 +212,53 @@ $MAIN { workgroupBarrier(); } +#if has_qlora_up + for (var idx = local_idx; idx < tile_rows * tile_k; idx += workgroup_size_x) { + let row = idx / tile_k; + let col = idx % tile_k; + let global_row = a_global_base + row; + if (global_row < uniforms.M) { + tile_A[idx] = compute_precision(lora_t[global_row * 8u + col / 4u][col % 4u]); + } + } + for (var idx = local_idx; idx < tile_cols * tile_k; idx += workgroup_size_x) { + let row = idx / tile_k; + let col = idx % tile_k; + let global_col = b_global_base + row; + if (global_col < uniforms.N) { + tile_B[idx] = compute_precision(qlora_weight(global_col, col / 4u)[col % 4u]); + } + } + workgroupBarrier(); + + for (var step: u32 = 0; step < tile_k; step += 8u) { + let matrix_a_offset = subtile_idy * subtile_rows * tile_k + step; + var matA0: subgroup_matrix_left = subgroupMatrixLoad, row_major>( + &tile_A, matrix_a_offset, tile_k); + var matA1: subgroup_matrix_left = subgroupMatrixLoad, row_major>( + &tile_A, matrix_a_offset + 8u * tile_k, tile_k); + let matrix_b_offset = subtile_idx * subtile_cols * tile_k + step; + var matB0: subgroup_matrix_right = subgroupMatrixLoad, col_major>( + &tile_B, matrix_b_offset, tile_k); + var matB1: subgroup_matrix_right = subgroupMatrixLoad, col_major>( + &tile_B, matrix_b_offset + 8u * tile_k, tile_k); + var matB2: subgroup_matrix_right = subgroupMatrixLoad, col_major>( + &tile_B, matrix_b_offset + 16u * tile_k, tile_k); + var matB3: subgroup_matrix_right = subgroupMatrixLoad, col_major>( + &tile_B, matrix_b_offset + 24u * tile_k, tile_k); + + matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00); + matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01); + matC02 = subgroupMatrixMultiplyAccumulate(matA0, matB2, matC02); + matC03 = subgroupMatrixMultiplyAccumulate(matA0, matB3, matC03); + matC10 = subgroupMatrixMultiplyAccumulate(matA1, matB0, matC10); + matC11 = subgroupMatrixMultiplyAccumulate(matA1, matB1, matC11); + matC12 = subgroupMatrixMultiplyAccumulate(matA1, matB2, matC12); + matC13 = subgroupMatrixMultiplyAccumulate(matA1, matB3, matC13); + } + workgroupBarrier(); +#endif + // Write out // Write out top block subgroupMatrixStore(&scratch[subtile_id][0], 0, matC00, 8); diff --git a/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc b/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc index d1cf54166e658..34ed6bb9e300e 100644 --- a/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc @@ -69,6 +69,7 @@ static const BuildKernelCreateInfoFn build_kernel_create_info_function_table[] = BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index fe19adb003419..a53cdca974176 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -4155,6 +4155,106 @@ Empty packed weights are currently supported by the CPU and WebGPU execution pro } }); + ONNX_CONTRIB_OPERATOR_SCHEMA(MatMulNBitsQLoRA) + .SetDomain(kMSDomain) + .SinceVersion(1) + .SetDoc( + "Computes MatMulNBits(X,W,scales) + MatMulNBits(MatMulNBits(X,lora_A,lora_A_scales),lora_B,lora_B_scales). " + "Each projection follows MatMulNBits block quantization and bit packing, with implicit zero points " + "2^(bits-1) for the base and 2^(lora_bits-1) for the adapters. No bias or explicit zero points are applied. " + "The adapter rank R may be any positive integer and need not be a multiple of lora_block_size. " + "Partial final blocks are padded; only the logical K or R elements participate in each projection. " + "An empty rank dimension in either packed adapter disables the adapter, while scale defaults are retained. " + "Intermediate projections and block scales use the activation type before the final addition. " + "Adapters remain replaceable. Execution providers may support a subset of the defined types, " + "quantization settings, and ranks.") + .Attr("K", "Positive input feature dimension of the base and down-projection weight matrices.", + AttributeProto::INT) + .Attr("N", "Positive output feature dimension of the base and up-projection weight matrices.", + AttributeProto::INT) + .Attr("bits", "Base weight bit width, following MatMulNBits (2, 4, or 8).", AttributeProto::INT, int64_t{4}) + .Attr("block_size", "Base quantization block size along K, following MatMulNBits: a power of two, at least 16.", + AttributeProto::INT, int64_t{32}) + .Attr("accuracy_level", + "Accuracy policy following MatMulNBits: 0 (unset), 1 (fp32), 2 (fp16), 3 (bf16), or 4 (int8). " + "Applies to the base and eligible adapter down projections; the up projection uses level 0.", + AttributeProto::INT, int64_t{0}) + .Attr("lora_bits", "Adapter weight bit width, following MatMulNBits (2, 4, or 8).", + AttributeProto::INT, int64_t{8}) + .Attr("lora_block_size", + "Adapter quantization block size along K for lora_A and R for lora_B: a power of two, at least 16.", + AttributeProto::INT, int64_t{32}) + .Input(0, "X", "Input of rank at least one, with last dimension K, following MatMulNBits.", "T") + .Input(1, "W", "Packed base weights [N,ceil(K/block_size),block_size*bits/8].", "TQ") + .Input(2, "scales", "Base block scales [N,ceil(K/block_size)]; may be flattened.", "T") + .Input(3, "lora_A", + "Packed down weights [R_A,ceil(K/lora_block_size),lora_block_size*lora_bits/8], R_A=R or 0.", "TQ") + .Input(4, "lora_A_scales", + "Down block scales [R,ceil(K/lora_block_size)]; retained when the adapter is empty.", "T") + .Input(5, "lora_B", + "Packed up weights [N,R_blocks,lora_block_size*lora_bits/8], R_blocks=ceil(R/lora_block_size) or 0.", "TQ") + .Input(6, "lora_B_scales", + "Up block scales [N,ceil(R/lora_block_size)]; retained when the adapter is empty.", "T") + .Output(0, "Y", "Input shape with the last dimension replaced by N.", "T") + .TypeConstraint("T", {"tensor(float16)", "tensor(float)"}, + "Activations and block scales share the same floating-point type (FP16 or FP32).") + .TypeConstraint("TQ", {"tensor(uint8)"}, "Packed quantized weights.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + const auto k = getAttribute(ctx, "K", -1); + const auto n = getAttribute(ctx, "N", -1); + const auto lora_bits = getAttribute(ctx, "lora_bits", 8); + const auto lora_block_size = getAttribute(ctx, "lora_block_size", 32); + const auto valid_quantization = [](int64_t bits, int64_t block_size) { + return (bits == 2 || bits == 4 || bits == 8) && block_size >= 16 && + (block_size & (block_size - 1)) == 0; + }; + if (k <= 0 || n <= 0 || + !valid_quantization(getAttribute(ctx, "bits", 4), getAttribute(ctx, "block_size", 32)) || + !valid_quantization(lora_bits, lora_block_size)) { + fail_shape_inference( + "MatMulNBitsQLoRA requires positive K and N, bit widths 2, 4, or 8, " + "and power-of-two block sizes of at least 16."); + } + propagateElemTypeFromInputToOutput(ctx, 0, 0); + MatmulWithQuantWeightShapeInference(ctx, k, n, true); + int64_t rank = -1; + for (size_t index : {size_t{3}, size_t{4}, size_t{5}, size_t{6}}) { + if (!hasInputShape(ctx, index)) continue; + const auto& shape = getInputShape(ctx, index); + const bool packed = index == 3 || index == 5; + if (shape.dim_size() != (packed ? 3 : 2)) fail_shape_inference("Invalid quantized adapter rank."); + if ((index == 3 || index == 4) && shape.dim(0).has_dim_value()) { + const auto value = shape.dim(0).dim_value(); + if (index == 3 && value == 0) continue; // Empty packed defaults retain the scale tensor's rank. + if (value <= 0 || (rank > 0 && value != rank)) { + fail_shape_inference("Inconsistent adapter rank between lora_A and lora_A_scales."); + } + rank = value; + } + } + const auto k_blocks = k / lora_block_size + (k % lora_block_size != 0); + const auto r_blocks = rank < 0 ? -1 : rank / lora_block_size + (rank % lora_block_size != 0); + const auto blob_size = (lora_block_size / 8) * lora_bits; + const int64_t dims[4][3] = { + {rank, k_blocks, blob_size}, + {rank, k_blocks, 0}, + {n, r_blocks, blob_size}, + {n, r_blocks, 0}, + }; + for (size_t index : {size_t{3}, size_t{4}, size_t{5}, size_t{6}}) { + if (!hasInputShape(ctx, index)) continue; + const auto& shape = getInputShape(ctx, index); + for (int axis = 0; axis < shape.dim_size(); ++axis) { + const auto& dim = shape.dim(axis); + const bool empty_axis = (index == 3 && axis == 0) || (index == 5 && axis == 1); + if (dim.has_dim_value() && dims[index - 3][axis] >= 0 && dim.dim_value() != dims[index - 3][axis] && + !(empty_axis && dim.dim_value() == 0)) { + fail_shape_inference("Invalid quantized adapter dimensions."); + } + } + } + }); + static const char* MatMulNBitsMlp_ver1_doc = R"DOC( MatMulNBitsMlp fuses two MatMulNBits projections that share the same input and computes diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index cc2c49dbc49e3..64b456a5c9f77 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -60,6 +60,7 @@ #include "core/optimizer/group_query_attention_pre_norm_fusion.h" #include "core/optimizer/matmul_nbits_qkv_fusion.h" #include "core/optimizer/matmul_nbits_mlp_fusion.h" +#include "core/optimizer/matmul_nbits_qlora_fusion.h" #include "core/optimizer/matmul_bn_fusion.h" #include "core/optimizer/matmul_integer_to_float.h" #include "core/optimizer/matmul_scale_fusion.h" @@ -465,6 +466,7 @@ InlinedVector> GenerateTransformers( transformers.emplace_back(std::make_unique( InlinedHashSet{onnxruntime::kCudaExecutionProvider, onnxruntime::kWebGpuExecutionProvider})); + bool has_matmul_nbits_qlora_kernel = false; bool has_matmul_nbits_mlp_kernel = false; bool has_matmul_nbits_qkv_kernel = false; if (execution_providers != nullptr) { @@ -482,11 +484,15 @@ InlinedVector> GenerateTransformers( nullptr) .IsOK(); }; + has_matmul_nbits_qlora_kernel = has_webgpu_kernel("MatMulNBitsQLoRA"); has_matmul_nbits_mlp_kernel = has_webgpu_kernel("MatMulNBitsMlp"); has_matmul_nbits_qkv_kernel = has_webgpu_kernel("MatMulNBitsQkv"); } } } + if (has_matmul_nbits_qlora_kernel) { + transformers.emplace_back(std::make_unique()); + } if (has_matmul_nbits_mlp_kernel) { transformers.emplace_back(std::make_unique( InlinedHashSet{onnxruntime::kWebGpuExecutionProvider})); diff --git a/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc b/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc new file mode 100644 index 0000000000000..971e13221d6fe --- /dev/null +++ b/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.cc @@ -0,0 +1,144 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/optimizer/matmul_nbits_qlora_fusion.h" + +#include "core/graph/graph_utils.h" +#include "core/graph/node_attr_utils.h" +#include "core/optimizer/matmul_nbits_fusion_utils.h" +#include "core/optimizer/utils.h" + +namespace onnxruntime { +namespace { + +const Node* Producer(const Graph& graph, const Node& node, int input) { + const auto* edge = graph_utils::GetInputEdge(node, input); + return edge ? graph.GetNode(edge->GetNode().Index()) : nullptr; +} + +bool Private(const Graph& graph, const Node& node) { + return !graph.NodeProducesGraphOutput(node) && optimizer_utils::CheckOutputEdges(graph, node, 1); +} + +bool Shape(const NodeArg& arg, std::initializer_list expected, int empty_axis = -1) { + const auto* shape = arg.Shape(); + if (!shape || shape->dim_size() != static_cast(expected.size())) return false; + int index = 0; + for (auto value : expected) { + const auto& dim = shape->dim(index); + if (index++ == empty_axis && (!dim.has_dim_value() || dim.dim_value() == 0)) continue; + if (!dim.has_dim_value() || dim.dim_value() != value) return false; + } + return true; +} + +bool AdapterShape(const Graph& graph, const NodeArg& arg, std::initializer_list expected, int empty_axis) { + // K/N attributes establish the logical rank even when an overridable default + // is empty. Only the axis used by the QLoRA empty-adapter convention may vary. + if (!Shape(arg, expected, empty_axis)) return false; + const ONNX_NAMESPACE::TensorProto* init = nullptr; + if (graph.GetInitializedTensor(arg.Name(), init)) { + if (init->dims_size() != static_cast(expected.size())) return false; + int index = 0; + for (auto value : expected) { + const auto dim = init->dims(index); + if (dim != value && !(index == empty_axis && dim == 0)) return false; + ++index; + } + } + return true; +} + +bool Projection(const Node* node, int64_t bits) { + using namespace matmul_nbits_fusion_utils; + return node && graph_utils::IsSupportedOptypeVersionAndDomain(*node, "MatMulNBits", {1}, kMSDomain) && + node->GetExecutionProviderType() == kWebGpuExecutionProvider && GetIntAttr(*node, "bits", 4) == bits && + GetIntAttr(*node, "block_size", -1, true) == 32 && GetIntAttr(*node, "weight_prepacked", 0) == 0 && + !HasInput(*node, 3) && !HasInput(*node, 4) && !HasInput(*node, 5) && + (bits != 8 || GetIntAttr(*node, "accuracy_level", 0) == 0); +} + +} // namespace + +Status MatMulNBitsQLoRAFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, + const logging::Logger& logger) const { + using namespace matmul_nbits_fusion_utils; + GraphViewer viewer(graph); + for (auto index : viewer.GetNodesInTopologicalOrder()) { + auto* add = graph.GetNode(index); + if (!add) continue; + ORT_RETURN_IF_ERROR(Recurse(*add, modified, graph_level, logger)); + if (!graph_utils::IsSupportedOptypeVersionAndDomain(*add, "Add", {7, 13, 14}) || + add->GetExecutionProviderType() != kWebGpuExecutionProvider) + continue; + auto* base = Producer(graph, *add, 0); + auto* up = Producer(graph, *add, 1); + if (!Projection(base, 4)) std::swap(base, up); + if (!Projection(base, 4) || !Projection(up, 8)) continue; + auto* down = Producer(graph, *up, 0); + if (!Projection(down, 8) || down->InputDefs()[0] != base->InputDefs()[0] || !Private(graph, *base) || + !Private(graph, *down) || !Private(graph, *up)) + continue; + const int64_t k = GetIntAttr(*base, "K", -1, true); + const int64_t n = GetIntAttr(*base, "N", -1, true); + if (k <= 0 || k % 32 != 0 || n <= 0 || GetIntAttr(*down, "K", -1, true) != k || + GetIntAttr(*down, "N", -1, true) != 32 || GetIntAttr(*up, "K", -1, true) != 32 || + GetIntAttr(*up, "N", -1, true) != n) + continue; + InlinedVector inputs{ + const_cast(base->InputDefs()[0]), const_cast(base->InputDefs()[1]), + const_cast(base->InputDefs()[2]), const_cast(down->InputDefs()[1]), + const_cast(down->InputDefs()[2]), const_cast(up->InputDefs()[1]), + const_cast(up->InputDefs()[2])}; + const auto* x_type = inputs[0]->TypeAsProto(); + if (!x_type || !x_type->has_tensor_type()) continue; + const auto float_type = x_type->tensor_type().elem_type(); + if (float_type != ONNX_NAMESPACE::TensorProto::FLOAT16 && float_type != ONNX_NAMESPACE::TensorProto::FLOAT) + continue; + bool types_ok = true; + for (size_t i = 0; i < inputs.size(); ++i) { + const auto* type = inputs[i]->TypeAsProto(); + const auto expected = + (i == 1 || i == 3 || i == 5) ? ONNX_NAMESPACE::TensorProto::UINT8 : float_type; + types_ok &= type && type->has_tensor_type() && type->tensor_type().elem_type() == expected; + } + if (!types_ok || !inputs[0]->Shape() || inputs[0]->Shape()->dim_size() < 2 || !Shape(*inputs[1], {n, k / 32, 16}) || + !AdapterShape(graph, *inputs[3], {32, k / 32, 32}, 0) || + !AdapterShape(graph, *inputs[4], {32, k / 32}, 0) || + !AdapterShape(graph, *inputs[5], {n, 1, 32}, 1) || !AdapterShape(graph, *inputs[6], {n, 1}, 1)) + continue; + const auto* s = inputs[2]->Shape(); + const bool flat_scales = s && s->dim_size() == 1 && s->dim(0).has_dim_value() && + s->dim(0).dim_value() / n == k / 32 && s->dim(0).dim_value() % n == 0; + if (!flat_scales && !Shape(*inputs[2], {n, k / 32})) continue; + auto attrs = base->GetAttributes(); + attrs.erase("weight_prepacked"); + utils::SetNodeAttribute(utils::MakeAttribute("lora_bits", int64_t{8}), attrs); + utils::SetNodeAttribute(utils::MakeAttribute("lora_block_size", int64_t{32}), attrs); + auto& fused = graph.AddNode(graph.GenerateNodeName("MatMulNBitsQLoRA"), "MatMulNBitsQLoRA", + "fused unmerged rank-32 INT8 LoRA projection", inputs, add->MutableOutputDefs(), &attrs, + kMSDomain); + fused.SetExecutionProviderType(kWebGpuExecutionProvider); + const Node* sources[] = {base, base, base, down, down, up, up}; + const int slots[] = {0, 1, 2, 1, 2, 1, 2}; + for (int i = 0; i < 7; ++i) { + if (const auto* edge = graph_utils::GetInputEdge(*sources[i], slots[i])) { + graph.AddEdge(edge->GetNode().Index(), fused.Index(), edge->GetSrcArgIndex(), i); + } + } + const auto outputs = graph_utils::GraphEdge::GetNodeOutputEdges(*add); + for (const Node* node : {base, down, up, static_cast(add)}) { + graph_utils::RemoveNodeOutputEdges(graph, *graph.GetNode(node->Index())); + } + const NodeIndex removed[] = {base->Index(), down->Index(), up->Index(), add->Index()}; + for (auto node : removed) graph.RemoveNode(node); + for (const auto& edge : outputs) graph.AddEdge(fused.Index(), edge.dst_node, 0, edge.dst_arg_index); + // Preserve initializer shapes and data. Replacing full-size zero weights with + // empty adapters would bypass arithmetic, including NaN/Inf propagation. + modified = true; + LOGS(logger, VERBOSE) << "MatMulNBitsQLoRAFusion: " << fused.Name(); + } + return Status::OK(); +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.h b/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.h new file mode 100644 index 0000000000000..943fe431c61e3 --- /dev/null +++ b/onnxruntime/core/optimizer/matmul_nbits_qlora_fusion.h @@ -0,0 +1,19 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/optimizer/graph_transformer.h" + +namespace onnxruntime { + +class MatMulNBitsQLoRAFusion final : public GraphTransformer { + public: + MatMulNBitsQLoRAFusion() : GraphTransformer("MatMulNBitsQLoRAFusion", {kWebGpuExecutionProvider}) {} + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(MatMulNBitsQLoRAFusion); + + private: + Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_shape_inference_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_shape_inference_test.cc new file mode 100644 index 0000000000000..d6c39aed2d4e7 --- /dev/null +++ b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_shape_inference_test.cc @@ -0,0 +1,139 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include + +#include "core/graph/model.h" +#include "test/test_environment.h" +#include "test/unittest_util/graph_transform_test_builder.h" +#include "test/util/include/asserts.h" +#include "test/util/include/test_utils.h" + +namespace onnxruntime::test { +namespace { + +struct QLoRAShapeConfig { + int64_t k = 64; + int64_t n = 128; + int64_t rank = 32; + int64_t bits = 4; + int64_t block_size = 32; + int64_t lora_bits = 8; + int64_t lora_block_size = 32; +}; + +template +Status ResolveQLoRAShape(const QLoRAShapeConfig& c, const std::function& edit = nullptr) { + Model model("qlora_shape", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, DefaultLoggingManager().DefaultLogger()); + ModelTestBuilder builder(model.MainGraph()); + const auto k_blocks = (c.k + c.block_size - 1) / c.block_size; + const auto lora_k_blocks = (c.k + c.lora_block_size - 1) / c.lora_block_size; + const auto r_blocks = (c.rank + c.lora_block_size - 1) / c.lora_block_size; + const auto lora_blob = c.lora_block_size * c.lora_bits / 8; + auto* x = builder.MakeInput(std::vector{2, 3, c.k}); + auto* w = builder.MakeInput(std::vector{c.n, k_blocks, c.block_size * c.bits / 8}); + auto* s = builder.MakeInput(std::vector{c.n, k_blocks}); + auto* a = builder.MakeInput(std::vector{c.rank, lora_k_blocks, lora_blob}); + auto* as = builder.MakeInput(std::vector{c.rank, lora_k_blocks}); + auto* b = builder.MakeInput(std::vector{c.n, r_blocks, lora_blob}); + auto* bs = builder.MakeInput(std::vector{c.n, r_blocks}); + auto* y = builder.MakeOutput(); + auto& node = builder.AddNode("MatMulNBitsQLoRA", {x, w, s, a, as, b, bs}, {y}, kMSDomain); + node.AddAttribute("K", c.k); + node.AddAttribute("N", c.n); + node.AddAttribute("bits", c.bits); + node.AddAttribute("block_size", c.block_size); + node.AddAttribute("lora_bits", c.lora_bits); + node.AddAttribute("lora_block_size", c.lora_block_size); + if (edit) edit(node); + builder.SetGraphOutputs(); + auto status = model.MainGraph().Resolve(); + if (status.IsOK()) { + EXPECT_EQ(y->TypeAsProto()->tensor_type().elem_type(), utils::ToTensorProtoElementType()); + auto expected = *x->Shape(); + expected.mutable_dim(expected.dim_size() - 1)->set_dim_value(c.n); + CheckShapeEquality(y->Shape(), &expected); + } + return status; +} + +void SetInputDimension(Node& node, size_t input, int axis, int64_t value) { + auto* arg = node.MutableInputDefs()[input]; + auto shape = *arg->Shape(); + auto* dim = shape.mutable_dim(axis); + dim->Clear(); + if (value >= 0) dim->set_dim_value(value); + arg->SetShape(shape); +} + +} // namespace + +TEST(MatMulNBitsQLoRASchema, GeneralQuantizationAndArbitraryRank) { + for (const QLoRAShapeConfig& c : { + QLoRAShapeConfig{63, 65, 7, 8, 64, 4, 16}, + {65, 33, 64, 2, 16, 8, 128}, + {256, 128, 129, 4, 128, 2, 64}, + {64, 128, 1}, + {64, 128, 16}, + {64, 128, 32}}) { + SCOPED_TRACE(testing::Message() << "K=" << c.k << " N=" << c.n << " R=" << c.rank); + ASSERT_STATUS_OK(ResolveQLoRAShape(c)); + ASSERT_STATUS_OK(ResolveQLoRAShape(c)); + } +} + +TEST(MatMulNBitsQLoRASchema, EmptyAdaptersRetainScaleDefaults) { + for (int mask : {1, 2, 3}) { + ASSERT_STATUS_OK(ResolveQLoRAShape({63, 65, 7, 8, 64, 4, 16}, [mask](Node& node) { + if (mask & 1) SetInputDimension(node, 3, 0, 0); + if (mask & 2) SetInputDimension(node, 5, 1, 0); + })); + } +} + +TEST(MatMulNBitsQLoRASchema, DynamicAdapterRankAndVectorInput) { + ASSERT_STATUS_OK(ResolveQLoRAShape({}, [](Node& node) { + SetInputDimension(node, 3, 0, -1); + SetInputDimension(node, 4, 0, -1); + SetInputDimension(node, 5, 1, -1); + SetInputDimension(node, 6, 1, -1); + ONNX_NAMESPACE::TensorShapeProto x_shape; + x_shape.add_dim()->set_dim_value(64); + node.MutableInputDefs()[0]->SetShape(x_shape); + })); +} + +#ifndef ORT_NO_EXCEPTIONS +TEST(MatMulNBitsQLoRASchema, RejectsInconsistentAdapterShapes) { + for (const auto& [input, axis, value] : {std::tuple{4, 0, 16}, + {3, 1, 3}, + {3, 2, 16}, + {5, 1, 2}, + {6, 1, 2}}) { + const auto status = ResolveQLoRAShape({}, [=](Node& node) { SetInputDimension(node, input, axis, value); }); + EXPECT_FALSE(status.IsOK()); + } +} + +TEST(MatMulNBitsQLoRASchema, RejectsInvalidQuantizationAttributes) { + const std::pair invalid_attributes[] = { + {"bits", 3}, + {"lora_bits", 16}, + {"block_size", 8}, + {"block_size", 48}, + {"lora_block_size", 0}, + {"lora_block_size", 24}, + {"K", 0}, + {"N", -1}, + }; + for (const auto& [name, value] : invalid_attributes) { + SCOPED_TRACE(name); + const auto status = ResolveQLoRAShape({}, [=](Node& node) { node.AddAttribute(name, value); }); + EXPECT_FALSE(status.IsOK()); + } +} +#endif // ORT_NO_EXCEPTIONS + +} // namespace onnxruntime::test diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc new file mode 100644 index 0000000000000..daa9a064aab3f --- /dev/null +++ b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc @@ -0,0 +1,504 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "test/providers/provider_test_utils.h" +#include "test/util/include/default_providers.h" + +#if defined(USE_WEBGPU) && !defined(DISABLE_CONTRIB_OPS) +#if !defined(ORT_USE_EP_API_ADAPTERS) +#include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" +#include "core/providers/webgpu/math/subgroup_matrix_config.h" +#include "core/providers/webgpu/webgpu_context.h" +#if !defined(__wasm__) && !defined(USE_EXTERNAL_DAWN) +#include "dawn/native/DawnNative.h" +#endif +#endif +#include "core/providers/webgpu/webgpu_provider_options.h" +#endif + +namespace onnxruntime::test { +#if defined(USE_WEBGPU) && !defined(DISABLE_CONTRIB_OPS) + +template +void RunEmptyAndActiveAdapters() { + // Scalar output, a two-column tail, and a partial wide tile. + for (const auto& [m, n] : {std::pair{1, 1}, {2, 2}, {37, 129}}) { + SCOPED_TRACE(MakeString("M=", m, ", N=", n)); + for (auto ranks : {std::pair{0, 0}, {0, 1}, {32, 0}, {32, 1}}) { + auto ep = DefaultWebGpuExecutionProvider(); + ASSERT_NE(ep, nullptr); + OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); + test.AddAttribute("K", 64); + test.AddAttribute("N", n); + test.AddInput("X", {1, m, 64}, std::vector(m * 64, T(0.5f))); + test.AddInput("W", {n, 2, 16}, std::vector(n * 32, 0x99), true); + test.AddInput("S", {n, 2}, std::vector(n * 2, T(0.03125f)), true); + test.AddInput("AQ", {ranks.first, 2, 32}, std::vector(ranks.first * 64, 129)); + test.AddInput("AS", {32, 2}, std::vector(64, T(0.03125f))); + test.AddInput("BQ", {n, ranks.second, 32}, std::vector(n * ranks.second * 32, 129)); + test.AddInput("BS", {n, 1}, std::vector(n, T(0.0625f))); + const float expected = ranks.first && ranks.second ? 3.0f : 1.0f; + test.AddOutput("Y", {1, m, n}, std::vector(m * n, T(expected))); + std::vector> eps; + eps.push_back(std::move(ep)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); + } + } +} + +TEST(MatMulNBitsQLoRA, EmptyAndActiveAdapters) { RunEmptyAndActiveAdapters(); } +TEST(MatMulNBitsQLoRA, EmptyAndActiveAdaptersFp32) { RunEmptyAndActiveAdapters(); } + +#ifndef ORT_NO_EXCEPTIONS +TEST(MatMulNBitsQLoRA, RejectsUnsupportedWebGpuConfigurations) { + struct Config { + int64_t k = 64, n = 128, rank = 32, bits = 4, block_size = 32, lora_bits = 8, lora_block_size = 32; + const char* error = "WebGPU MatMulNBitsQLoRA supports only bits=4"; + }; + const Config configs[] = { + {64, 128, 32, 8}, + {64, 128, 32, 4, 64}, + {64, 128, 32, 4, 32, 4}, + {64, 128, 32, 4, 32, 8, 64}, + {64, 128, 7, 4, 32, 8, 32, "WebGPU MatMulNBitsQLoRA supports only adapter rank 32"}, + {64, 128, 64, 4, 32, 8, 32, "WebGPU MatMulNBitsQLoRA supports only adapter rank 32"}, + {63, 128, 32, 4, 32, 8, 32, "WebGPU MatMulNBitsQLoRA requires positive K"}, + }; + for (const Config& c : configs) { + auto ep = DefaultWebGpuExecutionProvider(); + ASSERT_NE(ep, nullptr); + OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); + test.AddAttribute("K", c.k); + test.AddAttribute("N", c.n); + test.AddAttribute("bits", c.bits); + test.AddAttribute("block_size", c.block_size); + test.AddAttribute("lora_bits", c.lora_bits); + test.AddAttribute("lora_block_size", c.lora_block_size); + const auto k_blocks = (c.k + c.block_size - 1) / c.block_size; + const auto lora_k_blocks = (c.k + c.lora_block_size - 1) / c.lora_block_size; + const auto r_blocks = (c.rank + c.lora_block_size - 1) / c.lora_block_size; + const auto blob = c.block_size * c.bits / 8; + const auto lora_blob = c.lora_block_size * c.lora_bits / 8; + test.AddInput("X", {1, 1, c.k}, std::vector(c.k, 0.5f)); + test.AddInput("W", {c.n, k_blocks, blob}, std::vector(c.n * k_blocks * blob), true); + test.AddInput("S", {c.n, k_blocks}, std::vector(c.n * k_blocks, 1.0f), true); + test.AddInput("AQ", {c.rank, lora_k_blocks, lora_blob}, + std::vector(c.rank * lora_k_blocks * lora_blob)); + test.AddInput("AS", {c.rank, lora_k_blocks}, std::vector(c.rank * lora_k_blocks, 1.0f)); + test.AddInput("BQ", {c.n, r_blocks, lora_blob}, std::vector(c.n * r_blocks * lora_blob)); + test.AddInput("BS", {c.n, r_blocks}, std::vector(c.n * r_blocks, 1.0f)); + test.AddOutput("Y", {1, 1, c.n}, std::vector(c.n)); + std::vector> eps; + eps.push_back(std::move(ep)); + test.Run(OpTester::ExpectResult::kExpectFailure, c.error, {}, nullptr, &eps); + } +} +#endif // ORT_NO_EXCEPTIONS + +TEST(MatMulNBitsQLoRA, PrefillSubgroupMatrixNumerics) { +#if defined(ORT_USE_EP_API_ADAPTERS) + GTEST_SKIP() << "Hardware-specific dispatch checks require direct adapter capability inspection."; +#else + auto ep = DefaultWebGpuExecutionProvider(); + ASSERT_NE(ep, nullptr); + auto& gpu = webgpu::WebGpuContextFactory::GetContext(ep->GetDeviceId()); + constexpr auto kF16 = wgpu::SubgroupMatrixComponentType::F16; + const auto& adapter = gpu.AdapterInfo(); + const auto& device_configs = gpu.SubgroupMatrixConfigs(); + const auto config = webgpu::detail::SelectSubgroupMatrixConfigFromAdapterConfigs( + {device_configs.configs, device_configs.configCount}, + adapter.subgroupMinSize, adapter.subgroupMaxSize, + gpu.DeviceHasFeature(wgpu::FeatureName::SubgroupSizeControl), + {{kF16, kF16, 16, 16, 16, 32, true}, + {kF16, kF16, 8, 16, 16, 32, true}, + {kF16, kF16, 8, 8, 8, 32, false}}); + if (!gpu.DeviceHasFeature(wgpu::FeatureName::ChromiumExperimentalSubgroupMatrix) || !config.has_value()) { + GTEST_SKIP() << "Requires a compatible FP16 subgroup-matrix configuration."; + } + // Include a partial row tile; the dispatch assertion below prevents a fallback + // from hiding numerical failures in the subgroup-matrix kernel. + constexpr int64_t m = 65, k = 256, n = 128; + OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); + test.AddAttribute("K", k); + test.AddAttribute("N", n); + test.AddAttribute("accuracy_level", 4); + test.AddInput("X", {1, m, k}, std::vector(m * k, MLFloat16(0.5f))); + test.AddInput("W", {n, k / 32, 16}, std::vector(n * k / 2, 0x99), true); + test.AddInput("S", {n, k / 32}, + std::vector(n * k / 32, MLFloat16(0.03125f)), true); + test.AddInput("AQ", {32, k / 32, 32}, std::vector(32 * k, 129)); + test.AddInput("AS", {32, k / 32}, std::vector(k, MLFloat16(0.03125f))); + test.AddInput("BQ", {n, 1, 32}, std::vector(n * 32, 129)); + test.AddInput("BS", {n, 1}, std::vector(n, MLFloat16(0.0625f))); + test.AddOutput("Y", {1, m, n}, std::vector(m * n, MLFloat16(12.0f))); + + SessionOptions session_options; + session_options.session_log_severity_level = static_cast(logging::Severity::kINFO); + std::vector> eps; + eps.push_back(std::move(ep)); + testing::internal::CaptureStderr(); + test.Run(session_options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); + const auto logs = testing::internal::GetCapturedStderr(); + EXPECT_NE(logs.find("SubgroupMatrixMatMulNBitsQLoRAUp"), std::string::npos) << logs; +#endif +} + +void RunFp32Numerics(const ConfigOptions& provider_options = {}) { + struct Shape { + int64_t m, k, n, batch; + }; + // Small-K decode, joint decode with a K tail, long reductions, and batched + // generic/wide/DP4A prefill. Odd N covers scalar output tails. + const Shape shapes[] = { + {1, 64, 1, 1}, + {1, 2176, 127, 1}, + {1, 8192, 64, 1}, + {2, 2048, 129, 2}, + {37, 2176, 127, 2}, + {37, 2176, 64, 2}, + }; + for (const auto& shape : shapes) { + const auto [m, k, n, batch] = shape; + SCOPED_TRACE(MakeString("M=", m, ", K=", k, ", N=", n, ", batch=", batch)); + std::vector x(batch * m * k), as(k), bs(n), expected(batch * m * n); + std::vector aq(32 * k), bq(n * 32); + for (int64_t row = 0; row < batch * m; ++row) { + for (int64_t col = 0; col < k; ++col) { + // Exactly representable by DP4A quantization, including zero rows. + const int value = col % 128 == 0 ? 127 : static_cast((row * 19 + col * 11) % 255) - 127; + x[row * k + col] = row % 7 == 6 ? 0.0f : value / 128.0f; + } + } + for (int64_t r = 0; r < 32; ++r) { + for (int64_t col = 0; col < k; ++col) + aq[r * k + col] = static_cast(96 + (r * 17 + col * 13) % 65); + for (int64_t block = 0; block < k / 32; ++block) + as[r * (k / 32) + block] = 0.001234567f * (1 + (r + block) % 4); + } + for (int64_t col = 0; col < n; ++col) { + bs[col] = 0.002345678f * (1 + col % 3); + for (int64_t r = 0; r < 32; ++r) + bq[col * 32 + r] = static_cast(112 + (col * 7 + r * 3) % 33); + } + // Independent dense FP64 reference, rounding weights and intermediate T + // to FP32. Non-half-representable scales detect accidental FP16 arithmetic. + for (int64_t row = 0; row < batch * m; ++row) { + double base = 0; + for (int64_t col = 0; col < k; ++col) base += x[row * k + col] * 0.03125; + std::vector t(32); + for (int64_t r = 0; r < 32; ++r) { + double sum = 0; + for (int64_t col = 0; col < k; ++col) { + const float weight = (static_cast(aq[r * k + col]) - 128) * as[r * (k / 32) + col / 32]; + sum += static_cast(x[row * k + col]) * weight; + } + t[r] = static_cast(sum); + } + for (int64_t col = 0; col < n; ++col) { + double delta = 0; + for (int64_t r = 0; r < 32; ++r) { + const float weight = (static_cast(bq[col * 32 + r]) - 128) * bs[col]; + delta += static_cast(t[r]) * weight; + } + expected[row * n + col] = static_cast(base) + static_cast(delta); + } + } + for (int64_t accuracy : {0, 4}) { + SCOPED_TRACE(accuracy); + OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); + test.AddAttribute("K", k); + test.AddAttribute("N", n); + test.AddAttribute("accuracy_level", accuracy); + test.AddInput("X", {batch, m, k}, x); + test.AddInput("W", {n, k / 32, 16}, std::vector(n * k / 2, 0x99), true); + test.AddInput("S", {n, k / 32}, std::vector(n * k / 32, 0.03125f), true); + test.AddInput("AQ", {32, k / 32, 32}, aq); + test.AddInput("AS", {32, k / 32}, as); + test.AddInput("BQ", {n, 1, 32}, bq); + test.AddInput("BS", {n, 1}, bs); + test.AddOutput("Y", {batch, m, n}, expected); + test.SetOutputAbsErr("Y", 0.00001f); + test.SetOutputRelErr("Y", 0.00001f); + std::vector> eps; + eps.push_back(WebGpuExecutionProviderWithOptions(provider_options)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); + } + } +} + +TEST(MatMulNBitsQLoRA, Fp32Numerics) { RunFp32Numerics(); } + +TEST(MatMulNBitsQLoRA, Fp32NumericsWithoutSubgroups) { +#if defined(ORT_USE_EP_API_ADAPTERS) || defined(__wasm__) || defined(USE_EXTERNAL_DAWN) + GTEST_SKIP() << "Requires native Dawn to create a device without subgroups."; +#else + // Keep the default context alive to initialize Dawn's proc table. + auto initial_ep = DefaultWebGpuExecutionProvider(); + ASSERT_NE(initial_ep, nullptr); + const auto instance_feature = wgpu::InstanceFeatureName::TimedWaitAny; + wgpu::InstanceDescriptor instance_desc{}; + instance_desc.requiredFeatureCount = 1; + instance_desc.requiredFeatures = &instance_feature; + dawn::native::Instance instance{&instance_desc}; + wgpu::RequestAdapterOptions adapter_options{}; + adapter_options.backendType = static_cast(webgpu::WebGpuContextConfig{}.backend_type); + auto adapters = instance.EnumerateAdapters(&adapter_options); + ASSERT_FALSE(adapters.empty()); + + const auto device_feature = wgpu::FeatureName::ImplicitDeviceSynchronization; + wgpu::DeviceDescriptor device_desc{}; + device_desc.requiredFeatureCount = 1; + device_desc.requiredFeatures = &device_feature; + auto device = wgpu::Device::Acquire(adapters.front().CreateDevice(&device_desc)); + ASSERT_NE(device, nullptr); + ASSERT_FALSE(device.HasFeature(wgpu::FeatureName::Subgroups)); + + ConfigOptions options; + ASSERT_STATUS_OK(options.AddConfigEntry(webgpu::options::kDeviceId, "1")); + ASSERT_STATUS_OK(options.AddConfigEntry( + webgpu::options::kWebGpuInstance, std::to_string(reinterpret_cast(instance.Get())).c_str())); + ASSERT_STATUS_OK(options.AddConfigEntry( + webgpu::options::kWebGpuDevice, std::to_string(reinterpret_cast(device.Get())).c_str())); + auto external_ep = WebGpuExecutionProviderWithOptions(options); + ASSERT_NE(external_ep, nullptr); + ASSERT_FALSE(webgpu::WebGpuContextFactory::GetContext(1).DeviceHasFeature(wgpu::FeatureName::Subgroups)); + // Exercise the joint decode workgroup-memory reduction, including K tails, + // and the prefill fallback against the same dense FP64 reference. + RunFp32Numerics(options); +#endif +} + +void RunFp32DP4AQuantization(int64_t m) { + SCOPED_TRACE(m); +#if defined(ORT_USE_EP_API_ADAPTERS) + GTEST_SKIP() << "Hardware-specific specialization checks require direct adapter capability inspection."; +#else + auto initial_ep = DefaultWebGpuExecutionProvider(); + ASSERT_NE(initial_ep, nullptr); + if (!contrib::webgpu::HasDP4ADeviceSupport(initial_ep->GetDeviceId())) { + GTEST_SKIP() << "Requires DP4A device support."; + } +#endif + constexpr int64_t k = 2176, n = 64, batch = 3; + const int64_t rows = batch * m; + std::vector x(rows * k), quantized_x(rows * k), as(k), expected(rows * n); + std::vector aq(32 * k), bq(n * 32, 128); + for (int64_t row = 0; row < rows; ++row) { + for (int64_t col = 0; col < k; ++col) { + const float scale = 1.0f / static_cast(1 << (8 + (row + col / 128) % 3)); + const int value = col % 128 == 0 ? 127 : static_cast((row * 19 + col * 11) % 255) - 127; + quantized_x[row * k + col] = row == 1 ? 0.0f : value * scale; + // Stay inside the rounding interval of the known quantization grid. + const float offset = value == 127 || value == -127 ? 0.0f : 0.2f; + x[row * k + col] = row == 1 ? 0.0f : (value + offset) * scale; + } + } + for (int64_t r = 0; r < 32; ++r) { + for (int64_t col = 0; col < k; ++col) + aq[r * k + col] = static_cast((r * 17 + col * 13) % 256); + for (int64_t block = 0; block < k / 32; ++block) + as[r * (k / 32) + block] = (1 + (r + block) % 4) / 2048.0f; + } + for (int64_t col = 0; col < n; ++col) bq[col * 32 + col % 32] = 129; + // An identity up projection exposes every T value. The reference uses + // dense FP64 products on the known grid, including a zero row and a K tail. + for (int64_t row = 0; row < rows; ++row) { + double base = 0; + for (int64_t col = 0; col < k; ++col) base += quantized_x[row * k + col] * 0.03125; + for (int64_t r = 0; r < 32; ++r) { + double sum = 0; + for (int64_t col = 0; col < k; ++col) + sum += static_cast(quantized_x[row * k + col]) * + (static_cast(aq[r * k + col]) - 128) * as[r * (k / 32) + col / 32]; + expected[row * n + r] = expected[row * n + r + 32] = static_cast(base) + static_cast(sum); + } + } + OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); + test.AddAttribute("K", k); + test.AddAttribute("N", n); + test.AddAttribute("accuracy_level", 4); + test.AddInput("X", {batch, m, k}, x); + test.AddInput("W", {n, k / 32, 16}, std::vector(n * k / 2, 0x99), true); + test.AddInput("S", {n, k / 32}, std::vector(n * k / 32, 0.03125f), true); + test.AddInput("AQ", {32, k / 32, 32}, aq); + test.AddInput("AS", {32, k / 32}, as); + test.AddInput("BQ", {n, 1, 32}, bq); + test.AddInput("BS", {n, 1}, std::vector(n, 1.0f)); + test.AddOutput("Y", {batch, m, n}, expected); + test.SetOutputAbsErr("Y", 0.00001f); + test.SetOutputRelErr("Y", 0.00001f); + std::vector> eps; + eps.push_back(DefaultWebGpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); +} + +TEST(MatMulNBitsQLoRA, DecodeJointDP4AQuantization) { RunFp32DP4AQuantization(1); } + +TEST(MatMulNBitsQLoRA, SmallMPrefillDP4AQuantization) { + RunFp32DP4AQuantization(2); +} + +TEST(MatMulNBitsQLoRA, NarrowKParallelReduction) { +#if defined(ORT_USE_EP_API_ADAPTERS) + GTEST_SKIP() << "Hardware-specific specialization checks require direct adapter capability inspection."; +#else + auto initial_ep = DefaultWebGpuExecutionProvider(); + ASSERT_NE(initial_ep, nullptr); + if (!contrib::webgpu::HasDP4ADeviceSupport(initial_ep->GetDeviceId())) { + GTEST_SKIP() << "Requires a device that selects DP4A for the rank-32 down projection."; + } +#endif + struct Shape { + int64_t m; + int64_t k; + int64_t batch; + }; + // Cover both sides of the K-parallel threshold and a batched K/M tail. + for (const auto shape : {Shape{4, 1920, 1}, Shape{4, 2048, 1}, Shape{37, 2176, 2}}) { + const auto [m, k, batch] = shape; + SCOPED_TRACE(MakeString("M=", m, ", K=", k, ", batch=", batch)); + constexpr int64_t n = 64; + std::vector x(batch * m * k); + std::vector aq(32 * k); + std::vector as(k); + std::vector bq(n * 32, 128); + std::vector bs(n, MLFloat16(1.0f)); + std::vector expected(batch * m * n); + for (int64_t row = 0; row < batch * m; ++row) { + for (int64_t col = 0; col < k; ++col) { + // The maximum in every 128-element block is 127 * scale, so DP4A's + // input quantization preserves these signed integer values exactly. + const float scale = 1.0f / static_cast(1 << (7 + (row + col / 128) % 4)); + const int value = col % 128 == 0 ? 127 : static_cast((row * 19 + col * 11) % 255) - 127; + x[row * k + col] = MLFloat16(row % 7 == 0 ? 0.0f : value * scale); + } + } + for (int64_t r = 0; r < 32; ++r) { + for (int64_t col = 0; col < k; ++col) { + aq[r * k + col] = static_cast((r * 17 + col * 13) % 256); + } + for (int64_t block = 0; block < k / 32; ++block) { + as[r * (k / 32) + block] = MLFloat16(1.0f / static_cast(1 << (8 + (r + block) % 4))); + } + } + // A zero base and an identity up projection expose every element of T. + // Compute the reference as a dense dequantized dot product in double. + for (int64_t col = 0; col < n; ++col) { + bq[col * 32 + col % 32] = 129; + } + for (int64_t row = 0; row < batch * m; ++row) { + for (int64_t r = 0; r < 32; ++r) { + double sum = 0; + for (int64_t col = 0; col < k; ++col) { + sum += static_cast(x[row * k + col].ToFloat()) * + (static_cast(aq[r * k + col]) - 128) * as[r * (k / 32) + col / 32].ToFloat(); + } + expected[row * n + r] = expected[row * n + r + 32] = MLFloat16(static_cast(sum)); + } + } + OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); + test.AddAttribute("K", k); + test.AddAttribute("N", n); + test.AddAttribute("accuracy_level", 4); + test.AddInput("X", {batch, m, k}, x); + test.AddInput("W", {n, k / 32, 16}, std::vector(n * k / 2, 0x88), true); + test.AddInput("S", {n, k / 32}, std::vector(n * k / 32, MLFloat16(1.0f)), true); + test.AddInput("AQ", {32, k / 32, 32}, aq); + test.AddInput("AS", {32, k / 32}, as); + test.AddInput("BQ", {n, 1, 32}, bq); + test.AddInput("BS", {n, 1}, bs); + test.AddOutput("Y", {batch, m, n}, expected); + test.SetOutputAbsErr("Y", 0.002f); + test.SetOutputRelErr("Y", 0.001f); + ConfigOptions options; + // Short-K fallbacks do not force FP32 accumulation like the narrow DP4A + // kernel. Enable it explicitly for this tight dense-reference comparison. + if (k < 2048) { + ASSERT_STATUS_OK(options.AddConfigEntry(webgpu::options::kEnableMatmulFp32Accumulation, + webgpu::options::kEnableMatmulFp32Accumulation_ON)); + } + std::vector> eps; + eps.push_back(WebGpuExecutionProviderWithOptions(options)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); + } +} + +TEST(MatMulNBitsQLoRA, DecodeSubgroupReduction) { +#if defined(ORT_USE_EP_API_ADAPTERS) + GTEST_SKIP() << "Hardware-specific specialization checks require direct adapter capability inspection."; +#else + auto initial_ep = DefaultWebGpuExecutionProvider(); + ASSERT_NE(initial_ep, nullptr); + auto& gpu = webgpu::WebGpuContextFactory::GetContext(initial_ep->GetDeviceId()); + const auto& adapter = gpu.AdapterInfo(); + const bool subgroup_32 = (adapter.subgroupMinSize == 32 && adapter.subgroupMaxSize == 32) || + (gpu.DeviceHasFeature(wgpu::FeatureName::SubgroupSizeControl) && + adapter.subgroupMinSize <= 32 && adapter.subgroupMaxSize >= 32); + if (!gpu.DeviceHasFeature(wgpu::FeatureName::Subgroups) || !subgroup_32) { + GTEST_SKIP() << "The decode specialization requires fixed or selectable 32-lane subgroups."; + } +#endif + for (bool acc_f32 : {false, true}) { + for (const auto& [k, n] : {std::pair{2176, 127}, {8192, 64}}) { + SCOPED_TRACE(MakeString("K=", k, ", N=", n, ", acc_f32=", acc_f32)); + std::vector x(k); + std::vector aq(32 * k); + std::vector as(k); + std::vector bq(n * 32, 128); + std::vector expected(n); + double base = 0; + for (int64_t col = 0; col < k; ++col) { + x[col] = MLFloat16((static_cast(col % 17) - 8) / 64.0f); + base += x[col].ToFloat() * 0.03125; + } + std::vector t(32); + for (int64_t rank = 0; rank < 32; ++rank) { + for (int64_t block = 0; block < k / 32; ++block) { + as[rank * (k / 32) + block] = MLFloat16(0.001f * (1 + (rank + block) % 4)); + } + double sum = 0; + for (int64_t col = 0; col < k; ++col) { + const int q = rank == 0 ? 0 : static_cast((rank * 19 + col * 13) % 65) - 32; + aq[rank * k + col] = static_cast(q + 128); + const auto weight = MLFloat16(q * as[rank * (k / 32) + col / 32].ToFloat()); + sum += static_cast(x[col].ToFloat()) * weight.ToFloat(); + } + t[rank] = MLFloat16(static_cast(sum)); + } + // Identity rows expose all 32 T values; the nonzero base also exercises + // ordinary workgroups in the same dispatch and partial output tiles. + const auto rounded_base = MLFloat16(static_cast(base)); + for (int64_t col = 0; col < n; ++col) { + bq[col * 32 + col % 32] = 129; + expected[col] = MLFloat16(rounded_base.ToFloat() + t[col % 32].ToFloat()); + } + OpTester test("MatMulNBitsQLoRA", 1, kMSDomain); + test.AddAttribute("K", k); + test.AddAttribute("N", n); + test.AddInput("X", {1, 1, k}, x); + test.AddInput("W", {n, k / 32, 16}, std::vector(n * k / 2, 0x99), true); + test.AddInput("S", {n, k / 32}, std::vector(n * k / 32, MLFloat16(0.03125f)), true); + test.AddInput("AQ", {32, k / 32, 32}, aq); + test.AddInput("AS", {32, k / 32}, as); + test.AddInput("BQ", {n, 1, 32}, bq); + test.AddInput("BS", {n, 1}, std::vector(n, MLFloat16(1.0f))); + test.AddOutput("Y", {1, 1, n}, expected); + // The long-K half-precision lane accumulation can drift before the + // FP32 subgroup reduction (up to 0.00238 on RTX 4070 at K=8192). + // Keep the tighter bound for shorter reductions and FP32 accumulation. + test.SetOutputAbsErr("Y", !acc_f32 && k == 8192 ? 0.003f : 0.002f); + test.SetOutputRelErr("Y", 0.001f); + ConfigOptions options; + if (acc_f32) { + ASSERT_STATUS_OK(options.AddConfigEntry(webgpu::options::kEnableMatmulFp32Accumulation, + webgpu::options::kEnableMatmulFp32Accumulation_ON)); + } + std::vector> eps; + eps.push_back(WebGpuExecutionProviderWithOptions(options)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); + } + } +} + +#endif +} // namespace onnxruntime::test diff --git a/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc new file mode 100644 index 0000000000000..de62dbee98119 --- /dev/null +++ b/onnxruntime/test/optimizer/matmul_nbits_qlora_fusion_test.cc @@ -0,0 +1,465 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/optimizer/matmul_nbits_qlora_fusion.h" +#include "core/providers/webgpu/webgpu_provider_options.h" +#include "core/session/IOBinding.h" +#include "core/session/onnxruntime_run_options_config_keys.h" +#include "core/session/onnxruntime_session_options_config_keys.h" +#include "test/optimizer/graph_transform_test_fixture.h" +#include "test/optimizer/webgpu_fusion_test_util.h" +#include "test/util/include/default_providers.h" + +namespace onnxruntime::test { +#ifndef DISABLE_CONTRIB_OPS +namespace { + +template +void BuildQLoRA(ModelTestBuilder& builder, int64_t m = 1, int64_t k = 2048, int64_t n = 128, int64_t batch = 1, + bool expose = false, bool different_x = false, bool zero_base = false, int64_t adapter_accuracy = 0, + bool empty_defaults = false) { + auto* x = builder.MakeInput({batch, m, k}, T(-0.5f), T(0.5f)); + // Integer random ranges exclude the upper bound, even for constant weights. + const uint8_t base_weight_min = zero_base ? uint8_t{0x88} : (ExactBase ? uint8_t{0x99} : uint8_t{0}); + const uint8_t base_weight_max = zero_base ? uint8_t{0x89} : (ExactBase ? uint8_t{0x9a} : uint8_t{255}); + auto* w = builder.MakeInitializer({n, k / 32, 16}, base_weight_min, base_weight_max); + auto* s = builder.MakeInitializer({n, k / 32}, T(ExactBase ? 0.0625f : 0.03125f), T(0.0625f)); + auto* aq = builder.MakeInput({32, k / 32, 32}, uint8_t{100}, uint8_t{156}); + auto* as = builder.MakeInput({32, k / 32}, T(0.001f), T(0.002f)); + auto* bq = builder.MakeInput({n, 1, 32}, uint8_t{100}, uint8_t{156}); + auto* bs = builder.MakeInput({n, 1}, T(0.001f), T(0.002f)); + for (auto* arg : {aq, as, bq, bs}) { + ONNX_NAMESPACE::TensorProto init; + init.set_name(arg->Name()); + const bool packed = arg == aq || arg == bq; + init.set_data_type(packed ? ONNX_NAMESPACE::TensorProto::UINT8 : utils::ToTensorProtoElementType()); + int64_t count = 1; + for (const auto& dim : arg->Shape()->dim()) { + init.add_dims(dim.dim_value()); + count *= dim.dim_value(); + } + if (packed) { + init.set_raw_data(std::string(static_cast(count), static_cast(128))); + } else { + const std::vector ones(static_cast(count), T(1.0f)); + init.set_raw_data(ones.data(), ones.size() * sizeof(T)); + } + if (empty_defaults) { + const int axis = arg == aq || arg == as ? 0 : 1; + init.set_dims(axis, 0); + init.clear_raw_data(); + auto shape = *arg->Shape(); + shape.mutable_dim(axis)->Clear(); + arg->SetShape(shape); + } + builder.graph_.AddInitializedTensor(init); + } + auto* base = builder.MakeIntermediate(std::vector{batch, m, n}); + auto* t = expose ? builder.MakeOutput(std::vector{batch, m, 32}) + : builder.MakeIntermediate(std::vector{batch, m, 32}); + auto* delta = builder.MakeIntermediate(std::vector{batch, m, n}); + auto* y = builder.MakeOutput(std::vector{batch, m, n}); + auto* adapter_x = different_x ? builder.MakeInput({batch, m, k}, T(-0.5f), T(0.5f)) : x; + if (different_x) { + builder.graph_.SetInputs({x, aq, as, bq, bs, adapter_x}); + } else { + builder.graph_.SetInputs({x, aq, as, bq, bs}); + } + auto& base_node = builder.AddNode("MatMulNBits", {x, w, s}, {base}, kMSDomain); + auto& down = builder.AddNode("MatMulNBits", {adapter_x, aq, as}, {t}, kMSDomain); + auto& up = builder.AddNode("MatMulNBits", {t, bq, bs}, {delta}, kMSDomain); + for (auto* node : {&base_node, &down, &up}) { + node->AddAttribute("K", node == &up ? int64_t{32} : k); + node->AddAttribute("N", node == &down ? int64_t{32} : n); + node->AddAttribute("bits", node == &base_node ? int64_t{4} : int64_t{8}); + node->AddAttribute("block_size", int64_t{32}); + node->AddAttribute("accuracy_level", node == &base_node ? int64_t{4} : adapter_accuracy); + node->SetExecutionProviderType(kWebGpuExecutionProvider); + } + builder.AddNode("Add", {base, delta}, {y}).SetExecutionProviderType(kWebGpuExecutionProvider); +} + +Status CheckQLoRACount(const Graph& graph, size_t expected) { + size_t count = 0; + for (const auto& node : graph.Nodes()) count += node.OpType() == "MatMulNBitsQLoRA"; + ORT_RETURN_IF_NOT(count == expected, "Unexpected quantized LoRA fusion count."); + return Status::OK(); +} + +} // namespace + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionPreservesDefaults) { + for (bool empty_defaults : {false, true}) { + SCOPED_TRACE(empty_defaults); + InlinedHashMap initializers; + ASSERT_STATUS_OK(TestGraphTransformer( + [&](ModelTestBuilder& b) { + BuildQLoRA(b, 1, 64, 128, 1, false, false, false, 0, empty_defaults); + for (const auto& [name, init] : b.graph_.GetAllInitializedTensors()) { + initializers.emplace(name, init->SerializeAsString()); + } + }, + 21, *logger_, std::make_unique(), TransformerLevel::Level2, 1, nullptr, + [&](const Graph& graph) { + ORT_RETURN_IF_ERROR(CheckQLoRACount(graph, 1)); + ORT_RETURN_IF_NOT(graph.GetOverridableInitializers().size() == 4, "Adapter parameters were removed."); + for (const auto& [name, original] : initializers) { + const ONNX_NAMESPACE::TensorProto* init = nullptr; + ORT_RETURN_IF_NOT(graph.GetInitializedTensor(name, init), "Missing initializer."); + ORT_RETURN_IF_NOT(init->SerializeAsString() == original, "Initializer shape or data changed."); + } + return Status::OK(); + })); + } +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionRequiresWebGpu) { + Model model("CPU quantized LoRA", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, *logger_); + ModelTestBuilder builder(model.MainGraph()); + BuildQLoRA(builder, 1, 64); + builder.SetGraphOutputs(); + ASSERT_STATUS_OK(model.MainGraph().Resolve()); + const auto data = model.ToProto().SerializeAsString(); + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Level2; + InferenceSessionWrapper session(options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(DefaultCpuExecutionProvider())); + ASSERT_STATUS_OK(session.Load(data.data(), static_cast(data.size()))); + ASSERT_STATUS_OK(session.Initialize()); + ASSERT_STATUS_OK(CheckQLoRACount(session.GetGraph(), 0)); +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionNumerics) { + if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; + struct ShapeCase { + int64_t m, k, n, batch; + bool zero_base; + }; + const ShapeCase shapes[] = { + {1, 64, 1, 1, false}, + {1, 2176, 127, 1, true}, + {2, 2048, 129, 2, true}, + {37, 2048, 128, 1, false}, + {37, 2176, 127, 2, true}, + }; + for (const auto& shape : shapes) { + SCOPED_TRACE(shape.m); + SCOPED_TRACE(shape.k); + SCOPED_TRACE(shape.n); + SCOPED_TRACE(shape.batch); + // FP16 fusion can change rounding on Metal even without DP4A or subgroup + // reduction. Use the FP16 error budget for every dispatch path. + RunWebGpuFusionTransformerTest( + [shape](ModelTestBuilder& b) { + BuildQLoRA(b, shape.m, shape.k, shape.n, shape.batch, false, false, shape.zero_base); + }, + [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckQLoRACount(s.GetGraph(), 1)); }, + TransformerLevel::Level1, TransformerLevel::Level2, 21, + 0.002, 0.001, std::make_unique(), + [] { return DefaultWebGpuExecutionProvider(); }); + } +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionFp32Numerics) { + if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; + for (const auto& [m, n] : {std::pair{1, 128}, {2, 128}, {37, 128}, {2, 129}}) { + SCOPED_TRACE(MakeString("M=", m, ", N=", n)); + // Eligible FP32 decode and short prefill share quantized X with the base. + const bool quantized_down = m >= 4 || n % 16 == 0; + RunWebGpuFusionTransformerTest( + [m, n](ModelTestBuilder& b) { BuildQLoRA(b, m, 2048, n); }, + [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckQLoRACount(s.GetGraph(), 1)); }, + TransformerLevel::Level1, TransformerLevel::Level2, 21, + quantized_down ? 0.002 : 0.00001, quantized_down ? 0.001 : 0.00001, + std::make_unique(), [] { return DefaultWebGpuExecutionProvider(); }); + } +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionRejectsObservableIntermediate) { + ASSERT_STATUS_OK(TestGraphTransformer([](ModelTestBuilder& b) { BuildQLoRA(b, 1, 2048, 128, 1, true); }, 21, *logger_, + std::make_unique(), TransformerLevel::Level2, 1, + nullptr, [](const Graph& g) { return CheckQLoRACount(g, 0); })); +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionFp32Accumulation) { + if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; + for (int64_t m : {1, 2, 37}) { + SCOPED_TRACE(m); + // FP32 accumulation still produces FP16 intermediates and output, so + // fused and unfused execution need not be bitwise identical. + RunWebGpuFusionTransformerTest( + [m](ModelTestBuilder& b) { BuildQLoRA(b, m, 2048, 132, 1, false, false, true); }, + [](InferenceSessionWrapper& s) { ASSERT_STATUS_OK(CheckQLoRACount(s.GetGraph(), 1)); }, + TransformerLevel::Level1, TransformerLevel::Level2, 21, + 0.002, 0.001, std::make_unique(), + [] { + ConfigOptions options; + ORT_THROW_IF_ERROR(options.AddConfigEntry(webgpu::options::kEnableMatmulFp32Accumulation, + webgpu::options::kEnableMatmulFp32Accumulation_ON)); + return WebGpuExecutionProviderWithOptions(options); + }); + } +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionKeepsNonzeroOrSharedDefaults) { + for (bool shared : {false, true}) { + ASSERT_STATUS_OK(TestGraphTransformer( + [shared](ModelTestBuilder& b) { + BuildQLoRA(b); + const NodeArg* input = nullptr; + for (const auto& [name, tensor] : b.graph_.GetAllInitializedTensors()) { + if (tensor->data_type() == ONNX_NAMESPACE::TensorProto::UINT8 && tensor->dims(0) == 32) { + input = b.graph_.GetNodeArg(name); + break; + } + } + ORT_ENFORCE(input != nullptr); + if (shared) { + auto* observable = b.MakeOutput(std::vector{32, 64, 32}); + b.AddNode("Identity", {const_cast(input)}, {observable}); + } else { + const ONNX_NAMESPACE::TensorProto* original = nullptr; + ORT_ENFORCE(b.graph_.GetInitializedTensor(input->Name(), original)); + auto replacement = *original; + replacement.mutable_raw_data()->at(0) = static_cast(129); + b.graph_.RemoveInitializedTensor(input->Name()); + b.graph_.AddInitializedTensor(replacement); + } + }, + 21, *logger_, std::make_unique(), TransformerLevel::Level2, 1, nullptr, + [](const Graph& graph) { + ORT_RETURN_IF_ERROR(CheckQLoRACount(graph, 1)); + for (const auto* input : graph.GetOverridableInitializers()) { + const ONNX_NAMESPACE::TensorProto* init = nullptr; + ORT_RETURN_IF_NOT(graph.GetInitializedTensor(input->Name(), init), "Missing default."); + for (const auto dim : init->dims()) + ORT_RETURN_IF_NOT(dim > 0, "Observable or nonzero default was emptied."); + } + return Status::OK(); + })); + } +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionRejectsDifferentInput) { + ASSERT_STATUS_OK(TestGraphTransformer([](ModelTestBuilder& b) { BuildQLoRA(b, 1, 2048, 128, 1, false, true); }, 21, + *logger_, std::make_unique(), TransformerLevel::Level2, + 1, nullptr, [](const Graph& g) { return CheckQLoRACount(g, 0); })); +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionRejectsDifferentAdapterAccuracy) { + ASSERT_STATUS_OK( + TestGraphTransformer([](ModelTestBuilder& b) { BuildQLoRA(b, 1, 2048, 128, 1, false, false, false, 4); }, 21, + *logger_, std::make_unique(), TransformerLevel::Level2, 1, nullptr, + [](const Graph& g) { return CheckQLoRACount(g, 0); })); +} + +template +void RunQLoRASerializedDefaults(const logging::Logger& logger, int64_t k = 64, bool raw_empty_defaults = false) { + if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU unavailable."; + for (int64_t m : {1, 2, 37}) { + SCOPED_TRACE(m); + Model model("quantized LoRA defaults", false, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, logger); + auto& graph = model.MainGraph(); + ModelTestBuilder builder(graph); + // A random base can round differently when the subgroup kernel accumulates + // adapter-up into it. Keep this state/capture test exact; Numerics covers rounding. + BuildQLoRA(builder, m, k, 128, 1, false, false, false, 0, raw_empty_defaults); + builder.SetGraphOutputs(); + ASSERT_STATUS_OK(graph.Resolve()); + const auto adapters = graph.GetOverridableInitializers(); + ASSERT_EQ(adapters.size(), 4u); + std::string a_name, b_name; + // Dyadic inputs keep the switching checks exact on all reduction paths. + for (auto& [name, value] : builder.feeds_) { + Tensor* tensor = value.template GetMutable(); + if (tensor->IsDataType()) { + auto bytes = tensor->MutableDataAsSpan(); + std::fill(bytes.begin(), bytes.end(), uint8_t{129}); + (tensor->Shape()[0] == 32 ? a_name : b_name) = name; + } else { + auto values = tensor->MutableDataAsSpan(); + const float fill = tensor->Shape().NumDimensions() == 3 ? 0.25f : 0.03125f; + std::fill(values.begin(), values.end(), T(fill)); + } + } + ASSERT_FALSE(a_name.empty()); + ASSERT_FALSE(b_name.empty()); + std::vector expected(7); + // Compare the standard optimizer opt-out, automatic fusion, and serialized + // fusion using the same partial-override states and exact reference values. + for (int mode : {0, 1, 2}) { + const bool fused = mode != 0; + if (mode == 2) { + bool modified = false; + ASSERT_STATUS_OK(MatMulNBitsQLoRAFusion().Apply(graph, modified, logger)); + ASSERT_TRUE(modified); + ASSERT_STATUS_OK(graph.Resolve()); + } + const auto proto = model.ToProto(); + if (mode == 2) { + ASSERT_EQ(proto.graph().node_size(), 1); + ASSERT_EQ(proto.graph().node(0).op_type(), "MatMulNBitsQLoRA"); + int packed_count = 0; + for (const auto& init : proto.graph().initializer()) { + if (init.name() != a_name && init.name() != b_name) continue; + ++packed_count; + ASSERT_EQ(init.dims_size(), 3); + const int axis = init.name() == a_name ? 0 : 1; + EXPECT_EQ(init.dims(axis), raw_empty_defaults ? 0 : (axis == 0 ? 32 : 1)); + const int64_t size = init.dims(0) * init.dims(1) * init.dims(2); + EXPECT_EQ(init.raw_data(), std::string(static_cast(size), static_cast(128))); + EXPECT_EQ(init.int32_data_size(), 0); + EXPECT_EQ(init.external_data_size(), 0); + } + ASSERT_EQ(packed_count, 2); + } + std::string data; + ASSERT_TRUE(proto.SerializeToString(&data)); + SessionOptions options; + options.graph_optimization_level = mode == 2 ? TransformerLevel::Level1 : TransformerLevel::Level2; + if (mode == 0) { + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsDisableSpecifiedOptimizers, + "MatMulNBitsQLoRAFusion")); + } + ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.disable_cpu_ep_fallback", "1")); + InferenceSessionWrapper session(options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(DefaultWebGpuExecutionProvider())); + ASSERT_STATUS_OK(session.Load(data.data(), static_cast(data.size()))); + // Mode 2 loads the serialized operator without a LoRA transformer. + ASSERT_STATUS_OK(session.Initialize()); + ASSERT_STATUS_OK(CheckQLoRACount(session.GetGraph(), fused ? 1 : 0)); + // Off, A only, B only, scales only, packed weights only, all, off again. + for (int state = 0; state < 7; ++state) { + // Empty scale defaults cannot be used with nonempty packed overrides. + if (raw_empty_defaults && (state == 1 || state == 2 || state == 4)) continue; + SCOPED_TRACE(state); + auto feeds = builder.feeds_; + for (const auto* input : adapters) { + const auto& name = input->Name(); + const bool packed = name == a_name || name == b_name; + const bool supply = (state == 1 && name == a_name) || (state == 2 && name == b_name) || + (state == 3 && !packed) || (state == 4 && packed) || state == 5; + if (!supply) feeds.erase(name); + } + std::vector output; + ASSERT_STATUS_OK(session.Run(RunOptions{}, feeds, builder.output_names_, &output)); + ASSERT_EQ(output.size(), 1u); + // X=1/4 and base weights=1/16 give K/64. Unit adapter weights give + // 8*K with default scales, or K/128 when both scales are 1/32. + const float base_value = static_cast(k) / 64.0f; + const float delta = state == 4 ? 8.0f * static_cast(k) + : (state == 5 ? static_cast(k) / 128.0f : 0.0f); + const auto values = output[0].Get().DataAsSpan(); + EXPECT_TRUE(std::all_of(values.begin(), values.end(), + [base_value, delta](T value) { return value == T(base_value + delta); })) + << "Expected every output to equal " << base_value + delta; + if (!fused) { + expected[state] = std::move(output[0]); + } else { + const auto comparison = CompareOrtValue(output[0], expected[state], 0.0, 0.0, false); + EXPECT_EQ(comparison.first, COMPARE_RESULT::SUCCESS) << comparison.second; + } + } + } + EXPECT_NE(CompareOrtValue(expected[0], expected[5], 0.0, 0.0, false).first, COMPARE_RESULT::SUCCESS); + + ConfigOptions config; + ASSERT_STATUS_OK(config.AddConfigEntry(webgpu::options::kEnableGraphCapture, + webgpu::options::kEnableGraphCapture_ON)); + auto ep = WebGpuExecutionProviderWithOptions(config); + ASSERT_NE(ep, nullptr); + auto* capture_provider = ep.get(); + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Level1; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry("session.disable_cpu_ep_fallback", "1")); + InferenceSession session(options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(std::move(ep))); + const auto data = model.ToProto().SerializeAsString(); + ASSERT_STATUS_OK(session.Load(data.data(), static_cast(data.size()))); + ASSERT_STATUS_OK(session.Initialize()); + const auto allocators = capture_provider->CreatePreferredAllocators(); + ASSERT_FALSE(allocators.empty()); + const auto allocator = allocators.back(); + ASSERT_NE(allocator->Info().device.Type(), OrtDevice::CPU); + std::unique_ptr bindings[2]; + OrtValue outputs[2]; + for (int state : {0, 1}) { + ASSERT_STATUS_OK(session.NewIOBinding(&bindings[state])); + for (const auto& [name, value] : builder.feeds_) { + const bool adapter = std::any_of(adapters.begin(), adapters.end(), + [&name](const NodeArg* arg) { return arg->Name() == name; }); + if (state == 0 && adapter) continue; + const Tensor& source = value.template Get(); + OrtValue input; + Tensor::InitOrtValue(source.DataType(), source.Shape(), allocator, input); + ASSERT_STATUS_OK(session.GetDataTransferManager().CopyTensor(source, *input.GetMutable())); + ASSERT_STATUS_OK(bindings[state]->BindInput(name, input)); + } + const auto& reference = expected[state == 0 ? 0 : 5].Get(); + Tensor::InitOrtValue(reference.DataType(), reference.Shape(), allocator, outputs[state]); + ASSERT_STATUS_OK(bindings[state]->BindOutput(builder.output_names_[0], outputs[state])); + ASSERT_STATUS_OK(bindings[state]->SynchronizeInputs()); + } + OrtValue actual; + const auto& reference = expected[0].Get(); + Tensor::InitOrtValue(reference.DataType(), reference.Shape(), CPUAllocator::DefaultInstance(), actual); + // Each adapter state owns a capture ID and stable GPU bindings. Replay + // skips branch selection, so one ID cannot be reused across on/off states. + for (int state : {0, 0, 0, 1, 1, 1, 0, 1}) { + SCOPED_TRACE(MakeString("capture state=", state)); + auto* cpu_output = actual.GetMutable(); + auto values = cpu_output->MutableDataAsSpan(); + std::fill(values.begin(), values.end(), T(-99.0f)); + ASSERT_STATUS_OK(session.GetDataTransferManager().CopyTensor(*cpu_output, *outputs[state].GetMutable())); + RunOptions run_options; + ASSERT_STATUS_OK(run_options.config_options.AddConfigEntry(kOrtRunOptionsConfigCudaGraphAnnotation, + state == 0 ? "1" : "2")); + ASSERT_STATUS_OK(session.Run(run_options, *bindings[state])); + ASSERT_STATUS_OK(bindings[state]->SynchronizeOutputs()); + ASSERT_STATUS_OK(session.GetDataTransferManager().CopyTensor(outputs[state].Get(), *cpu_output)); + const auto comparison = CompareOrtValue(actual, expected[state == 0 ? 0 : 5], 0.0, 0.0, false); + EXPECT_EQ(comparison.first, COMPARE_RESULT::SUCCESS) << comparison.second; + } + EXPECT_TRUE(capture_provider->IsGraphCaptured(1)); + EXPECT_TRUE(capture_provider->IsGraphCaptured(2)); + } +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionSerializedDefaults) { + RunQLoRASerializedDefaults(*logger_); +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionSerializedDefaultsFp32) { + RunQLoRASerializedDefaults(*logger_); +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionJointDP4ACapture) { + RunQLoRASerializedDefaults(*logger_, 2048); +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionFromEmptyRawDefaults) { + RunQLoRASerializedDefaults(*logger_, 64, true); + RunQLoRASerializedDefaults(*logger_, 64, true); +} + +TEST_F(GraphTransformationTests, MatMulNBitsQLoRAFusionRejectsInvalidDefaultRank) { + ASSERT_STATUS_OK(TestGraphTransformer( + [](ModelTestBuilder& b) { + BuildQLoRA(b, 1, 64, 128, 1, false, false, false, 0, true); + auto* arg = b.graph_.GetInputsIncludingInitializers()[1]; + const ONNX_NAMESPACE::TensorProto* original = nullptr; + ORT_ENFORCE(b.graph_.GetInitializedTensor(arg->Name(), original)); + auto invalid = *original; + invalid.set_dims(0, 16); + invalid.set_raw_data(std::string(16 * 2 * 32, static_cast(128))); + b.graph_.RemoveInitializedTensor(arg->Name()); + b.graph_.AddInitializedTensor(invalid); + }, + 21, *logger_, std::make_unique(), TransformerLevel::Level2, 1, nullptr, + [](const Graph& g) { return CheckQLoRACount(g, 0); })); +} + +#endif +} // namespace onnxruntime::test diff --git a/onnxruntime/test/providers/webgpu/dispatch_normalization_test.cc b/onnxruntime/test/providers/webgpu/dispatch_normalization_test.cc index a197551015f3e..a77a7ae791c1d 100644 --- a/onnxruntime/test/providers/webgpu/dispatch_normalization_test.cc +++ b/onnxruntime/test/providers/webgpu/dispatch_normalization_test.cc @@ -26,6 +26,11 @@ #include "test/test_environment.h" #include "test/util/include/asserts.h" +#if !defined(DISABLE_CONTRIB_OPS) +#include "contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h" +#include "contrib_ops/webgpu/quantization/matmul_nbits.h" +#endif + namespace onnxruntime { namespace test { namespace { @@ -218,6 +223,116 @@ void CheckResultsAndPadding(gsl::span output_data, size_t output_si << "Normalization-added workgroups modified output padding."; } +#if !defined(DISABLE_CONTRIB_OPS) +void CheckQLoRADownPadding(WebGpuExecutionProvider& ep, bool dp4a, bool subgroups) { + auto& context = WebGpuContextFactory::GetContext(0); + const auto& adapter = context.AdapterInfo(); + const bool fixed_subgroup_32 = adapter.subgroupMinSize == 32 && adapter.subgroupMaxSize == 32; + if (subgroups && + (!context.DeviceHasFeature(wgpu::FeatureName::Subgroups) || + !context.Instance().HasWGSLLanguageFeature(wgpu::WGSLLanguageFeatureName::SubgroupId) || + (!fixed_subgroup_32 && + (!context.DeviceHasFeature(wgpu::FeatureName::SubgroupSizeControl) || + adapter.subgroupMinSize > 32 || adapter.subgroupMaxSize < 32)))) { + GTEST_SKIP() << "QLoRA subgroup decode requires 32-lane subgroups and subgroup_id."; + } + + constexpr uint32_t N = 24, K = 512, rank = 32, tile_n = 8; + const uint32_t batches = dp4a ? 2 : 1; + const uint32_t base_groups = batches * (N / tile_n); + const uint32_t down_groups = dp4a ? batches * 8 : 4; + const OrtMemoryInfo memory_info(WEBGPU_BUFFER, OrtDeviceAllocator, webgpu::WebGpuDevice(0), OrtMemTypeDefault); + // Exercise both cube and square normalization with small physical buffers. + for (uint32_t limit : {dp4a ? 3u : 2u, 5u}) { + SCOPED_TRACE(limit); + std::array physical{base_groups + down_groups, 1, 1}; + ASSERT_STATUS_OK(webgpu::detail::NormalizeDispatchGroupSize(physical[0], physical[1], physical[2], limit)); + const uint32_t physical_count = physical[0] * physical[1] * physical[2]; + ASSERT_GT(physical_count, base_groups + down_groups); + const uint32_t padded_batches = dp4a ? CeilDiv(physical_count - base_groups, 8u) : 1; + const uint32_t padded_ranks = dp4a ? rank : (physical_count - base_groups) * 8; + // Oversize the backing buffers so an unguarded shader writes the canary + // instead of accessing unallocated GPU memory, even with robustness enabled. + std::vector x_data(padded_batches * K, 0.5f); + std::vector xq_data(padded_batches * K / 4, 0x40404040u); + std::vector xs_data(padded_batches * K / 128, 1.0f / 128.0f); + std::vector w_data(N * K / 2, 0x99); + std::vector s_data(N * K / 32, 0.03125f); + std::vector aq_data(padded_ranks * K, 129); + std::vector as_data(padded_ranks * K / 32, 0.0625f); + std::vector output_data(batches * N + rank, kOutputPaddingCanary); + std::vector down_data(padded_batches * padded_ranks, kOutputPaddingCanary); + auto x_buffer = CreateStorageBuffer(ep, x_data); + auto xq_buffer = CreateStorageBuffer(ep, xq_data); + auto xs_buffer = CreateStorageBuffer(ep, xs_data); + auto w_buffer = CreateStorageBuffer(ep, w_data); + auto s_buffer = CreateStorageBuffer(ep, s_data); + auto aq_buffer = CreateStorageBuffer(ep, aq_data); + auto as_buffer = CreateStorageBuffer(ep, as_data); + auto output_buffer = CreateStorageBuffer(ep, output_data); + auto down_buffer = CreateStorageBuffer(ep, down_data); + Tensor x(DataTypeImpl::GetType(), {batches, 1, K}, x_buffer.Get(), memory_info); + Tensor xq(DataTypeImpl::GetType(), {batches, 1, K / 4}, xq_buffer.Get(), memory_info); + Tensor xs(DataTypeImpl::GetType(), {batches, 1, 1, K / 128}, xs_buffer.Get(), memory_info); + Tensor w(DataTypeImpl::GetType(), {N, K / 32, 16}, w_buffer.Get(), memory_info); + Tensor s(DataTypeImpl::GetType(), {N, K / 32}, s_buffer.Get(), memory_info); + Tensor aq(DataTypeImpl::GetType(), {rank, K / 32, 32}, aq_buffer.Get(), memory_info); + Tensor as(DataTypeImpl::GetType(), {rank, K / 32}, as_buffer.Get(), memory_info); + Tensor output(DataTypeImpl::GetType(), {batches, 1, N}, output_buffer.Get(), memory_info); + Tensor down(DataTypeImpl::GetType(), {batches, 1, rank}, down_buffer.Get(), memory_info); + if (dp4a) { + contrib::webgpu::DP4AMatMulNBitsKParallelProgram program{ + 32, tile_n, 4, false, false, false, false, false, false, true, false, true}; + program.AddInputs({{&xq, ProgramTensorMetadataDependency::TypeAndRank, 4}, + {&xs, ProgramTensorMetadataDependency::TypeAndRank}, + {&w, ProgramTensorMetadataDependency::TypeAndRank, 16}, + {&s, ProgramTensorMetadataDependency::TypeAndRank}, + {&aq, ProgramTensorMetadataDependency::Type, 16}, + {&as, ProgramTensorMetadataDependency::Type}}) + .AddOutputs({{&output, ProgramTensorMetadataDependency::TypeAndRank}, + {&down, ProgramTensorMetadataDependency::Type}}) + .AddUniformVariables({{batches}, {1u}, {N}, {K}, {K / 16}, {K / 32}, {32u}, {N / tile_n}, {0u}, {0u}, {1u}}) + .SetWorkgroupSize(128) + .SetDispatchGroupSize(physical[0], physical[1], physical[2]); + ASSERT_NO_FATAL_FAILURE(RunAndReadProgram(ep, program, output_buffer.Get(), output_data)); + } else { + contrib::webgpu::MatMulNBitsProgram program{ + tile_n, 4, false, false, false, false, false, 32, false, false, false, true, subgroups}; + program.AddInputs({{&x, ProgramTensorMetadataDependency::TypeAndRank, 4}, + {&w, ProgramTensorMetadataDependency::TypeAndRank, 16}, + {&s, ProgramTensorMetadataDependency::TypeAndRank}, + {&aq, ProgramTensorMetadataDependency::Type, 16}, + {&as, ProgramTensorMetadataDependency::Type}}) + .AddOutputs({{&output, ProgramTensorMetadataDependency::TypeAndRank}, + {&down, ProgramTensorMetadataDependency::Type}}) + .AddUniformVariables({{1u}, {N}, {K}, {K / 4}, {K / 32}, {32u}, {K / 32}, {0u}, {N / tile_n}, {1u}, {0u}, {1u}}) + .SetWorkgroupSize(128) + .SetUseSubgroupId(subgroups) + .SetDispatchGroupSize(physical[0], physical[1], physical[2]); + if (subgroups && !fixed_subgroup_32) { + program.SetSubgroupSize(32); + } + ASSERT_NO_FATAL_FAILURE(RunAndReadProgram(ep, program, output_buffer.Get(), output_data)); + } + ep.BufferManager().Download(ep.Recording(), down_buffer.Get(), down_data.data(), down_data.size() * sizeof(float)); + CheckResultsAndPadding(output_data, batches * N, 8.0f); + CheckResultsAndPadding(down_data, batches * rank, 16.0f); + } +} + +TEST_F(WebGpuDispatchExecutionTest, QLoRASharedDownPreservesPaddingAfterNormalization) { + CheckQLoRADownPadding(Ep(), false, false); +} + +TEST_F(WebGpuDispatchExecutionTest, QLoRASubgroupDownPreservesPaddingAfterNormalization) { + CheckQLoRADownPadding(Ep(), false, true); +} + +TEST_F(WebGpuDispatchExecutionTest, QLoRADP4ADownPreservesPaddingAfterNormalization) { + CheckQLoRADownPadding(Ep(), true, false); +} +#endif // !defined(DISABLE_CONTRIB_OPS) + class LogicalDispatchProbeProgram final : public Program { public: LogicalDispatchProbeProgram() : Program{"LogicalDispatchProbe"} {}