Pull requests / #2070

#2070 sm_75: a hyper-connection up kernel for 3+ tokens and the staged down kernel in token pairs (bitwise)

open · @rafatxf · 0 Kommentare · Auf GitHub

BenchmarksNVIDIA / CUDAWindows

Beschreibung

The hyper-connection read is ~19 % of an RTX 2080 Ti's decode GPU time (96 reads per window, 13 MB of BF16 weights each). Two Turing-only changes, both bitwise equal to the kernels they replace (`gr_parity`, `gr_multi_parity`, and the same `fused_gr_bench` hashes):

1. **`gr_up75_kernel`** for 3+ tokens: two batches of 4 rows of `w_up` loaded before the dots (was: one row's load after the previous row's epilogue; ncu showed ~9k of a warp's ~22k cycles waiting on memory), each token's activations read from shared memory once per batch, the 4 xor reductions interleaved, and the rows' last 8 chunks held 4 rows to a `uint4` and shuffled to their lane. Every row's lane sums, xor tree and epilogue are `gr_up_multi_kernel`'s. At 1-2 tokens `gr_up_multi_kernel` is as fast or faster, so it stays there. `STRATA_GR_UP75=0` turns it off.
2. **The staged down kernel takes a token group from `blockIdx.y`**: on sm_75 a launch of 4 or 6 tokens runs as pairs, more blocks for the 41-block grid; every (row, token) is still summed by one warp exactly as before. `STRATA_GR_DOWN_TG=n` forces groups of n (0 = off). One token per group is slower at every T (the weights are read from L2 once per group), so only pairs.

## `fused_gr_bench 300 1 6 8` (cold: 8 weight sets rotated), `STRATA_HC_SPLIT=2`, RTX 2080 Ti at 1650 MHz

Whole read (norm + down + up), µs:

| tokens | before | after |
|---:|---:|---:|
| 1 | 37.1 | 37.1 |
| 2 | 43.6 | 43.7 |
| 3 | 53.2 | 50.2 |
| 4 | 64.1 | 57.6 |
| 5 | 76.2 | 69.3 |
| 6 | 86.1 | 70.9 |

## Engine: same binary, switches off / on

One-shot, 256 new tokens, greedy, runs interleaved off / on / on / off (off = `STRATA_GR_UP75=0 STRATA_GR_DOWN_TG=0`; measured with the Turing tables of #2069 on, both arms), decode tok/s:

| prompt | off | on | |
|---|---:|---:|---:|
| 4K | 55.79 / 55.74 | 56.09 / 56.07 | +0.6 % |
| 32K | 53.52 / 53.16 | 53.98 / 55.15 | +2.3 % |

At 4K all four runs wrote the same tokens. A verify window here averages ~2.2 tokens, where the kernels change little; the gain grows with the window (T 4-6 above), so it is larger with `--spec` above 4 or a layer split's larger windows.

## Tests

`gr_parity --selftest`: OK. `gr_multi_parity`: 32 cases, 0 failures. `fused_gr_bench`: bitwise equal, the same hashes with the switches on and off.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Mehr auf der Site

Links zu Install, Modellen, Releases.