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 commentaires · Sur GitHub
BenchmarksNVIDIA / CUDAWindows
Description
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)
Sur le site
Liens install, modèles, releases.