Initial commit

2026-09-09 13:45:04 +02:00
commit 0f9b410151
5 changed files with 327 additions and 0 deletions
@@ -0,0 +1 @@
+__pycache__/
@@ -0,0 +1,113 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "869892b8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import triton\n",
+    "import triton.language as tl\n",
+    "\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "\n",
+    "torch.set_default_device(\"cuda\")\n",
+    "torch.set_default_dtype(torch.bfloat16)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1a4b0ef4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "@triton.jit\n",
+    "def add_kernel(x_ptr,  # *Pointer* to first input vector.\n",
+    "               y_ptr,  # *Pointer* to second input vector.\n",
+    "               output_ptr,  # *Pointer* to output vector.\n",
+    "               n_elements,  # Size of the vector.\n",
+    "               BLOCK_SIZE: tl.constexpr,  # Number of elements each program should process.\n",
+    "               # NOTE: `constexpr` so it can be used as a shape value.\n",
+    "               ):\n",
+    "    # There are multiple 'programs' processing different data. We identify which program\n",
+    "    # we are here:\n",
+    "    pid = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.\n",
+    "    # This program will process inputs that are offset from the initial data.\n",
+    "    # For instance, if you had a vector of length 256 and block_size of 64, the programs\n",
+    "    # would each access the elements [0:64, 64:128, 128:192, 192:256].\n",
+    "    # Note that offsets is a list of pointers:\n",
+    "    block_start = pid * BLOCK_SIZE\n",
+    "    offsets = block_start + tl.arange(0, BLOCK_SIZE)\n",
+    "    # Create a mask to guard memory operations against out-of-bounds accesses.\n",
+    "    mask = offsets < n_elements\n",
+    "    # Load x and y from DRAM, masking out any extra elements in case the input is not a\n",
+    "    # multiple of the block size.\n",
+    "    x = tl.load(x_ptr + offsets, mask=mask)\n",
+    "    y = tl.load(y_ptr + offsets, mask=mask)\n",
+    "    output = x + y\n",
+    "    # Write x + y back to DRAM.\n",
+    "    tl.store(output_ptr + offsets, output, mask=mask)\n",
+    "\n",
+    "def add(x: torch.Tensor, y: torch.Tensor):\n",
+    "    # We need to preallocate the output.\n",
+    "    output = torch.empty_like(x)\n",
+    "    n_elements = output.numel()\n",
+    "    # The SPMD launch grid denotes the number of kernel instances that run in parallel.\n",
+    "    # It is analogous to CUDA launch grids. It can be either Tuple[int], or Callable(metaparameters) -> Tuple[int].\n",
+    "    # In this case, we use a 1D grid where the size is the number of blocks:\n",
+    "    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )\n",
+    "    # NOTE:\n",
+    "    #  - Each torch.tensor object is implicitly converted into a pointer to its first element.\n",
+    "    #  - `triton.jit`'ed functions can be indexed with a launch grid to obtain a callable GPU kernel.\n",
+    "    #  - Don't forget to pass meta-parameters as keywords arguments.\n",
+    "    add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)\n",
+    "    # We return a handle to z but, since `torch.cuda.synchronize()` hasn't been called, the kernel is still\n",
+    "    # running asynchronously at this point.\n",
+    "    return output"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7bcab582",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "torch.manual_seed(0)\n",
+    "size = 98432\n",
+    "x = torch.rand(size)\n",
+    "y = torch.rand(size)\n",
+    "output_torch = x + y\n",
+    "output_triton = add(x, y)\n",
+    "print(output_torch)\n",
+    "print(output_triton)\n",
+    "print(f'The maximum difference between torch and triton is '\n",
+    "      f'{torch.max(torch.abs(output_torch - output_triton))}')"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "pyvenv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.13.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
@@ -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
@@ -0,0 +1,18 @@
+QK4_1 = 32
+
+
+def dequantize_q4_1_triton():
+    pass
+def dequantize_q4_1_pytorch():
+    pass
+
+def quantize_q4_1_triton():
+    pass
+def quantize_q4_1_pytorch():
+    pass
+
+def q4_1_fused_gemm_triton():
+    pass
+def q4_1_fused_gemm_pytorch():
+    pass
+
@@ -0,0 +1,61 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "448d884d",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor(0.0140)\n",
+      "20\n",
+      "cpu\n",
+      "cuda\n"
+     ]
+    }
+   ],
+   "source": [
+    "import lib\n",
+    "\n",
+    "import torch\n",
+    "\n",
+    "T = torch.rand(8, 4)\n",
+    "Q = lib.quantize_q4_1(T)\n",
+    "\n",
+    "dif = T - Q.dequantize()\n",
+    "print(abs(dif.sum()))\n",
+    "\n",
+    "print(Q.nbytes)\n",
+    "\n",
+    "print(Q.device)\n",
+    "\n",
+    "Q.to(\"cuda\")\n",
+    "print(Q.device)\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "pyvenv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.13.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}