From 3c6f22f9dc0c3c08788b70b8b70b35cc64496e6e Mon Sep 17 00:00:00 2001
From: Shivani767 <shivanibhandari767@gmail.com>
Date: Mon, 15 Jun 2026 21:10:10 +0530
Subject: [PATCH] Add 2-bit quantization support to WebGPU GatherBlockQuantized
 operator

---
 .../webgpu/quantization/gather_block_quantized.cc   | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/onnxruntime/contrib_ops/webgpu/quantization/gather_block_quantized.cc b/onnxruntime/contrib_ops/webgpu/quantization/gather_block_quantized.cc
index b7f213a57398a..f1c70868af15e 100755
--- a/onnxruntime/contrib_ops/webgpu/quantization/gather_block_quantized.cc
+++ b/onnxruntime/contrib_ops/webgpu/quantization/gather_block_quantized.cc
@@ -74,6 +74,10 @@ Status GatherBlockQuantizedProgram::GenerateShaderCode(ShaderHelper& shader) con
         << "  let byte_in_word_2b = byte_idx_2b % 4;\n"
         << "  let unpacked_bytes_2b = " << unpack << "(u32(packed_word_2b));\n"
         << "  var quantized_data = (unpacked_bytes_2b[byte_in_word_2b] >> bit_shift_2b) & 0x3;\n";
+    if (is_signed_) {
+      shader.MainFunctionBody()
+          << "  if((quantized_data & 0x2) != 0) { quantized_data = quantized_data - 4 ;};\n";
+    }
   } else if (is_4bit) {
     shader.MainFunctionBody()
         << "  let data_index = data_offset % 8;\n"
@@ -149,8 +153,13 @@ Status GatherBlockQuantizedProgram::GenerateShaderCode(ShaderHelper& shader) con
           << "  var zero_point = zero_point_vec[zero_point_index];\n";
     }
     if (is_signed_) {
-      shader.MainFunctionBody()
-          << "  if((zero_point & 0x8) != 0) { zero_point = zero_point - 16 ;};\n";
+      if (is_2bit) {
+        shader.MainFunctionBody()
+            << "  if((zero_point & 0x2) != 0) { zero_point = zero_point - 4 ;};\n";
+      } else if (is_4bit) {
+        shader.MainFunctionBody()
+            << "  if((zero_point & 0x8) != 0) { zero_point = zero_point - 16 ;};\n";
+      }
     }
   }
   shader.MainFunctionBody()