bug fox

microsoft · Aug 20, 2024 · 0ae3953 · 0ae3953
1 parent fc63748
commit 0ae3953
Showing 1 changed file with 26 additions and 28 deletions.
diff --git a/bitblas/gpu/gemv_dequantize.py b/bitblas/gpu/gemv_dequantize.py
@@ -158,20 +158,19 @@ def get_idx(weight_decode_info: Dict):
 
         skip_blocks = [block_shared_local_B]
 
-        if "zeros_mode" in weight_decode_info and weight_decode_info["zeros_mode"] == "quantized":
-            if "with_scaling" in weight_decode_info and weight_decode_info["with_scaling"]:
-                block_local_scales = sch.cache_read(block_decode_B,
-                                                    get_idx(weight_decode_info) + 1, "local")
-                sch.compute_at(block_local_scales, tx, preserve_unit_loops=True)
-                auto_inline_producers(sch, block_local_scales)
-                skip_blocks.append(block_local_scales)
-
-            if "with_zeros" in weight_decode_info and weight_decode_info["with_zeros"]:
-                block_local_zeros = sch.cache_read(block_decode_B,
-                                                   get_idx(weight_decode_info) + 2, "local")
-                sch.compute_at(block_local_zeros, tx, preserve_unit_loops=True)
-                auto_inline_producers(sch, block_local_zeros)
-                skip_blocks.append(block_local_zeros)
+        if "with_scaling" in weight_decode_info and weight_decode_info["with_scaling"]:
+            block_local_scales = sch.cache_read(block_decode_B,
+                                                get_idx(weight_decode_info) + 1, "local")
+            sch.compute_at(block_local_scales, tx, preserve_unit_loops=True)
+            auto_inline_producers(sch, block_local_scales)
+            skip_blocks.append(block_local_scales)
+
+        if "with_zeros" in weight_decode_info and weight_decode_info["with_zeros"]:
+            block_local_zeros = sch.cache_read(block_decode_B,
+                                               get_idx(weight_decode_info) + 2, "local")
+            sch.compute_at(block_local_zeros, tx, preserve_unit_loops=True)
+            auto_inline_producers(sch, block_local_zeros)
+            skip_blocks.append(block_local_zeros)
 
         auto_inline_producers(sch, block_decode_B, skip_blocks)
 
@@ -329,20 +328,19 @@ def get_idx(weight_decode_info: Dict):
 
         skip_blocks = [block_shared_local_B]
 
-        if "zeros_mode" in weight_decode_info and weight_decode_info["zeros_mode"] == "quantized":
-            if "with_scaling" in weight_decode_info and weight_decode_info["with_scaling"]:
-                block_local_scales = sch.cache_read(block_decode_B,
-                                                    get_idx(weight_decode_info) + 1, "local")
-                sch.compute_at(block_local_scales, tx, preserve_unit_loops=True)
-                auto_inline_producers(sch, block_local_scales)
-                skip_blocks.append(block_local_scales)
-
-            if "with_zeros" in weight_decode_info and weight_decode_info["with_zeros"]:
-                block_local_zeros = sch.cache_read(block_decode_B,
-                                                   get_idx(weight_decode_info) + 2, "local")
-                sch.compute_at(block_local_zeros, tx, preserve_unit_loops=True)
-                auto_inline_producers(sch, block_local_zeros)
-                skip_blocks.append(block_local_zeros)
+        if "with_scaling" in weight_decode_info and weight_decode_info["with_scaling"]:
+            block_local_scales = sch.cache_read(block_decode_B,
+                                                get_idx(weight_decode_info) + 1, "local")
+            sch.compute_at(block_local_scales, tx, preserve_unit_loops=True)
+            auto_inline_producers(sch, block_local_scales)
+            skip_blocks.append(block_local_scales)
+
+        if "with_zeros" in weight_decode_info and weight_decode_info["with_zeros"]:
+            block_local_zeros = sch.cache_read(block_decode_B,
+                                               get_idx(weight_decode_info) + 2, "local")
+            sch.compute_at(block_local_zeros, tx, preserve_unit_loops=True)
+            auto_inline_producers(sch, block_local_zeros)
+            skip_blocks.append(block_local_zeros)
 
         auto_inline_producers(sch, block_decode_B, skip_blocks)