Initial commit

lib.py

@@ -0,0 +1,134 @@
+import json
+from pathlib import Path
+
+CONFIG_FILE = Path(".QuantTensor_config.json")
+__version__ = "0.0.1"
+
+# === Chargement ou création de la config ===
+if CONFIG_FILE.exists():
+    with open(CONFIG_FILE, "r") as f:
+        config = json.load(f)
+else:
+    config = None
+
+# Vérification si on doit relancer les tests
+if config is None or config.get("version") != __version__:
+    print("[QuantTensor] Première exécution ou mise à jour détectée : test des capacités Triton...")
+
+    # Test réel de Triton (import conditionnel)
+    try:
+        import triton
+        import triton.language as tl
+        import torch
+        print(f"[QuantTensor] Triton est fonctionnel.")
+
+        try:
+            pass # Tester le premier kernel de Quant
+            
+            print(f" [X] Kernels de quantifications Q4_1")
+            Q4_1__triton_quant = True
+        except Exception as e :
+            
+            print(f" [ ] Kernels de quantifications Q4_1 :\n{e}")
+            Q4_1__triton_quant = False
+        
+        try:
+            pass # Tester le premier kernel de DeQuant
+
+            print(f" [X] Kernels de déquantifications Q4_1")
+            Q4_1__triton_dequant = True
+        except Exception as e :
+
+            print(f" [ ] Kernels de déquantifications Q4_1 :\n{e}")
+            Q4_1__triton_dequant = False
+        
+        try:
+            pass # Tester kernel de GEMM même si les précédents passent pas
+
+            print(f" [X] Kernels de fused-GEMM Q4_1")
+            Q4_1__triton_fused_gemm = True
+        except Exception as e :
+
+            print(f" [ ] Kernels de fused-GEMM Q4_1 :\n{e}")
+            Q4_1__triton_fused_gemm = False
+        
+        # Et Etc....
+
+    except Exception as e :
+        print(f"[QuantTensor] La bibliothèque Triton ne fonctionne pas :\n{e}")
+        Q4_1__triton_dequant = False
+        Q4_1__triton_quant = False
+        Q4_1__triton_fused_gemm = False
+        
+
+    config = {
+        "version": __version__,
+        "Q4_1__triton_dequant": Q4_1__triton_dequant,
+        "Q4_1__triton_quant": Q4_1__triton_quant,
+        "Q4_1__fused_gemm": Q4_1__triton_fused_gemm
+    }
+
+    with open(CONFIG_FILE, "w") as f:
+        json.dump(config, f, indent=2)
+
+
+
+# === Binding des fonctions au moment de l'import ===
+if config["triton_dequant"]:
+    from ops import dequantize_q4_1_triton as dequantize_q4_1
+else:
+    from ops import dequantize_q4_1_pytorch as dequantize_q4_1
+
+if config["triton_quant"]:
+    from ops import quantize_q4_1_triton as dequantize_q4_1
+else:
+    from ops import quantize_q4_1_pytorch as dequantize_q4_1
+
+if config["q4_1_fused_gemm"]:
+    from ops import q4_1_fused_gemm_triton as q4_1_fused_gemm
+else:
+    from ops import q4_1_fused_gemm_pytorch as q4_1_fused_gemm
+
+# =================
+
+class QuantTensor:
+    def __init__(self, format: str, qs: torch.Tensor, d: torch.Tensor, 
+                 m: torch.Tensor, original_shape: torch.Size):
+        self.format = format
+        self.qs = qs
+        self.d = d
+        self.m = m
+        self.original_shape = original_shape
+
+        self.dtype = d.dtype
+        self.device = d.device
+
+    def numel(self) -> int:
+        return self.original_shape.numel()
+
+    @property
+    def nbytes(self) -> int:
+        total = 0
+        for attr in self.__dict__.values():
+            if isinstance(attr, torch.Tensor):
+                total += attr.numel() * attr.element_size()
+        return total
+
+    def dequantize(self):
+        if self.format == "Q4_1":
+            w = dequantize_q4_1(self.qs, self.d, self.m)
+            return w.view(self.original_shape)
+        else:
+            raise NotImplementedError(f"Format {self.format} non supporté")
+
+    def to(self, device=None, dtype=None):
+        if device is not None:
+            self.device = device
+            self.qs = self.qs.to(device=device)
+            self.d = self.d.to(device=device)
+            self.m = self.m.to(device=device)
+
+        if dtype is not None:
+            self.dtype = dtype
+
+        return self