Pull requests / #3498

#3498 [Engine] hip: STRATA_SELECT_SPLIT2=1 - the selection scorer's hi/mid bf16 pieces (3 wmma instead of 6)

open · @vlupilin · 0 Kommentare · Auf GitHub

BenchmarksMulti-GPUAMD / HIPNVIDIA / CUDAModels & quantsLinux

Beschreibung

Problem and result: The QSA selection scorer on gfx11/gfx12 (the WMMA path, STRATA_SELECT_WMMA=1) computes 6 wmma products per block. With STRATA_SELECT_SPLIT2=1 the three products that carry the lo piece are dropped (3 wmma instead of 6): the scorer runs 1.5-2.2x faster at the kernel level, with a rounding-level accuracy change (rel err ~1.5e-5, the class upstream's own WMMA-vs-warp gate accepts at 254/256).

Related discussion: #3163 (the original PR, closed in the 0.2.0 tracker reset; the check passed but GitHub would not reopen it - bot said to open a new PR and link this one). Companion to the accepted wide top-k (#1749): together they optimize both stages of the QSA selection pipeline.

Approach: Template parameter bool S2 on block_scores_wmma_kernel; opt-in env STRATA_SELECT_SPLIT2=1 on top of STRATA_SELECT_WMMA=1. Default off: the 6-product kernel's bits are unchanged. One file changed: src/kernels/cuda/qsa_select.cu (+35/-14).

Validation:
  Passed: Kernel benchmark (v0.2.1.3 @ c093ac26, gfx1151 Radeon 8060S Strix Halo, ROCm 10.1, Linux 7.0.0-38-generic): qsa_select_bench at 32768/131072/262144/524288 context, 256 queries, 5 reps — scorer 0.809/2.674/8.187/12.409 ms (6-product) vs 0.525/1.551/3.722/7.488 ms (split2), 1.54x/1.72x/2.20x/1.66x. Selections identical 256/256 at 32K, 253/256 at 262K (cells 0.0013%). Engine interleaved 2 rounds at 131072/262144/524288: 946.8/1002.8/928.1 tok/s (mainline) vs 940.7/1033.9/978.7 tok/s (+split2), depth-dependent +3.1%/+5.4% at 262K/524K.
  Failed: none.
  Not run: CUDA backend (the kernel is HIP-only, guarded by #if defined(__HIPCC__)); gfx12 hardware (maintainer's R9700 can verify); long-prompt quality battery.

Limits: HIP only; requires STRATA_SELECT_WMMA=1; accuracy 253/256 (upstream gate class 254/256).

Mehr auf der Site

Links zu Install, Modellen, Releases.