Skip to content

[WebGPU] Reuse MatMulNBits MLP/QKV for QLoRA adapters (alternative to #33901) #33912

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Jiajia Qin (qjia7) wants to merge 4 commits into
base: main
Choose a base branch
from
Draft

[WebGPU] Reuse MatMulNBits MLP/QKV for QLoRA adapters (alternative to #33901) #33912

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
1c86c61
[WebGPU] Reuse MatMulNBits MLP and QKV for runtime QLoRA adapters
qjia7 Nov 5, 2026
4febae8
[WebGPU] Preserve base fusion compatibility and fix plugin build
qjia7 Nov 5, 2026
87f75ca
Fix optional adapter detection in fused WebGPU kernels
qjia7 Nov 5, 2026
079608a
Account for ORT-format resolver growth in Android size check
qjia7 Nov 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion .github/workflows/android.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}"

Expand Down
64 changes: 62 additions & 2 deletions docs/ContribOperators.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4445,6 +4445,10 @@ This version of the operator has been available since version 1 of the 'com.micr

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)
Expand Down Expand Up @@ -4473,9 +4477,13 @@ This version of the operator has been available since version 1 of the 'com.micr
<dd>Size of each quantization block along the K dimension. Currently only block_size=32 is supported by the WebGPU kernel.</dd>
<dt><tt>epsilon</tt> : float</dt>
<dd>Epsilon used by the optional fused (Skip)SimplifiedLayerNormalization. Defaults to 1e-5.</dd>
<dt><tt>lora_bits</tt> : int</dt>
<dd>Bit-width of optional adapter weights.</dd>
<dt><tt>lora_block_size</tt> : int</dt>
<dd>Quantization block size of optional adapter weights.</dd>
</dl>

#### Inputs (8 - 9)
#### Inputs (8 - 17)

<dl>
<dt><tt>A</tt> : T1</dt>
Expand All @@ -4496,6 +4504,22 @@ This version of the operator has been available since version 1 of the 'com.micr
<dd>Per-block scaling factors for the up projection.</dd>
<dt><tt>up_bias</tt> (optional) : T1</dt>
<dd>Optional bias for the up projection with shape [N].</dd>
<dt><tt>gate_lora_A</tt> (optional) : TAdapter</dt>
<dd>Optional packed gate adapter-down weights.</dd>
<dt><tt>gate_lora_A_scales</tt> (optional) : T1</dt>
<dd>Optional gate adapter-down scales.</dd>
<dt><tt>gate_lora_B</tt> (optional) : TAdapter</dt>
<dd>Optional packed gate adapter-up weights; may be empty.</dd>
<dt><tt>gate_lora_B_scales</tt> (optional) : T1</dt>
<dd>Optional gate adapter-up scales.</dd>
<dt><tt>up_lora_A</tt> (optional) : TAdapter</dt>
<dd>Optional packed up adapter-down weights.</dd>
<dt><tt>up_lora_A_scales</tt> (optional) : T1</dt>
<dd>Optional up adapter-down scales.</dd>
<dt><tt>up_lora_B</tt> (optional) : TAdapter</dt>
<dd>Optional packed up adapter-up weights; may be empty.</dd>
<dt><tt>up_lora_B_scales</tt> (optional) : T1</dt>
<dd>Optional up adapter-up scales.</dd>
</dl>

#### Outputs (1 - 2)
Expand All @@ -4514,6 +4538,8 @@ This version of the operator has been available since version 1 of the 'com.micr
<dd>Constrain input and output types to float tensors.</dd>
<dt><tt>T2</tt> : tensor(uint8)</dt>
<dd>Constrain quantized weight types to uint8.</dd>
<dt><tt>TAdapter</tt> : tensor(uint8)</dt>
<dd>Constrain packed adapter weights to uint8.</dd>
</dl>


Expand Down Expand Up @@ -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

Expand All @@ -4617,9 +4645,13 @@ This version of the operator has been available since version 1 of the 'com.micr
<dd>Size of each quantization block along the K dimension. Currently only block_size=32 is supported by the WebGPU kernel.</dd>
<dt><tt>epsilon</tt> : float</dt>
<dd>Epsilon used by the simplified layer norm reduction.</dd>
<dt><tt>lora_bits</tt> : int</dt>
<dd>Bit-width of optional adapter weights.</dd>
<dt><tt>lora_block_size</tt> : int</dt>
<dd>Quantization block size of optional adapter weights.</dd>
</dl>

#### Inputs (11 - 12)
#### Inputs (11 - 24)

<dl>
<dt><tt>A</tt> : T1</dt>
Expand All @@ -4646,6 +4678,30 @@ This version of the operator has been available since version 1 of the 'com.micr
<dd>Per-block scaling factors for the V projection.</dd>
<dt><tt>v_bias</tt> (optional) : T1</dt>
<dd>Optional bias for the V projection with shape [Nkv].</dd>
<dt><tt>q_lora_A</tt> (optional) : TAdapter</dt>
<dd>Optional packed Q adapter-down weights.</dd>
<dt><tt>q_lora_A_scales</tt> (optional) : T1</dt>
<dd>Optional Q adapter-down scales.</dd>
<dt><tt>q_lora_B</tt> (optional) : TAdapter</dt>
<dd>Optional packed Q adapter-up weights; may be empty.</dd>
<dt><tt>q_lora_B_scales</tt> (optional) : T1</dt>
<dd>Optional Q adapter-up scales.</dd>
<dt><tt>k_lora_A</tt> (optional) : TAdapter</dt>
<dd>Optional packed K adapter-down weights.</dd>
<dt><tt>k_lora_A_scales</tt> (optional) : T1</dt>
<dd>Optional K adapter-down scales.</dd>
<dt><tt>k_lora_B</tt> (optional) : TAdapter</dt>
<dd>Optional packed K adapter-up weights; may be empty.</dd>
<dt><tt>k_lora_B_scales</tt> (optional) : T1</dt>
<dd>Optional K adapter-up scales.</dd>
<dt><tt>v_lora_A</tt> (optional) : TAdapter</dt>
<dd>Optional packed V adapter-down weights.</dd>
<dt><tt>v_lora_A_scales</tt> (optional) : T1</dt>
<dd>Optional V adapter-down scales.</dd>
<dt><tt>v_lora_B</tt> (optional) : TAdapter</dt>
<dd>Optional packed V adapter-up weights; may be empty.</dd>
<dt><tt>v_lora_B_scales</tt> (optional) : T1</dt>
<dd>Optional V adapter-up scales.</dd>
</dl>

#### Outputs (3 - 4)
Expand All @@ -4668,6 +4724,8 @@ This version of the operator has been available since version 1 of the 'com.micr
<dd>Constrain input and output types to float tensors.</dd>
<dt><tt>T2</tt> : tensor(uint8)</dt>
<dd>Constrain quantized weight types to uint8.</dd>
<dt><tt>TAdapter</tt> : tensor(uint8)</dt>
<dd>Constrain packed adapter weights to uint8.</dd>
</dl>


Expand Down Expand Up @@ -9328,3 +9386,5 @@ No versioning maintained for experimental ops.
<dt><tt>T</tt> : tensor(float)</dt>
<dd>Constrain input and output types to float32 tensors.</dd>
</dl>


12 changes: 6 additions & 6 deletions docs/OperatorKernels.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1777,9 +1777,9 @@ The **OpSet Version** column uses the following notation:
|||[6, 12]|**T** = tensor(float), tensor(float16)|
|GRU|*in* X:**T**<br> *in* W:**T**<br> *in* R:**T**<br> *in* B:**T**<br> *in* sequence_lens:**T1**<br> *in* initial_h:**T**<br> *out* Y:**T**<br> *out* Y_h:**T**|14+|**T** = tensor(float), tensor(float16)<br/> **T1** = tensor(int32)|
|||[7, 13]|**T** = tensor(float), tensor(float16)<br/> **T1** = tensor(int32)|
|Gather|*in* data:**T**<br> *in* indices:**Tind**<br> *out* output:**T**|13+|**T** = tensor(bool), tensor(float), tensor(float16), tensor(int32), tensor(uint32), tensor(uint8)<br/> **Tind** = tensor(int32), tensor(int64)|
|||[11, 12]|**T** = tensor(bool), tensor(float), tensor(float16), tensor(int32), tensor(uint32), tensor(uint8)<br/> **Tind** = tensor(int32), tensor(int64)|
|||[1, 10]|**T** = tensor(bool), tensor(float), tensor(float16), tensor(int32), tensor(uint32), tensor(uint8)<br/> **Tind** = tensor(int32), tensor(int64)|
|Gather|*in* data:**T**<br> *in* indices:**Tind**<br> *out* output:**T**|13+|**T** = tensor(bool), tensor(float), tensor(float16), tensor(int32), tensor(int8), tensor(uint32), tensor(uint8)<br/> **Tind** = tensor(int32), tensor(int64)|
|||[11, 12]|**T** = tensor(bool), tensor(float), tensor(float16), tensor(int32), tensor(int8), tensor(uint32), tensor(uint8)<br/> **Tind** = tensor(int32), tensor(int64)|
|||[1, 10]|**T** = tensor(bool), tensor(float), tensor(float16), tensor(int32), tensor(int8), tensor(uint32), tensor(uint8)<br/> **Tind** = tensor(int32), tensor(int64)|
|GatherElements|*in* data:**T**<br> *in* indices:**Tind**<br> *out* output:**T**|13+|**T** = tensor(float), tensor(float16)|
|||[11, 12]|**T** = tensor(float), tensor(float16)|
|GatherND|*in* data:**T**<br> *in* indices:**tensor(int64)**<br> *out* output:**T**|13+|**T** = tensor(float), tensor(float16), tensor(int32), tensor(uint32)<br/> **indices** = tensor(int64)|
Expand Down Expand Up @@ -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**<br> *out* Y:**T**|14+|**T** = tensor(float), tensor(float16)|
|Relu|*in* X:**T**<br> *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**<br> *in* shape:**tensor(int64)**<br> *out* reshaped:**T**<br><br>or<br><br>*in* data:**T**<br> *out* reshaped:**T**|25+|**T** = tensor(bool), tensor(float), tensor(float16), tensor(int32), tensor(uint32), tensor(uint8)<br/> **shape** = tensor(int64)|
Expand Down Expand Up @@ -2056,9 +2056,9 @@ The **OpSet Version** column uses the following notation:
|MatMulBlockQuantizedFp8Weight|*in* A:**T**<br> *in* B:**T1**<br> *in* b_scale:**T2**<br> *in* a_scale:**T2**<br> *in* bias:**T**<br> *out* Y:**T**|1+|**T** = tensor(float16)<br/> **T1** = tensor(float8e4m3fn)<br/> **T2** = tensor(float)|
|MatMulBnb4|*in* A:**T1**<br> *in* B:**T2**<br> *in* absmax:**T1**<br> *out* Y:**T1**|1+|**T1** = tensor(float), tensor(float16)<br/> **T2** = tensor(uint8)|
|MatMulNBits|*in* A:**T1**<br> *in* B:**T2**<br> *in* scales:**T1**<br> *in* zero_points:**T3**<br> *in* g_idx:**T4**<br> *in* bias:**T1**<br> *out* Y:**T1**|1+|**T1** = tensor(float), tensor(float16)<br/> **T2** = tensor(uint8)<br/> **T3** = tensor(uint8)<br/> **T4** = tensor(int32)|
|MatMulNBitsMlp|*in* A:**T1**<br> *in* skip:**T1**<br> *in* norm_scale:**T1**<br> *in* gate_B:**T2**<br> *in* gate_scales:**T1**<br> *in* gate_bias:**T1**<br> *in* up_B:**T2**<br> *in* up_scales:**T1**<br> *in* up_bias:**T1**<br> *out* Y:**T1**<br> *out* input_skip_bias_sum:**T1**|1+|**T1** = tensor(float), tensor(float16)<br/> **T2** = tensor(uint8)|
|MatMulNBitsMlp|*in* A:**T1**<br> *in* skip:**T1**<br> *in* norm_scale:**T1**<br> *in* gate_B:**T2**<br> *in* gate_scales:**T1**<br> *in* gate_bias:**T1**<br> *in* up_B:**T2**<br> *in* up_scales:**T1**<br> *in* up_bias:**T1**<br> *in* gate_lora_A:**TAdapter**<br> *in* gate_lora_A_scales:**T1**<br> *in* gate_lora_B:**TAdapter**<br> *in* gate_lora_B_scales:**T1**<br> *in* up_lora_A:**TAdapter**<br> *in* up_lora_A_scales:**T1**<br> *in* up_lora_B:**TAdapter**<br> *in* up_lora_B_scales:**T1**<br> *out* Y:**T1**<br> *out* input_skip_bias_sum:**T1**|1+|**T1** = tensor(float), tensor(float16)<br/> **T2** = tensor(uint8)<br/> **TAdapter** = tensor(uint8)|
|MatMulNBitsQLoRA|*in* X:**T**<br> *in* W:**TQ**<br> *in* scales:**T**<br> *in* lora_A:**TQ**<br> *in* lora_A_scales:**T**<br> *in* lora_B:**TQ**<br> *in* lora_B_scales:**T**<br> *out* Y:**T**|1+|**T** = tensor(float), tensor(float16)<br/> **TQ** = tensor(uint8)|
|MatMulNBitsQkv|*in* A:**T1**<br> *in* skip:**T1**<br> *in* norm_scale:**T1**<br> *in* q_B:**T2**<br> *in* q_scales:**T1**<br> *in* q_bias:**T1**<br> *in* k_B:**T2**<br> *in* k_scales:**T1**<br> *in* k_bias:**T1**<br> *in* v_B:**T2**<br> *in* v_scales:**T1**<br> *in* v_bias:**T1**<br> *out* Q:**T1**<br> *out* K:**T1**<br> *out* V:**T1**<br> *out* input_skip_bias_sum:**T1**|1+|**T1** = tensor(float), tensor(float16)<br/> **T2** = tensor(uint8)|
|MatMulNBitsQkv|*in* A:**T1**<br> *in* skip:**T1**<br> *in* norm_scale:**T1**<br> *in* q_B:**T2**<br> *in* q_scales:**T1**<br> *in* q_bias:**T1**<br> *in* k_B:**T2**<br> *in* k_scales:**T1**<br> *in* k_bias:**T1**<br> *in* v_B:**T2**<br> *in* v_scales:**T1**<br> *in* v_bias:**T1**<br> *in* q_lora_A:**TAdapter**<br> *in* q_lora_A_scales:**T1**<br> *in* q_lora_B:**TAdapter**<br> *in* q_lora_B_scales:**T1**<br> *in* k_lora_A:**TAdapter**<br> *in* k_lora_A_scales:**T1**<br> *in* k_lora_B:**TAdapter**<br> *in* k_lora_B_scales:**T1**<br> *in* v_lora_A:**TAdapter**<br> *in* v_lora_A_scales:**T1**<br> *in* v_lora_B:**TAdapter**<br> *in* v_lora_B_scales:**T1**<br> *out* Q:**T1**<br> *out* K:**T1**<br> *out* V:**T1**<br> *out* input_skip_bias_sum:**T1**|1+|**T1** = tensor(float), tensor(float16)<br/> **T2** = tensor(uint8)<br/> **TAdapter** = tensor(uint8)|
|MoE|*in* input:**T**<br> *in* router_probs:**T**<br> *in* fc1_experts_weights:**T**<br> *in* fc1_experts_bias:**T**<br> *in* fc2_experts_weights:**T**<br> *in* fc2_experts_bias:**T**<br> *in* fc3_experts_weights:**T**<br> *in* fc3_experts_bias:**T**<br> *out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
|MultiHeadAttention|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* bias:**T**<br> *in* key_padding_mask:**M**<br> *in* attention_bias:**T**<br> *in* past_key:**T**<br> *in* past_value:**T**<br> *in* past_sequence_length:**M**<br> *in* cache_indirection:**M**<br> *out* output:**T**<br> *out* present_key:**T**<br> *out* present_value:**T**<br> *out* qk:**QK**|1+|**T** = tensor(float), tensor(float16)|
|NGramHashMapping|*in* input_ids:**M**<br> *in* multipliers:**M**<br> *in* vocab_sizes:**M**<br> *in* past_ids:**M**<br> *in* head_offsets:**M**<br> *in* eos_token_id:**M**<br> *in* segment_ids:**tensor(int32)**<br> *out* hash_ids:**M**<br> *out* present_ids:**M**|1+|**M** = tensor(int32)|
Expand Down
80 changes: 79 additions & 1 deletion onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
#include <optional>

#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"
Expand Down Expand Up @@ -240,10 +241,18 @@ ONNX_OPERATOR_KERNEL_EX(
kWebGpuExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T1", WebGpuSupportedFloatTypes())
.TypeConstraint("T2", DataTypeImpl::GetTensorType<uint8_t>()),
.TypeConstraint("T2", DataTypeImpl::GetTensorType<uint8_t>())
.TypeConstraint("TAdapter", DataTypeImpl::GetTensorType<uint8_t>()),
Comment thread
qjia7 marked this conversation as resolved.
Resolved
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<Tensor>(0);
const Tensor* skip = context.Input<Tensor>(1);
const Tensor* norm_scale = context.Input<Tensor>(2);
Expand Down Expand Up @@ -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<Tensor>(5) != nullptr || context.Input<Tensor>(8) != nullptr,
"MatMulNBitsMlp with adapters does not support projection biases.");
const Tensor* gate_lora_b = context.Input<Tensor>(11);
const Tensor* up_lora_b = context.Input<Tensor>(15);
if (gate_lora_b->Shape().Size() == 0 && up_lora_b->Shape().Size() == 0) {
ORT_RETURN_IF_ERROR(ValidateQLoRABaseWeights(context.Input<Tensor>(3), context.Input<Tensor>(4), K_, N_));
ORT_RETURN_IF_ERROR(ValidateQLoRABaseWeights(context.Input<Tensor>(6), context.Input<Tensor>(7), K_, N_));
return ComputeBase(context);
}

const Tensor* a = context.Input<Tensor>(0);
const Tensor* skip = context.Input<Tensor>(1);
const Tensor* norm_scale = context.Input<Tensor>(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<Tensor> 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<uint32_t>(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<Tensor>(3), context.Input<Tensor>(4),
context.Input<Tensor>(9), context.Input<Tensor>(10),
gate_lora_b, context.Input<Tensor>(12),
K_, N_, accuracy_level_, context, &gate));
ORT_RETURN_IF_ERROR(ApplyMatMulNBitsQLoRA(projection_input, context.Input<Tensor>(6), context.Input<Tensor>(7),
context.Input<Tensor>(13), context.Input<Tensor>(14),
up_lora_b, context.Input<Tensor>(16),
K_, N_, accuracy_level_, context, &up));

const uint32_t vec_size = (narrow<uint32_t>(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
Loading
Loading