Pull requests / #1821

#1821 qsa_prompt_attn: int8 K/V prefetched in registers and sharing one buffer (Turing / sm_75)

open · @xrip · 0 Kommentare · Auf GitHub

Multi-GPUNVIDIA / CUDAModels & quants

Beschreibung

## What

Turing (sm_75) only, `--kv int8` only: the prompt attention kernel that Turing runs (v1, `qsa_prompt_attn.cu`) gets
faster with the same output bits. Other GPUs and other KV modes run the same code as before.

## Why

On Turing the v1 kernel used 37.9 KB of shared memory per 128-thread block, so only one block fit on an SM (4 warps,
12.5 % of what an SM can hold). Each 32-cell chunk also loaded its scattered K and V rows only after the previous
chunk had finished. Nsight Compute on a CMP 50HX (TU102, RTX 2080 Ti class) showed `long_scoreboard` (waiting on
global loads) as the top stall.

## Change

1. **The next chunk's int8 K/V rows and scales go into registers while the current chunk computes**, then are stored
   where the old gather stored them. (sm_80+ uses `cp.async` for this; Turing does not have it.)
2. **K and V share one shared-memory buffer.** K is read only by the scores and V only by `p·v`, so V is written over
   K after the scores' barrier; the softmax's barrier keeps it before `p·v`. 37.9 -> 29.2 KB per block: **two blocks
   per SM** instead of one.

The same bytes land in the same places and the arithmetic is unchanged, so the results are bitwise the same.

`qsa_prompt_attn_parity` now also prints a checksum of the new kernel's output, so two builds can be compared bit
for bit.

## Measured

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

`qsa_prompt_attn_parity`, int8, 2,048 queries:

| context | before | after | change |
|---:|---:|---:|---:|
| 40K | 21.45 ms | 12.91 ms | -40 % |
| 140K | 21.94 ms | 13.31 ms | -39 % |

All output checksums (int8 and fp16, at 4K / 40K / 140K) match the old kernel.

Whole engine: Flash-Next IQ3_XXS, 2x CMP 50HX with `--layer-split auto`, 2x Xeon E5-2670 v2. A coding-agent session
grew from 34K to 190K tokens over 36 turns (+~4.4K tokens per turn, greedy, fresh server per run). The text of all 36
turns was identical to the old build. Incremental prompt read time went down 3.5 % (median per-turn ratio 0.965), and
decode did not change.

Note: these numbers were 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.

## How to check

On an sm_75 card, before and after:

```
qsa_prompt_attn_parity 40000 2048
qsa_prompt_attn_parity 140000 2048
```

The `new output checksum` lines should be equal and the time per chunk lower.

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

Mehr auf der Site

Links zu Install, Modellen, Releases.