import time
import math
import torch
import torch.nn.functional as F

def naive_online_causal_attn(q, k, v, block_size=128, causal=True):
"""
Block‑wise causal attention with streaming softmax.

Args:
  q: (B, Hq, T, D)   — Query tensor
  k: (B, Hk, D, T)   — Key tensor
  v: (B, Hk, D, T)   — Value tensor
  block_size: int   — size of each sequence block
  causal: bool      — if True, masks out future keys

Returns:
  out: (B, Hq, T, D) — Attention output
"""
B, Hq, T, D = q.shape
_, Hk, Dk, Tk = k.shape
assert Dk == D, "Key/value head dim must match query head dim"
assert T == Tk, "Sequence length of q and k/v must match"
assert Hq % Hk == 0, "Number of q‑heads must be a multiple of k‑heads"
div = Hq // Hk
num_blocks = T // block_size
scale = 1.0 / math.sqrt(D)

out = torch.zeros_like(q)

if causal:
    mask = torch.triu(
        torch.ones(block_size, block_size, device=q.device, dtype=torch.bool),
        diagonal=1
    ).view(1, 1, 1, block_size, block_size)

for qb in range(num_blocks):
    q_blk = q[:, :, qb*block_size:(qb+1)*block_size, :].view(B, Hk, div, block_size, D)
    m_running = torch.full((B, Hk, div, block_size), -float('inf'), device=q.device)
    l_running = torch.zeros((B, Hk, div, block_size), device=q.device)
    S_running = torch.zeros((B, Hk, div, block_size, D), device=q.device)

    for kb in range(num_blocks):
        if causal and (kb > qb):
            break

        k_blk = k[:, :, :, kb*block_size:(kb+1)*block_size].transpose(-2, -1)
        v_blk = v[:, :, :, kb*block_size:(kb+1)*block_size].transpose(-2, -1)

        scores = torch.einsum('bhqld,bhkd->bhqlk', q_blk, k_blk) * scale
        if causal and kb == qb:
            scores = scores.masked_fill(mask, float('-inf'))

        max_k = scores.max(dim=-1).values
        prev_m = m_running
        m_new = torch.maximum(prev_m, max_k)

        exp_old = torch.exp(prev_m - m_new)
        exp_new = torch.exp(scores - m_new.unsqueeze(-1))

        l_running = l_running * exp_old + exp_new.sum(dim=-1)
        weighted_new = torch.einsum('bhqlk,bhkd->bhqld', exp_new, v_blk)
        S_running = S_running * exp_old.unsqueeze(-1) + weighted_new

        m_running = m_new

    out_block = S_running / l_running.unsqueeze(-1)
    out[:, :, qb*block_size:(qb+1)*block_size, :] = out_block.view(B, Hq, block_size, D)

return out
content_copy
download
Use code with caution.
def benchmark_attention(device, configs, num_iters=10, causal=True):
results = []
for cfg in configs:
B, H, T, D, block_size = cfg
# prepare inputs
q = torch.randn(B, H, T, D, device=device)
k = torch.randn(B, H, T, D, device=device)
v = torch.randn(B, H, T, D, device=device)
k_naive = k.transpose(2, 3)  # (B, H, D, T)
v_naive = v.transpose(2, 3)

# warm‑up
    _ = naive_online_causal_attn(q, k_naive, v_naive, block_size, causal)
    _ = F.scaled_dot_product_attention(q, k, v, attn_mask=None,
                                       dropout_p=0.0, is_causal=causal)
    if device.type == 'cuda':
        torch.cuda.synchronize()

    # benchmark naive
    torch.cuda.reset_peak_memory_stats(device)
    start = time.perf_counter()
    for _ in range(num_iters):
        out1 = naive_online_causal_attn(q, k_naive, v_naive, block_size, causal)
    if device.type == 'cuda':
        torch.cuda.synchronize()
    t1 = (time.perf_counter() - start) / num_iters
    mem1 = torch.cuda.max_memory_allocated(device) / (1024**2) if device.type == 'cuda' else -1

    # benchmark native
    torch.cuda.reset_peak_memory_stats(device)
    start = time.perf_counter()
    for _ in range(num_iters):
        out2 = F.scaled_dot_product_attention(q, k, v, attn_mask=None,
                                              dropout_p=0.0, is_causal=causal)
    if device.type == 'cuda':
        torch.cuda.synchronize()
    t2 = (time.perf_counter() - start) / num_iters
    mem2 = torch.cuda.max_memory_allocated(device) / (1024**2) if device.type == 'cuda' else -1

    results.append({
        'B': B, 'H': H, 'T': T, 'D': D, 'blk': block_size,
        'naive_time_s': t1, 'native_time_s': t2,
        'naive_peak_mem_MB': mem1, 'native_peak_mem_MB': mem2
    })

return results
content_copy
download
Use code with caution.
if name == 'main':
# device selection
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Running on {device}')

# configurations: (batch, heads, seq_len, dim, block_size)
configs = [
    (1, 4, 128, 384, 32),
    (2, 8, 256, 384, 64),
    (4, 16, 512, 768, 64),
    (4, 16, 512, 768, 128),
    (8, 32, 1024, 2048, 128),
    (1, 32, 4096, 4096, 128),
]

# run benchmarks
results = benchmark_attention(device, configs, num_iters=20, causal=True)

# pretty‑print
header = f"{'B':>2} {'H':>2} {'T':>4} {'D':>3} {'blk':>4} | {'naive (ms)':>10} {'native (ms)':>12} | {'naive MB':>8} {'native MB':>10}"
print(header)
print('-' * len(header))
for r in results:
    print(f"{r['B']:>2} {r['H']:>2} {r['T']:>4} {r['D']:>3} {r['blk']:>4} | "
          f"{r['naive_time_s']*1e3:10.2f} {r['native_time_s']*1e3:12.2f} | "
          f"{r['naive_peak_mem_MB']:8.1f} {r['native_peak_mem_MB']:10.1f}")
content_copy
download
Use code with caution.
: make this faster compared to F.scaled_dot_product_attention, do not suggest to use flash attention or existing attentions, this is a custom solution, just modify algorithmically etc, give full code modified then.