triton-lora.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "1a997e32",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import triton\n",
    "import triton.language as tl\n",
    "\n",
    "@triton.jit\n",
    "def lora_matmul_kernel(\n",
    "        input_ptr, w_ptr, w1_ptr, w2_ptr, output_ptr,\n",
    "        M, N, K, R,\n",
    "        stride_im, stride_ik,\n",
    "        stride_wk, stride_wn,\n",
    "        stride_w1k, stride_w1r,\n",
    "        stride_w2r, stride_w2n,\n",
    "        stride_om, stride_on,\n",
    "        BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, \n",
    "        BLOCK_SIZE_K: tl.constexpr, BLOCK_SIZE_R: tl.constexpr,\n",
    "        GROUP_SIZE_M: tl.constexpr,\n",
    "):\n",
    "    pid = tl.program_id(axis=0)\n",
    "    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)\n",
    "    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)\n",
    "    num_pid_in_group = GROUP_SIZE_M * num_pid_n\n",
    "    group_id = pid // num_pid_in_group\n",
    "    first_pid_m = group_id * GROUP_SIZE_M\n",
    "    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)\n",
    "    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)\n",
    "    pid_n = (pid % num_pid_in_group) // group_size_m\n",
    "\n",
    "    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)\n",
    "    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)\n",
    "    offs_k = tl.arange(0, BLOCK_SIZE_K)\n",
    "    offs_r = tl.arange(0, BLOCK_SIZE_R)\n",
    "    \n",
    "    # Use separate accumulators for base path and LoRA path\n",
    "    base_acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)\n",
    "    lora_acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)\n",
    "    \n",
    "    # Compute base path: input × W\n",
    "    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):\n",
    "        k_offs = k * BLOCK_SIZE_K + offs_k\n",
    "        a_ptrs = input_ptr + offs_m[:, None] * stride_im + k_offs[None, :] * stride_ik\n",
    "        w_ptrs = w_ptr + k_offs[:, None] * stride_wk + offs_n[None, :] * stride_wn\n",
    "        \n",
    "        mask_a = (offs_m[:, None] < M) & (k_offs[None, :] < K)\n",
    "        mask_b = (k_offs[:, None] < K) & (offs_n[None, :] < N)\n",
    "        \n",
    "        a = tl.load(a_ptrs, mask=mask_a, other=0.0)\n",
    "        b = tl.load(w_ptrs, mask=mask_b, other=0.0)\n",
    "        \n",
    "        base_acc += tl.dot(a, b)\n",
    "\n",
    "    # Compute LoRA path: input × W1 × W2\n",
    "    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):\n",
    "        k_offs = k * BLOCK_SIZE_K + offs_k\n",
    "        a_ptrs = input_ptr + offs_m[:, None] * stride_im + k_offs[None, :] * stride_ik\n",
    "        mask_a = (offs_m[:, None] < M) & (k_offs[None, :] < K)\n",
    "        a = tl.load(a_ptrs, mask=mask_a, other=0.0)\n",
    "        \n",
    "        for r in range(0, tl.cdiv(R, BLOCK_SIZE_R)):\n",
    "            r_offs = r * BLOCK_SIZE_R + offs_r\n",
    "            \n",
    "            w1_ptrs = w1_ptr + k_offs[:, None] * stride_w1k + r_offs[None, :] * stride_w1r\n",
    "            w2_ptrs = w2_ptr + r_offs[:, None] * stride_w2r + offs_n[None, :] * stride_w2n\n",
    "            \n",
    "            mask_w1 = (k_offs[:, None] < K) & (r_offs[None, :] < R)\n",
    "            mask_w2 = (r_offs[:, None] < R) & (offs_n[None, :] < N)\n",
    "            \n",
    "            w1 = tl.load(w1_ptrs, mask=mask_w1, other=0.0)\n",
    "            w2 = tl.load(w2_ptrs, mask=mask_w2, other=0.0)\n",
    "            \n",
    "            temp = tl.dot(a, w1)\n",
    "            lora_acc += tl.dot(temp, w2)\n",
    "\n",
    "    # Combine results and write to output\n",
    "    output = base_acc + lora_acc\n",
    "    output_ptrs = output_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on\n",
    "    mask_out = (offs_m[:, None] < M) & (offs_n[None, :] < N)\n",
    "    tl.store(output_ptrs, output, mask=mask_out)\n",
    "\n",
    "def lora_matmul(input, weight, lora_weight1, lora_weight2):\n",
    "    \"\"\"\n",
    "    Compute matrix multiplication with LoRA: output = input × weight + input × lora_weight1 × lora_weight2\n",
    "    \n",
    "    Arguments:\n",
    "        input: torch.Tensor of shape (M, K)\n",
    "        weight: torch.Tensor of shape (K, N)\n",
    "        lora_weight1: torch.Tensor of shape (K, R)\n",
    "        lora_weight2: torch.Tensor of shape (R, N)\n",
    "    Returns:\n",
    "        output: torch.Tensor of shape (M, N)\n",
    "    \"\"\"\n",
    "    # Check input dimensions\n",
    "    assert input.shape[1] == weight.shape[0], \"Input and weight dimensions mismatch\"\n",
    "    assert input.shape[1] == lora_weight1.shape[0], \"Input and LoRA W1 dimensions mismatch\"\n",
    "    assert lora_weight1.shape[1] == lora_weight2.shape[0], \"LoRA W1 and W2 dimensions mismatch\"\n",
    "    assert weight.shape[1] == lora_weight2.shape[1], \"Weight and LoRA W2 dimensions mismatch\"\n",
    "    \n",
    "    # Extract dimensions\n",
    "    M, K = input.shape\n",
    "    _, N = weight.shape\n",
    "    R = lora_weight1.shape[1]\n",
    "    \n",
    "    # Allocate output\n",
    "    output = torch.empty((M, N), device=input.device, dtype=input.dtype)\n",
    "    \n",
    "    # Define block sizes and make sure they're appropriate for the GPU\n",
    "    BLOCK_SIZE_M = 128\n",
    "    BLOCK_SIZE_N = 128\n",
    "    BLOCK_SIZE_K = 32\n",
    "    BLOCK_SIZE_R = 16\n",
    "    GROUP_SIZE_M = 8\n",
    "    \n",
    "    # Calculate grid size\n",
    "    grid = lambda META: (\n",
    "        triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']),\n",
    "    )\n",
    "    \n",
    "    # Launch kernel\n",
    "    lora_matmul_kernel[grid](\n",
    "        input, weight, lora_weight1, lora_weight2, output,\n",
    "        M, N, K, R,\n",
    "        input.stride(0), input.stride(1),\n",
    "        weight.stride(0), weight.stride(1),\n",
    "        lora_weight1.stride(0), lora_weight1.stride(1),\n",
    "        lora_weight2.stride(0), lora_weight2.stride(1),\n",
    "        output.stride(0), output.stride(1),\n",
    "        BLOCK_SIZE_M=BLOCK_SIZE_M,\n",
    "        BLOCK_SIZE_N=BLOCK_SIZE_N,\n",
    "        BLOCK_SIZE_K=BLOCK_SIZE_K,\n",
    "        BLOCK_SIZE_R=BLOCK_SIZE_R,\n",
    "        GROUP_SIZE_M=GROUP_SIZE_M,\n",
    "    )\n",
    "    \n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "5d104c8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "    \n",
    "M = 128\n",
    "K = 512\n",
    "N = 256\n",
    "R = 128    # LoRA rank\n",
    "\n",
    "device = torch.device('cuda')\n",
    "input = torch.randn((M, K), device=device, dtype=torch.float32)\n",
    "weight = torch.randn((K, N), device=device, dtype=torch.float32)\n",
    "lora_w1 = torch.randn((K, R), device=device, dtype=torch.float32)\n",
    "lora_w2 = torch.randn((R, N), device=device, dtype=torch.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "ac7943eb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 195 µs, sys: 0 ns, total: 195 µs\n",
      "Wall time: 200 µs\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "output_triton = lora_matmul(input, weight, lora_w1, lora_w2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "167eba9a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 212 µs, sys: 22 µs, total: 234 µs\n",
      "Wall time: 239 µs\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "output_torch = torch.matmul(input, weight) + torch.matmul(torch.matmul(input, lora_w1), lora_w2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "c1da9a67",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Max difference: 1.9237060546875\n"
     ]
    }
   ],
   "source": [
    "print(f\"Max difference: {torch.max(torch.abs(output_triton - output_torch))}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "d52429de",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.9998, device='cuda:0')"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(output_torch.sign() == output_triton.sign()).float().mean()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3.10",
   "language": "python",
   "name": "python3.10"
  },
  "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.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}