Sequence Packing & Variable-Length Attention: The Trick That Boosts Throughput by 40%

# === Variable-length packing — flash_attn_varlen_func ===
import torch
from flash_attn import flash_attn_varlen_func
 
# Pack 3 example into one sequence of length 4096
example_lens = [1500, 1200, 1100]   # toplam 3800; rest padding (296 token)
cu_seqlens = torch.tensor([0, 1500, 2700, 3800], dtype=torch.int32, device="cuda")
# cu_seqlens[i+1] - cu_seqlens[i] = i-th example length
# cu_seqlens shape: [batch_size + 1]
 
max_seqlen = max(example_lens)
total_tokens = cu_seqlens[-1].item()
 
# Q, K, V shape: [total_tokens, n_heads, head_dim]
q = torch.randn(total_tokens, 32, 128, device="cuda", dtype=torch.bfloat16)
k = torch.randn(total_tokens, 32, 128, device="cuda", dtype=torch.bfloat16)
v = torch.randn(total_tokens, 32, 128, device="cuda", dtype=torch.bfloat16)
 
# Block-diagonal attention — example A can't attend to B
out = flash_attn_varlen_func(
    q, k, v,
    cu_seqlens_q=cu_seqlens, cu_seqlens_k=cu_seqlens,
    max_seqlen_q=max_seqlen, max_seqlen_k=max_seqlen,
    causal=True,
)
# out shape: [total_tokens, n_heads, head_dim]