Update README.md for correct LaTeX.

Author: kpoeppel

.gitignore

@@ -0,0 +1 @@
+outputs_speed_exps*

.pre-commit-config.yaml

@@ -0,0 +1,18 @@
+# See https://pre-commit.com for more information
+# See https://pre-commit.com/hooks.html for more hooks
+repos:
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v3.2.0
+    hooks:
+    -   id: trailing-whitespace
+    -   id: end-of-file-fixer
+    -   id: check-yaml
+    -   id: check-added-large-files
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  # Ruff version.
+  rev: v0.6.2
+  hooks:
+    # Run the linter.
+    - id: ruff
+    # Run the formatter.
+    - id: ruff-format

README.md

@@ -3,19 +3,19 @@ Korbinian Pöppel, Maximilian Beck, Sepp Hochreiter
 
 ## Intro
 
-FlashRNN implements traditional RNNs like LSTMs, GRUs and Elman networks as well as the recent sLSTM architecture in CUDA and Triton. In contrary to common modern sequence models they have state tracking capabilities (Merrill et al., 2024). All of them are of the basic recurrent structure with input $\mathbf{x}^{(n)}_t$, bias $\mathbf{b}^{(n)}$, recurrent matrix $\mathbf{R}^{(n)}$ :
-$$ 
-\mathbf{g}^{(n)}_{t} = \mathbf{R}^{(n)} \ \mathbf{s}^{(0)}_{t-1} + \mathbf{x}^{(n)}_{t} + \mathbf{b}^{(n)} \\
-\mathbf{y}^{(m)}_t = \mathcal{P}^{(m)}\left( \left( \mathbf{s}^{(m')}_{t-1} \right)_{m'} , \left( \mathbf{g}^{(n)}_{t}  \right)_{n} \right)
+FlashRNN implements traditional RNNs like LSTMs, GRUs and Elman networks as well as the recent sLSTM architecture in CUDA and Triton. In contrary to common modern sequence models they have state tracking capabilities (Merrill et al., 2024). All of them are of the basic recurrent structure with input $\mathbf{x}^{(n)}\_t$, bias $\mathbf{b}^{(n)}$, recurrent matrix $\mathbf{R}^{(n)}$ :
+
+$$
+\mathbf{g}^{(n)}\_{t} = \mathbf{R}^{(n)} \ \mathbf{s}^{(0)}\_{t-1} + \mathbf{x}^{(n)}_{t} + \mathbf{b}^{(n)} \\
+\mathbf{y}^{(m)}\_t = \mathcal{P}^{(m)}\left( \left( \mathbf{s}^{(m')}\_{t-1} \right)\_{m'} , \left( \mathbf{g}^{(n)}\_{t}  \right)\_{n} \right)
 $$
 
-Typically the inputs are modified with a linear layer which is omitted here for flexibility (it would look like $\mathbf{x}^{n}_t = \mathbf{W}^{n} \mathbf{x'}_t$). This operation can be parallelized along the sequence dimension in contrary to the recurrent part, \\
-It employs a multi-head structure, which is equivalent to having a block-diagonal recurrent matrix. The hidden state and gate vectors of hidden dimension $d$ are split into heads of head dimension $d_{head}$. 
+Typically the inputs are modified with a linear layer which is omitted here for flexibility (it would look like $\mathbf{x}^{n}\_t = \mathbf{W}^{n} \mathbf{x'}\_t$). This operation can be parallelized along the sequence dimension in contrary to the recurrent part, \\
+It employs a multi-head structure, which is equivalent to having a block-diagonal recurrent matrix. The hidden state and gate vectors of hidden dimension $d$ are split into heads of head dimension $d\_{head}$. 
 
 For the fused `triton` backend, heads are limited to small head dimensions $d_{head} \leq 64$. For the CUDA backend there are two versions. The basic `cuda` one that alternates between recurrent matrix multiplication the non-linear pointwise function $\mathcal{P}$ application. This version is not limited in head dimension $d_{head}$. The second is a `cuda_fused` version, which fuses matrix multiplication with point-wise non-linearity into one CUDA kernel using `wmma` instructions and custom caching on SRAM / registers (similar to FlashAttention (Dao et al., 2022), but with a different focus here). Since the recurrent matrix $\mathbf{R}$ and biases $\mathbf{b}$ are used for for every time step, they are customly cached in registers and SRAM, enabling a $2 \times$ to $5 \times$
 speedup over the alternating option. 
 
-
 ## Speed comparison
 
 ![speed_comparison](head_dim--lstm.svg)