Pull requests / #1525

#1525 prefill: dense IQ4_XS / Q8_0 to FP16 with 16-byte stores, swiglu + q8_1 and the MoE combine + hyper-connection write fused (32K prompt -4.0%, the same bits)

open · @sergqwer · 0 コメント · GitHub で見る

AMD / HIPNVIDIA / CUDAModels & quants

本文

A 32K prompt on the IQ2_XS model spends ~420 ms of its GPU time in three places. A faster kernel does not help any of them: the time goes to how they use memory.

```
IQ2_XS, 32K prompt (4 x 8192-token chunks), nsys, RTX 5090, 6674a006
dequant_kernel<IQ4_XS, H16>          141.1 ms   716 calls   the dense IQ4_XS weights to FP16 before each cuBLAS GEMM
moe_combine4 + gr_write_norm_rs      234.4 ms   192 + 184   the MoE half's combine, then the hyper-connection write
swiglu_kernel + quantize_mmq_q8_1     43.3 ms  5551 x 2     each MMQ group's SwiGLU, then H to q8_1 for the down product
```

- **The dequant** writes 4.67 GB of FP16 per chunk (179 matrices) in 35.3 ms: ~170 GB/s, under 10% of the card's bandwidth. Each thread writes its 32-value group's halves one by one, so one store instruction of a warp touches 32 places 64 bytes apart.
- **The combine and the write** both run at bandwidth: moe_combine4 reads ~1.0 GB per layer and chunk in 622 us (~1.6 TB/s), and gr_write_norm_rs ~0.9 GB in 625 us. Between the two, `bo` (84 MB at 8192 tokens) is written and read back once for each of the 4 hyper-connection rows.
- **swiglu + quantize** are two small launches per group of 16 experts (~31 groups per layer and chunk, 4.4 + 3.4 us), with H in FP32 between them.

**Change**, all on CUDA with the same bits:

- `dequant_f16` / `dequant_f16_ld` of Q8_0 and IQ4_XS take `dequant_h16x8_kernel` when the output is 16-byte aligned and `ld % 8 == 0`: one thread per 8 outputs, one 16-byte store. Each value is group32's expression, rounded by `__float2half_rn`. IQ4_XS reads its 16-entry table from shared memory, because the constant bank serializes divergent indices. Other types and the BF16 output keep `dequant_kernel`.
- `mmq::swiglu_quant`: swiglu, then `quantize_mmq_q8_1<D4>`'s rows of H, in one pass, and H is never stored. It is taken where the down type reads D4 rows (Q2_0, IQ4_NL, Q8_0, the i-quants), with gate/up split (GGUF) or interleaved (Strata pack). Each value is swiglu_kernel's expression, and each 32-value block is the quantizer's (same unit, same flags). `STRATA_PREFILL_SWIGLU_QUANT=0` keeps the two kernels. So does `STRATA_DBG_NAN`, which reads H.
- `moe_combine_write_norm_rs`: one block per token. It computes moe_combine4's sum in registers (the same thread-to-element mapping and k order), then gr_write_norm_rs's four rows, and `bo` is never stored. It is used exactly where the write would be gr_write_norm_rs: not steered, not the last layer, not before layer 1's PLE block, not `STRATA_GR_UNFUSED`, and with Dm / shared 16-byte aligned (where moe_combine takes moe_combine4). `STRATA_PREFILL_COMBINE_WRITE=0` keeps the two kernels, and so does `STRATA_DBG_NAN`.
- Tests (exit 77 without a GPU):
  - `prefill_combine_write_test`: R, rs, xn16 and xn16_lo against the two kernels, bit for bit; T 1 / 333 / 1000 / 517, padded token stride, with and without xn16_lo.
  - `prefill_swiglu_quant_test`: the q8_1 bytes against swiglu + quantize; Q2_0 / IQ4_NL / Q8_0, split and interleaved, 1 / 37 / 4113 rows, a zero row and a 3e4 outlier.
  - `dequant_f16_test`: every element against the exact value rounded to nearest-even FP16; zero, subnormal and overflowing scales, row slices, padded `ld`, and an unaligned output that takes the generic kernel.
- HIP builds compile the new code. Both fusions report themselves off there (not checked against the HIP kernels). The dequant kernel runs there unmeasured; its values do not depend on contraction.

After the change, in the same profile: the IQ4_XS dequant 141.1 -> 11.2 ms, the combine and write 234.4 -> 203.0 ms (one kernel at ~1.6 TB/s), and swiglu + quantize 43.3 -> 19.6 ms.

**Identity** (first-token logits, sha256, `--expert-cache 12000 --pcie-frac 0.25`, 6674a006 vs this branch):

| prompt | 6674a006 | this PR |
|---|---|---|
| 2K | a40b6fea84356938 | a40b6fea84356938 |
| 8K | d349749a07af4cc5 | d349749a07af4cc5 |
| 32K | 050aa170500d5340 | 050aa170500d5340 |

Every A/B run below had the same logits within its prompt as well.

**A/B** on current main: RTX 5090, Ryzen 9 9950X3D, IQ2_XS, `--expert-cache 12000 --pcie-frac 0.25 --vram-reserve-mib 1500`, interleaved pairs, order alternating. The prompt times are from the engine's `prefill N tokens ... ms` line.

| prompt | chunks | pairs | 6674a006 | this PR | paired diff |
|---|---|---|---|---|---|
| 32K | 4 x 8192 (`auto`) | 3 | 4874.8 ± 7.9 ms | 4679.4 ± 4.6 ms | **-195.4 ± 7.3 ms (-4.0%)** |
| 8K | 1 x 8034 (`auto`) | 3 | 1328.2 ± 3.8 ms | 1290.9 ± 16.7 ms | **-37.3 ± 16.9 ms (-2.8%)** |
| 2K | 1 x 2099 | 1 | 669.4 ms | 628.2 ms | -41.2 ms (one pair) |

The kernel savings in the profile add up to ~186 ms per 32K prompt, which matches the 195 ms A/B.

**Not measured:** HIP (the fusions are off there), GPUs other than sm_120, the Q2_0 Strata pack (its interleaved rows are covered only by the test), `--prefill auto:32768`, `--peer-device`, the CPU share, and 96K prompts. Decode does not call these kernels.

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

https://claude.ai/code/session_01VZy1yKaDDiA8a7svdwaHio

関連リンク

インストール・モデル・リリースへの站内リンク。