Modular collision-free sparse attention

Integrate the construction at the attention-mask stage, before computing attention logits. Reshape N=A times B tokens into an A by A grid of blocks, each block containing B within-block positions. For each head h, choose a shift s_h in {0,...,B-1} and define p_h(i)=1+((2(i-1)+s_h) modulo B). For every block row R, block column C, and within-block query position i, add one edge from r=B(R-1)+i to c=B(C-1)+p_h(i). Store the edge list once as a compressed sparse row structure. At runtime, gather Q[r], K[c], and V[c] for all edges, compute z[r,c]=Q[r] dot K[c] divided by the square root of the head dimension, apply a segmented softmax over keys sharing the same query r, and scatter-add the weighted values into each query output. Use different shifts across heads, for example s_h=h modulo B, and optionally add a small fixed local-window edge set so nearby tokens remain directly connected. The modular formulas and injectivity claims are computed exactly; attention quality, useful shift assignments, and the benefit of local edges must be measured empirically. Start with ViT-Tiny on CIFAR-100 or ImageNet-100, using N=120 tokens, A=8, B=15, and four shifted heads. Compare dense attention, standard window attention, and the modular mask at equal hidden size and batch size. The first success criterion is at least 2 times lower attention memory and FLOPs with no more than one percentage point top-1 accuracy loss. A stronger result is matching dense-attention accuracy at larger token counts or outperforming random sparse attention at the same number of edges.