Undergraduate Honors Research Thesis, The Ohio State University, 2026

FlashStieltjes: Fast and Memory-Efficient Flash-Style Attention with Algebraic Normalization

Alex Gulko · Advised by Sachin Kumar

TL;DR

FlashStieltjes is my honors thesis at Ohio State. It is a GPU kernel for Stieltjes attention, which weighs keys with a power law where softmax uses an exponential. It is written in Triton in the style of FlashAttention and never builds the N×N attention matrix. On an A100 its forward pass ran 3.37× faster (geometric mean over the benchmarked shapes) than a PyTorch reference implementation of the same attention.

Why touch softmax

Nearly every Transformer turns attention scores into weights with softmax. But softmax is only one of many functions that take a list of scores and return positive numbers that add up to 1. It does so with an exponential: every step down in score shrinks the weight by the same factor.

Stieltjes attention uses a different rule:

p_j = (\lambda - s_j)^{-q}

Here $s_j$ is the score of key $j$ , the order $q$ controls how sharp the attention is, and $\lambda$ (the “pole”) sits above the best score. Every key is weighed by how far its score is from the pole, and the pole is placed so that all the weights add up to exactly 1. The weights now fall off like a power law. In this example the top two keys get more weight than under softmax, and so does the far tail: the lowest key gets about 380 times more.

The same 64 scores, each key 0.25 below the previous one, turned into weights two ways. On a log scale softmax's exponential is a straight line down. The algebraic weights bend, ending up heavier at the very top and in the far tail.

Data table

Score gap	Softmax	Stieltjes, q = 4
0	0.221	0.457
1	0.081	0.041
2	0.030	0.009
4	0.004	0.001
8	7.4 × 10⁻⁵	1.4 × 10⁻⁴
12	1.4 × 10⁻⁶	3.3 × 10⁻⁵
15.75	3.2 × 10⁻⁸	1.2 × 10⁻⁵

The challenge

Softmax can be computed in a single pass over the keys: keep a running maximum and a running sum of exponentials, and divide at the end. This is the online softmax of Milakov and Gimelshein, and FlashAttention is built on it. Stieltjes attention has no such decomposition. The pole $\lambda$ depends on every score in the row at once, so it has to be found by solving an equation for each row.

I solve it with Newton-Raphson. After subtracting each row’s largest score, every row starts from the same guess of 1.1. On synthetic scores it reaches single-precision accuracy in 2-3 steps for $q \le 4$ , and the kernel runs three steps by default.

The kernel

The reference way to compute attention builds a table of all N×N scores, and at long sequences that table takes up most of the memory. A flash-style kernel keeps one block of queries in the GPU’s on-chip memory and streams the keys past it in tiles instead. Since Stieltjes attention cannot finish in one pass, FlashStieltjes streams them five times:

One block of queries, five sweeps over the keys. Scores are recomputed on the fly in every sweep, so the only things that survive between sweeps are two numbers per row.

Memory now grows with the sequence length and not with its square.

Results

3.37× faster forward pass than the PyTorch reference on an A100, geometric mean (3.78× on an H100)

4.9× less peak memory at batch 4, 8 heads, 2,048 tokens (329 MB instead of 1.6 GB)

320× less memory at 16,384 tokens, batch 1 (about 65 MB instead of 21 GB)

To see what swapping softmax out does to accuracy, I trained six-layer nanoGPT models (about 23M parameters) on four algorithmic tasks: binary search, finding the maximum, BFS, and sorting. In my runs with $q$ set to 2 or 4, Stieltjes attention scored within 0.005 of softmax on every task.

See more

FlashStieltjes is one piece of a larger ongoing project on Stieltjes attention with Jack Taylor, Tarun Kathuria, and my advisor Sachin Kumar.