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 comentarios · En GitHub
BenchmarksSetup & installAMD / HIPNVIDIA / CUDAModels & quantsLinux
Descripción
## 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)
En el sitio
Enlaces a install, modelos, releases.