[Dev] Implement ScheduleUnsafeInjectCallArgument Primitive to Hack decoding

bitblas/base/utils.py

@@ -210,7 +210,11 @@ def tvm_callback_cuda_postproc(code, _):
             code = tensor_remove_make_int2(code)
             return code
 
-        with tvm.transform.PassContext(config={"tir.use_async_copy": True, **config.pass_context}):
+        with tvm.transform.PassContext(config={
+                "tir.use_async_copy": True,
+                "tir.disable_cse_tir": True,
+                **config.pass_context
+        }):
             rt_mod = tvm.build(mod, target=arch.target)
 
         from tvm.contrib.tar import tar  # pylint: disable=import-outside-toplevel

bitblas/ops/general_matmul/tirscript/matmul_dequantize_impl.py

@@ -14,6 +14,7 @@
     _tir_u8_to_f8_e4m3_to_f16,
     _tir_packed_to_unsigned_convert_with_zeros,
 )
+from typing import Union
 
 
 class MatMulNTDequantizeEmitter:
@@ -514,8 +515,11 @@ def matmul_nt_dequantize_b_propagate_b(
     fast_decoding=False,
     with_bias=False,
     zeros_mode="original",
-    transform_kind: TransformKind = TransformKind.IntraWarpTransform,
+    transform_kind: Union[int, TransformKind] = TransformKind.NonTransform,
 ):
+    if isinstance(transform_kind, int):
+        transform_kind = TransformKind(transform_kind)
+
     assert bit in [1, 2, 4, 8], "Unsupported bit: {}".format(bit)
     if not isinstance(M, int):
         M = tvm.te.var("m")
@@ -695,9 +699,14 @@ def matmul_nt_dequantize_b_propagate_a_propagate_b(
     fast_decoding=False,
     with_bias=False,
     zeros_mode="original",
-    transform_kind_input: TransformKind = TransformKind.IntraWarpTransform,
-    transform_kind_weight: TransformKind = TransformKind.IntraWarpTransform,
+    transform_kind_input: Union[int, TransformKind] = TransformKind.NonTransform,
+    transform_kind_weight: Union[int, TransformKind] = TransformKind.NonTransform,
 ):
+    if isinstance(transform_kind_input, int):
+        transform_kind_input = TransformKind(transform_kind_input)
+    if isinstance(transform_kind_weight, int):
+        transform_kind_weight = TransformKind(transform_kind_weight)
+
     assert bit in [1, 2, 4, 8], "Unsupported bit: {}".format(bit)
     if not isinstance(M, int):
         M = tvm.te.var("m")

bitblas/ops/ladder_permutate/ladder_permutate_impl.py

@@ -63,8 +63,6 @@ def select_implementation(
             None,
         )
 
-    ladder_stage3_map_inverse = ladder_stage3_map.inverse([l, r])
-
     inp = te.placeholder((M, N // scaling_factor), name="inp", dtype=storage_dtype)
     args = [inp]
 
@@ -97,6 +95,8 @@ def fcompute(*args):
     if transform_kind >= 3:
         arg = args[-1]
 
+        ladder_stage3_map_inverse = ladder_stage3_map.inverse([l, r])
+
         def fcompute(*args):
             warp_i, warp_j = args[-2:]
             spatial_args = args[:-2]

bitblas/ops/operator.py

@@ -100,6 +100,7 @@ def tvm_callback_cuda_postproc(code, _):
                 # Use a specific TVM pass context for CUDA platforms
                 with tvm.transform.PassContext(config={
                         "tir.use_async_copy": True,
+                        "tir.disable_cse_tir": True,
                         **self.pass_context
                 }):
                     rt_mod = tvm.build(self.optimized_func, target=target, name=self.name)

testing/python/module/test_repack_from_gptq.py

@@ -4,7 +4,7 @@
 import torch
 
 try:
-    import auto_gptq  # noqa
+    import auto_gptq  # noqa: F401
 except ImportError as e:
     raise ImportError("Please install auto-gptq by running `pip install auto-gptq`") from e
 
@@ -16,6 +16,8 @@
 
 
 def assert_output_with_gptq(m, in_features, out_features, group_size):
+    if group_size == -1:
+        group_size = in_features
     _, linear, s, _ = bitblas.quantization.gen_quant4(in_features, out_features, group_size)
 
     zeros = torch.full((in_features // group_size, out_features), 7, dtype=torch.int32)