Pull requests / #655

#655 sm_75: run the prompt path's BF16 products on the FP16 tensor cores (+15-18 % prefill on RTX 2080 Ti)

closed · @rafatxf · 0 评论 · 在 GitHub 查看

BenchmarksSetup & installAMD / HIPNVIDIA / CUDAModels & quantsLinux

描述

## What

Turing (sm_75) and Volta (sm_70) have FP16 tensor cores but no BF16 ones, so `cublasGemmEx` on `CUDA_R_16BF` inputs falls back to a SIMT fp32 kernel (`magma_sgemmEx`). On an RTX 2080 Ti that was the largest kernel of the prompt path, **20 % of a 4K prompt** (nsys). It runs every `Gemm::bf16` product: the hyper-connection down / up projections, the router and indexer rows, and PLE key / value.

On compute capability 7.x, `Gemm::bf16` now:
- converts the weight and the activations BF16 → FP16, the activations in row slices of at most 32 MiB;
- runs the existing `Gemm::f16` product (fp32 accumulate) on the tensor cores.

BF16 → FP16 is exact for every value inside FP16's range, because the 7-bit mantissa fits in 10 bits.

- **Scope.** Only compute capability 7.x (Volta, Turing) by default; every other card runs exactly what it ran before. HIP is untouched (`#if !defined(__HIPCC__)`).
- **Override.** `STRATA_BF16_TC=0/1` forces the path off or on.
- **Single-row products** (N = 1) stay on cuBLAS's GEMV, which is faster than the conversions.
- **Scratch.** It belongs to the `Gemm` instance (there is one, in `Prefill::Impl`) and is freed with it. With Qwen3.8-Flash-Next it peaks at about 45 MiB: the 2560×2560 PLE value weight plus one activation slice.
- **Allocation failure.** If the scratch allocation fails, the call falls back to the cuBLAS BF16 path.

## Accuracy

**New test `gemm_bf16_parity`.** It compares `Gemm::bf16` with cuBLAS's own BF16 product on the same inputs. GPU, synthetic, no model; `--bench` adds timings. It covers:
- the prompt path's shapes;
- a T large enough to slice the activations;
- `beta = 1` accumulation (the bf16x2 low parts);
- an output row stride wider than N.

On the 2080 Ti the worst relative difference is 1.7e-5. With `STRATA_BF16_TC=0` it is 0, so the test also passes unchanged on other cards.

**Real products.** I also checked every `beta = 0` product of a real 32K prompt (400 products) against the cuBLAS BF16 result, with a temporary diagnostic that is not in this PR:
- worst relative difference 3.5e-5;
- no activation outside FP16 range (max |X| 75.5);
- no non-finite values.

**Greedy output.** It matches the unpatched engine for the first 52 (4K) and 105 (32K) tokens, then diverges the usual way for a change in summation order.

`gr_parity` still passes.

## Numbers

Setup:
- RTX 2080 Ti 22 GB (PCIe 3.0 x16), i7-6700, 32 GB DDR4-2133, Linux.
- Qwen3.8-Flash-Next Q2_0, `--resident-experts`, `--spec 4`.
- Fixed-token prompts, greedy, same binary with `STRATA_BF16_TC=0` vs `1`.

| prompt | prefill t/s, before | after | |
|---|---:|---:|---:|
| 4K | 877 | 1020 | +16 % |
| 16K | 1134 | 1337 | +18 % |
| 32K | 1153 | 1366 | +18 % |
| 64K | 1116 | 1312 | +18 % |
| 128K | 1009 | 1166 | +15 % |

Decode does not use `Gemm::bf16` and is unchanged within noise.

Per product (`gemm_bf16_parity --bench`):

| product | cuBLAS BF16 | this PR |
|---|---:|---:|
| hc down, T 4096 × N 320 × K 10240 | 4428 µs | 1668 µs |
| hc up, T 4096 × N 10240 × K 320 | 3468 µs | 613 µs |
| PLE value, T 777 × N 2560 × K 2560 | 1440 µs | 325 µs |

After the change, the conversion kernel is 3.8 % of a 32K prefill. Writing FP16 directly from the producers would take that back, but it touches more code, so I left it out of this PR.

Measured on sm_75 only (Volta was added in review, unmeasured). Other cards keep the default off, so their path is the same code as before.

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

站内延伸阅读

链到安装、模型与版本说明,便于 SEO/GEO,非官方 issue 正文。