Linear attention without writing a CUDA kernel

Most linear-attention work ships with a custom CUDA kernel. That is a reasonable engineering decision and it is also a wall: you now need the right toolchain, the right GPU, and a willingness to debug someone else's kernel when it fails on your hardware.

I wanted to know what the architecture is worth without that. MicroRWKV is a 130M parameter RWKV implementation in plain PyTorch, no custom kernels, small enough to train on a CPU. The point was never a state-of-the-art model. The point was to find out which of the efficiency claims survive contact with an ordinary implementation.

The recurrence

A transformer computes attention over all previous positions for every new token, so per-token cost grows with context. RWKV replaces that with a recurrence carrying a fixed-size state, which makes each step cost the same no matter how long the context is.

The weighted key-value operator is where the mixing happens:

$$wkv_t = \frac{\sum_{i=1}^{t-1} e^{-(t-1-i)w + k_i} v_i + e^{u + k_t} v_t}{\sum_{i=1}^{t-1} e^{-(t-1-i)w + k_i} + e^{u + k_t}}$$

Two learned vectors do the work. $w$ is a per-channel decay, so a channel can choose how fast the past fades. $u$ is a bonus applied to the current token so the present is not drowned by history.

Written as a sum it looks like it needs the whole sequence. It does not. Because the decay is exponential, the numerator and denominator can each be carried forward as a running pair, which is what makes generation a constant-cost step:

import torch


def wkv_step(state, k, v, w, u):
    """One RWKV time-mixing step.

    state is (numerator, denominator, max_exponent), each shaped like k. The
    max_exponent term keeps the exponentials from overflowing, the same trick
    a stable softmax uses.
    """
    num, den, max_e = state

    # Compare the running maximum against this token's exponent.
    e = torch.maximum(max_e, k + u)
    a = torch.exp(max_e - e)
    b = torch.exp(k + u - e)

    # Output uses the bonus term, so the current token is weighted directly.
    out = (a * num + b * v) / (a * den + b)

    # Then advance the state, decaying the past by w.
    e_next = torch.maximum(max_e - w, k)
    a_next = torch.exp(max_e - w - e_next)
    b_next = torch.exp(k - e_next)

    return out, (a_next * num + b_next * v, a_next * den + b_next, e_next)

That is the entire mechanism. No kernel, no flash-attention dependency, nothing that cares which GPU you have.

Configuration

Parameter	Value
Layers	8
Heads	8
Embedding dim	768
MoE experts	4
Block size	1024
Vocab size	50304
Parameters	130M

Grouped-query attention and a sliding window sit alongside the recurrence, and a small mixture-of-experts layer replaces the dense feed-forward block. The MoE is the piece I would defend least: at this scale the routing overhead eats most of the gain, and I kept it because I wanted the plumbing, not because it paid for itself.

What held up, and what did not

The efficiency claims held up, and they are the honest reason to care about this architecture:

• Per-token generation latency is flat as context grows, where the transformer baseline climbs quadratically. This is the whole point and it is very visible past about 512 tokens.
• Memory during generation is constant, because the state is fixed size. No KV cache growing with every token.
• The model trains on CPU. Slowly, but it trains, which makes the architecture something you can actually study without renting a GPU.

The quality claim did not hold up, and I want to be direct about it because the loss table is published in the repository and anyone can read it:

Model	Params	Train loss	Val loss
MicroRWKV	130M	2.85	2.88
GPT-2	124M	2.82	2.86

That is a slightly worse validation loss than GPT-2 at a slightly larger parameter count. Close enough to say the architecture is competitive at this scale, not close enough to claim it wins. An earlier write-up of mine described this as outperforming GPT-2 and that was wrong; the table above is what the runs actually produced.

Worth it or not

If you need the best quality per parameter at small scale, take the transformer. It is better understood, better tooled, and in this comparison marginally better.

If you need generation cost that does not grow with context, on hardware you do not control, without a compile step that can fail in six different ways, the recurrence is a genuinely different trade and pure PyTorch is enough to get it.

I would use it for long-context streaming on an edge device. I would not use it because it is novel.

Code and weights are linked from the project page if you want to reproduce any of this.