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 comentários · No GitHub
AMD / HIPNVIDIA / CUDAModels & quants
Descrição
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
No site
Links install, modelos, releases.