Fast Attention For Short Sequences

Nowadays, there is a big spotlight on big transformer decoder models with big context windows and billions (even trillions) of parameters.

But machine learning/deep learning is not limited to just LLMs; there are other applications of transformers too.

And sometimes we need to use attention on short sequences: 32-128.
For some tasks, like classification or representation learning, we also want to have a Transformer Encoder with one CLS token.

If we only need an embedding of a cls token from the transformer, then it's a good optimization to skip computing self-attention for non-cls tokens at the last layer.
First, lets benchmark different pytorch backends (MemEff,FlashAttention2,CUDNN) on a full self-attention setup:
All times below are in ms, e2e in torch.
Setup: H100, torch 2.11.0+cu128, bf16, non-causal, w/o dropout, batch=16000, heads=8, headdim=64

Interesting, that the more advanced algorithm (Flash Attention 2), works slower. But why? What is the difference between MemEfficient and Flash Attention 2? First, lets examine what kernels are actually being run on seq=32

MemEfficient

fmha_cutlassF_bf16_aligned_64x64_rf_sm80(PyTorchMemEffAttention::AttentionKernel<cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64, true, true>::Params)
fmha_cutlassB_bf16_aligned_64x64_k64_sm80(PyTorchMemEffAttention::AttentionBackwardKernel<cutlass::arch::Sm80, cutlass::bfloat16_t, true, false, true, 64, 64, 64, false>::Params)

FlashAttention 2

void pytorch_flash::flash_fwd_kernel<Flash_fwd_kernel_traits<64, 128, 128, 4, false, false, cutlass::bfloat16_t, Flash_kernel_traits<64, 128, 128, 4, cutlass::bfloat16_t> >, false, false, false, false, false, true, false, false>(pytorch_flash::Flash_fwd_params)  
void pytorch_flash::flash_bwd_dq_dk_dv_loop_seqk_parallel_kernel<Flash_bwd_kernel_traits<64, 128, 128, 8, 4, 4, 4, false, false, cutlass::bfloat16_t, Flash_kernel_traits<64, 128, 128, 8, cutlass::bfloat16_t> >, false, false, false, false, false, true, false>(pytorch_flash::Flash_bwd_params)

Well, that explains why 32vs64 seqlen had a smaller increase in fwd/bwd time than 128. MemEff uses a local QK^T of 64 query rows by 64 key rows, while FlashAttention2 uses 128 query rows by 128 key rows. It is underutilized on seq=32.

If seq=32 was already using the tile efficiently, we would expect kernel time to grow much more from 32 -> 64. It does not. What actually grows sharply is the memory-side throughput shown by Memory %. Now we will profile only the main attention kernel with ncu, not the full end-to-end time (sdpa also launches kernels besides the main attention kernel).

FlashAttention2 main attention kernel

MemEff main attention kernel

So seq=32 is underutilized for both kernels, and much more so for FlashAttention2. The useful work grows fast from 32 -> 64, but the attention kernel barely gets slower.

At seq=64, MemEff forward is almost a perfect 64x64 one-tile problem, so memory throughput gets much higher. At seq=128, the same kernel has to iterate over multiple K/V tiles and merge online-softmax state, which adds more on-chip work and lowers the reported Memory %. Looking at these results, is it possible to write a much simpler triton kernel that outperforms MemEfficient/Flash Attention 2 on short sequences?

For seq=32 or seq=64, the whole attention problem is small enough that one kernel instance can handle the entire sequence tile directly. In that regime, we do not need extra complexity, but this kernel will only be suitable for short sequences.

Here is the forward of our new triton kernel:

    pid_b = tl.program_id(0).to(tl.int64)
    pid_h = tl.program_id(1).to(tl.int64)
    offs_s = tl.arange(0, block_s)
    row_mask = offs_s < seq_len
    qkv_base = qkv_ptr + pid_b * stride_qkv_b + pid_h * stride_qkv_h

    q_blk = tl.make_block_ptr(
        qkv_base + 0 * stride_qkv_c,
        shape=[seq_len, head_dim],
        strides=[stride_qkv_s, stride_qkv_d],
        offsets=[0, 0],
        block_shape=[block_s, block_d],
        order=(1, 0),
    )
    k_blk = tl.make_block_ptr(
        qkv_base + 1 * stride_qkv_c,
        shape=[seq_len, head_dim],
        strides=[stride_qkv_s, stride_qkv_d],
        offsets=[0, 0],
        block_shape=[block_s, block_d],
        order=(1, 0),
    )
    v_blk = tl.make_block_ptr(
        qkv_base + 2 * stride_qkv_c,
        shape=[seq_len, head_dim],
        strides=[stride_qkv_s, stride_qkv_d],
        offsets=[0, 0],
        block_shape=[block_s, block_d],
        order=(1, 0),
    )

    q = tl.load(q_blk, boundary_check=(0, 1), padding_option="zero")
    k = tl.load(k_blk, boundary_check=(0, 1), padding_option="zero")
    v = tl.load(v_blk, boundary_check=(0, 1), padding_option="zero")

    scores = tl.dot(q, tl.trans(k), out_dtype=tl.float32) * scale
    scores = tl.where(row_mask[None, :], scores, float("-inf"))
    scores_max = tl.max(scores, axis=1)

    probs = tl.exp(scores - scores_max[:, None])
    probs = tl.where(row_mask[:, None], probs, 0.0)
    probs_sum = tl.sum(probs, axis=1)
    probs_sum_safe = tl.where(row_mask, probs_sum, 1.0)
    lse = tl.where(row_mask, tl.log(probs_sum_safe) + scores_max, 0.0)
    probs = probs / probs_sum_safe[:, None]

    out = tl.dot(probs.to(q.dtype), v, out_dtype=tl.float32)
    o_blk = tl.make_block_ptr(
        o_ptr + pid_b * stride_o_b + pid_h * stride_o_h,
        shape=[seq_len, head_dim],
        strides=[stride_o_s, stride_o_d],
        offsets=[0, 0],
        block_shape=[block_s, block_d],
        order=(1, 0),
    )
    tl.store(o_blk, out.to(q.dtype), boundary_check=(0, 1))
    tl.store(lse_ptr + pid_b * stride_lse_b + pid_h * stride_lse_h + offs_s, lse, mask=row_mask)

Backward has the same idea - recalculate scores using lse, calculate whole sequence in one go.
There are also separate fwd/bwd kernels for a case where len(q)=1 e.g. just a cls token, they work basically the same way.

Lets benchmark our new kernels For Triton vs best, values are best non-Triton time / Triton time; >1x favors Triton and <1x favors the best non-Triton backend.

q=1

Well, seems that our kernels do pretty well on small seq length and headdim.

At seq_len=128, head_dim=256, the current backward kernel gets close to the SM resource limits. Shared memory usage becomes very high, and many intermediate values have to go through local memory, which is private to each thread but backed by global memory. That makes the kernel extremely slow. The table below compares seq_len=64 and 128 at head_dim=256 and shows how shared memory usage, warp activity, and local-memory traffic change (measured using ncu).

To fix it, we have to design a better kernel, but it will not be as simple as the current one.
In this blog post we showed that a simple triton attention kernel can be faster than default torch attention backends, if the seq_len and headdim are sufficiently small
https://github.com/qwertyforce/fast_attn_short_seq/tree/main