Pull requests / #1825

#1825 qsa_select: prompt-path block scores with a thread per query head (below sm_80, same bits)

open · @xrip · 0 Kommentare · Auf GitHub

Setup & installMulti-GPUAMD / HIPNVIDIA / CUDAModels & quants

Beschreibung

## What

The prompt path's QSA block scores on cards without the TF32 scorer (below sm_80; measured on Turing, sm_75) run on a
new kernel that is 4.6-5.2x faster per call, with bitwise the same scores. sm_80 and newer, HIP, and decode run the
same code as before. `STRATA_SCORES_THREAD=0` turns it off.

Two commits:

1. **A thread per (query, head)** (`block_scores_thread_kernel`). The old kernel scored every (query, key block) with
   a whole warp: 20 shuffles and five 512-byte loads per score. Now each lane keeps one 128-value query head in
   registers, reads each key once from shared memory as a broadcast, and adds its 32 four-value partials in the
   same tree that lane 0 of the old xor-shuffle sum used. Same pairs, same order, so the same bits. The heads are
   summed 0..3 from 0.0f as before, and the tail block keeps `block_scores_tail_kernel`.
2. **Four lanes per query.** Nsight Compute showed the first version bound by the shared-memory pipe
   (`mio_throttle`): one key load per 4 FMAs. The old shuffle tree splits into four quarters that only meet at the
   end, so now four lanes share a query, each holds one quarter of all four heads, one key load feeds 16 FMAs, and
   two xor shuffles finish the same tree. Still the same bits.

New test: `qsa_scores_parity` checks every score of the new kernel against `block_scores_kernel` (normal values, wide
magnitudes with zeros, small integers with ties; up to 262K context), then times both.

## Measured

CMP 50HX (TU102, sm_75, unlocked to RTX 2080 speed), CUDA 13.3, GCC 14.2.

- `qsa_scores_parity`: 141 M scores bitwise identical on both cards.
- Time per call, old warp kernel -> this PR:

  | queries | context | before | after |
  |---:|---:|---:|---:|
  | 1,024 | 140K | 35.9 ms | 6.2 ms |
  | 1,024 | 200K | 52.4 ms | 8.5 ms |

- Why it matters: in an nsys profile of one +4.4K-token prompt at 140K (Flash-Next IQ3_XXS, 2x CMP 50HX,
  `--layer-split auto`), `block_scores_kernel` was the largest prompt kernel: 1,045 ms per card, 17 % of the time.
- Whole engine, same setup: a coding-agent session from 34K to 190K tokens (36 turns, +~4.4K tokens each, greedy,
  fresh server per run) gave word-for-word the same text as before. Time to first token at 190K went from 10.49 s to
  7.13 s, but that number also includes the tokenizer change in #1823; the scorer's own share was not split out.

Note: measured on a fork based on 0.1.38. The patch applies to current `main` without changes; re-measuring on 0.1.41
is still to do. Volta and Pascal (the experimental build) also take this kernel and were not tested.

## How to check

On an sm_75 card: `qsa_scores_parity` (also in ctest). It prints `PASS` when every score matches, and the time of both
kernels.

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

Mehr auf der Site

Links zu Install, Modellen, Releases.