From 0f9b41015154ccf642198dcbeb7996b1f377640c Mon Sep 17 00:00:00 2001 From: Logikisto Date: Wed, 9 Sep 2026 13:45:04 +0200 Subject: [PATCH] Initial commit --- .gitignore | 1 + apprendre.ipynb | 113 ++++++++++++++++++++++++++++++++++++++++ lib.py | 134 ++++++++++++++++++++++++++++++++++++++++++++++++ ops.py | 18 +++++++ test.ipynb | 61 ++++++++++++++++++++++ 5 files changed, 327 insertions(+) create mode 100644 .gitignore create mode 100644 apprendre.ipynb create mode 100644 lib.py create mode 100644 ops.py create mode 100644 test.ipynb diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ba0430d --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +__pycache__/ \ No newline at end of file diff --git a/apprendre.ipynb b/apprendre.ipynb new file mode 100644 index 0000000..f480d6f --- /dev/null +++ b/apprendre.ipynb @@ -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 +} diff --git a/lib.py b/lib.py new file mode 100644 index 0000000..6f901b9 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/ops.py b/ops.py new file mode 100644 index 0000000..73e1a1b --- /dev/null +++ b/ops.py @@ -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 + diff --git a/test.ipynb b/test.ipynb new file mode 100644 index 0000000..25e7689 --- /dev/null +++ b/test.ipynb @@ -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 +}