Kernels for higher order gradients of Flash Attention.
Project description
Flash Hog
This repo contains the code for Flash Higher-Order-Gradients, aka. Flash Hog. This kernel achieves around a 3.7x speedup over an XLA optimized kernel, with linear memory scaling instead of quadratic scaling.
Installation
uv add flash-hog
Method
Flash Hog does 4 recomputation passes to avoid any atomics or saving any intermediary tensors of shape (N_Q, N_K).
This shakes out to be thread-wise tiling across Q in 3 passes first, once to compute dd, then once for b, then once for both dQ' and ddO.
Finally we do another pass tiled over K, producing dK' and dV'.
The equations we implement are the following:
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file flash_hog-0.6.0.tar.gz.
File metadata
- Download URL: flash_hog-0.6.0.tar.gz
- Upload date:
- Size: 20.8 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: uv/0.11.21 {"installer":{"name":"uv","version":"0.11.21","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"macOS","version":null,"id":null,"libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
23641f16935f1b7aa04d62ed6a82c2b9b40ea5497d0c5201e5072ae31d32960c
|
|
| MD5 |
833a23200ca572c9620431c7586e0389
|
|
| BLAKE2b-256 |
13ac1c1c34491a0e5b5284413a587a030c5c80e2846878d71c168572718fb8b0
|
File details
Details for the file flash_hog-0.6.0-py3-none-any.whl.
File metadata
- Download URL: flash_hog-0.6.0-py3-none-any.whl
- Upload date:
- Size: 27.7 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: uv/0.11.21 {"installer":{"name":"uv","version":"0.11.21","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"macOS","version":null,"id":null,"libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
760acea57d630edb8412f5fabf295b5fcd0e2cfdcfe9002430e3c8fd7b643cf7
|
|
| MD5 |
cedfb94d3b0e3e86005be66a87e0f40f
|
|
| BLAKE2b-256 |
3a3f98d5f450dc17b0abd5a219634f4de440176287bac6eb13c6d8a65790d0a9
|
