code wiki / (root) / nx_attn_window.nx

nx_attn_window.nx

buildroot/runtime/nx_attn_window.nx

15593 B391 linesdepth 5pulls 7 transitivereach 0 importersview sourcekind tool
docsdependenciesstructsconstsfunctions

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

nx_syscalls.nx nx_tier.nx nx_loop.nx nx_tensor.nx nx_attention.nx nx_attn_window.nx

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

main nx_attn_w_band_start nx_attn_w_band_len nx_attn_w_band_start ↻ nx_attn_w_memory_ratio_q10 nx_attn_w_verdict_is_valid sys_mmap nx_t_alloc nx_dt_is_valid nx_dt_element_bytes sys_mmap ↻ nx_t_compute_strides_rowma nx_attn_w_score_matrix nx_t_is_contiguous nx_attn_w_band_start ↻ nx_attn_w_band_len ↻

structs

none

consts

64const NX_MAGIC_1024: i64 = 1024
65const NX_MAGIC_8000: i64 = 8000
66const NX_MAGIC_8200: i64 = 8200
67const NX_MAGIC_4096: i64 = 4096
77const NX_ATTN_W_NEG_INF: nx_int = -1000000000
81const NX_ATTN_W_OK: nx_int = 0
82const NX_ATTN_W_ERR_BAD_DTYPE: nx_int = 1
83const NX_ATTN_W_ERR_BAD_NDIM: nx_int = 2
84const NX_ATTN_W_ERR_SHAPE_MISMATCH: nx_int = 3
85const NX_ATTN_W_ERR_NOT_CONTIGUOUS: nx_int = 4
86const NX_ATTN_W_ERR_BAD_RADIUS: nx_int = 5
87const NX_ATTN_W_N_VERDICTS: nx_int = 6

functions

89func nx_attn_w_verdict_is_valid(v: nx_int) -> nx_int
called by 1: main
100func nx_attn_w_band_start(i: nx_int, w: nx_int) -> nx_int
106func nx_attn_w_band_len(i: nx_int, n_kv: nx_int, w: nx_int) -> nx_int
123func nx_attn_w_score_matrix(q: *NxTensor, k: *NxTensor,
216func nx_attn_w_apply_to_v(weights_band: *NxTensor, v: *NxTensor,
281func nx_attn_w_memory_ratio_q10(n_kv: nx_int, w_radius: nx_int) -> nx_int
called by 1: main
298func main() -> i64