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 comments · View on GitHub

BenchmarksMulti-GPUAMD / HIPNVIDIA / CUDAWindows

Description

`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

Related on strata.com

Editorial links to help you install, pick models, or read release notes — not part of the upstream thread.