Flex RNN compiles linear attention mechanisms like RWKV and Mamba from simple pytorch code. For example, the Delta Rule kernel can be implemented as

@flex_rnn.jit
def delta_rule(q, k, v, beta, S):
    S = S + beta * k * (v - (S*k).sum(-1,True))
    return (S*q).sum(-1,True), S

See tests/examples.py for example implementations of RWVK-7, RWKV-6, RWKV-4, Mamba-2, Mamba, Delta Rule, Gated Delta Rule, Retention, GLA, GSA, HGRN, Longhorn, mLSTM and S4D.

The compiler traces the code, generates the backward pass, compiles it to cuda code, and returns a callable function equivalent to naive defined below. Regardless of input and output dtypes, all internal calculations are performed using 32-bit floats. The backward pass fully recalculates all states, so precision should be similar to torch autograd. See below for how the speed compares to other kernels.

For batch size B, sequence length T, number of heads H and state dimensions MxN, the resulting function expects inputs with shapes [B,T,H,M,N] followed by initial state(s) with shape(s) [B,H,M,N]. Inputs with dimensions of size 1 will be broadcast, so for the Delta Rule, q should have shape [B,T,H,1,N] and v have shape [B,T,H,M,1]. Mamba-2's A is a headwise parameter with no batch or time dependence, so it would have shape [1,1,H,1,1]. Initial states are not broadcast.

Let's see how to implement an equivalent to fla.ops.delta_rule.chunk_delta_rule(q, k, v, beta, scale = 1, initial_state = S, output_final_state = True). That function expects q, k and v to have shape [B,T,H,M], beta to have shape B,T,H and the initial state to have shape [B,T,H,M,N]. Hence, we can implement it by unsqueezing as follows:

def matching_delta_rule(q, k, v, beta, S):
    q, k, v, beta = [i.unsqueeze(-2) for i in (q, k, v, beta.unsqueeze(-1))]
    y, S = delta_rule(q, k, v.mT, beta, S.mT)
    return y, S.mT 

In general, the compiled function behaves like jit.naive from flex_rnn/jit.py. It is often helpful to first debug using the naive function (for example, use y, S = delta_rule.naive(q, k, v.mT, beta, S.mT) instead of y, S = delta_rule(q, k, v.mT, beta, S.mT) in the code above), before working with the fast compiled kernels.

When there's a single state variable and a single output variable, and disregarding strides of outputs, the compiled function behaves like

def simplified_naive(step, *args):
    inputs, state = args[:-1], args[-1]
    T = max(i.shape[1] for i in inputs)
    output = []
    for t in range(T):
        inputs_t = [i.expand(-1,T,-1,-1,-1)[:,t].float() for i in inputs]
        output_t, state = step(*inputs_t, state)
        output.append(output_t)
    return torch.stack(output, dim=1).squeeze(-1), state

git clone https://github.com/johanwind/flex_rnn
cd flex_rnn
pip install .

You should then be able to run examples like

python tests/examples.py --op rwkv7 --check-flex

Batch size B = 8, sequence length T = 1024, number of heads H = 4096 / 128 = 32 and state dimensions 128x128. NVIDIA 4070 mobile GPU.

All inputs and initial states in bfloat16 except S4D and Mamba (because their reference implementations only support float32).

• Compatible with torch.compile without graph breaks
• Handles non-contiguous input strides without copying
• float32, bfloat16, float16 input and output types
• High precision because of float32 internal calculations

• The backend computes the state row-wise. Specifically, that means we only support reductions in the last dimension, i.e. torch.sum(x, dim=-1, keepdim=True).
• The supported operations are essentially +, -, *, /, <, >, exp, log, sqrt and x.sum(dim=-1, keepdim=True).
• Input dimensions have only been tested for powers of 2.
• Speed and memory usage are often worse than specialized implementations (see speed comparison above).