مقایسه‌ی انواع مختلف اتنشن

def swa_mask(n, window):
    M = [[NEG_INF] * n for _ in range(n)]
    for i in range(n):
        lo = max(0, i - window + 1)   # اگر پنجره >= n باشد اتنشن علّی کامل حاصل می‌شود
        for j in range(lo, i + 1):
            M[i][j] = 0.0             # 0.0 یعنی اتنشن دارد، NEG_INF ماسک می‌کند
    return M

def strided_mask(n, window, stride):  # پنجره‌ی محلی، سپس:
        for j in range(0, i + 1, stride):
            M[i][j] = 0.0             # هر توکن به فاصله‌ی استراید به عقب تا 0

render(swa_mask(8, window=4), "sliding window (W=4)")
render(strided_mask(8, window=2, stride=3), "local (W=2) + strided (stride=3)")

from torch.nn.attention.flex_attention import flex_attention, create_block_mask

def swa_pattern(b, h, q_idx, kv_idx):
    return (q_idx - kv_idx < 1024) & (q_idx >= kv_idx)  # پنجره‌ی 1024، علّی

mask = create_block_mask(swa_pattern, B=batch, H=heads, Q_LEN=n, KV_LEN=n)
out = flex_attention(q, k, v, block_mask=mask)          # کامپایل به یک کرنل Triton