Pull requests / #1370
#1370 prefill: opt-in STRATA_HC_UPMIX=1 on CUDA - the hyper-connection up projection with gr_mix_r as its epilogue (sm_80+)
open · @sergqwer · 0 comentários · No GitHub
BenchmarksMulti-GPUAMD / HIPNVIDIA / CUDAWindows
Descrição
`STRATA_HC_UPMIX=1` (31b2472a) fuses the hyper-connection read's up projection with `gr_mix_r` on gfx11; on CUDA `gr_upmix` returns false and the pair runs. This adds a CUDA kernel (sm_80+) behind the same switch. Unset (the default), nothing changes: no new code runs, and the output is byte-identical to v0.1.40.2. ### Why Each layer half's read runs the up GEMM, which writes `gated` (T x 10240 FP32, 336 MB for an 8,192-token chunk). `gr_mix_r` then reads it back with R. `gr_mix_r` already runs at ~96% of the bandwidth (102 KB per token in 0.489 ms at 8K on an RTX 5090), so only fewer bytes make the read faster. Fused, `gated` is never written or read. ### What it does - **The kernel.** `gr_upmix_cuda_kernel` computes a block of 128 tokens x 32 columns x the 4 streams. - K = 320 is staged in 32-deep slices (`cp.async`, double-buffered). - 8 warps run `mma.m16n8k16` BF16 with FP32 accumulation (`ldmatrix`). - A lane ends with its (token, column) pairs for all 4 streams, so the mix runs in registers, in `gr_mix_r_kernel`'s order, and writes `mixed` with its BF16 / FP16 images. - The epilogue's R rows are prefetched into L2 when the block starts. The grid is column-tile fastest, as on gfx11. - **Where it runs.** On sm_80+ (`cp.async` and the BF16 `mma`), checked per device; `STRATA_EMULATE_CC` is honoured. On Turing, `gr_upmix` returns false and the pair runs. On CUDA, chunks under 33 tokens also keep the pair (see the parity results below). - **Not changed.** The gfx11 kernel, HIP and SYCL; `gr_upmix`'s signature; `STRATA_HC_UPMIX_CHECK`. - **Test.** `gr_upmix_parity` (ctest; GPU, synthetic, no model) compares the fused read with cuBLAS + `gr_mix_r` and with FP64. `--bench` adds timings. ### Measured RTX 5090, CUDA 13.3, Windows. **Parity** (`gr_upmix_parity`, 21 sizes from 1 to 32,768 tokens, 0 failures): - **From 33 tokens up, bitwise.** `mixed`, the BF16 image and the FP16 image all equal the pair's. cuBLAS's up GEMM sums K in the kernel's order at these sizes. - **Below 33 tokens, rounding-level.** cuBLAS picks another kernel there: `mixed` max |diff| / max |value| is at most 1.8e-7, and at most 0.010% of the BF16 and 0.058% of the FP16 values differ. Both paths are equally close to FP64. These chunks keep the pair anyway. **Kernel** (`--bench`, medians of 20 alternating runs, ms): | T | up GEMM + `gr_mix_r` | fused | speedup | |---|---|---|---| | 2,048 | 0.100 + 0.101 = 0.201 | 0.141 | 1.42x | | 8,192 | 0.338 + 0.489 = 0.828 | 0.458 | 1.81x | | 32,768 | 1.814 + 4.103 = 5.919 | 2.707 | 2.19x | **Prompt** (IQ2_XS, `generate`, `--prefill auto`, 3 interleaved pairs, ms; "hc read" is `STRATA_PREFILL_TIMING`'s phase): | prompt | off: prompt / hc read | `=1`: prompt / hc read | |---|---|---| | 2,099 tokens | 1,004 / 114, 901 / 103, 896 / 104 | 955 / 101, 863 / 97, 898 / 102 | | 8,034 tokens | 1,646 / 202, 1,713 / 206, 1,726 / 210 | 1,571 / 164, 1,697 / 166, 1,637 / 173 | | 32,036 tokens (4 chunks) | 5,968 / 567, 6,012 / 575, 5,874 / 549 | 5,752 / 429, 5,736 / 431, 5,519 / 407 | The hc read takes 19% less time at 8K and 24% less at 32K, which makes the prompt 4% faster. At 2K the read is a small share of the prompt and the difference is within the noise. **Identity** (IQ2_XS, first-token logits and 32 greedy tokens against v0.1.40.2's engine, same expert cache): - **Unset.** Identical after 600, 2K and 32K tokens. - **`=1`.** The logits are identical at all three lengths, and so are the tokens at 600 and 2K. `STRATA_HC_UPMIX_CHECK=100000` found every fused read of the 32K prompt bitwise equal to the pair (all 384 of them, 4 chunks). - **The 32K tokens with `=1`.** They differ, but not because of the numbers. - The suffix-draft policy (#1252) learns from measured time, so a faster prompt changes which drafts it tries and how the verify windows are batched. The log shows 2 suffix-draft windows (2 of 3 drafts accepted) instead of 4 (6 of 9). - With `--suffix-draft 0`, both engines give the same 32 tokens. - That is why the switch stays opt-in on CUDA too. **Builds.** `75-real;75-virtual;86-real;89-real;120a-real` (the release list plus sm_75 PTX) builds `strata` and the test. `gr_upmix_parity` and `gemm_bf16_parity` pass. The kernel ran only on sm_120; the bitwise result is cuBLAS 13.3's on that card. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01VZy1yKaDDiA8a7svdwaHio
No site
Links install, modelos, releases.