Regular Hypergraph Masks for Balanced Sparse Attention

Integrate this at the sparse-attention mask generator, either during training for block-sparse attention or at inference for a fixed routing pattern. Let Q queries select w of S keys each, represented by a binary matrix A in {0,1}^{Q times S}. Partition queries into K score or locality classes using their current top-w attention scores. Repair each class so key incidences are regular or almost regular. The simplest method sorts queries by preferred key lists, maintains key deficits d_j, and replaces the least costly selected key whenever a key exceeds its target incidence. A stronger method solves a min-cost bipartite b-matching: each query has degree w, each key has degree floor(Q_k w/S) or ceil(Q_k w/S), and edge cost is -q_i^T k_j/sqrt(d) plus a locality penalty. Keep the repaired binary mask for the forward pass; use original attention scores only to calculate reassignment costs, and update masks every few steps. The hypergraph equation determines total incidences and the almost-regular degree targets. Begin with Longformer-style local-plus-global attention on a small Transformer using S=1024, w=64, and 8 heads. Compare learned top-w, random, and balanced masks at equal nonzeros. Measure maximum key fan-in, activation-memory peaks, kernel occupancy, throughput, and validation perplexity. Success is lower fan-in variance and memory spikes with equal or better perplexity at identical attention FLOPs.