nx_attn_window.nx
buildroot/runtime/nx_attn_window.nx
about
nx_attn_window.nx -- sliding-window (local) attention kernel.
Ships VRAM-track W-001 per docs/VRAM_OPTIMIZATION_REALISTIC_TRACKING.md:
sub-quadratic attention via a fixed-width local receptive field.
Each query token i attends only to K/V positions within
[i - W, i + W], so the score matrix shrinks from [n_q, n_kv] to
[n_q, 2W+1]. For n=4096 tokens, W=256 -> 8x smaller scores
buffer (4096 * 513 vs 4096 * 4096). Working memory savings
dominate at long context.
Composition (bits-up cardinal):
* Operates on `*NxTensor` -- canonical L1 container from
nx_tensor.nx (the same surface nx_attention.nx uses).
* Reuses nx_attn_softmax_row_q10 from nx_attention.nx -- the
softmax kernel is row-wise so it works on the banded
[n_q, 2W+1] layout unchanged.
* Sentinel-mask for invalid slots: NX_ATTN_W_NEG_INF
(chosen so the existing _attn_exp_q10 clamps it to 0).
===== Math =======================================================
Dense attention (existing):
scores[i, j] = Q[i] . K[j] for all i, j
weights[i, *] = softmax(scores[i, *])
out[i, d] = sum_j weights[i, j] * V[j, d]
Sliding-window:
band_start(i) = max(0, i - W)
band_len(i) = min(2W + 1, n_kv - band_start(i))
scores[i, j] = Q[i] . K[band_start(i) + j] for j in [0, band_len(i))
scores[i, j] = NEG_INF for j in [band_len(i), 2W+1)
weights[i, *] = softmax(scores[i, *]) (NEG_INF slots -> 0 weight)
out[i, d] = sum_{j < band_len(i)} weights[i, j] * V[band_start(i) + j, d]
Identity check: with W >= n_kv - 1, every query has band_len == n_kv,
no invalid slots, and the result equals dense attention exactly.
===== Quality envelope (the honest measurement basis) =============
Sliding-window matches dense attention quality within ~1% on
dependencies 5 imports · 0 importers
imports: nx_syscalls.nxnx_tier.nxnx_loop.nxnx_tensor.nxnx_attention.nx
imported by: nobody (leaf or entry point)
call flow from main pre-order; caps 40 nodes / depth 6 declared; ↻ = already shown
structs
| none |
consts
| 64 | const NX_MAGIC_1024: i64 = 1024 |
| 65 | const NX_MAGIC_8000: i64 = 8000 |
| 66 | const NX_MAGIC_8200: i64 = 8200 |
| 67 | const NX_MAGIC_4096: i64 = 4096 |
| 77 | const NX_ATTN_W_NEG_INF: nx_int = -1000000000 |
| 81 | const NX_ATTN_W_OK: nx_int = 0 |
| 82 | const NX_ATTN_W_ERR_BAD_DTYPE: nx_int = 1 |
| 83 | const NX_ATTN_W_ERR_BAD_NDIM: nx_int = 2 |
| 84 | const NX_ATTN_W_ERR_SHAPE_MISMATCH: nx_int = 3 |
| 85 | const NX_ATTN_W_ERR_NOT_CONTIGUOUS: nx_int = 4 |
| 86 | const NX_ATTN_W_ERR_BAD_RADIUS: nx_int = 5 |
| 87 | const NX_ATTN_W_N_VERDICTS: nx_int = 6 |
functions
| 89 | func nx_attn_w_verdict_is_valid(v: nx_int) -> nx_int called by 1: main |
| 100 | func nx_attn_w_band_start(i: nx_int, w: nx_int) -> nx_int |
| 106 | func nx_attn_w_band_len(i: nx_int, n_kv: nx_int, w: nx_int) -> nx_int |
| 123 | func nx_attn_w_score_matrix(q: *NxTensor, k: *NxTensor, |
| 216 | func nx_attn_w_apply_to_v(weights_band: *NxTensor, v: *NxTensor, |
| 281 | func nx_attn_w_memory_ratio_q10(n_kv: nx_int, w_radius: nx_int) -> nx_int called by 1: main |
| 298 | func main() -> i64 |