Pull requests / #2815

#2815 strata-q36 on Volta: K-quant prompt experts on a register-fragment tensor-core kernel - prompt +22-25% on a V100, same tokens

open · @merbanan · 0 commentaires · Sur GitHub

BenchmarksNVIDIA / CUDAModels & quants

Description

On a Volta GPU, strata-q36's prompt-path MoE experts in Q4_K / Q5_K / Q6_K now use a new FP16 tensor-core kernel (`moe_tc_kq`) instead of MMQ's dp4a path. This covers the gate/up and down products of a prompt chunk.

**How it works**
- A block handles 64 routed rows against 128 weight rows, with 8 warps.
- In each warp, the four `mma.m8n8k4` quad pairs each take a 32-value chunk of a 128-value k step.
- Each lane decodes its row's 32 values straight into A fragments in registers, with the same arithmetic as ggml's `dequantize_row_q*_K`. Weights are never staged in shared memory.
- The q8_1 activations MMQ already produces (DS4 for Q4_K/Q5_K, D4 for Q6_K) are staged as FP16 in shared memory, double-buffered, with the next step's loads in flight under the current step's products.
- The four chunks' partial sums are added with shuffles. Accumulation is FP32.

**Scope**
- Used only when the device is sm_70. Turing and newer already run MMQ on int8 tensor cores.
- `STRATA_PREFILL_TCKQ=0` keeps MMQ.
- The source is compiled only when `CMAKE_CUDA_ARCHITECTURES` includes 70 (or `native`/`all`). It is about 140 KB of device code, so default builds (`120`) are unchanged. Outside that build the header's inline no-ops leave the call site free of `#ifdef`s.

**Measured** on one Tesla V100 PCIe 32 GB with Qwen3.6-35B-A3B-MTP-UD-Q4_K_M. The build is this branch on main 296b59d2 with `-DCMAKE_CUDA_ARCHITECTURES=70 -DSTRATA_MMQ_KQUANTS=ON`. Each run is `strata-q36 --bench-pp` (`--max-context 131072`), plus a 16K-token greedy generation compared by md5. Two cooled rounds, alternating the two settings:

| | `STRATA_PREFILL_TCKQ=0` (MMQ) | default (this kernel) |
|---|---:|---:|
| pp512 | 1,685.9 / 1,684.4 | **2,109.0 / 2,109.3** (+25%) |
| pp16384 | 1,518.8 / 1,518.3 | **1,854.6 / 1,854.2** (+22%) |
| MTP2 pp512 | 1,640.5 / 1,640.6 | 2,057.1 / 2,056.4 |
| tg256 | 100.6 / 100.3 | 99.7 / 100.8 |
| 16K greedy, tokens md5 | 75a302c75b7c | 75a302c75b7c (identical) |

The products use FP16 operands with FP32 sums, so they are not bit-identical to MMQ. The 16K greedy tokens came out identical in all four runs. The same kernel design in a Maya-based engine also matched llama.cpp's greedy output for 128 of 128 tokens on this model.

Sur le site

Liens install, modèles, releases.