SDPA fallback: sliding window doesn't reduce memory (full mask built regardless of window size)

In _sdpa_attention (nanochat/flash_attention.py), the sliding-window branch (lines 99-110) builds a full (Tq, Tk) boolean mask before applying the window restriction:

row_idx = (Tk - Tq) + torch.arange(Tq, device=device).unsqueeze(1)
col_idx = torch.arange(Tk, device=device).unsqueeze(0)
mask = col_idx <= row_idx

if window >= 0 and window < Tk:
    mask = mask & ((row_idx - col_idx) <= window)

return F.scaled_dot_product_attention(q, k, v, attn_mask=mask, enable_gqa=enable_gqa)

The mask tensor is always shape (Tq, Tk), regardless of window. The window only narrows which entries inside that full-size mask are True after it's built. Passing an arbitrary attn_mask= tensor also forces SDPA onto its unfused math backend, which materializes the full Tq x Tk score matrix before masking.

I measured this directly on a Tesla T4 (no FA2/FA3 available, so this is the actual path used). At T=2048, batch=1, 6 query heads:

window=16    -> peak memory ~296 MB
window=64    -> peak memory ~301 MB
window=256   -> peak memory ~300 MB
window=1024  -> peak memory ~300 MB
unlimited    -> peak memory ~305 MB

Memory is effectively flat regardless of window size - a window=16 layer costs about the same as unlimited attention on this fallback path.

I tried an alternative that slices k/v per query block before calling SDPA, so SDPA never sees more than the window actually needs:

for start in range(0, Tq, block_size):
    end = min(start + block_size, Tq)
    q_block = q[:, :, start:end, :]
    k_start = max(0, (offset + start) - window)
    k_end = offset + end
    k_block = k[:, :, k_start:k_end, :]
    v_block = v[:, :, k_start:k_end, :]
    # build a small mask sized to the block, not (Tq, Tk)
    ...
    out_block = F.scaled_dot_product_attention(q_block, k_block, v_block, attn_mask=mask, enable_gqa=enable_gqa)
    outputs.append(out_block)
return torch.cat(outputs, dim=2)

Same T=2048/batch=1 shape, output verified identical to the original (max abs diff 0.0 across several window sizes):

window=16    -> ~41 MB   (296 -> 41)
window=64    -> ~43 MB
window=256   -> ~47 MB
window=1024  -> ~64 MB

Memory now scales with window size instead of staying flat. At the larger batch size where I originally hit an OOM (batch=8), the blocked version used roughly 10x less memory than the masked version.

Happy to open a PR with this if it's useful, or if there's a reason this tradeoff was intentional (e.g. simplicity, or block_size tuning concerns) I'd be glad to hear it.