[BUG Fix] Fix the NVCC Comple options for CUDA Version >= 12.5

bitblas/builder/lib_generator/__init__.py

@@ -43,7 +43,8 @@ def compile_lib(self, timeout: float = None):
             "--shared",
             src.name,
             "-lcuda",
-            f"-gencode=arch=compute_{compute_version},code=compute_{compute_version}",
+            "-gencode",
+            f"arch=compute_{compute_version},code=sm_{compute_version}",
             "-o",
             libpath,
         ]

bitblas/gpu/intrin/lop3.py

@@ -1495,6 +1495,7 @@ def fast_decode_impl(
     (1, "int8", "float16", 8, "local", "uint", True, True, "rescale"),
     (4, "int8", "int8", 8, "local", "uint", False, False, "original"),
     (4, "int8", "int8", 16, "local", "uint", False, False, "original"),
+    (4, "int8", "int8", 16, "local", "int", False, False, "original"),
     (2, "int8", "int8", 16, "local", "uint", False, False, "original"),
     (2, "int8", "int8", 16, "local", "int", False, False, "original"),
     (1, "int8", "int8", 16, "local", "uint", False, False, "original"),
@@ -1523,6 +1524,7 @@ def fast_decode_impl(
     (1, "int8", "float16", 8, "warp", "uint", True, True, "rescale"),
     (4, "int8", "int8", 8, "warp", "uint", False, False, "original"),
     (4, "int8", "int8", 16, "warp", "uint", False, False, "original"),
+    (4, "int8", "int8", 16, "warp", "int", False, False, "original"),
     (2, "int8", "int8", 16, "warp", "uint", False, False, "original"),
     (2, "int8", "int8", 16, "warp", "int", False, False, "original"),
     (1, "int8", "int8", 16, "warp", "uint", False, False, "original"),

bitblas/gpu/matmul_analysis.py

@@ -623,7 +623,8 @@ def check_last_trait(region: List[Range]):
         # Currently, we only support block reduction depth 2 for small M
         # When the func is a dequantize like ops, we should consider the M
         require_block_reduce = False
-        if hasattr(func.attrs, "dequantize_info"):
+        # And we only support float16 for now
+        if hasattr(func.attrs, "dequantize_info") and in_dtype == "float16":
             for arg in func.params:
                 inp_shape = func.buffer_map[arg].shape
                 M = inp_shape[0]

bitblas/ops/general_matmul/__init__.py

@@ -85,6 +85,8 @@ class MatmulConfig(OperatorConfig):
         None  # propagate_b is a flag to control the ladder permutation
     )
 
+    # TODO: This is a temporary solution to legalize the dynamic symbolic.
+    # Maybe we should remove this in the future.
     # optimize strategy, default is SingleBatchDecodeOnly
     optimize_stratety: Union[int, OptimizeStrategy] = OptimizeStrategy.SingleBatchDecodeOnly
 

bitblas/wrapper/general.py

@@ -160,7 +160,8 @@ def compile_lib(self, timeout: float = None):
             "--shared",
             src.name,
             "-lcuda",
-            f"-gencode=arch=compute_{compute_version},code=compute_{compute_version}",
+            "-gencode",
+            f"arch=compute_{compute_version},code=sm_{compute_version}",
             "-o",
             libpath,
         ]

integration/BitNet/eval_correctness.py

@@ -72,18 +72,19 @@ def get_runtime(num_repeats=1):
 def main():
     model = BitnetForCausalLM.from_pretrained(
         model_path,
-        use_flash_attention_2=True,
+        use_flash_attention_2=False,
         torch_dtype=torch.float16,
     ).cuda().half()
-    with torch.no_grad():
-        model._post_process_weights()
 
     tokenizer = BitnetTokenizer.from_pretrained(model_path, use_fast=False)
     input_id = tokenizer("Hello")['input_ids']
     input_id = torch.tensor(input_id).unsqueeze(0).cuda()
-    output = model(input_id)
-    print(output)
 
+    print("original model generated text:")
+    print(generate_text(model, tokenizer, "Hello", max_length=100))
+
+    model.quantize()
+    print("quantized model generated text:")
     print(generate_text(model, tokenizer, "Hello", max_length=100))
 
 

integration/BitNet/maint/create_bitblas_ckpt.py

@@ -68,7 +68,7 @@ def main():
     model = (
         BitnetForCausalLM.from_pretrained(
             model_name_or_path,
-            use_flash_attention_2=True,
+            use_flash_attention_2=False,
             torch_dtype=torch.float16,
         ).cuda().half())
     tokenizer = BitnetTokenizer.from_pretrained(model_name_or_path, use_fast=False)

integration/BitNet/maint/generate_bitnet_model_bitblas_format.sh

@@ -18,6 +18,9 @@ fi
 if [ -z "$SAVED_MODEL_DIR" ]; then
   python ./maint/create_bitblas_ckpt.py --model_name_or_path $MODEL_DIR
 else
+  if [ ! -d "$SAVED_MODEL_DIR" ]; then
+    mkdir -p $SAVED_MODEL_DIR
+  fi
   python ./maint/create_bitblas_ckpt.py --model_name_or_path $MODEL_DIR --saved_model_path $SAVED_MODEL_DIR
 fi
 

integration/BitNet/requirements.txt

@@ -1,2 +1,3 @@
 lm_eval==0.3.0
 flash_attn
+transformers==4.40