feat: add attention notebook

ex3ndr · Jul 10, 2024 · dea4c61 · dea4c61
1 parent 6b9e1b6
commit dea4c61
Showing 1 changed file with 12 additions and 12 deletions.
diff --git a/attention.ipynb b/attention.ipynb
@@ -45,8 +45,8 @@
       "With padding\n",
       "torch 0.0\n",
       "direct 0.00048828125\n",
-      "xformers 0.000244140625\n",
-      "flash 0.000244140625\n"
+      "xformers 0.0\n",
+      "flash 0.0\n"
      ]
     }
    ],
@@ -81,15 +81,15 @@
      "output_type": "stream",
      "text": [
       "Without padding\n",
-      "torch 0.4039938449859619\n",
-      "direct 0.9630444049835205\n",
-      "xformers 0.9124343395233154\n",
-      "flash 0.3562278747558594\n",
+      "torch 4.43998122215271\n",
+      "direct 10.454708576202393\n",
+      "xformers 8.981621742248535\n",
+      "flash 3.2445619106292725\n",
       "With padding\n",
-      "torch 1.6574337482452393\n",
-      "direct 2.504969835281372\n",
-      "xformers 1.5768730640411377\n",
-      "flash 1.4189743995666504\n"
+      "torch 16.73130464553833\n",
+      "direct 25.77487087249756\n",
+      "xformers 16.095849990844727\n",
+      "flash 15.358261585235596\n"
      ]
     }
    ],
@@ -98,14 +98,14 @@
     "print(\"Without padding\")\n",
     "for a in attentions:\n",
     "    start = time.time()\n",
-    "    for i in range(10000):\n",
+    "    for i in range(100000):\n",
     "        a(query, key, value)\n",
     "    print(a.engine, time.time() - start)\n",
     "\n",
     "print(\"With padding\")\n",
     "for a in attentions:\n",
     "    start = time.time()\n",
-    "    for i in range(10000):\n",
+    "    for i in range(100000):\n",
     "        a(query, key, value, lenghts = lengths)\n",
     "    print(a.engine, time.time() - start)"
    ]