Pull requests / #139
#139 perf(prefill): add shape-selected pre-Ampere BF16 SGEMM with current-main A/B
closed · @rluisr · 0 comments · View on GitHub
BenchmarksServer & APIMulti-GPUAMD / HIPNVIDIA / CUDAModels & quantsLinux
Description
## Summary
Add an optional BF16-to-FP32 expansion + tiled SGEMM path for Volta/Turing (`sm_70` / `sm_75`). Current main already handles BF16 with `cublasGemmEx` on these cards; this is an alternative for selected projection shapes, not an enablement fix.
**The final current-main request A/B is essentially unchanged: 65.590 s → 65.478 s (−0.17%).** This PR is not demonstrated to improve end-to-end performance on the measured workload; the shape-level results and tradeoffs are reported below.
Revised after review:
- Merged current main at `bbaaabb4643bb7873d4cef9d48b5dcf96e6cbff4` (engine 0.1.28), including #87 and the HIP backend.
- Reserve the optional **64 MiB per Gemm at init**, in both `init` and `init_external`. There is no allocation or retry in `bf16()`.
- If the optional allocation cannot fit, use the existing `cublasGemmEx` path. Other CUDA errors are not treated as an OOM.
- Preserve the same fallback for narrow products, deep/narrow projections, and dimensions that do not fit a tile. Removed the custom SIMT fallback.
- Expand BF16 exactly to FP32, preserving its exponent range. Tile prompt rows and output columns; preserve `beta` and output stride.
- Leave sm_80+ and HIP/hipBLASLt unchanged.
- Moved the independent `STRATA_PREFILL_RING=8` staging change to **#225**. It is not in this diff or the current-main A/B.
- Added the repository-owned `prefill_gemm_test` CUDA/CTest target, including deterministic allocation-failure injection on Linux.
## V100 A/B against current main
| Warm median | pp s | pp tok/s | tg tok/s | Wall s |
|---|---:|---:|---:|---:|
| Current main `bbaaabb` / GemmEx | 60.151 | 527.50 | 50.15 | 65.590 |
| Revised PR / shape-selected SGEMM | 60.131 | 527.65 | 51.15 | 65.478 |
Observed warm wall-time difference: **-0.17%**. With only two warm samples per arm, this is not evidence of a reliable general end-to-end speedup.
<details>
<summary>All request measurements</summary>
| Arm | Run | pp ms | pp tok/s | tg ms | tg tok/s | Wall s |
|---|---|---:|---:|---:|---:|---:|
| main | cold | 60264.7 | 526.5 | 7137.8 | 35.9 | 67.731 |
| main | warm 1 | 60152.5 | 527.5 | 5352.7 | 47.8 | 65.829 |
| main | warm 2 | 60149.3 | 527.5 | 4872.3 | 52.5 | 65.350 |
| pr | cold | 60194.0 | 527.1 | 7264.7 | 35.2 | 67.796 |
| pr | warm 1 | 60044.0 | 528.4 | 5288.7 | 48.4 | 65.660 |
| pr | warm 2 | 60217.2 | 526.9 | 4751.7 | 53.9 | 65.297 |
All six requests: zero reused prompt tokens, zero cgroup CPU throttling events. No OOM kills were recorded.
</details>
### GEMM-only comparison (exact 64 MiB tiled production path)
Three warmups, then seven CUDA-event-timed calls on an otherwise idle GPU; median below. Initialization is excluded, but both BF16 expansion and SGEMM are included. No staging or server is involved.
The two arms were run sequentially with normal GPU boost behavior, not locked clocks. In particular, `4096/320/10240` uses the same GemmEx path in both final arms; its GPU 0 timing difference should not be interpreted as a kernel improvement.
| GPU | T / N / K | Main GemmEx ms | Final PR ms |
|---|---|---:|---:|
| 0 | 4096 / 320 / 10240 | 2.744 | 2.483 |
| 0 | 4096 / 10240 / 320 | 3.186 | 2.241 |
| 0 | 4096 / 512 / 2560 | 1.283 | 1.054 |
| 0 | 4096 / 10240 / 2560 | 20.711 | 17.745 |
| 1 | 4096 / 320 / 10240 | 2.481 | 2.484 |
| 1 | 4096 / 10240 / 320 | 2.890 | 2.242 |
| 1 | 4096 / 512 / 2560 | 1.153 | 1.059 |
| 1 | 4096 / 10240 / 2560 | 20.609 | 17.636 |
| 2 | 4096 / 320 / 10240 | 2.483 | 2.483 |
| 2 | 4096 / 10240 / 320 | 2.884 | 2.241 |
| 2 | 4096 / 512 / 2560 | 1.159 | 1.061 |
| 2 | 4096 / 10240 / 2560 | 20.606 | 17.631 |
The initial revised candidate regressed `4096/320/10240` by about 19% (3.27 ms vs 2.75 ms on GPU 0). The final branch therefore retains GemmEx when `N < ceil(K/16)`, as well as when `N < 32`. The regression test includes that selection. This heuristic is based on this V100 experiment, not a claim of universal optimality.
### Matched conditions
- 3 × Tesla V100-PCIE-16GB, driver 580.173.02, CUDA Toolkit 12.9.86; dual Xeon E5-2650 v4, 12-core container quota.
- Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S, context 32768, INT8 KV, MTP spec=4, pool-workers 8, temperature 0, no thinking.
- **Both arms:** prefill 4096, fixed layer split `18,34` (18/16/14 layers), `STRATA_PREFILL_RING` unset, `STRATA_PREFILL_TIMING=1`, prompt-cache/root/every disabled.
- Identical expert-profile file and cache allocation: 4332 / 3469 / 2675 slots on GPUs 0 / 1 / 2.
- Restart per arm, one cold and two warm requests, no concurrent requests. Each request has 31,730 input and 256 generated tokens; warm means warm expert cache, not reused input.
- Both engines built from the same current-main source with only the PR GEMM/header substitution in the candidate. **Current main still rejects sm_70:** for this V100 experiment only, both arms identically lower the CMake/runtime acceptance checks from 75 to 70. Those two local edits are not proposed by this PR; no older local kernel or staging patches were carried over.
This is a small, repetitive-input experiment, not a model-quality benchmark or a statistically established speedup. SGEMM changes reduction order and rounding; greedy outputs need not be bitwise identical.
### Request reproduction
For run indices 0, 1, 2, construct a user message from:
1. Prefix `Benchmark 31744-{run}-coldwarm-20260930.\n`.
2. Repetitions of `The expedition recorded observations of the landscape, atmosphere, rocks, and distant stars.\n`.
3. Suffix `\nWrite a detailed original science-fiction story of at least 2000 words about an expedition to a distant planet. Start the story directly, continue at length, and do not discuss the notes above.`
Choose the largest repetition count whose model-tokenized rendered chat template (`enable_thinking=False`, special-token parsing enabled) fits 31,744 tokens. The rendered input here was 31,730 tokens. POST `/v1/chat/completions` with `max_tokens: 256`, `temperature: 0`, and `reasoning_effort: "none"`. Measure wall time until the complete JSON response is received; use the engine-reported prompt/decode timings. Report run 0 separately and the median of runs 1 and 2 as warm.
## Validation
- [x] Exact revised GEMM compiled for sm_70, sm_75, and sm_80 with CUDA 12.9; sm_75 and sm_80 are compile-only here.
- [x] Full current-main and candidate experimental V100 engines built cleanly.
- [x] Repository `prefill_gemm_test` built through CMake and passed through CTest on all three V100s.
- [x] CPU double-reference checks for owned/external modes and normal/injected-OOM initialization; allocation is attempted exactly once at init and never retried by requests; external workspace remains owned by the caller after destruction and rebind.
- [x] Shapes `(1,35,47)`, `(17,35,47)`, `(33,35,47)`, `(1031,35,47)`, `(3,4103,4097)`, `(17,4,47)`, `(3,35,577)`; beta 0/0.5/1, output stride `N+4`, row/column tile boundaries, narrow/deep-narrow fallback, and BF16 65536 outside FP16 range. Tolerance: `1e-5 * (1 + sum(abs(products)))` against CPU double.
- [x] The existing user server was restored after the isolated experiments, with its original 262144-context configuration and arithmetic smoke check. Existing uncommitted work was not overwritten.
On the experimental V100 build:
```sh
cmake --build build --target prefill_gemm_test
ctest --test-dir build -R '^prefill_gemm_test$' --output-on-failure
```
## sm_75 evidence
@Adamyno reported carrying the original changes onto main `a790805` and rebuilding on an RTX 2080 Ti, CUDA 13.0. No measured numbers were included in that comment, so this PR does not treat it as quantified runtime evidence. The revised scratch reservation and shape selection still need an sm_75 runtime A/B.
<details>
<summary>Historical evidence from the original submission (not a current-main comparison)</summary>
The original combined experiment used a local sm_70 SIMT baseline, routed-only staging, and larger chunks. It reported warm wall time 90.82 s → 56.52 s (−37.8%), while decode throughput fell from 43.75 to 35.05 tok/s. The baseline was an experimental V100 build based on `f1b1d961537fd66d37fee68a60015701375b7b5a`, with #87 changes and local sm_70 fixes.
Those measurements remain historical evidence for that combined local tuning. They do **not** establish this GEMM-only PR's advantage over main's existing `cublasGemmEx`, and the original standalone diagnostic used full expansion buffers rather than the exact 64 MiB tiled path. The comparison above supersedes those claims for this PR.
</details>
Related on strata.com
Editorial links to help you install, pick models, or read release notes — not part of the upstream thread.