From 1c86c61fec4c963ae798c26db048c8702ac31f19 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Thu, 5 Nov 2026 13:56:26 +0800 Subject: [PATCH 1/4] [WebGPU] Reuse MatMulNBits MLP and QKV for runtime QLoRA adapters --- docs/ContribOperators.md | 62 +++++- docs/OperatorKernels.md | 4 +- .../webgpu/quantization/matmul_nbits_mlp.cc | 80 ++++++- .../webgpu/quantization/matmul_nbits_mlp.h | 17 +- .../webgpu/quantization/matmul_nbits_qkv.cc | 75 ++++++- .../webgpu/quantization/matmul_nbits_qkv.h | 17 +- .../webgpu/quantization/matmul_nbits_qlora.cc | 73 ++++--- .../webgpu/quantization/matmul_nbits_qlora.h | 10 + .../core/graph/contrib_ops/contrib_defs.cc | 48 ++++- .../core/optimizer/graph_transformer_utils.cc | 25 ++- .../core/optimizer/matmul_nbits_mlp_fusion.cc | 45 +++- .../core/optimizer/matmul_nbits_mlp_fusion.h | 8 +- .../core/optimizer/matmul_nbits_qkv_fusion.cc | 42 +++- .../core/optimizer/matmul_nbits_qkv_fusion.h | 8 +- .../contrib_ops/matmul_nbits_qlora_test.cc | 65 ++++++ .../optimizer/matmul_nbits_mlp_fusion_test.cc | 121 ++++++++++- .../optimizer/matmul_nbits_qkv_fusion_test.cc | 181 ++++++++++++++-- .../matmul_nbits_qlora_fused_switch_test.cc | 197 ++++++++++++++++++ 18 files changed, 997 insertions(+), 81 deletions(-) create mode 100644 onnxruntime/test/optimizer/matmul_nbits_qlora_fused_switch_test.cc diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 102e6d0653d13..bc8da37eb8745 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -4444,6 +4444,10 @@ This version of the operator has been available since version 1 of the 'com.micr derived from the runtime shape of A and the shared attributes K and N. The operator contract includes a string attribute describing the fused gate activation. + + Optional gate and up LoRA inputs apply MatMulNBitsQLoRA to each projection before the + activation. Empty adapter-up weights use the base projection. The adapter inputs are + runtime tensors and may change between calls. When fused from SkipSimplifiedLayerNormalization, the optional residual-sum output may also be materialized: @@ -4473,9 +4477,13 @@ This version of the operator has been available since version 1 of the 'com.micr
Size of each quantization block along the K dimension. Currently only block_size=32 is supported by the WebGPU kernel.
epsilon : float
Epsilon used by the optional fused (Skip)SimplifiedLayerNormalization. Defaults to 1e-5.
+
lora_bits : int
+
Bit-width of optional adapter weights.
+
lora_block_size : int
+
Quantization block size of optional adapter weights.
-#### Inputs (8 - 9) +#### Inputs (8 - 17)
A : T1
@@ -4496,6 +4504,22 @@ This version of the operator has been available since version 1 of the 'com.micr
Per-block scaling factors for the up projection.
up_bias (optional) : T1
Optional bias for the up projection with shape [N].
+
gate_lora_A (optional) : TAdapter
+
Optional packed gate adapter-down weights.
+
gate_lora_A_scales (optional) : T1
+
Optional gate adapter-down scales.
+
gate_lora_B (optional) : TAdapter
+
Optional packed gate adapter-up weights; may be empty.
+
gate_lora_B_scales (optional) : T1
+
Optional gate adapter-up scales.
+
up_lora_A (optional) : TAdapter
+
Optional packed up adapter-down weights.
+
up_lora_A_scales (optional) : T1
+
Optional up adapter-down scales.
+
up_lora_B (optional) : TAdapter
+
Optional packed up adapter-up weights; may be empty.
+
up_lora_B_scales (optional) : T1
+
Optional up adapter-up scales.
#### Outputs (1 - 2) @@ -4514,6 +4538,8 @@ This version of the operator has been available since version 1 of the 'com.micr
Constrain input and output types to float tensors.
T2 : tensor(uint8)
Constrain quantized weight types to uint8.
+
TAdapter : tensor(uint8)
+
Constrain packed adapter weights to uint8.
@@ -4595,6 +4621,8 @@ This version of the operator has been available since version 1 of the 'com.micr and may also return the input+skip residual sum as output 3. This operator is intended as a decode-oriented QKV fusion primitive. + Optional Q, K, and V LoRA inputs apply MatMulNBitsQLoRA to each projection. Empty + adapter-up weights use the base projection, and runtime adapters may change between calls. #### Version @@ -4617,9 +4645,13 @@ This version of the operator has been available since version 1 of the 'com.micr
Size of each quantization block along the K dimension. Currently only block_size=32 is supported by the WebGPU kernel.
epsilon : float
Epsilon used by the simplified layer norm reduction.
+
lora_bits : int
+
Bit-width of optional adapter weights.
+
lora_block_size : int
+
Quantization block size of optional adapter weights.
-#### Inputs (11 - 12) +#### Inputs (11 - 24)
A : T1
@@ -4646,6 +4678,30 @@ This version of the operator has been available since version 1 of the 'com.micr
Per-block scaling factors for the V projection.
v_bias (optional) : T1
Optional bias for the V projection with shape [Nkv].
+
q_lora_A (optional) : TAdapter
+
Optional packed Q adapter-down weights.
+
q_lora_A_scales (optional) : T1
+
Optional Q adapter-down scales.
+
q_lora_B (optional) : TAdapter
+
Optional packed Q adapter-up weights; may be empty.
+
q_lora_B_scales (optional) : T1
+
Optional Q adapter-up scales.
+
k_lora_A (optional) : TAdapter
+
Optional packed K adapter-down weights.
+
k_lora_A_scales (optional) : T1
+
Optional K adapter-down scales.
+
k_lora_B (optional) : TAdapter
+
Optional packed K adapter-up weights; may be empty.
+
k_lora_B_scales (optional) : T1
+
Optional K adapter-up scales.
+
v_lora_A (optional) : TAdapter
+
Optional packed V adapter-down weights.
+
v_lora_A_scales (optional) : T1
+
Optional V adapter-down scales.
+
v_lora_B (optional) : TAdapter
+
Optional packed V adapter-up weights; may be empty.
+
v_lora_B_scales (optional) : T1
+
Optional V adapter-up scales.
#### Outputs (3 - 4) @@ -4668,6 +4724,8 @@ This version of the operator has been available since version 1 of the 'com.micr
Constrain input and output types to float tensors.
T2 : tensor(uint8)
Constrain quantized weight types to uint8.
+
TAdapter : tensor(uint8)
+
Constrain packed adapter weights to uint8.
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index ee0f01826c69b..c0e1da179691d 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -2056,9 +2056,9 @@ The **OpSet Version** column uses the following notation: |MatMulBlockQuantizedFp8Weight|*in* A:**T**
*in* B:**T1**
*in* b_scale:**T2**
*in* a_scale:**T2**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(float16)
**T1** = tensor(float8e4m3fn)
**T2** = tensor(float)| |MatMulBnb4|*in* A:**T1**
*in* B:**T2**
*in* absmax:**T1**
*out* Y:**T1**|1+|**T1** = tensor(float), tensor(float16)
**T2** = tensor(uint8)| |MatMulNBits|*in* A:**T1**
*in* B:**T2**
*in* scales:**T1**
*in* zero_points:**T3**
*in* g_idx:**T4**
*in* bias:**T1**
*out* Y:**T1**|1+|**T1** = tensor(float), tensor(float16)
**T2** = tensor(uint8)
**T3** = tensor(uint8)
**T4** = tensor(int32)| -|MatMulNBitsMlp|*in* A:**T1**
*in* skip:**T1**
*in* norm_scale:**T1**
*in* gate_B:**T2**
*in* gate_scales:**T1**
*in* gate_bias:**T1**
*in* up_B:**T2**
*in* up_scales:**T1**
*in* up_bias:**T1**
*out* Y:**T1**
*out* input_skip_bias_sum:**T1**|1+|**T1** = tensor(float), tensor(float16)
**T2** = tensor(uint8)| +|MatMulNBitsMlp|*in* A:**T1**
*in* skip:**T1**
*in* norm_scale:**T1**
*in* gate_B:**T2**
*in* gate_scales:**T1**
*in* gate_bias:**T1**
*in* up_B:**T2**
*in* up_scales:**T1**
*in* up_bias:**T1**
*in* gate_lora_A:**TAdapter**
*in* gate_lora_A_scales:**T1**
*in* gate_lora_B:**TAdapter**
*in* gate_lora_B_scales:**T1**
*in* up_lora_A:**TAdapter**
*in* up_lora_A_scales:**T1**
*in* up_lora_B:**TAdapter**
*in* up_lora_B_scales:**T1**
*out* Y:**T1**
*out* input_skip_bias_sum:**T1**|1+|**T1** = tensor(float), tensor(float16)
**T2** = tensor(uint8)
**TAdapter** = tensor(uint8)| |MatMulNBitsQLoRA|*in* X:**T**
*in* W:**TQ**
*in* scales:**T**
*in* lora_A:**TQ**
*in* lora_A_scales:**T**
*in* lora_B:**TQ**
*in* lora_B_scales:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)
**TQ** = tensor(uint8)| -|MatMulNBitsQkv|*in* A:**T1**
*in* skip:**T1**
*in* norm_scale:**T1**
*in* q_B:**T2**
*in* q_scales:**T1**
*in* q_bias:**T1**
*in* k_B:**T2**
*in* k_scales:**T1**
*in* k_bias:**T1**
*in* v_B:**T2**
*in* v_scales:**T1**
*in* v_bias:**T1**
*out* Q:**T1**
*out* K:**T1**
*out* V:**T1**
*out* input_skip_bias_sum:**T1**|1+|**T1** = tensor(float), tensor(float16)
**T2** = tensor(uint8)| +|MatMulNBitsQkv|*in* A:**T1**
*in* skip:**T1**
*in* norm_scale:**T1**
*in* q_B:**T2**
*in* q_scales:**T1**
*in* q_bias:**T1**
*in* k_B:**T2**
*in* k_scales:**T1**
*in* k_bias:**T1**
*in* v_B:**T2**
*in* v_scales:**T1**
*in* v_bias:**T1**
*in* q_lora_A:**TAdapter**
*in* q_lora_A_scales:**T1**
*in* q_lora_B:**TAdapter**
*in* q_lora_B_scales:**T1**
*in* k_lora_A:**TAdapter**
*in* k_lora_A_scales:**T1**
*in* k_lora_B:**TAdapter**
*in* k_lora_B_scales:**T1**
*in* v_lora_A:**TAdapter**
*in* v_lora_A_scales:**T1**
*in* v_lora_B:**TAdapter**
*in* v_lora_B_scales:**T1**
*out* Q:**T1**
*out* K:**T1**
*out* V:**T1**
*out* input_skip_bias_sum:**T1**|1+|**T1** = tensor(float), tensor(float16)
**T2** = tensor(uint8)
**TAdapter** = tensor(uint8)| |MoE|*in* input:**T**
*in* router_probs:**T**
*in* fc1_experts_weights:**T**
*in* fc1_experts_bias:**T**
*in* fc2_experts_weights:**T**
*in* fc2_experts_bias:**T**
*in* fc3_experts_weights:**T**
*in* fc3_experts_bias:**T**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| |MultiHeadAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* bias:**T**
*in* key_padding_mask:**M**
*in* attention_bias:**T**
*in* past_key:**T**
*in* past_value:**T**
*in* past_sequence_length:**M**
*in* cache_indirection:**M**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**
*out* qk:**QK**|1+|**T** = tensor(float), tensor(float16)| |NGramHashMapping|*in* input_ids:**M**
*in* multipliers:**M**
*in* vocab_sizes:**M**
*in* past_ids:**M**
*in* head_offsets:**M**
*in* eos_token_id:**M**
*in* segment_ids:**tensor(int32)**
*out* hash_ids:**M**
*out* present_ids:**M**|1+|**M** = tensor(int32)| diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.cc index ab500c7742bb4..323359ef016f1 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.cc @@ -6,6 +6,7 @@ #include #include "contrib_ops/webgpu/quantization/matmul_nbits.h" +#include "contrib_ops/webgpu/quantization/matmul_nbits_qlora.h" #include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" #include "contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h" #include "contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h" @@ -240,10 +241,18 @@ ONNX_OPERATOR_KERNEL_EX( kWebGpuExecutionProvider, (*KernelDefBuilder::Create()) .TypeConstraint("T1", WebGpuSupportedFloatTypes()) - .TypeConstraint("T2", DataTypeImpl::GetTensorType()), + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) + .TypeConstraint("TAdapter", DataTypeImpl::GetTensorType()), MatMulNBitsMlp); Status MatMulNBitsMlp::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const { + if (has_adapter_inputs_) { + return ComputeWithAdapters(context); + } + return ComputeBase(context); +} + +Status MatMulNBitsMlp::ComputeBase(onnxruntime::webgpu::ComputeContext& context) const { const Tensor* a = context.Input(0); const Tensor* skip = context.Input(1); const Tensor* norm_scale = context.Input(2); @@ -542,6 +551,75 @@ Status MatMulNBitsMlp::ComputeInternal(onnxruntime::webgpu::ComputeContext& cont y); } +// Empty adapter B tensors reuse the base decode shader, including normalization. +Status MatMulNBitsMlp::ComputeWithAdapters(onnxruntime::webgpu::ComputeContext& context) const { + ORT_RETURN_IF_NOT(context.InputCount() == 17 && bits_ == 4 && block_size_ == 32 && + K_ > 0 && K_ % 32 == 0 && N_ > 0 && + lora_bits_ == 8 && lora_block_size_ == 32, + "WebGPU MatMulNBitsMlp requires 4-bit block-32 base and 8-bit block-32 adapters."); + ORT_RETURN_IF(context.Input(5) != nullptr || context.Input(8) != nullptr, + "MatMulNBitsMlp with adapters does not support projection biases."); + const Tensor* gate_lora_b = context.Input(11); + const Tensor* up_lora_b = context.Input(15); + if (gate_lora_b->Shape().Size() == 0 && up_lora_b->Shape().Size() == 0) { + ORT_RETURN_IF_ERROR(ValidateQLoRABaseWeights(context.Input(3), context.Input(4), K_, N_)); + ORT_RETURN_IF_ERROR(ValidateQLoRABaseWeights(context.Input(6), context.Input(7), K_, N_)); + return ComputeBase(context); + } + + const Tensor* a = context.Input(0); + const Tensor* skip = context.Input(1); + const Tensor* norm_scale = context.Input(2); + ORT_RETURN_IF(skip != nullptr && norm_scale == nullptr, + "MatMulNBitsMlp with adapters requires norm_scale when skip is present."); + ORT_RETURN_IF(norm_scale != nullptr && norm_scale->Shape().Size() != K_, + "MatMulNBitsMlp with adapters requires norm_scale to contain K elements."); + + MatMulComputeHelper helper; + ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), TensorShape({N_, K_}), false, true)); + Tensor* y = context.Output(0, helper.OutputShape()); + Tensor* skip_sum = skip != nullptr && context.OutputCount() > 1 ? context.Output(1, a->Shape()) : nullptr; + if (y->Shape().Size() == 0) return Status::OK(); + + std::optional normalized; + const Tensor* projection_input = a; + if (norm_scale != nullptr) { + normalized.emplace(context.CreateGPUTensor(a->DataType(), a->Shape())); + if (skip != nullptr) { + ORT_RETURN_IF_ERROR(RunSkipLayerNormProgram(context, a, skip, norm_scale, + nullptr, nullptr, epsilon_, true, &*normalized, skip_sum)); + } else { + const auto& shape = a->Shape(); + const int64_t norm_size = shape[shape.NumDimensions() - 1]; + const uint32_t norm_count = narrow(shape.Size() / norm_size); + ORT_RETURN_IF_ERROR(onnxruntime::webgpu::RunLayerNormProgram( + context, a, norm_scale, nullptr, epsilon_, norm_count, norm_size, + true, &*normalized, nullptr, nullptr)); + } + projection_input = &*normalized; + } + + Tensor gate = context.CreateGPUTensor(a->DataType(), helper.OutputShape()); + Tensor up = context.CreateGPUTensor(a->DataType(), helper.OutputShape()); + ORT_RETURN_IF_ERROR(ApplyMatMulNBitsQLoRA(projection_input, context.Input(3), context.Input(4), + context.Input(9), context.Input(10), + gate_lora_b, context.Input(12), + K_, N_, accuracy_level_, context, &gate)); + ORT_RETURN_IF_ERROR(ApplyMatMulNBitsQLoRA(projection_input, context.Input(6), context.Input(7), + context.Input(13), context.Input(14), + up_lora_b, context.Input(16), + K_, N_, accuracy_level_, context, &up)); + + const uint32_t vec_size = (narrow(y->Shape().Size()) + 3u) / 4u; + MatMulNBitsMlpProgram program{activation_kind_}; + program.AddInputs({{&gate, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten, 4}, + {&up, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten, 4}}) + .AddOutput({y, ProgramTensorMetadataDependency::Type, {vec_size}, 4}) + .SetDispatchGroupSize((vec_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) + .AddUniformVariables({vec_size}); + return context.RunProgram(program); +} + } // namespace webgpu } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.h index 002ebca4c0e54..4f2c038a4ea7a 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.h @@ -28,7 +28,7 @@ enum class MlpActivationKind : uint32_t { // Status for unsupported activations so the kernel rejects unknown values up front. Status ParseMlpActivation(std::string_view name, MlpActivationKind* out); -class MatMulNBitsMlp final : public WebGpuKernel { +class MatMulNBitsMlp : public WebGpuKernel { public: explicit MatMulNBitsMlp(const OpKernelInfo& info) : WebGpuKernel(info) { K_ = info.GetAttr("K"); @@ -37,6 +37,15 @@ class MatMulNBitsMlp final : public WebGpuKernel { bits_ = info.GetAttr("bits"); accuracy_level_ = info.GetAttrOrDefault("accuracy_level", 4); epsilon_ = info.GetAttrOrDefault("epsilon", 1e-5f); + lora_bits_ = info.GetAttrOrDefault("lora_bits", 8); + lora_block_size_ = info.GetAttrOrDefault("lora_block_size", 32); + const auto& inputs = info.node().InputDefs(); + auto input_exists = [&](size_t i) { return inputs[i] != nullptr && inputs[i]->Exists(); }; + has_adapter_inputs_ = inputs.size() > 9 && input_exists(9); + for (size_t i = 9; i < inputs.size(); ++i) { + ORT_ENFORCE(input_exists(i) == has_adapter_inputs_, + "MatMulNBitsMlp requires all adapter inputs when any is present."); + } std::string activation; ORT_ENFORCE(info.GetAttr("activation", &activation).IsOK(), "MatMulNBitsMlp requires the 'activation' attribute."); @@ -49,6 +58,9 @@ class MatMulNBitsMlp final : public WebGpuKernel { Status ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const override; private: + Status ComputeBase(onnxruntime::webgpu::ComputeContext& context) const; + Status ComputeWithAdapters(onnxruntime::webgpu::ComputeContext& context) const; + int64_t K_; int64_t N_; int64_t block_size_; @@ -56,6 +68,9 @@ class MatMulNBitsMlp final : public WebGpuKernel { int64_t bits_; float epsilon_; MlpActivationKind activation_kind_; + int64_t lora_bits_; + int64_t lora_block_size_; + bool has_adapter_inputs_; }; } // namespace webgpu diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.cc index 8c62a0d2c5fc3..36a859139ffb8 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.cc @@ -7,6 +7,7 @@ #include "contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h" #include "contrib_ops/webgpu/quantization/matmul_nbits.h" +#include "contrib_ops/webgpu/quantization/matmul_nbits_qlora.h" #include "contrib_ops/webgpu/quantization/matmul_nbits_common.h" #include "contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h" #include "contrib_ops/webgpu/bert/skip_layer_norm.h" @@ -140,10 +141,18 @@ ONNX_OPERATOR_KERNEL_EX( kWebGpuExecutionProvider, (*KernelDefBuilder::Create()) .TypeConstraint("T1", WebGpuSupportedFloatTypes()) - .TypeConstraint("T2", DataTypeImpl::GetTensorType()), + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) + .TypeConstraint("TAdapter", DataTypeImpl::GetTensorType()), MatMulNBitsQkv); Status MatMulNBitsQkv::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const { + if (has_adapter_inputs_) { + return ComputeWithAdapters(context); + } + return ComputeBase(context); +} + +Status MatMulNBitsQkv::ComputeBase(onnxruntime::webgpu::ComputeContext& context) const { const Tensor* a = context.Input(0); const Tensor* skip = context.Input(1); const Tensor* norm_scale = context.Input(2); @@ -389,6 +398,70 @@ Status MatMulNBitsQkv::ComputeInternal(onnxruntime::webgpu::ComputeContext& cont return context.RunProgram(program); } +// Reuse the base QKV decode shader when all three up-adapter tensors are empty. +Status MatMulNBitsQkv::ComputeWithAdapters(onnxruntime::webgpu::ComputeContext& context) const { + ORT_RETURN_IF_NOT(context.InputCount() == 24 && block_size_ == 32 && K_ > 0 && K_ % 32 == 0 && + Nq_ > 0 && Nkv_ > 0 && lora_bits_ == 8 && lora_block_size_ == 32, + "WebGPU MatMulNBitsQkv requires 4-bit block-32 base and 8-bit block-32 adapters."); + ORT_RETURN_IF(context.Input(5) != nullptr || context.Input(8) != nullptr || + context.Input(11) != nullptr, + "MatMulNBitsQkv with adapters does not support projection biases."); + if (context.Input(14)->Shape().Size() == 0 && + context.Input(18)->Shape().Size() == 0 && + context.Input(22)->Shape().Size() == 0) { + ORT_RETURN_IF_ERROR(ValidateQLoRABaseWeights(context.Input(3), context.Input(4), K_, Nq_)); + ORT_RETURN_IF_ERROR(ValidateQLoRABaseWeights(context.Input(6), context.Input(7), K_, Nkv_)); + ORT_RETURN_IF_ERROR(ValidateQLoRABaseWeights(context.Input(9), context.Input(10), K_, Nkv_)); + return ComputeBase(context); + } + + const Tensor* a = context.Input(0); + const Tensor* skip = context.Input(1); + const Tensor* norm_scale = context.Input(2); + ORT_RETURN_IF_NOT(norm_scale != nullptr, "MatMulNBitsQkv with adapters requires norm_scale."); + ORT_RETURN_IF_NOT(norm_scale->Shape().Size() == K_, + "MatMulNBitsQkv with adapters requires norm_scale to contain K elements."); + MatMulComputeHelper helper; + ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), TensorShape({Nq_, K_}), false, true)); + auto q_shape = helper.OutputShape(); + auto kv_dims = q_shape.AsShapeVector(); + kv_dims.back() = Nkv_; + TensorShape kv_shape(kv_dims); + Tensor* q = context.Output(0, q_shape); + Tensor* k = context.Output(1, kv_shape); + Tensor* v = context.Output(2, kv_shape); + Tensor* skip_sum = skip != nullptr && context.OutputCount() > 3 ? context.Output(3, a->Shape()) : nullptr; + if (q->Shape().Size() == 0) return Status::OK(); + + Tensor normalized = context.CreateGPUTensor(a->DataType(), a->Shape()); + if (skip != nullptr) { + ORT_RETURN_IF_ERROR(RunSkipLayerNormProgram(context, a, skip, norm_scale, + nullptr, nullptr, epsilon_, true, &normalized, skip_sum)); + } else { + const auto& shape = a->Shape(); + const int64_t norm_size = shape[shape.NumDimensions() - 1]; + const uint32_t norm_count = narrow(shape.Size() / norm_size); + ORT_RETURN_IF_ERROR(onnxruntime::webgpu::RunLayerNormProgram( + context, a, norm_scale, nullptr, epsilon_, norm_count, norm_size, + true, &normalized, nullptr, nullptr)); + } + + Tensor* outputs[] = {q, k, v}; + const int64_t sizes[] = {Nq_, Nkv_, Nkv_}; + const int base_slots[] = {3, 6, 9}; + const int adapter_slots[] = {12, 16, 20}; + for (int i = 0; i < 3; ++i) { + const int base = base_slots[i]; + const int adapter = adapter_slots[i]; + ORT_RETURN_IF_ERROR(ApplyMatMulNBitsQLoRA( + &normalized, context.Input(base), context.Input(base + 1), + context.Input(adapter), context.Input(adapter + 1), + context.Input(adapter + 2), context.Input(adapter + 3), + K_, sizes[i], accuracy_level_, context, outputs[i])); + } + return Status::OK(); +} + } // namespace webgpu } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.h index 4d57ab5ac2b3c..50637ee024e5d 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.h @@ -11,7 +11,7 @@ namespace webgpu { using namespace onnxruntime::webgpu; -class MatMulNBitsQkv final : public WebGpuKernel { +class MatMulNBitsQkv : public WebGpuKernel { public: explicit MatMulNBitsQkv(const OpKernelInfo& info) : WebGpuKernel(info) { K_ = info.GetAttr("K"); @@ -21,6 +21,15 @@ class MatMulNBitsQkv final : public WebGpuKernel { bits_ = info.GetAttr("bits"); accuracy_level_ = info.GetAttrOrDefault("accuracy_level", 4); epsilon_ = info.GetAttrOrDefault("epsilon", 1e-6f); + lora_bits_ = info.GetAttrOrDefault("lora_bits", 8); + lora_block_size_ = info.GetAttrOrDefault("lora_block_size", 32); + const auto& inputs = info.node().InputDefs(); + auto input_exists = [&](size_t i) { return inputs[i] != nullptr && inputs[i]->Exists(); }; + has_adapter_inputs_ = inputs.size() > 12 && input_exists(12); + for (size_t i = 12; i < inputs.size(); ++i) { + ORT_ENFORCE(input_exists(i) == has_adapter_inputs_, + "MatMulNBitsQkv requires all adapter inputs when any is present."); + } ORT_ENFORCE(bits_ == 4, "MatMulNBitsQkv currently supports 4-bit weights only."); } @@ -28,6 +37,9 @@ class MatMulNBitsQkv final : public WebGpuKernel { Status ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const override; private: + Status ComputeBase(onnxruntime::webgpu::ComputeContext& context) const; + Status ComputeWithAdapters(onnxruntime::webgpu::ComputeContext& context) const; + int64_t K_; int64_t Nq_; int64_t Nkv_; @@ -35,6 +47,9 @@ class MatMulNBitsQkv final : public WebGpuKernel { int64_t accuracy_level_; int64_t bits_; float epsilon_; + int64_t lora_bits_; + int64_t lora_block_size_; + bool has_adapter_inputs_; }; } // namespace webgpu diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc index 4c8c4c172a0d3..b2c119829b0bc 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.cc @@ -89,6 +89,14 @@ Status ApplyQLoRAOutput(onnxruntime::webgpu::ComputeContext& context, const Tens return context.RunProgram(program); } +Status ValidateQLoRABaseWeights(const Tensor* w, const Tensor* scales, int64_t k, int64_t n) { + ORT_RETURN_IF_NOT(w != nullptr && scales != nullptr && + w->Shape() == TensorShape({n, k / 32, 16}) && + scales->Shape().Size() == SafeInt(n) * (k / 32), + "Invalid base weight/scales shape."); + return Status::OK(); +} + class MatMulNBitsQLoRA final : public WebGpuKernel { public: ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(MatMulNBitsQLoRA); @@ -113,30 +121,36 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { const auto* as = context.Input(4); const auto* bq = context.Input(5); const auto* bs = context.Input(6); + MatMulComputeHelper helper; + ORT_RETURN_IF_ERROR(helper.Compute(x->Shape(), TensorShape({n_, k_}), false, true)); + auto* y = context.Output(0, helper.OutputShape()); + return ComputeProjection(x, w, s, aq, as, bq, bs, k_, n_, accuracy_level_, helper, context, y); + } + + static Status ComputeProjection(const Tensor* x, const Tensor* w, const Tensor* s, + const Tensor* aq, const Tensor* as, const Tensor* bq, const Tensor* bs, + int64_t k, int64_t n, int64_t accuracy_level, + const MatMulComputeHelper& helper, + onnxruntime::webgpu::ComputeContext& context, Tensor* y) { 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."); + ORT_RETURN_IF_ERROR(ValidateQLoRABaseWeights(w, s, k, n)); // 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, + 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 (blocks_b == 0) { - return ApplyMatMulNBits(x, w, s, nullptr, nullptr, k_, n_, 32, accuracy_level_, 4, context, y); + return ApplyMatMulNBits(x, w, s, nullptr, nullptr, k, n, 32, accuracy_level, 4, context, y); } - ORT_RETURN_IF_NOT((rank_a == 0 || as->Shape() == TensorShape({32, k_ / 32})) && - bs->Shape() == TensorShape({n_, 1}), + ORT_RETURN_IF_NOT((rank_a == 0 || 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). @@ -149,14 +163,14 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { 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); + 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(), + context, accuracy_level, 32, narrow(helper.OutputOffsets().size()), + narrow(n), narrow(k), 4, x->IsDataType(), base_subgroup_matrix_config, subgroup_m, false); if (rank_a != 0 && base_subgroup_matrix && base_subgroup_matrix_config->Is(8, 16, 16)) { const auto down_config = SelectSubgroupMatrixConfig( @@ -164,24 +178,24 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { 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_), + x, w, s, aq, as, bq, bs, subgroup_m, narrow(n), narrow(k), *base_subgroup_matrix_config, *down_config, context, &t, y); } } - if (rank_a != 0 && !fuse_up && m == 1 && x->IsDataType() && base_dp4a && k_ >= 2048 && k_ <= 8192) { + if (rank_a != 0 && !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, + 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 = rank_a != 0 && !fuse_up && x->Shape().Size() == k_ && k_ >= 2048 && k_ <= 8192 && !base_dp4a; + const bool joint_down = rank_a != 0 && !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, + 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); } @@ -190,19 +204,19 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { const auto down_accuracy = fuse_up && (m >= kMinMForTileOptimization || (x->IsDataType() && base_dp4a)) - ? accuracy_level_ + ? accuracy_level : 0; DP4AQuantizedInput quantized_x; // Empty down weights fill T with zeros. The active up projection must still // run: its dequantized weights may overflow, and 0 * Inf must not be skipped. - ORT_RETURN_IF_ERROR(ApplyMatMulNBits(x, aq, as, nullptr, nullptr, k_, 32, 32, down_accuracy, 8, context, &t, 0, + 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, + 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, + 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); } @@ -213,6 +227,15 @@ class MatMulNBitsQLoRA final : public WebGpuKernel { int64_t accuracy_level_; }; +Status ApplyMatMulNBitsQLoRA(const Tensor* x, const Tensor* w, const Tensor* s, + const Tensor* aq, const Tensor* as, const Tensor* bq, const Tensor* bs, + int64_t k, int64_t n, int64_t accuracy_level, + onnxruntime::webgpu::ComputeContext& context, Tensor* y) { + MatMulComputeHelper helper; + ORT_RETURN_IF_ERROR(helper.Compute(x->Shape(), TensorShape({n, k}), false, true)); + return MatMulNBitsQLoRA::ComputeProjection(x, w, s, aq, as, bq, bs, k, n, accuracy_level, helper, context, y); +} + ONNX_OPERATOR_KERNEL_EX(MatMulNBitsQLoRA, kMSDomain, 1, kWebGpuExecutionProvider, (*KernelDefBuilder::Create()) .TypeConstraint("T", WebGpuSupportedFloatTypes()) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h index 72cdea459d857..8b12dda721894 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qlora.h @@ -3,6 +3,8 @@ #pragma once +#include + #include "core/providers/webgpu/program.h" namespace onnxruntime::webgpu { @@ -38,4 +40,12 @@ Status GenerateQLoRAUp(onnxruntime::webgpu::ShaderHelper& shader, Status ApplyQLoRAOutput(onnxruntime::webgpu::ComputeContext& context, const Tensor* base, const Tensor* t, const Tensor* b, const Tensor* scales, Tensor* y); +// Match the packed base-weight contract checked by the standalone QLoRA kernel. +Status ValidateQLoRABaseWeights(const Tensor* w, const Tensor* scales, int64_t k, int64_t n); + +Status ApplyMatMulNBitsQLoRA(const Tensor* x, const Tensor* w, const Tensor* s, + const Tensor* aq, const Tensor* as, const Tensor* bq, const Tensor* bs, + int64_t k, int64_t n, int64_t accuracy_level, + onnxruntime::webgpu::ComputeContext& context, Tensor* y); + } // namespace onnxruntime::contrib::webgpu diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index f98390a7b5555..8973db11dab50 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -4409,6 +4409,10 @@ derived from the runtime shape of A and the shared attributes K and N. The operator contract includes a string attribute describing the fused gate activation. +Optional gate and up LoRA inputs apply MatMulNBitsQLoRA to each projection before the +activation. Empty adapter-up weights use the base projection. The adapter inputs are +runtime tensors and may change between calls. + When fused from SkipSimplifiedLayerNormalization, the optional residual-sum output may also be materialized: A_norm, input_skip_bias_sum = SkipSimplifiedLayerNormalization(A, skip, norm_scale, epsilon) @@ -4436,6 +4440,9 @@ When fused from SkipSimplifiedLayerNormalization, the optional residual-sum outp .Attr("epsilon", "Epsilon used by the optional fused (Skip)SimplifiedLayerNormalization. Defaults to 1e-5.", AttributeProto::FLOAT, 1e-5f) + .Attr("lora_bits", "Bit-width of optional adapter weights.", AttributeProto::INT, int64_t{8}) + .Attr("lora_block_size", "Quantization block size of optional adapter weights.", AttributeProto::INT, + int64_t{32}) .Input(0, "A", "The shared input tensor.", "T1") .Input(1, "skip", "Optional skip input used by SkipSimplifiedLayerNormalization.", "T1", OpSchema::Optional) .Input(2, "norm_scale", "Optional RMSNorm scale with shape [K] used by SimplifiedLayerNormalization or SkipSimplifiedLayerNormalization.", "T1", OpSchema::Optional) @@ -4445,12 +4452,27 @@ When fused from SkipSimplifiedLayerNormalization, the optional residual-sum outp .Input(6, "up_B", "Packed uint8 tensor for the up projection weights.", "T2") .Input(7, "up_scales", "Per-block scaling factors for the up projection.", "T1") .Input(8, "up_bias", "Optional bias for the up projection with shape [N].", "T1", OpSchema::Optional) + .Input(9, "gate_lora_A", "Optional packed gate adapter-down weights.", "TAdapter", OpSchema::Optional) + .Input(10, "gate_lora_A_scales", "Optional gate adapter-down scales.", "T1", OpSchema::Optional) + .Input(11, "gate_lora_B", "Optional packed gate adapter-up weights; may be empty.", "TAdapter", OpSchema::Optional) + .Input(12, "gate_lora_B_scales", "Optional gate adapter-up scales.", "T1", OpSchema::Optional) + .Input(13, "up_lora_A", "Optional packed up adapter-down weights.", "TAdapter", OpSchema::Optional) + .Input(14, "up_lora_A_scales", "Optional up adapter-down scales.", "T1", OpSchema::Optional) + .Input(15, "up_lora_B", "Optional packed up adapter-up weights; may be empty.", "TAdapter", OpSchema::Optional) + .Input(16, "up_lora_B_scales", "Optional up adapter-up scales.", "T1", OpSchema::Optional) .Output(0, "Y", "The fused gated MLP output tensor.", "T1") .Output(1, "input_skip_bias_sum", "Optional residual-sum output for SkipSimplifiedLayerNormalization.", "T1", OpSchema::Optional) .TypeConstraint("T1", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") .TypeConstraint("T2", {"tensor(uint8)"}, "Constrain quantized weight types to uint8.") + .TypeConstraint("TAdapter", {"tensor(uint8)"}, "Constrain packed adapter weights to uint8.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + const bool has_adapter_inputs = ctx.hasInput(9); + for (int i = 9; i < 17; ++i) { + if (ctx.hasInput(i) != has_adapter_inputs) { + fail_shape_inference("all gate and up adapter inputs must be present together"); + } + } propagateElemTypeFromInputToOutput(ctx, 0, 0); if (ctx.getNumOutputs() > 1) { propagateElemTypeFromInputToOutput(ctx, 0, 1); @@ -4522,6 +4544,8 @@ If skip is provided, the operator computes the SkipSimplifiedLayerNormalization and may also return the input+skip residual sum as output 3. This operator is intended as a decode-oriented QKV fusion primitive. +Optional Q, K, and V LoRA inputs apply MatMulNBitsQLoRA to each projection. Empty +adapter-up weights use the base projection, and runtime adapters may change between calls. )DOC"; ONNX_CONTRIB_OPERATOR_SCHEMA(MatMulNBitsQkv) @@ -4539,6 +4563,9 @@ This operator is intended as a decode-oriented QKV fusion primitive. "The minimum accuracy level of input A. It follows the same semantics as MatMulNBits.", AttributeProto::INT, static_cast(0)) .Attr("epsilon", "Epsilon used by the simplified layer norm reduction.", AttributeProto::FLOAT, 1e-6f) + .Attr("lora_bits", "Bit-width of optional adapter weights.", AttributeProto::INT, int64_t{8}) + .Attr("lora_block_size", "Quantization block size of optional adapter weights.", AttributeProto::INT, + int64_t{32}) .Input(0, "A", "The shared input tensor.", "T1") .Input(1, "skip", "Optional residual input for SkipSimplifiedLayerNormalization.", "T1", OpSchema::Optional) .Input(2, "norm_scale", "Scale input for the simplified layer norm with shape [K].", "T1") @@ -4551,6 +4578,18 @@ This operator is intended as a decode-oriented QKV fusion primitive. .Input(9, "v_B", "Packed uint8 tensor for the V projection weights.", "T2") .Input(10, "v_scales", "Per-block scaling factors for the V projection.", "T1") .Input(11, "v_bias", "Optional bias for the V projection with shape [Nkv].", "T1", OpSchema::Optional) + .Input(12, "q_lora_A", "Optional packed Q adapter-down weights.", "TAdapter", OpSchema::Optional) + .Input(13, "q_lora_A_scales", "Optional Q adapter-down scales.", "T1", OpSchema::Optional) + .Input(14, "q_lora_B", "Optional packed Q adapter-up weights; may be empty.", "TAdapter", OpSchema::Optional) + .Input(15, "q_lora_B_scales", "Optional Q adapter-up scales.", "T1", OpSchema::Optional) + .Input(16, "k_lora_A", "Optional packed K adapter-down weights.", "TAdapter", OpSchema::Optional) + .Input(17, "k_lora_A_scales", "Optional K adapter-down scales.", "T1", OpSchema::Optional) + .Input(18, "k_lora_B", "Optional packed K adapter-up weights; may be empty.", "TAdapter", OpSchema::Optional) + .Input(19, "k_lora_B_scales", "Optional K adapter-up scales.", "T1", OpSchema::Optional) + .Input(20, "v_lora_A", "Optional packed V adapter-down weights.", "TAdapter", OpSchema::Optional) + .Input(21, "v_lora_A_scales", "Optional V adapter-down scales.", "T1", OpSchema::Optional) + .Input(22, "v_lora_B", "Optional packed V adapter-up weights; may be empty.", "TAdapter", OpSchema::Optional) + .Input(23, "v_lora_B_scales", "Optional V adapter-up scales.", "T1", OpSchema::Optional) .Output(0, "Q", "The Q projection output tensor.", "T1") .Output(1, "K", "The K projection output tensor.", "T1") .Output(2, "V", "The V projection output tensor.", "T1") @@ -4558,7 +4597,14 @@ This operator is intended as a decode-oriented QKV fusion primitive. .TypeConstraint("T1", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") .TypeConstraint("T2", {"tensor(uint8)"}, "Constrain quantized weight types to uint8.") + .TypeConstraint("TAdapter", {"tensor(uint8)"}, "Constrain packed adapter weights to uint8.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + const bool has_adapter_inputs = ctx.hasInput(12); + for (int i = 12; i < 24; ++i) { + if (ctx.hasInput(i) != has_adapter_inputs) { + fail_shape_inference("all Q, K, and V adapter inputs must be present together"); + } + } for (size_t output_index = 0; output_index < ctx.getNumOutputs(); ++output_index) { propagateElemTypeFromInputToOutput(ctx, 0, output_index); } @@ -4584,7 +4630,7 @@ This operator is intended as a decode-oriented QKV fusion primitive. set_output_shape(0, q_out_features); set_output_shape(1, kv_out_features); set_output_shape(2, kv_out_features); - if (ctx.getNumOutputs() > 3) { + if (ctx.hasOutput(3)) { auto* output_shape = getOutputShape(ctx, 3); *output_shape = input_shape; } diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index bfbc6fa7aaca0..a37b91e02d8b7 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -471,24 +471,33 @@ InlinedVector> GenerateTransformers( bool has_matmul_nbits_qlora_kernel = false; bool has_matmul_nbits_mlp_kernel = false; bool has_matmul_nbits_qkv_kernel = false; + bool has_matmul_nbits_mlp_adapter_kernel = false; + bool has_matmul_nbits_qkv_adapter_kernel = false; if (execution_providers != nullptr) { const auto* webgpu_ep = execution_providers->Get(onnxruntime::kWebGpuExecutionProvider); if (webgpu_ep != nullptr) { auto registry = webgpu_ep->GetKernelRegistry(); if (registry) { - auto has_webgpu_kernel = [&](std::string_view op_type) { + auto has_webgpu_kernel = [&](std::string_view op_type, + const KernelRegistry::TypeConstraintMap& type_constraints) { return registry->TryFindKernel(onnxruntime::kWebGpuExecutionProvider, op_type, kMSDomain, 1, - KernelRegistry::TypeConstraintMap{}, + type_constraints, logger, 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"); + const KernelRegistry::TypeConstraintMap no_constraints{}; + const KernelRegistry::TypeConstraintMap adapter_constraint{ + {"TAdapter", DataTypeImpl::GetTensorType()}}; + has_matmul_nbits_qlora_kernel = has_webgpu_kernel("MatMulNBitsQLoRA", no_constraints); + has_matmul_nbits_mlp_kernel = has_webgpu_kernel("MatMulNBitsMlp", no_constraints); + has_matmul_nbits_qkv_kernel = has_webgpu_kernel("MatMulNBitsQkv", no_constraints); + // Older WebGPU plugins expose the base kernels without the adapter type constraint. + has_matmul_nbits_mlp_adapter_kernel = has_webgpu_kernel("MatMulNBitsMlp", adapter_constraint); + has_matmul_nbits_qkv_adapter_kernel = has_webgpu_kernel("MatMulNBitsQkv", adapter_constraint); } } } @@ -497,11 +506,13 @@ InlinedVector> GenerateTransformers( } if (has_matmul_nbits_mlp_kernel) { transformers.emplace_back(std::make_unique( - InlinedHashSet{onnxruntime::kWebGpuExecutionProvider})); + InlinedHashSet{onnxruntime::kWebGpuExecutionProvider}, + has_matmul_nbits_mlp_adapter_kernel)); } if (has_matmul_nbits_qkv_kernel) { transformers.emplace_back(std::make_unique( - InlinedHashSet{onnxruntime::kWebGpuExecutionProvider})); + InlinedHashSet{onnxruntime::kWebGpuExecutionProvider}, + has_matmul_nbits_qkv_adapter_kernel)); } #endif // !defined(DISABLE_CONTRIB_OPS) diff --git a/onnxruntime/core/optimizer/matmul_nbits_mlp_fusion.cc b/onnxruntime/core/optimizer/matmul_nbits_mlp_fusion.cc index 510fe3b8f7fa2..cc1bcbd54a43e 100644 --- a/onnxruntime/core/optimizer/matmul_nbits_mlp_fusion.cc +++ b/onnxruntime/core/optimizer/matmul_nbits_mlp_fusion.cc @@ -78,6 +78,10 @@ bool HasExpectedNormConsumers(const Graph& graph, const Node& node) { } bool IsMatMulNBitsWithoutZeroPointOrGroupIdx(const Node& node) { + if (graph_utils::IsSupportedOptypeVersionAndDomain(node, "MatMulNBitsQLoRA", {1}, kMSDomain)) { + return node.InputDefs().size() == 7 && GetIntAttr(node, "lora_bits", 8) == 8 && + GetIntAttr(node, "lora_block_size", 32) == 32; + } return graph_utils::IsSupportedOptypeVersionAndDomain(node, "MatMulNBits", {1}, kMSDomain) && !HasInput(node, 3) && !HasInput(node, 4); } @@ -137,6 +141,7 @@ bool ValidateMatMulNBitsPair(const Graph& graph, if (!IsMatMulNBitsWithoutZeroPointOrGroupIdx(gate_matmul) || !IsMatMulNBitsWithoutZeroPointOrGroupIdx(up_matmul)) { return false; } + if (gate_matmul.OpType() != up_matmul.OpType()) return false; if (!HasSingleNonGraphConsumer(graph, up_matmul)) { return false; @@ -161,6 +166,11 @@ bool ValidateMatMulNBitsPair(const Graph& graph, const int64_t up_block_size = GetIntAttr(up_matmul, "block_size", -1, true); const int64_t gate_accuracy_level = GetIntAttr(gate_matmul, "accuracy_level", 0); const int64_t up_accuracy_level = GetIntAttr(up_matmul, "accuracy_level", 0); + if (gate_matmul.OpType() == "MatMulNBitsQLoRA" && + (GetIntAttr(gate_matmul, "lora_bits", 8) != GetIntAttr(up_matmul, "lora_bits", 8) || + GetIntAttr(gate_matmul, "lora_block_size", 32) != GetIntAttr(up_matmul, "lora_block_size", 32))) { + return false; + } // Fusion intentionally narrower than the kernel: although the MatMulNBitsMlp // kernel itself accepts {2, 4, 8} bits and any block_size, the fused decode @@ -336,6 +346,10 @@ Status MatMulNBitsMlpFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l matched_shape = "silu"; } + if (gate_matmul->OpType() == "MatMulNBitsQLoRA" && !enable_qlora_fusion_) { + continue; + } + LOGS(logger, VERBOSE) << "MatMulNBitsMlpFusion: matched candidate shape='" << matched_shape << "' output_mul='" << node.Name() << "' gate='" << gate_matmul->Name() << "' up='" << up_matmul->Name() @@ -378,6 +392,11 @@ Status MatMulNBitsMlpFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l utils::SetNodeAttribute(utils::MakeAttribute("accuracy_level", GetIntAttr(*gate_matmul, "accuracy_level", 0)), attrs); utils::SetNodeAttribute(utils::MakeAttribute(kActivationAttrName, std::string{kSupportedActivation}), attrs); utils::SetNodeAttribute(utils::MakeAttribute("epsilon", GetFloatAttr(*norm, "epsilon", 1e-5f)), attrs); + const bool is_qlora = gate_matmul->OpType() == "MatMulNBitsQLoRA"; + if (is_qlora) { + utils::SetNodeAttribute(utils::MakeAttribute("lora_bits", GetIntAttr(*gate_matmul, "lora_bits", 8)), attrs); + utils::SetNodeAttribute(utils::MakeAttribute("lora_block_size", GetIntAttr(*gate_matmul, "lora_block_size", 32)), attrs); + } NodeArg& empty_arg = graph.GetOrCreateNodeArg("", nullptr); // `norm` is guaranteed non-null here: the GetNormProducer() / continue guard @@ -390,11 +409,18 @@ Status MatMulNBitsMlpFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l const_cast(norm->InputDefs()[is_skip_sln ? 2 : 1]), const_cast(gate_matmul->InputDefs()[1]), const_cast(gate_matmul->InputDefs()[2]), - HasInput(*gate_matmul, 5) ? const_cast(gate_matmul->InputDefs()[5]) : &empty_arg, + !is_qlora && HasInput(*gate_matmul, 5) ? const_cast(gate_matmul->InputDefs()[5]) : &empty_arg, const_cast(up_matmul->InputDefs()[1]), const_cast(up_matmul->InputDefs()[2]), - HasInput(*up_matmul, 5) ? const_cast(up_matmul->InputDefs()[5]) : &empty_arg, + !is_qlora && HasInput(*up_matmul, 5) ? const_cast(up_matmul->InputDefs()[5]) : &empty_arg, }; + if (is_qlora) { + for (const Node* projection : {gate_matmul, up_matmul}) { + for (size_t input = 3; input < 7; ++input) { + fused_inputs.push_back(const_cast(projection->InputDefs()[input])); + } + } + } InlinedVector fused_outputs{const_cast(node.OutputDefs()[0])}; const bool preserve_skip_output = is_skip_sln && HasProducedOutput(*norm, 3); @@ -424,8 +450,9 @@ Status MatMulNBitsMlpFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l graph_utils::RemoveNodeOutputEdges(graph, node); graph.RemoveNode(node.Index()); - Node& fused_node = graph.AddNode(graph.GenerateNodeName("MatMulNBitsMlp"), - "MatMulNBitsMlp", + const char* fused_op = "MatMulNBitsMlp"; + Node& fused_node = graph.AddNode(graph.GenerateNodeName(fused_op), + fused_op, "fused MatMulNBits gated MLP projections", fused_inputs, fused_outputs, @@ -457,10 +484,16 @@ Status MatMulNBitsMlpFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l add_input_edge_if_present(gate_input_edges, 1, 3); add_input_edge_if_present(gate_input_edges, 2, 4); - add_input_edge_if_present(gate_input_edges, 5, 5); + if (!is_qlora) add_input_edge_if_present(gate_input_edges, 5, 5); add_input_edge_if_present(up_input_edges, 1, 6); add_input_edge_if_present(up_input_edges, 2, 7); - add_input_edge_if_present(up_input_edges, 5, 8); + if (!is_qlora) add_input_edge_if_present(up_input_edges, 5, 8); + if (is_qlora) { + for (int input = 3; input < 7; ++input) { + add_input_edge_if_present(gate_input_edges, input, 9 + input - 3); + add_input_edge_if_present(up_input_edges, input, 13 + input - 3); + } + } for (const auto& output_edge : final_mul_output_edges) { graph.AddEdge(fused_node.Index(), output_edge.dst_node, 0, output_edge.dst_arg_index); diff --git a/onnxruntime/core/optimizer/matmul_nbits_mlp_fusion.h b/onnxruntime/core/optimizer/matmul_nbits_mlp_fusion.h index 9698c467e84ed..bb08db9fdba6b 100644 --- a/onnxruntime/core/optimizer/matmul_nbits_mlp_fusion.h +++ b/onnxruntime/core/optimizer/matmul_nbits_mlp_fusion.h @@ -30,11 +30,15 @@ namespace onnxruntime { class MatMulNBitsMlpFusion : public GraphTransformer { public: explicit MatMulNBitsMlpFusion( - const InlinedHashSet& compatible_execution_providers = {}) noexcept - : GraphTransformer("MatMulNBitsMlpFusion", compatible_execution_providers) {} + const InlinedHashSet& compatible_execution_providers = {}, + bool enable_qlora_fusion = false) noexcept + : GraphTransformer("MatMulNBitsMlpFusion", compatible_execution_providers), + enable_qlora_fusion_(enable_qlora_fusion) {} private: Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override; + + const bool enable_qlora_fusion_; }; } // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/matmul_nbits_qkv_fusion.cc b/onnxruntime/core/optimizer/matmul_nbits_qkv_fusion.cc index d322964c278fa..b844018efd930 100644 --- a/onnxruntime/core/optimizer/matmul_nbits_qkv_fusion.cc +++ b/onnxruntime/core/optimizer/matmul_nbits_qkv_fusion.cc @@ -28,6 +28,10 @@ bool IsSupportedNormForFusion(const Node& node) { } bool IsMatMulNBitsWithoutOptionalInputs(const Node& node) { + if (graph_utils::IsSupportedOptypeVersionAndDomain(node, "MatMulNBitsQLoRA", {1}, kMSDomain)) { + return node.InputDefs().size() == 7 && GetIntAttr(node, "lora_bits", 8) == 8 && + GetIntAttr(node, "lora_block_size", 32) == 32; + } return graph_utils::IsSupportedOptypeVersionAndDomain(node, "MatMulNBits", {1}, kMSDomain) && !HasInput(node, 3) && !HasInput(node, 4) && !HasInput(node, 5); } @@ -149,6 +153,15 @@ bool IsFuseCandidate(const Node& norm, const QkvNodes& qkv) { !HasSupportedExecutionProvider(*qkv.k) || !HasSupportedExecutionProvider(*qkv.v)) { return false; } + if (qkv.q->OpType() != qkv.k->OpType() || qkv.q->OpType() != qkv.v->OpType()) return false; + if (qkv.q->OpType() == "MatMulNBitsQLoRA") { + for (const Node* projection : {qkv.k, qkv.v}) { + if (GetIntAttr(*projection, "lora_bits", 8) != GetIntAttr(*qkv.q, "lora_bits", 8) || + GetIntAttr(*projection, "lora_block_size", 32) != GetIntAttr(*qkv.q, "lora_block_size", 32)) { + return false; + } + } + } const size_t min_norm_inputs = IsSupportedSkipSimplifiedLayerNormalization(norm) ? 3u : 2u; if (norm.InputDefs().size() < min_norm_inputs || qkv.q->InputDefs().empty() || qkv.k->InputDefs().empty() || qkv.v->InputDefs().empty()) { @@ -210,6 +223,9 @@ Status MatMulNBitsQkvFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l if (!qkv_nodes || !IsFuseCandidate(node, *qkv_nodes)) { continue; } + if (qkv_nodes->q->OpType() == "MatMulNBitsQLoRA" && !enable_qlora_fusion_) { + continue; + } const int64_t K = GetIntAttr(*qkv_nodes->q, "K", -1, true); const int64_t Nq = GetIntAttr(*qkv_nodes->q, "N", -1, true); @@ -236,6 +252,11 @@ Status MatMulNBitsQkvFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l utils::SetNodeAttribute(utils::MakeAttribute("block_size", block_size), attrs); utils::SetNodeAttribute(utils::MakeAttribute("accuracy_level", accuracy_level), attrs); utils::SetNodeAttribute(utils::MakeAttribute("epsilon", epsilon), attrs); + const bool is_qlora = qkv_nodes->q->OpType() == "MatMulNBitsQLoRA"; + if (is_qlora) { + utils::SetNodeAttribute(utils::MakeAttribute("lora_bits", GetIntAttr(*qkv_nodes->q, "lora_bits", 8)), attrs); + utils::SetNodeAttribute(utils::MakeAttribute("lora_block_size", GetIntAttr(*qkv_nodes->q, "lora_block_size", 32)), attrs); + } NodeArg& empty_arg = graph.GetOrCreateNodeArg("", nullptr); @@ -253,6 +274,13 @@ Status MatMulNBitsQkvFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l const_cast(qkv_nodes->v->InputDefs()[2]), &empty_arg, }; + if (is_qlora) { + for (const Node* projection : {qkv_nodes->q, qkv_nodes->k, qkv_nodes->v}) { + for (size_t input = 3; input < 7; ++input) { + fused_inputs.push_back(const_cast(projection->InputDefs()[input])); + } + } + } InlinedVector fused_outputs{ const_cast(qkv_nodes->q->OutputDefs()[0]), @@ -288,8 +316,9 @@ Status MatMulNBitsQkvFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l graph_utils::RemoveNodeOutputEdges(graph, node); graph.RemoveNode(node.Index()); - Node& fused_node = graph.AddNode(graph.GenerateNodeName("MatMulNBitsQkv"), - "MatMulNBitsQkv", + const char* fused_op = "MatMulNBitsQkv"; + Node& fused_node = graph.AddNode(graph.GenerateNodeName(fused_op), + fused_op, "fused SimplifiedLayerNormalization with Q/K/V MatMulNBits projections", fused_inputs, fused_outputs, @@ -328,6 +357,15 @@ Status MatMulNBitsQkvFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l add_input_edge_if_present(k_input_edges, 2, 7); // k_scales add_input_edge_if_present(v_input_edges, 1, 9); // v_weight add_input_edge_if_present(v_input_edges, 2, 10); // v_scales + if (is_qlora) { + const std::array*, 3> input_edges{ + &q_input_edges, &k_input_edges, &v_input_edges}; + for (int projection = 0; projection < 3; ++projection) { + for (int input = 3; input < 7; ++input) { + add_input_edge_if_present(*input_edges[projection], input, 12 + projection * 4 + input - 3); + } + } + } for (const auto& output_edge : q_output_edges) { graph.AddEdge(fused_node.Index(), output_edge.dst_node, 0, output_edge.dst_arg_index); diff --git a/onnxruntime/core/optimizer/matmul_nbits_qkv_fusion.h b/onnxruntime/core/optimizer/matmul_nbits_qkv_fusion.h index fcbbb78457f52..2e30438830131 100644 --- a/onnxruntime/core/optimizer/matmul_nbits_qkv_fusion.h +++ b/onnxruntime/core/optimizer/matmul_nbits_qkv_fusion.h @@ -26,11 +26,15 @@ namespace onnxruntime { class MatMulNBitsQkvFusion : public GraphTransformer { public: explicit MatMulNBitsQkvFusion( - const InlinedHashSet& compatible_execution_providers = {}) noexcept - : GraphTransformer("MatMulNBitsQkvFusion", compatible_execution_providers) {} + const InlinedHashSet& compatible_execution_providers = {}, + bool enable_qlora_fusion = false) noexcept + : GraphTransformer("MatMulNBitsQkvFusion", compatible_execution_providers), + enable_qlora_fusion_(enable_qlora_fusion) {} private: Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override; + + const bool enable_qlora_fusion_; }; } // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc index 8bedf6de6e7cb..c479e1fe0a109 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc @@ -52,6 +52,71 @@ void RunEmptyAndActiveAdapters() { TEST(MatMulNBitsQLoRA, EmptyAndActiveAdapters) { RunEmptyAndActiveAdapters(); } TEST(MatMulNBitsQLoRA, EmptyAndActiveAdaptersFp32) { RunEmptyAndActiveAdapters(); } +TEST(MatMulNBitsQLoRA, FusedEmptyAdaptersRejectMalformedBaseWeights) { + for (bool qkv : {false, true}) { + const int projection_count = qkv ? 3 : 2; + for (int bad_projection = 0; bad_projection < projection_count; ++bad_projection) { + for (bool bad_weight : {false, true}) { + SCOPED_TRACE(MakeString(qkv ? "QKV" : "MLP", ", projection=", bad_projection, + ", bad_weight=", bad_weight)); + auto ep = DefaultWebGpuExecutionProvider(); + ASSERT_NE(ep, nullptr); + OpTester test(qkv ? "MatMulNBitsQkv" : "MatMulNBitsMlp", 1, kMSDomain); + test.AddAttribute("K", 32); + if (qkv) { + test.AddAttribute("Nq", 8); + test.AddAttribute("Nkv", 4); + } else { + test.AddAttribute("N", 8); + test.AddAttribute("activation", "silu"); + } + test.AddAttribute("bits", 4); + test.AddAttribute("block_size", 32); + test.AddAttribute("lora_bits", 8); + test.AddAttribute("lora_block_size", 32); + test.AddInput("A", {1, 32}, std::vector(32, MLFloat16(0.5f))); + test.AddOptionalInputEdge(); // skip + if (qkv) { + test.AddInput("norm_scale", {32}, std::vector(32, MLFloat16(1.0f))); + } else { + test.AddOptionalInputEdge(); // norm_scale + } + for (int projection = 0; projection < projection_count; ++projection) { + const int64_t n = qkv && projection > 0 ? 4 : 8; + const bool invalid = projection == bad_projection; + const int64_t weight_n = invalid && bad_weight ? n - 1 : n; + const int64_t scale_n = invalid && !bad_weight ? n / 2 : n; + test.AddInput(MakeString("weight_", projection).c_str(), {weight_n, 1, 16}, + std::vector(weight_n * 16, 0x99), true); + test.AddInput(MakeString("scales_", projection).c_str(), {scale_n, 1}, + std::vector(scale_n, MLFloat16(0.03125f)), true); + test.AddOptionalInputEdge(); // bias + } + for (int projection = 0; projection < projection_count; ++projection) { + const int64_t n = qkv && projection > 0 ? 4 : 8; + test.AddInput(MakeString("lora_A_", projection).c_str(), {32, 1, 32}, + std::vector(1024, 129)); + test.AddInput(MakeString("lora_A_scales_", projection).c_str(), {32, 1}, + std::vector(32, MLFloat16(0.03125f))); + test.AddInput(MakeString("lora_B_", projection).c_str(), {n, 0, 32}, {}); + test.AddInput(MakeString("lora_B_scales_", projection).c_str(), {n, 1}, + std::vector(n, MLFloat16(0.0625f))); + } + if (qkv) { + test.AddOutput("Q", {1, 8}, std::vector(8)); + test.AddOutput("K", {1, 4}, std::vector(4)); + test.AddOutput("V", {1, 4}, std::vector(4)); + } else { + test.AddOutput("Y", {1, 8}, std::vector(8)); + } + std::vector> eps; + eps.push_back(std::move(ep)); + test.Run(OpTester::ExpectResult::kExpectFailure, "Invalid base weight/scales shape.", {}, nullptr, &eps); + } + } + } +} + #ifndef ORT_NO_EXCEPTIONS TEST(MatMulNBitsQLoRA, RejectsUnsupportedWebGpuConfigurations) { struct Config { diff --git a/onnxruntime/test/optimizer/matmul_nbits_mlp_fusion_test.cc b/onnxruntime/test/optimizer/matmul_nbits_mlp_fusion_test.cc index a249eb6ef4b6b..2f4c0f22cd42a 100644 --- a/onnxruntime/test/optimizer/matmul_nbits_mlp_fusion_test.cc +++ b/onnxruntime/test/optimizer/matmul_nbits_mlp_fusion_test.cc @@ -187,7 +187,8 @@ void BuildMatMulNBitsMlpWebGpuPatternImpl(ModelTestBuilder& builder, SkipOutputKind skip_output_kind = SkipOutputKind::kNone, BiasKind bias_kind = BiasKind::kWithBias, ActivationShape activation_shape = ActivationShape::kSilu, - SkipNormVariant skip_norm_variant = SkipNormVariant::kDefault) { + SkipNormVariant skip_norm_variant = SkipNormVariant::kDefault, + bool qlora = false, bool empty_adapter = false) { constexpr int64_t k = 32; constexpr int64_t n = 8; constexpr int64_t block_size = 32; @@ -226,6 +227,10 @@ void BuildMatMulNBitsMlpWebGpuPatternImpl(ModelTestBuilder& builder, NodeArg* output = builder.MakeOutput(std::vector{1, n}); NodeAttributes matmul_attrs = MakeMatMulNBitsAttrs(k, n, block_size, bits, accuracy_level); + if (qlora) { + utils::SetNodeAttribute(utils::MakeAttribute("lora_bits", int64_t{8}), matmul_attrs); + utils::SetNodeAttribute(utils::MakeAttribute("lora_block_size", int64_t{32}), matmul_attrs); + } Node* norm = nullptr; if (norm_anchor_kind == NormAnchorKind::kSkipSimplified) { NodeArg* skip_input = builder.MakeInput( @@ -260,14 +265,25 @@ void BuildMatMulNBitsMlpWebGpuPatternImpl(ModelTestBuilder& builder, } } - Node& gate_matmul = builder.AddNode("MatMulNBits", - {normalized_input, gate_weight, gate_scale, optional_tensor, optional_tensor, - gate_bias}, - {gate_out}, kMSDomain, &matmul_attrs); - Node& up_matmul = builder.AddNode("MatMulNBits", - {normalized_input, up_weight, up_scale, optional_tensor, optional_tensor, - up_bias}, - {up_out}, kMSDomain, &matmul_attrs); + std::vector gate_inputs{normalized_input, gate_weight, gate_scale}; + std::vector up_inputs{normalized_input, up_weight, up_scale}; + if (qlora) { + const std::vector down_shape{32, 1, 32}; + const std::vector up_shape{n, empty_adapter ? 0 : 1, 32}; + for (auto* inputs : {&gate_inputs, &up_inputs}) { + inputs->push_back(builder.MakeInitializer(down_shape, uint8_t{130}, uint8_t{131})); + inputs->push_back(builder.MakeInitializer({32, 1}, MLFloat16(0.01f), MLFloat16(0.01f))); + inputs->push_back(builder.MakeInitializer(up_shape, uint8_t{130}, uint8_t{131})); + inputs->push_back(builder.MakeInitializer({n, 1}, MLFloat16(0.01f), MLFloat16(0.01f))); + } + } else { + gate_inputs.insert(gate_inputs.end(), {optional_tensor, optional_tensor, gate_bias}); + up_inputs.insert(up_inputs.end(), {optional_tensor, optional_tensor, up_bias}); + } + Node& gate_matmul = builder.AddNode(qlora ? "MatMulNBitsQLoRA" : "MatMulNBits", + gate_inputs, {gate_out}, kMSDomain, &matmul_attrs); + Node& up_matmul = builder.AddNode(qlora ? "MatMulNBitsQLoRA" : "MatMulNBits", + up_inputs, {up_out}, kMSDomain, &matmul_attrs); Node* sigmoid = nullptr; Node* silu_mul = nullptr; @@ -337,6 +353,65 @@ void BuildMatMulNBitsMlpSkipQuickGeluWebGpuPattern(ModelTestBuilder& builder) { BiasKind::kWithBias, ActivationShape::kQuickGelu); } +void BuildQLoRAMlpActive(ModelTestBuilder& builder) { + BuildMatMulNBitsMlpWebGpuPatternImpl(builder, NormAnchorKind::kSimplified, SkipOutputKind::kNone, + BiasKind::kNoBias, ActivationShape::kSilu, SkipNormVariant::kDefault, + true, false); +} + +void BuildQLoRAMlpEmpty(ModelTestBuilder& builder) { + BuildMatMulNBitsMlpWebGpuPatternImpl(builder, NormAnchorKind::kSimplified, SkipOutputKind::kNone, + BiasKind::kNoBias, ActivationShape::kSilu, SkipNormVariant::kDefault, + true, true); +} + +void BuildQLoRAMlpSkipActive(ModelTestBuilder& builder) { + BuildMatMulNBitsMlpWebGpuPatternImpl(builder, NormAnchorKind::kSkipSimplified, SkipOutputKind::kGraphOutput, + BiasKind::kNoBias, ActivationShape::kSilu, SkipNormVariant::kDefault, + true, false); +} + +void BuildQLoRAMlpSkipEmpty(ModelTestBuilder& builder) { + BuildMatMulNBitsMlpWebGpuPatternImpl(builder, NormAnchorKind::kSkipSimplified, SkipOutputKind::kGraphOutput, + BiasKind::kNoBias, ActivationShape::kSilu, SkipNormVariant::kDefault, + true, true); +} + +void BuildQLoRAMlpQuickGeluActive(ModelTestBuilder& builder) { + BuildMatMulNBitsMlpWebGpuPatternImpl(builder, NormAnchorKind::kSkipSimplified, SkipOutputKind::kGraphOutput, + BiasKind::kNoBias, ActivationShape::kQuickGelu, SkipNormVariant::kDefault, + true, false); +} + +void BuildQLoRAMlpQuickGeluEmpty(ModelTestBuilder& builder) { + BuildMatMulNBitsMlpWebGpuPatternImpl(builder, NormAnchorKind::kSkipSimplified, SkipOutputKind::kGraphOutput, + BiasKind::kNoBias, ActivationShape::kQuickGelu, SkipNormVariant::kDefault, + true, true); +} + +Status CheckQLoRAMlpFusedGraph(const Graph& graph) { + const auto counts = CountOpsInGraph(graph); + ORT_RETURN_IF_NOT(OpCount(counts, "com.microsoft.MatMulNBitsMlp") == 1 && + OpCount(counts, "com.microsoft.MatMulNBitsQLoRA") == 0 && + OpCount(counts, "SimplifiedLayerNormalization") == 0 && + OpCount(counts, "Mul") == 0, + "Expected one fused QLoRA MLP."); + for (const auto& node : graph.Nodes()) { + if (node.OpType() == "MatMulNBitsMlp") { + ORT_RETURN_IF_NOT(node.InputDefs().size() == 17, "Fused QLoRA MLP input count differs."); + } + } + return Status::OK(); +} + +Status CheckQLoRAMlpUnfusedGraph(const Graph& graph) { + const auto counts = CountOpsInGraph(graph); + ORT_RETURN_IF_NOT(OpCount(counts, "com.microsoft.MatMulNBitsMlp") == 0 && + OpCount(counts, "com.microsoft.MatMulNBitsQLoRA") == 2, + "QLoRA MLP must remain unfused without its WebGPU kernel."); + return Status::OK(); +} + void BuildMatMulNBitsMlpSkipWebGpuPatternWithSkipNormBias(ModelTestBuilder& builder) { BuildMatMulNBitsMlpWebGpuPatternImpl(builder, NormAnchorKind::kSkipSimplified, SkipOutputKind::kNone, BiasKind::kWithBias, ActivationShape::kSilu, @@ -351,6 +426,34 @@ void BuildMatMulNBitsMlpSimplifiedWebGpuPatternWithNormAxisZero(ModelTestBuilder } // namespace +TEST_F(GraphTransformationTests, MatMulNBitsMlpFusionFusesQLoRA) { + ASSERT_STATUS_OK(TestGraphTransformer( + BuildQLoRAMlpActive, 21, *logger_, + std::make_unique(InlinedHashSet{kWebGpuExecutionProvider}, true), + TransformerLevel::Level2, 1, nullptr, CheckQLoRAMlpFusedGraph)); +} + +TEST_F(GraphTransformationTests, MatMulNBitsMlpFusionLeavesQLoRAWithoutKernel) { + ASSERT_STATUS_OK(TestGraphTransformer( + BuildQLoRAMlpActive, 21, *logger_, + std::make_unique(InlinedHashSet{kWebGpuExecutionProvider}), + TransformerLevel::Level2, 0, nullptr, CheckQLoRAMlpUnfusedGraph)); +} + +TEST_F(GraphTransformationTests, MatMulNBitsMlpFusionQLoRANumerics) { + if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU EP unavailable in this build."; + for (auto build : {BuildQLoRAMlpActive, BuildQLoRAMlpEmpty, + BuildQLoRAMlpSkipActive, BuildQLoRAMlpSkipEmpty, + BuildQLoRAMlpQuickGeluActive, BuildQLoRAMlpQuickGeluEmpty}) { + RunWebGpuFusionTransformerTest( + build, + [](InferenceSessionWrapper& session) { ASSERT_STATUS_OK(CheckQLoRAMlpFusedGraph(session.GetGraph())); }, + TransformerLevel::Level1, TransformerLevel::Level2, 21, 1e-3, 1e-3, + std::make_unique(InlinedHashSet{kWebGpuExecutionProvider}, true), + [] { return DefaultWebGpuExecutionProvider(); }); + } +} + TEST_F(GraphTransformationTests, MatMulNBitsMlpFusionFusesSimplifiedWebGpuPattern) { ASSERT_STATUS_OK(TestGraphTransformer( BuildMatMulNBitsMlpSimplifiedWebGpuPattern, diff --git a/onnxruntime/test/optimizer/matmul_nbits_qkv_fusion_test.cc b/onnxruntime/test/optimizer/matmul_nbits_qkv_fusion_test.cc index c01312f57ea3b..8364842671426 100644 --- a/onnxruntime/test/optimizer/matmul_nbits_qkv_fusion_test.cc +++ b/onnxruntime/test/optimizer/matmul_nbits_qkv_fusion_test.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include "core/graph/node_attr_utils.h" +#include "core/graph/model.h" #include "core/optimizer/graph_transformer_mgr.h" #include "core/optimizer/matmul_nbits_qkv_fusion.h" #include "core/optimizer/utils.h" @@ -9,6 +10,7 @@ #include "core/session/onnxruntime_session_options_config_keys.h" #include "test/util/include/asserts.h" +#include "test/test_environment.h" #include "test/util/include/default_providers.h" #include "test/unittest_util/framework_test_utils.h" #include "test/unittest_util/graph_transform_test_builder.h" @@ -125,8 +127,9 @@ Status CheckMatMulNBitsQkvPatternNotFusedGraph(Graph& graph) { void BuildMatMulNBitsQkvWebGpuPatternImpl(ModelTestBuilder& builder, bool with_skip_input, bool with_skip_output, SkipNormVariant skip_norm_variant = SkipNormVariant::kDefault, - MatMulBiasVariant matmul_bias_variant = MatMulBiasVariant::kNone) { - constexpr int64_t k = 16; + MatMulBiasVariant matmul_bias_variant = MatMulBiasVariant::kNone, + bool qlora = false, bool empty_adapter = false) { + const int64_t k = qlora ? 32 : 16; constexpr int64_t q_n = 8; constexpr int64_t kv_n = 4; constexpr int64_t block_size = 32; @@ -134,21 +137,22 @@ void BuildMatMulNBitsQkvWebGpuPatternImpl(ModelTestBuilder& builder, bool with_s constexpr int64_t accuracy_level = 4; constexpr int64_t blob_size = block_size * bits / 8; - NodeArg* input = builder.MakeInput( - std::vector{1, k}, - std::vector{ - MLFloat16(-1.0f), MLFloat16(-0.875f), MLFloat16(-0.75f), MLFloat16(-0.625f), - MLFloat16(-0.5f), MLFloat16(-0.375f), MLFloat16(-0.25f), MLFloat16(-0.125f), - MLFloat16(0.125f), MLFloat16(0.25f), MLFloat16(0.375f), MLFloat16(0.5f), - MLFloat16(0.625f), MLFloat16(0.75f), MLFloat16(0.875f), MLFloat16(1.0f)}); + NodeArg* input = qlora ? builder.MakeInput({1, k}, MLFloat16(-1.0f), MLFloat16(1.0f)) + : builder.MakeInput( + std::vector{1, k}, + std::vector{ + MLFloat16(-1.0f), MLFloat16(-0.875f), MLFloat16(-0.75f), MLFloat16(-0.625f), + MLFloat16(-0.5f), MLFloat16(-0.375f), MLFloat16(-0.25f), MLFloat16(-0.125f), + MLFloat16(0.125f), MLFloat16(0.25f), MLFloat16(0.375f), MLFloat16(0.5f), + MLFloat16(0.625f), MLFloat16(0.75f), MLFloat16(0.875f), MLFloat16(1.0f)}); + std::vector skip_values{ + MLFloat16(1.0f), MLFloat16(0.875f), MLFloat16(0.75f), MLFloat16(0.625f), + MLFloat16(0.5f), MLFloat16(0.375f), MLFloat16(0.25f), MLFloat16(0.125f), + MLFloat16(-0.125f), MLFloat16(-0.25f), MLFloat16(-0.375f), MLFloat16(-0.5f), + MLFloat16(-0.625f), MLFloat16(-0.75f), MLFloat16(-0.875f), MLFloat16(-1.0f)}; + skip_values.resize(static_cast(k), MLFloat16(0.125f)); NodeArg* skip_input = with_skip_input - ? builder.MakeInput( - std::vector{1, k}, - std::vector{ - MLFloat16(1.0f), MLFloat16(0.875f), MLFloat16(0.75f), MLFloat16(0.625f), - MLFloat16(0.5f), MLFloat16(0.375f), MLFloat16(0.25f), MLFloat16(0.125f), - MLFloat16(-0.125f), MLFloat16(-0.25f), MLFloat16(-0.375f), MLFloat16(-0.5f), - MLFloat16(-0.625f), MLFloat16(-0.75f), MLFloat16(-0.875f), MLFloat16(-1.0f)}) + ? builder.MakeInput(std::vector{1, k}, skip_values) : nullptr; NodeArg* norm_scale = builder.MakeInitializer({k}, MLFloat16(1.0f), MLFloat16(1.0f)); @@ -180,6 +184,12 @@ void BuildMatMulNBitsQkvWebGpuPatternImpl(ModelTestBuilder& builder, bool with_s NodeAttributes q_attrs = MakeMatMulNBitsAttrs(k, q_n, block_size, bits, accuracy_level); NodeAttributes kv_attrs = MakeMatMulNBitsAttrs(k, kv_n, block_size, bits, accuracy_level); + if (qlora) { + for (auto* attrs : {&q_attrs, &kv_attrs}) { + utils::SetNodeAttribute(utils::MakeAttribute("lora_bits", int64_t{8}), *attrs); + utils::SetNodeAttribute(utils::MakeAttribute("lora_block_size", int64_t{32}), *attrs); + } + } NodeArg* norm_bias = (with_skip_input && skip_norm_variant == SkipNormVariant::kWithBiasInput) ? builder.MakeInitializer({k}, MLFloat16(0.0f), MLFloat16(0.0f)) @@ -201,9 +211,25 @@ void BuildMatMulNBitsQkvWebGpuPatternImpl(ModelTestBuilder& builder, bool with_s norm.AddAttribute("axis", static_cast(0)); } - Node& q_matmul = builder.AddNode("MatMulNBits", {norm_out, q_weight, q_scale, optional_tensor, optional_tensor, q_bias}, {q_output}, kMSDomain, &q_attrs); - Node& k_matmul = builder.AddNode("MatMulNBits", {norm_out, k_weight, k_scale, optional_tensor, optional_tensor, k_bias}, {k_output}, kMSDomain, &kv_attrs); - Node& v_matmul = builder.AddNode("MatMulNBits", {norm_out, v_weight, v_scale, optional_tensor, optional_tensor, v_bias}, {v_output}, kMSDomain, &kv_attrs); + auto projection_inputs = [&](NodeArg* weight, NodeArg* scale, NodeArg* bias, int64_t n) { + std::vector inputs{norm_out, weight, scale}; + if (qlora) { + inputs.push_back(builder.MakeInitializer({32, 1, 32}, uint8_t{130}, uint8_t{131})); + inputs.push_back(builder.MakeInitializer({32, 1}, MLFloat16(0.01f), MLFloat16(0.01f))); + inputs.push_back(builder.MakeInitializer({n, empty_adapter ? 0 : 1, 32}, uint8_t{130}, uint8_t{131})); + inputs.push_back(builder.MakeInitializer({n, 1}, MLFloat16(0.01f), MLFloat16(0.01f))); + } else { + inputs.insert(inputs.end(), {optional_tensor, optional_tensor, bias}); + } + return inputs; + }; + const char* projection_op = qlora ? "MatMulNBitsQLoRA" : "MatMulNBits"; + Node& q_matmul = builder.AddNode(projection_op, projection_inputs(q_weight, q_scale, q_bias, q_n), + {q_output}, kMSDomain, &q_attrs); + Node& k_matmul = builder.AddNode(projection_op, projection_inputs(k_weight, k_scale, k_bias, kv_n), + {k_output}, kMSDomain, &kv_attrs); + Node& v_matmul = builder.AddNode(projection_op, projection_inputs(v_weight, v_scale, v_bias, kv_n), + {v_output}, kMSDomain, &kv_attrs); SetWebGpuProvider(norm); SetWebGpuProvider(q_matmul); @@ -220,6 +246,48 @@ void BuildMatMulNBitsQkvWebGpuPattern(ModelTestBuilder& builder) { BuildMatMulNBitsQkvWebGpuPatternImpl(builder, false, false); } +void BuildQLoRAQkvActive(ModelTestBuilder& builder) { + BuildMatMulNBitsQkvWebGpuPatternImpl(builder, false, false, SkipNormVariant::kDefault, + MatMulBiasVariant::kNone, true, false); +} + +void BuildQLoRAQkvEmpty(ModelTestBuilder& builder) { + BuildMatMulNBitsQkvWebGpuPatternImpl(builder, false, false, SkipNormVariant::kDefault, + MatMulBiasVariant::kNone, true, true); +} + +void BuildQLoRAQkvSkipActive(ModelTestBuilder& builder) { + BuildMatMulNBitsQkvWebGpuPatternImpl(builder, true, true, SkipNormVariant::kDefault, + MatMulBiasVariant::kNone, true, false); +} + +void BuildQLoRAQkvSkipEmpty(ModelTestBuilder& builder) { + BuildMatMulNBitsQkvWebGpuPatternImpl(builder, true, true, SkipNormVariant::kDefault, + MatMulBiasVariant::kNone, true, true); +} + +Status CheckQLoRAQkvFusedGraph(const Graph& graph) { + const auto counts = CountOpsInGraph(graph); + ORT_RETURN_IF_NOT(OpCount(counts, "com.microsoft.MatMulNBitsQkv") == 1 && + OpCount(counts, "com.microsoft.MatMulNBitsQLoRA") == 0 && + OpCount(counts, "SimplifiedLayerNormalization") == 0, + "Expected one fused QLoRA QKV node."); + for (const auto& node : graph.Nodes()) { + if (node.OpType() == "MatMulNBitsQkv") { + ORT_RETURN_IF_NOT(node.InputDefs().size() == 24, "Fused QLoRA QKV input count differs."); + } + } + return Status::OK(); +} + +Status CheckQLoRAQkvUnfusedGraph(const Graph& graph) { + const auto counts = CountOpsInGraph(graph); + ORT_RETURN_IF_NOT(OpCount(counts, "com.microsoft.MatMulNBitsQkv") == 0 && + OpCount(counts, "com.microsoft.MatMulNBitsQLoRA") == 3, + "QLoRA QKV must remain unfused without its WebGPU kernel."); + return Status::OK(); +} + void BuildMatMulNBitsQkvSkipWebGpuPattern(ModelTestBuilder& builder) { BuildMatMulNBitsQkvWebGpuPatternImpl(builder, true, false); } @@ -256,6 +324,81 @@ void BuildMatMulNBitsQkvWebGpuPatternWithVBias(ModelTestBuilder& builder) { } // namespace +TEST(MatMulNBitsQkvSchemaTest, EmptyResidualOutputWithoutSkip) { + for (const bool qlora : {false, true}) { + SCOPED_TRACE(qlora ? "QLoRA QKV" : "base QKV"); + std::unordered_map domain_to_version{{kOnnxDomain, 21}, {kMSDomain, 1}}; + Model model("matmul_nbits_qkv_optional_output", false, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), domain_to_version, {}, + DefaultLoggingManager().DefaultLogger()); + ModelTestBuilder builder(model.MainGraph()); + NodeArg* empty = builder.MakeOptionalTensor(); + NodeArg* a = builder.MakeInput({1, 32}, MLFloat16(1.0f), MLFloat16(1.0f)); + NodeArg* norm_scale = builder.MakeInitializer({32}, MLFloat16(1.0f), MLFloat16(1.0f)); + std::vector inputs{a, empty, norm_scale}; + for (int64_t n : {8, 4, 4}) { + inputs.push_back(builder.MakeInitializer({n, 1, 16}, uint8_t{0}, uint8_t{15})); + inputs.push_back(builder.MakeInitializer({n, 1}, MLFloat16(1.0f), MLFloat16(1.0f))); + inputs.push_back(empty); + } + if (qlora) { + for (int64_t n : {8, 4, 4}) { + inputs.push_back(builder.MakeInitializer({32, 1, 32}, uint8_t{0}, uint8_t{255})); + inputs.push_back(builder.MakeInitializer({32, 1}, MLFloat16(1.0f), MLFloat16(1.0f))); + inputs.push_back(builder.MakeInitializer({n, 1, 32}, uint8_t{0}, uint8_t{255})); + inputs.push_back(builder.MakeInitializer({n, 1}, MLFloat16(1.0f), MLFloat16(1.0f))); + } + } + + NodeArg* q = builder.MakeOutput(std::nullopt); + NodeArg* k = builder.MakeOutput(std::nullopt); + NodeArg* v = builder.MakeOutput(std::nullopt); + Node& node = builder.AddNode("MatMulNBitsQkv", + inputs, {q, k, v, empty}, kMSDomain); + node.AddAttribute("K", int64_t{32}); + node.AddAttribute("Nq", int64_t{8}); + node.AddAttribute("Nkv", int64_t{4}); + node.AddAttribute("bits", int64_t{4}); + node.AddAttribute("block_size", int64_t{32}); + if (qlora) { + node.AddAttribute("lora_bits", int64_t{8}); + node.AddAttribute("lora_block_size", int64_t{32}); + } + builder.SetGraphOutputs(); + ASSERT_STATUS_OK(model.MainGraph().Resolve()); + } +} + +TEST_F(GraphTransformationTests, MatMulNBitsQkvFusionFusesQLoRA) { + ASSERT_STATUS_OK(TestGraphTransformer( + BuildQLoRAQkvActive, 21, *logger_, + std::make_unique(InlinedHashSet{kWebGpuExecutionProvider}, true), + TransformerLevel::Level2, 1, nullptr, CheckQLoRAQkvFusedGraph)); +} + +TEST_F(GraphTransformationTests, MatMulNBitsQkvFusionLeavesQLoRAWithoutKernel) { + ASSERT_STATUS_OK(TestGraphTransformer( + BuildQLoRAQkvActive, 21, *logger_, + std::make_unique(InlinedHashSet{kWebGpuExecutionProvider}), + TransformerLevel::Level2, 0, nullptr, CheckQLoRAQkvUnfusedGraph)); +} + +TEST_F(GraphTransformationTests, MatMulNBitsQkvFusionQLoRANumerics) { + if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU EP unavailable in this build."; + int case_index = 0; + for (auto build : {BuildQLoRAQkvActive, BuildQLoRAQkvEmpty, + BuildQLoRAQkvSkipActive, BuildQLoRAQkvSkipEmpty}) { + SCOPED_TRACE(case_index++); + RunWebGpuFusionTransformerTest( + build, + [](InferenceSessionWrapper& session) { ASSERT_STATUS_OK(CheckQLoRAQkvFusedGraph(session.GetGraph())); }, + // The extra FP16 adapter projection can double the base QKV reduction-order difference on Metal. + TransformerLevel::Level1, TransformerLevel::Level2, 21, 2e-3, 1e-2, + std::make_unique(InlinedHashSet{kWebGpuExecutionProvider}, true), + [] { return DefaultWebGpuExecutionProvider(); }); + } +} + TEST_F(GraphTransformationTests, MatMulNBitsQkvFusionFusesWebGpuPattern) { ASSERT_STATUS_OK(TestGraphTransformer( BuildMatMulNBitsQkvWebGpuPattern, diff --git a/onnxruntime/test/optimizer/matmul_nbits_qlora_fused_switch_test.cc b/onnxruntime/test/optimizer/matmul_nbits_qlora_fused_switch_test.cc new file mode 100644 index 0000000000000..576172096a65b --- /dev/null +++ b/onnxruntime/test/optimizer/matmul_nbits_qlora_fused_switch_test.cc @@ -0,0 +1,197 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include + +#include "core/framework/tensor.h" +#include "core/graph/model.h" +#include "core/optimizer/matmul_nbits_mlp_fusion.h" +#include "core/optimizer/matmul_nbits_qkv_fusion.h" +#include "test/optimizer/webgpu_fusion_test_util.h" +#include "test/util/include/default_providers.h" +#include "test/unittest_util/framework_test_utils.h" + +#include "gtest/gtest.h" + +namespace onnxruntime::test { + +#if !defined(DISABLE_CONTRIB_OPS) +namespace { + +struct AdapterInputs { + std::string down_weights; + std::string down_scales; + std::string up_weights; + std::string up_scales; + int64_t n; +}; + +template +OrtValue FilledCpuTensor(const TensorShape& shape, T value) { + OrtValue result; + Tensor::InitOrtValue(DataTypeImpl::GetType(), shape, CPUAllocator::DefaultInstance(), result); + auto data = result.GetMutable()->MutableDataAsSpan(); + std::fill(data.begin(), data.end(), value); + return result; +} + +float FirstValue(const OrtValue& value) { + return static_cast(value.Get().Data()[0]); +} + +void RunAdapterSwitchTest(bool qkv) { + if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU EP unavailable in this build."; + + const int projection_count = qkv ? 3 : 2; + Model model("qlora_fused_adapter_switch", false, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 21}, {kMSDomain, 1}}, {}, + DefaultLoggingManager().DefaultLogger()); + ModelTestBuilder builder(model.MainGraph()); + NodeArg* x = builder.MakeInput({1, 32}, MLFloat16(0.5f), MLFloat16(0.5f)); + NodeArg* norm_scale = builder.MakeInitializer({32}, MLFloat16(1.0f), MLFloat16(1.0f)); + NodeArg* normalized = builder.MakeIntermediate(std::vector{1, 32}); + Node& norm = builder.AddNode("SimplifiedLayerNormalization", {x, norm_scale}, {normalized}); + norm.AddAttribute("epsilon", 1e-6f); + norm.SetExecutionProviderType(kWebGpuExecutionProvider); + + std::vector adapters; + std::vector projections; + for (int projection = 0; projection < projection_count; ++projection) { + const int64_t n = qkv && projection > 0 ? 4 : 8; + NodeArg* base_weights = builder.MakeInitializer({n, 1, 16}, std::vector(n * 16, 0x99)); + NodeArg* base_scales = builder.MakeInitializer({n, 1}, MLFloat16(0.03125f), MLFloat16(0.03125f)); + NodeArg* down_weights = builder.MakeInput({32, 1, 32}, std::vector(1024, 129)); + NodeArg* down_scales = builder.MakeInput({32, 1}, MLFloat16(0.03125f), MLFloat16(0.03125f)); + NodeArg* up_weights = builder.MakeInput({n, 1, 32}, std::vector(n * 32, 129)); + NodeArg* up_scales = builder.MakeInput({n, 1}, MLFloat16(0.0625f), MLFloat16(0.0625f)); + auto dynamic_shape = *up_weights->Shape(); + dynamic_shape.mutable_dim(1)->Clear(); + up_weights->SetShape(dynamic_shape); + adapters.push_back({down_weights->Name(), down_scales->Name(), up_weights->Name(), up_scales->Name(), n}); + + NodeArg* output = qkv ? builder.MakeOutput(std::vector{1, n}) + : builder.MakeIntermediate(std::vector{1, n}); + Node& node = builder.AddNode("MatMulNBitsQLoRA", + {normalized, base_weights, base_scales, + down_weights, down_scales, up_weights, up_scales}, + {output}, kMSDomain); + node.AddAttribute("K", int64_t{32}); + node.AddAttribute("N", n); + node.AddAttribute("bits", int64_t{4}); + node.AddAttribute("block_size", int64_t{32}); + node.AddAttribute("lora_bits", int64_t{8}); + node.AddAttribute("lora_block_size", int64_t{32}); + node.SetExecutionProviderType(kWebGpuExecutionProvider); + projections.push_back(output); + } + if (!qkv) { + NodeArg* sigmoid_output = builder.MakeIntermediate(std::vector{1, 8}); + NodeArg* activated = builder.MakeIntermediate(std::vector{1, 8}); + NodeArg* y = builder.MakeOutput(std::vector{1, 8}); + Node& sigmoid = builder.AddNode("Sigmoid", {projections[0]}, {sigmoid_output}); + Node& silu = builder.AddNode("Mul", {projections[0], sigmoid_output}, {activated}); + Node& output_mul = builder.AddNode("Mul", {activated, projections[1]}, {y}); + for (Node* node : {&sigmoid, &silu, &output_mul}) { + node->SetExecutionProviderType(kWebGpuExecutionProvider); + } + } + builder.SetGraphOutputs(); + ASSERT_STATUS_OK(model.MainGraph().Resolve()); + std::string model_data; + ASSERT_TRUE(model.ToProto().SerializeToString(&model_data)); + + // The graph and both sessions stay fixed while adapter B shapes and contents change. + std::vector> states{{0, false}}; + for (int projection = 0; projection < projection_count; ++projection) { + states.emplace_back(1 << projection, false); + } + const int all_on = (1 << projection_count) - 1; + states.emplace_back(all_on, false); + states.emplace_back(all_on, true); + states.emplace_back(0, false); + + std::vector> baseline_outputs; + for (bool fused : {false, true}) { + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Level1; + InferenceSessionWrapper session(options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(DefaultWebGpuExecutionProvider())); + ASSERT_STATUS_OK(session.Load(model_data.data(), static_cast(model_data.size()))); + if (fused) { + if (qkv) { + ASSERT_STATUS_OK(session.RegisterGraphTransformer( + std::make_unique(InlinedHashSet{kWebGpuExecutionProvider}, true), + TransformerLevel::Level2)); + } else { + ASSERT_STATUS_OK(session.RegisterGraphTransformer( + std::make_unique(InlinedHashSet{kWebGpuExecutionProvider}, true), + TransformerLevel::Level2)); + } + } + ASSERT_STATUS_OK(session.Initialize()); + size_t fused_count = 0; + for (const auto& node : session.GetGraph().Nodes()) { + fused_count += node.OpType() == (qkv ? "MatMulNBitsQkv" : "MatMulNBitsMlp"); + } + ASSERT_EQ(fused_count, fused ? 1u : 0u); + + std::vector> current_outputs; + for (size_t state_index = 0; state_index < states.size(); ++state_index) { + const auto [mask, changed] = states[state_index]; + SCOPED_TRACE(MakeString(qkv ? "QKV" : "MLP", ", fused=", fused, + ", mask=", mask, ", changed=", changed)); + NameMLValMap feeds = builder.feeds_; + for (int projection = 0; projection < projection_count; ++projection) { + const auto& adapter = adapters[projection]; + if ((mask & (1 << projection)) == 0) { + feeds[adapter.up_weights] = FilledCpuTensor({adapter.n, 0, 32}, 0); + } else if (changed) { + feeds[adapter.down_weights] = FilledCpuTensor({32, 1, 32}, 130); + feeds[adapter.down_scales] = FilledCpuTensor({32, 1}, MLFloat16(0.0625f)); + feeds[adapter.up_weights] = FilledCpuTensor({adapter.n, 1, 32}, 130); + feeds[adapter.up_scales] = FilledCpuTensor({adapter.n, 1}, MLFloat16(0.125f)); + } + } + std::vector outputs; + ASSERT_STATUS_OK(session.Run(RunOptions{}, feeds, builder.output_names_, &outputs)); + if (fused) { + ASSERT_EQ(outputs.size(), baseline_outputs[state_index].size()); + for (size_t output = 0; output < outputs.size(); ++output) { + const auto comparison = CompareOrtValue(outputs[output], baseline_outputs[state_index][output], + 2e-3, 1e-2, false); + EXPECT_EQ(comparison.first, COMPARE_RESULT::SUCCESS) << comparison.second; + } + } + current_outputs.push_back(std::move(outputs)); + } + + for (int projection = 0; projection < projection_count; ++projection) { + const size_t output = qkv ? static_cast(projection) : 0u; + EXPECT_GT(std::abs(FirstValue(current_outputs[projection + 1][output]) - + FirstValue(current_outputs[0][output])), + 0.1f); + } + const size_t changed_index = static_cast(projection_count) + 2; + EXPECT_GT(std::abs(FirstValue(current_outputs[changed_index][0]) - + FirstValue(current_outputs[changed_index - 1][0])), + 0.1f); + for (size_t output = 0; output < current_outputs.front().size(); ++output) { + const auto comparison = CompareOrtValue(current_outputs.back()[output], current_outputs.front()[output], + 1e-3, 1e-3, false); + EXPECT_EQ(comparison.first, COMPARE_RESULT::SUCCESS) << comparison.second; + } + if (!fused) baseline_outputs = std::move(current_outputs); + } +} + +} // namespace + +TEST(MatMulNBitsQLoRAFusedAdapterTest, MlpSwitchesInOneSession) { RunAdapterSwitchTest(false); } +TEST(MatMulNBitsQLoRAFusedAdapterTest, QkvSwitchesInOneSession) { RunAdapterSwitchTest(true); } + +#endif // !defined(DISABLE_CONTRIB_OPS) + +} // namespace onnxruntime::test From 4febae8397633cbfa1137a7417530daa7e5ec72e Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Thu, 5 Nov 2026 14:48:26 +0800 Subject: [PATCH 2/4] [WebGPU] Preserve base fusion compatibility and fix plugin build --- .../webgpu/quantization/matmul_nbits_mlp.h | 12 ++-- .../webgpu/quantization/matmul_nbits_qkv.h | 12 ++-- .../framework/kernel_type_str_resolver.cc | 28 +++++++- .../kernel_type_str_resolver_utils_test.cc | 68 +++++++++++++++++++ .../optimizer/matmul_nbits_mlp_fusion_test.cc | 18 +++++ .../optimizer/matmul_nbits_qkv_fusion_test.cc | 18 +++++ .../test/optimizer/webgpu_fusion_test_util.h | 48 +++++++++++++ 7 files changed, 187 insertions(+), 17 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.h index 4f2c038a4ea7a..db97da7b9a68e 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.h @@ -39,13 +39,11 @@ class MatMulNBitsMlp : public WebGpuKernel { epsilon_ = info.GetAttrOrDefault("epsilon", 1e-5f); lora_bits_ = info.GetAttrOrDefault("lora_bits", 8); lora_block_size_ = info.GetAttrOrDefault("lora_block_size", 32); - const auto& inputs = info.node().InputDefs(); - auto input_exists = [&](size_t i) { return inputs[i] != nullptr && inputs[i]->Exists(); }; - has_adapter_inputs_ = inputs.size() > 9 && input_exists(9); - for (size_t i = 9; i < inputs.size(); ++i) { - ORT_ENFORCE(input_exists(i) == has_adapter_inputs_, - "MatMulNBitsMlp requires all adapter inputs when any is present."); - } + // GetInputCount is available through the plugin EP API. The schema checks + // that the optional adapter inputs are either all present or all absent. + has_adapter_inputs_ = info.GetInputCount() > 9; + ORT_ENFORCE(!has_adapter_inputs_ || info.GetInputCount() == 17, + "MatMulNBitsMlp requires all adapter inputs when any is present."); std::string activation; ORT_ENFORCE(info.GetAttr("activation", &activation).IsOK(), "MatMulNBitsMlp requires the 'activation' attribute."); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.h index 50637ee024e5d..9bae2ec278bc1 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.h @@ -23,13 +23,11 @@ class MatMulNBitsQkv : public WebGpuKernel { epsilon_ = info.GetAttrOrDefault("epsilon", 1e-6f); lora_bits_ = info.GetAttrOrDefault("lora_bits", 8); lora_block_size_ = info.GetAttrOrDefault("lora_block_size", 32); - const auto& inputs = info.node().InputDefs(); - auto input_exists = [&](size_t i) { return inputs[i] != nullptr && inputs[i]->Exists(); }; - has_adapter_inputs_ = inputs.size() > 12 && input_exists(12); - for (size_t i = 12; i < inputs.size(); ++i) { - ORT_ENFORCE(input_exists(i) == has_adapter_inputs_, - "MatMulNBitsQkv requires all adapter inputs when any is present."); - } + // GetInputCount is available through the plugin EP API. The schema checks + // that the optional adapter inputs are either all present or all absent. + has_adapter_inputs_ = info.GetInputCount() > 12; + ORT_ENFORCE(!has_adapter_inputs_ || info.GetInputCount() == 24, + "MatMulNBitsQkv requires all adapter inputs when any is present."); ORT_ENFORCE(bits_ == 4, "MatMulNBitsQkv currently supports 4-bit weights only."); } diff --git a/onnxruntime/core/framework/kernel_type_str_resolver.cc b/onnxruntime/core/framework/kernel_type_str_resolver.cc index e92cf7f46a0fa..81be949828bed 100644 --- a/onnxruntime/core/framework/kernel_type_str_resolver.cc +++ b/onnxruntime/core/framework/kernel_type_str_resolver.cc @@ -274,9 +274,31 @@ static std::string LoadFromOrtFormatImpl(const fbs::KernelTypeStrResolver& fbs_k Status KernelTypeStrResolver::LoadFromOrtFormat(const fbs::KernelTypeStrResolver& fbs_kernel_type_str_resolver) { const auto error_message = LoadFromOrtFormatImpl(fbs_kernel_type_str_resolver, op_kernel_type_str_map_); - return error_message.empty() ? Status::OK() - : ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, error_message, " ", - fbs::utils::kInvalidOrtFormatModelMessage); + if (!error_message.empty()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, error_message, " ", + fbs::utils::kInvalidOrtFormatModelMessage); + } + + // These version-1 schemas gained optional QLoRA inputs. Older ORT-format + // models have only the base inputs and omit TAdapter from their saved + // resolver. Add the new mapping so the updated WebGPU kernel can still + // match those base-only nodes. Missing optional inputs are skipped by the + // kernel type matcher. + for (auto& [op_id, type_str_map] : op_kernel_type_str_map_) { + if (op_id.domain != kMSDomain || op_id.since_version != 1 || Contains(type_str_map, "TAdapter")) { + continue; + } + + if (op_id.op_type == "MatMulNBitsMlp") { + auto& args = type_str_map["TAdapter"]; + for (size_t i = 9; i < 17; i += 2) args.emplace_back(ArgType::kInput, i); + } else if (op_id.op_type == "MatMulNBitsQkv") { + auto& args = type_str_map["TAdapter"]; + for (size_t i = 12; i < 24; i += 2) args.emplace_back(ArgType::kInput, i); + } + } + + return Status::OK(); } void KernelTypeStrResolver::Merge(KernelTypeStrResolver src) { diff --git a/onnxruntime/test/framework/kernel_type_str_resolver_utils_test.cc b/onnxruntime/test/framework/kernel_type_str_resolver_utils_test.cc index 32720e42988f6..25e70ebdf7423 100644 --- a/onnxruntime/test/framework/kernel_type_str_resolver_utils_test.cc +++ b/onnxruntime/test/framework/kernel_type_str_resolver_utils_test.cc @@ -6,9 +6,11 @@ #include #include #include +#include #include "gtest/gtest.h" #include "core/flatbuffers/schema/ort.fbs.h" +#include "core/framework/kernel_registry.h" #include "core/graph/constants.h" #include "core/graph/model.h" #include "core/graph/schema_registry.h" @@ -61,6 +63,72 @@ TEST(KernelTypeStrResolverUtilsTest, VerifyLayoutTransformationRequiredOpsResolv } #if !defined(ORT_MINIMAL_BUILD) && !defined(DISABLE_CONTRIB_OPS) +TEST(KernelTypeStrResolverUtilsTest, LegacyMatMulNBitsFusionsResolveAdapterKernel) { + for (const auto& [op_type, first_adapter_input, adapter_input_count] : + {std::tuple{"MatMulNBitsMlp", size_t{9}, size_t{4}}, + std::tuple{"MatMulNBitsQkv", size_t{12}, size_t{6}}}) { + SCOPED_TRACE(op_type); + // Reproduce the type-string metadata written before these version-1 ops + // acquired optional adapter inputs. + ONNX_NAMESPACE::OpSchema old_schema; + old_schema.SetName(op_type) + .SetDomain(kMSDomain) + .SinceVersion(1) + .Output(0, "Y", "Result", "T1") + .TypeConstraint("T1", {"tensor(float)"}, "Activation type") + .TypeConstraint("T2", {"tensor(uint8)"}, "Packed weight type"); + for (size_t i = 0; i < first_adapter_input; ++i) { + old_schema.Input(static_cast(i), "input_" + std::to_string(i), "Base input", + i >= 3 && i % 3 == 0 ? "T2" : "T1"); + } + KernelTypeStrResolver old_resolver; + ASSERT_STATUS_OK(old_resolver.RegisterOpSchema(old_schema)); + + flatbuffers::DetachedBuffer buffer; + gsl::span buffer_span; + ASSERT_STATUS_OK(kernel_type_str_resolver_utils::SaveKernelTypeStrResolverToBuffer( + old_resolver, buffer, buffer_span)); + KernelTypeStrResolver loaded_resolver; + ASSERT_STATUS_OK(kernel_type_str_resolver_utils::LoadKernelTypeStrResolverFromBuffer( + loaded_resolver, buffer_span)); + + Model model("legacy_matmul_nbits_fusion", false, DefaultLoggingManager().DefaultLogger()); + auto& graph = model.MainGraph(); + ONNX_NAMESPACE::TypeProto float_type; + float_type.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + ONNX_NAMESPACE::TypeProto uint8_type; + uint8_type.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_UINT8); + std::vector base_inputs; + for (size_t i = 0; i < first_adapter_input; ++i) { + base_inputs.push_back(&graph.GetOrCreateNodeArg( + "input_" + std::to_string(i), i >= 3 && i % 3 == 0 ? &uint8_type : &float_type)); + } + auto& output = graph.GetOrCreateNodeArg("output", &float_type); + auto& node = graph.AddNode("base_fusion", op_type, "", base_inputs, {&output}, nullptr, kMSDomain); + node.SetSinceVersion(1); + + gsl::span resolved_args; + ASSERT_STATUS_OK(loaded_resolver.ResolveKernelTypeStr(node, "TAdapter", resolved_args)); + ASSERT_EQ(resolved_args.size(), adapter_input_count); + EXPECT_EQ(resolved_args.front(), (ArgTypeAndIndex{ArgType::kInput, first_adapter_input})); + + KernelRegistry registry; + KernelDefBuilder builder; + builder.SetName(op_type) + .SetDomain(kMSDomain) + .SinceVersion(1) + .Provider(kWebGpuExecutionProvider) + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) + .TypeConstraint("TAdapter", DataTypeImpl::GetTensorType()); + ASSERT_STATUS_OK(registry.Register(builder, {})); + const KernelCreateInfo* kernel = nullptr; + ASSERT_STATUS_OK(registry.TryFindKernel(node, kWebGpuExecutionProvider, loaded_resolver, + DefaultLoggingManager().DefaultLogger(), &kernel)); + ASSERT_NE(kernel, nullptr); + } +} + TEST(KernelTypeStrResolverUtilsTest, ResolveNhwcFusedConvFromLayoutTransformationRequiredOps) { KernelTypeStrResolver resolver; ASSERT_STATUS_OK(kernel_type_str_resolver_utils::AddLayoutTransformationRequiredOpsToKernelTypeStrResolver(resolver)); diff --git a/onnxruntime/test/optimizer/matmul_nbits_mlp_fusion_test.cc b/onnxruntime/test/optimizer/matmul_nbits_mlp_fusion_test.cc index 2f4c0f22cd42a..781c536a743c7 100644 --- a/onnxruntime/test/optimizer/matmul_nbits_mlp_fusion_test.cc +++ b/onnxruntime/test/optimizer/matmul_nbits_mlp_fusion_test.cc @@ -440,6 +440,24 @@ TEST_F(GraphTransformationTests, MatMulNBitsMlpFusionLeavesQLoRAWithoutKernel) { TransformerLevel::Level2, 0, nullptr, CheckQLoRAMlpUnfusedGraph)); } +TEST_F(GraphTransformationTests, MatMulNBitsMlpFusionProbesPluginAdapterSupport) { + for (bool adapter_aware : {false, true}) { + SCOPED_TRACE(adapter_aware ? "adapter-aware plugin" : "base-only plugin"); + auto qlora_fusion = GenerateWebGpuFusionTransformerForTest("MatMulNBitsMlpFusion", adapter_aware); + ASSERT_NE(qlora_fusion, nullptr); + ASSERT_STATUS_OK(TestGraphTransformer( + BuildQLoRAMlpActive, 21, *logger_, std::move(qlora_fusion), + TransformerLevel::Level2, adapter_aware ? 1 : 0, nullptr, + adapter_aware ? CheckQLoRAMlpFusedGraph : CheckQLoRAMlpUnfusedGraph)); + + auto base_fusion = GenerateWebGpuFusionTransformerForTest("MatMulNBitsMlpFusion", adapter_aware); + ASSERT_NE(base_fusion, nullptr); + ASSERT_STATUS_OK(TestGraphTransformer( + BuildMatMulNBitsMlpSimplifiedWebGpuPatternNoBias, 21, *logger_, std::move(base_fusion), + TransformerLevel::Level2, 1, nullptr, CheckMatMulNBitsMlpSimplifiedFusedGraph)); + } +} + TEST_F(GraphTransformationTests, MatMulNBitsMlpFusionQLoRANumerics) { if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU EP unavailable in this build."; for (auto build : {BuildQLoRAMlpActive, BuildQLoRAMlpEmpty, diff --git a/onnxruntime/test/optimizer/matmul_nbits_qkv_fusion_test.cc b/onnxruntime/test/optimizer/matmul_nbits_qkv_fusion_test.cc index 8364842671426..2ace0bfd95009 100644 --- a/onnxruntime/test/optimizer/matmul_nbits_qkv_fusion_test.cc +++ b/onnxruntime/test/optimizer/matmul_nbits_qkv_fusion_test.cc @@ -383,6 +383,24 @@ TEST_F(GraphTransformationTests, MatMulNBitsQkvFusionLeavesQLoRAWithoutKernel) { TransformerLevel::Level2, 0, nullptr, CheckQLoRAQkvUnfusedGraph)); } +TEST_F(GraphTransformationTests, MatMulNBitsQkvFusionProbesPluginAdapterSupport) { + for (bool adapter_aware : {false, true}) { + SCOPED_TRACE(adapter_aware ? "adapter-aware plugin" : "base-only plugin"); + auto qlora_fusion = GenerateWebGpuFusionTransformerForTest("MatMulNBitsQkvFusion", adapter_aware); + ASSERT_NE(qlora_fusion, nullptr); + ASSERT_STATUS_OK(TestGraphTransformer( + BuildQLoRAQkvActive, 21, *logger_, std::move(qlora_fusion), + TransformerLevel::Level2, adapter_aware ? 1 : 0, nullptr, + adapter_aware ? CheckQLoRAQkvFusedGraph : CheckQLoRAQkvUnfusedGraph)); + + auto base_fusion = GenerateWebGpuFusionTransformerForTest("MatMulNBitsQkvFusion", adapter_aware); + ASSERT_NE(base_fusion, nullptr); + ASSERT_STATUS_OK(TestGraphTransformer( + BuildMatMulNBitsQkvWebGpuPattern, 21, *logger_, std::move(base_fusion), + TransformerLevel::Level2, 1, nullptr, CheckMatMulNBitsQkvFusedGraph)); + } +} + TEST_F(GraphTransformationTests, MatMulNBitsQkvFusionQLoRANumerics) { if (!DefaultWebGpuExecutionProvider()) GTEST_SKIP() << "WebGPU EP unavailable in this build."; int case_index = 0; diff --git a/onnxruntime/test/optimizer/webgpu_fusion_test_util.h b/onnxruntime/test/optimizer/webgpu_fusion_test_util.h index 88e922519d537..555a24a37fd39 100644 --- a/onnxruntime/test/optimizer/webgpu_fusion_test_util.h +++ b/onnxruntime/test/optimizer/webgpu_fusion_test_util.h @@ -6,15 +6,20 @@ #include #include #include +#include #include #include #include #include "core/framework/execution_provider.h" +#include "core/framework/execution_providers.h" +#include "core/framework/kernel_registry.h" #include "core/framework/session_options.h" #include "core/graph/constants.h" #include "core/graph/model.h" #include "core/optimizer/graph_transformer.h" +#include "core/optimizer/graph_transformer_utils.h" +#include "core/providers/cpu/cpu_execution_provider.h" #include "core/session/inference_session.h" #include "test/compare_ortvalue.h" #include "test/test_environment.h" @@ -25,6 +30,49 @@ namespace onnxruntime { namespace test { +#if !defined(DISABLE_CONTRIB_OPS) +// Exercise the real GenerateTransformers kernel probe without a GPU device. +class RegistryOnlyWebGpuExecutionProvider final : public IExecutionProvider { + public: + explicit RegistryOnlyWebGpuExecutionProvider(bool adapter_aware) + : IExecutionProvider(kWebGpuExecutionProvider), registry_(std::make_shared()) { + for (const char* op_type : {"MatMulNBitsMlp", "MatMulNBitsQkv"}) { + KernelDefBuilder builder; + builder.SetName(op_type) + .SetDomain(kMSDomain) + .SinceVersion(1) + .Provider(kWebGpuExecutionProvider) + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) + .TypeConstraint("T2", DataTypeImpl::GetTensorType()); + if (adapter_aware) { + builder.TypeConstraint("TAdapter", DataTypeImpl::GetTensorType()); + } + ORT_THROW_IF_ERROR(registry_->Register(builder, {})); + } + } + + std::shared_ptr GetKernelRegistry() const override { return registry_; } + + private: + std::shared_ptr registry_; +}; + +inline std::unique_ptr GenerateWebGpuFusionTransformerForTest( + std::string_view name, bool adapter_aware) { + ExecutionProviders providers; + ORT_THROW_IF_ERROR(providers.Add( + kWebGpuExecutionProvider, std::make_shared(adapter_aware))); + CPUExecutionProvider cpu_ep{CPUExecutionProviderInfo{}}; + auto transformers = optimizer_utils::GenerateTransformers( + TransformerLevel::Level2, {}, cpu_ep, DefaultLoggingManager().DefaultLogger(), + {}, nullptr, &providers); + for (auto& transformer : transformers) { + if (transformer->Name() == name) return std::move(transformer); + } + return nullptr; +} +#endif // !defined(DISABLE_CONTRIB_OPS) + // Variant of TransformerTester for WebGPU fusion tests that creates a fresh execution provider // per session via the provided factory, instead of sharing one EP across the baseline and target // sessions. Sharing a single WebGPU EP across multiple InferenceSessions in series can leave the From 87f75ca6eb0478d29d44f7689f13695694f1ba23 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Thu, 5 Nov 2026 17:56:43 +0800 Subject: [PATCH 3/4] Fix optional adapter detection in fused WebGPU kernels --- docs/ContribOperators.md | 4 +- docs/OperatorKernels.md | 8 +-- .../webgpu/quantization/matmul_nbits_mlp.h | 18 +++++-- .../webgpu/quantization/matmul_nbits_qkv.h | 18 +++++-- .../contrib_ops/matmul_nbits_qlora_test.cc | 54 +++++++++++++++++++ 5 files changed, 89 insertions(+), 13 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index bc8da37eb8745..872ad1a0e8fab 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -4444,7 +4444,7 @@ This version of the operator has been available since version 1 of the 'com.micr derived from the runtime shape of A and the shared attributes K and N. The operator contract includes a string attribute describing the fused gate activation. - + Optional gate and up LoRA inputs apply MatMulNBitsQLoRA to each projection before the activation. Empty adapter-up weights use the base projection. The adapter inputs are runtime tensors and may change between calls. @@ -9386,3 +9386,5 @@ No versioning maintained for experimental ops.
T : tensor(float)
Constrain input and output types to float32 tensors.
+ + diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index c0e1da179691d..081fc41754f7d 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -1777,9 +1777,9 @@ The **OpSet Version** column uses the following notation: |||[6, 12]|**T** = tensor(float), tensor(float16)| |GRU|*in* X:**T**
*in* W:**T**
*in* R:**T**
*in* B:**T**
*in* sequence_lens:**T1**
*in* initial_h:**T**
*out* Y:**T**
*out* Y_h:**T**|14+|**T** = tensor(float), tensor(float16)
**T1** = tensor(int32)| |||[7, 13]|**T** = tensor(float), tensor(float16)
**T1** = tensor(int32)| -|Gather|*in* data:**T**
*in* indices:**Tind**
*out* output:**T**|13+|**T** = tensor(bool), tensor(float), tensor(float16), tensor(int32), tensor(uint32), tensor(uint8)
**Tind** = tensor(int32), tensor(int64)| -|||[11, 12]|**T** = tensor(bool), tensor(float), tensor(float16), tensor(int32), tensor(uint32), tensor(uint8)
**Tind** = tensor(int32), tensor(int64)| -|||[1, 10]|**T** = tensor(bool), tensor(float), tensor(float16), tensor(int32), tensor(uint32), tensor(uint8)
**Tind** = tensor(int32), tensor(int64)| +|Gather|*in* data:**T**
*in* indices:**Tind**
*out* output:**T**|13+|**T** = tensor(bool), tensor(float), tensor(float16), tensor(int32), tensor(int8), tensor(uint32), tensor(uint8)
**Tind** = tensor(int32), tensor(int64)| +|||[11, 12]|**T** = tensor(bool), tensor(float), tensor(float16), tensor(int32), tensor(int8), tensor(uint32), tensor(uint8)
**Tind** = tensor(int32), tensor(int64)| +|||[1, 10]|**T** = tensor(bool), tensor(float), tensor(float16), tensor(int32), tensor(int8), tensor(uint32), tensor(uint8)
**Tind** = tensor(int32), tensor(int64)| |GatherElements|*in* data:**T**
*in* indices:**Tind**
*out* output:**T**|13+|**T** = tensor(float), tensor(float16)| |||[11, 12]|**T** = tensor(float), tensor(float16)| |GatherND|*in* data:**T**
*in* indices:**tensor(int64)**
*out* output:**T**|13+|**T** = tensor(float), tensor(float16), tensor(int32), tensor(uint32)
**indices** = tensor(int64)| @@ -1933,7 +1933,7 @@ The **OpSet Version** column uses the following notation: |||[13, 17]|**T** = tensor(float), tensor(float16), tensor(int32), tensor(uint32)| |||[11, 12]|**T** = tensor(float), tensor(float16), tensor(int32), tensor(uint32)| |||[1, 10]|**T** = tensor(float), tensor(float16), tensor(int32), tensor(uint32)| -|Relu|*in* X:**T**
*out* Y:**T**|14+|**T** = tensor(float), tensor(float16)| +|Relu|*in* X:**T**
*out* Y:**T**|14+|**T** = tensor(float), tensor(float16), tensor(int32)| |||13|**T** = tensor(float), tensor(float16)| |||[6, 12]|**T** = tensor(float), tensor(float16)| |Reshape|*in* data:**T**
*in* shape:**tensor(int64)**
*out* reshaped:**T**

or

*in* data:**T**
*out* reshaped:**T**|25+|**T** = tensor(bool), tensor(float), tensor(float16), tensor(int32), tensor(uint32), tensor(uint8)
**shape** = tensor(int64)| diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.h index db97da7b9a68e..220d5b1a1bfd8 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.h @@ -39,10 +39,20 @@ class MatMulNBitsMlp : public WebGpuKernel { epsilon_ = info.GetAttrOrDefault("epsilon", 1e-5f); lora_bits_ = info.GetAttrOrDefault("lora_bits", 8); lora_block_size_ = info.GetAttrOrDefault("lora_block_size", 32); - // GetInputCount is available through the plugin EP API. The schema checks - // that the optional adapter inputs are either all present or all absent. - has_adapter_inputs_ = info.GetInputCount() > 9; - ORT_ENFORCE(!has_adapter_inputs_ || info.GetInputCount() == 17, + // Trailing optional slots may be present in the node even when their edges + // are omitted. Inspect each edge rather than relying on GetInputCount(). + bool all_adapter_inputs = true; + has_adapter_inputs_ = false; + for (size_t i = 9; i < 17; ++i) { +#ifdef ORT_USE_EP_API_ADAPTERS + const bool exists = info.node().InputExists(i); +#else + const bool exists = i < static_cast(info.GetInputCount()) && info.node().InputDefs()[i]->Exists(); +#endif + has_adapter_inputs_ |= exists; + all_adapter_inputs &= exists; + } + ORT_ENFORCE(!has_adapter_inputs_ || all_adapter_inputs, "MatMulNBitsMlp requires all adapter inputs when any is present."); std::string activation; ORT_ENFORCE(info.GetAttr("activation", &activation).IsOK(), diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.h b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.h index 9bae2ec278bc1..9049a7e2a5c83 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.h +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.h @@ -23,10 +23,20 @@ class MatMulNBitsQkv : public WebGpuKernel { epsilon_ = info.GetAttrOrDefault("epsilon", 1e-6f); lora_bits_ = info.GetAttrOrDefault("lora_bits", 8); lora_block_size_ = info.GetAttrOrDefault("lora_block_size", 32); - // GetInputCount is available through the plugin EP API. The schema checks - // that the optional adapter inputs are either all present or all absent. - has_adapter_inputs_ = info.GetInputCount() > 12; - ORT_ENFORCE(!has_adapter_inputs_ || info.GetInputCount() == 24, + // Trailing optional slots may be present in the node even when their edges + // are omitted. Inspect each edge rather than relying on GetInputCount(). + bool all_adapter_inputs = true; + has_adapter_inputs_ = false; + for (size_t i = 12; i < 24; ++i) { +#ifdef ORT_USE_EP_API_ADAPTERS + const bool exists = info.node().InputExists(i); +#else + const bool exists = i < static_cast(info.GetInputCount()) && info.node().InputDefs()[i]->Exists(); +#endif + has_adapter_inputs_ |= exists; + all_adapter_inputs &= exists; + } + ORT_ENFORCE(!has_adapter_inputs_ || all_adapter_inputs, "MatMulNBitsQkv requires all adapter inputs when any is present."); ORT_ENFORCE(bits_ == 4, "MatMulNBitsQkv currently supports 4-bit weights only."); diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc index c479e1fe0a109..3c3f29cb7cab4 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_qlora_test.cc @@ -52,6 +52,60 @@ void RunEmptyAndActiveAdapters() { TEST(MatMulNBitsQLoRA, EmptyAndActiveAdapters) { RunEmptyAndActiveAdapters(); } TEST(MatMulNBitsQLoRA, EmptyAndActiveAdaptersFp32) { RunEmptyAndActiveAdapters(); } +TEST(MatMulNBitsQLoRA, BaseOnlyFusedOpsWithPaddedAdapterInputs) { + for (bool qkv : {false, true}) { + const int projection_count = qkv ? 3 : 2; + const int adapter_count = qkv ? 12 : 8; + for (int padded_inputs : {1, adapter_count}) { + SCOPED_TRACE(MakeString(qkv ? "QKV" : "MLP", ", padded_inputs=", padded_inputs)); + auto ep = DefaultWebGpuExecutionProvider(); + ASSERT_NE(ep, nullptr); + OpTester test(qkv ? "MatMulNBitsQkv" : "MatMulNBitsMlp", 1, kMSDomain); + test.AddAttribute("K", 32); + test.AddAttribute(qkv ? "Nq" : "N", 8); + if (qkv) { + test.AddAttribute("Nkv", 4); + } else { + test.AddAttribute("activation", "silu"); + } + test.AddAttribute("bits", 4); + test.AddAttribute("block_size", 32); + test.AddInput("A", {1, 32}, std::vector(32, 0.5f)); + test.AddOptionalInputEdge(); // skip + if (qkv) { + test.AddInput("norm_scale", {32}, std::vector(32, 1.0f)); + } else { + test.AddOptionalInputEdge(); // norm_scale + } + for (int projection = 0; projection < projection_count; ++projection) { + const int64_t n = qkv && projection > 0 ? 4 : 8; + test.AddInput(MakeString("weight_", projection).c_str(), {n, 1, 16}, + std::vector(n * 16, 0x88), true); + test.AddInput(MakeString("scales_", projection).c_str(), {n, 1}, + std::vector(n, 0.03125f), true); + test.AddOptionalInputEdge(); // bias + } + for (int i = 0; i < padded_inputs; ++i) { + if (i % 4 == 0 || i % 4 == 2) { + test.AddOptionalInputEdge(); + } else { + test.AddOptionalInputEdge(); + } + } + if (qkv) { + test.AddOutput("Q", {1, 8}, std::vector(8)); + test.AddOutput("K", {1, 4}, std::vector(4)); + test.AddOutput("V", {1, 4}, std::vector(4)); + } else { + test.AddOutput("Y", {1, 8}, std::vector(8)); + } + std::vector> eps; + eps.push_back(std::move(ep)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); + } + } +} + TEST(MatMulNBitsQLoRA, FusedEmptyAdaptersRejectMalformedBaseWeights) { for (bool qkv : {false, true}) { const int projection_count = qkv ? 3 : 2; From 079608aa4dd16a2a222a0c29ef3e96cedb971a4c Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Thu, 5 Nov 2026 18:00:45 +0800 Subject: [PATCH 4/4] Account for ORT-format resolver growth in Android size check --- .github/workflows/android.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/android.yml b/.github/workflows/android.yml index a0975d1cf4abf..e86e0a26602b3 100644 --- a/.github/workflows/android.yml +++ b/.github/workflows/android.yml @@ -80,7 +80,8 @@ jobs: run: | set -e -x BINARY_SIZE_THRESHOLD_ARGS="" - BINARY_SIZE_THRESHOLD_IN_BYTES=1610752 + # ORT-format compatibility for legacy MatMulNBitsMlp/Qkv nodes adds 1,920 bytes. + BINARY_SIZE_THRESHOLD_IN_BYTES=1612672 echo "Binary size threshold in bytes: ${BINARY_SIZE_THRESHOLD_IN_BYTES}" BINARY_SIZE_THRESHOLD_ARGS="--threshold_size_in_bytes ${BINARY_SIZE_THRESHOLD_IN_BYTES}"