Pull requests / #2058
#2058 Verify window: the routed experts' gate/up with the group's q8_1 rows in shared memory (STRATA_EXPDEC=1, opt-in, bit-exact)
open · @imanu86 · 0 コメント · GitHub で見る
BenchmarksMulti-GPUAMD / HIPNVIDIA / CUDAModels & quantsDocumentationWindows
本文
## Title Issue: none (an opt-in speed change). ## Summary `STRATA_EXPDEC=1` (opt-in, off by default) runs the verify window's VRAM-resident routed experts' **gate/up** through a kernel that keeps the group's q8_1 activation rows in shared memory and computes several rows per lane. Same kernels otherwise, same launch order, and the output is **bit for bit** the default kernels' (`expdec_bench` checks it with `memcmp` on every cell). CUDA only; HIP and the prompt path are not touched. ## Why `native_gu_multi_kernel` already decodes each weight part once per block and applies it to the group's 1-4 entries, but the apply re-reads each entry's q8_1 block from global memory with 9 scattered 4-byte loads (36-byte blocks: a warp's load touches 16-32 sectors). Measured with nsys on a 131k decode (RTX 2080 Ti, stage 1 of a layer split), an extra row of the window costs about 60% of a whole expert - ~1.6 us against ~2.6 us per expert and layer - although it reads nothing from DRAM. That is most of the gate/up kernel's slope with the window's rows (+0.86 ms per row). ## What changed - `src/kernels/cuda/iq_kernels.cu`: `expdec_gu_kernel<TG, R, MINB>` (gate/up IQ2_XXS, IQ2_XS, IQ3_XXS, IQ3_S, IQ2_S at n_embd 2560) and `expdec_down_kernel<TD, R>` (down Q2_0 at n_ff 640; IQ4_NL with `STRATA_EXPDEC_D20=1`). The group's entries are staged 4 at a time in a layout where a call's 8 words are two conflict-free 16-byte loads (IQ4_NL: two 8-byte loads) and the fp16 block scale is already a float; R rows per lane share the activation words; a block takes `rb` rows (auto from the launch's groups) instead of 16; the IQ3_XXS / IQ3_S codebooks are staged in shared memory like the IQ2 ones. `Split<TG>::apply_u` is `apply` on the block's words; `Split<18>/<21>::load<true>` read the staged grid. `native_expert_grouped` takes the new kernels when the lever is on and the layout matches, else the default ones. - `include/strata/kernels/iq_kernels.hpp`: `native_expert_set_expdec`, `native_expert_set_expdec_params` / `native_expert_expdec_params` (`NativeExpertExpdecParams`), `native_expert_expdec_takes`. - `src/kernels/expdec_bench.cpp` (`EXCLUDE_FROM_ALL`): the default kernels against EXPDEC at T 1-4 on 10 format pairs, synthetic blobs read cold from a >= 96 MB pool, `memcmp` of the output, median of 5 x `reps` calls, per phase. - `docs/EXPERTDEC.md`. Knobs (environment, read at startup): `STRATA_EXPDEC_R` (gate/up rows per lane, 1 | 2, default 1), `STRATA_EXPDEC_DR` (down, 1 | 2 | 4, default 4), `STRATA_EXPDEC_MINB` (gate/up `__launch_bounds__` min blocks per SM, 4 | 2, default 4), `STRATA_EXPDEC_RB` / `_DRB` (rows per block, 0 = auto), `STRATA_EXPDEC_MINBLK` (the auto rule's minimum blocks, 800), `STRATA_EXPDEC_DOWN` (the down too, default 0), `STRATA_EXPDEC_D20` (IQ4_NL down too, default 0). The first launch of each (role, format) logs one line (`strata expdec: CUDA1 gate/up type 21: expdec_gu_kernel R 1 MINB 4, ... captured into a graph`). **The down stays separate on purpose.** In the micro-bench the Q2_0 down with EXPDEC is 0.65-0.85x of the default, but in the engine (nsys, two captures with two block geometries) it was about 10% *slower* than the default kernel, which already keeps the group's h rows in shared memory; IQ4_NL showed no gain in either. So `STRATA_EXPDEC=1` takes the gate/up only and `STRATA_EXPDEC_DOWN=1` is a separate opt-in for whoever wants to measure it on another card. ## Bit for bit For every (row, entry) the sum is the default kernel's tree (`row_dot_80_sub16`: lane t of a 16-lane sub-warp does calls t, t+32, t+64 into `s` and t+16, t+48 into `s16`, `s = __fadd_rn(s, s16)`, butterfly xor 8-4-2-1; `row_dot_40_sub8` / `row_dot_20_sub16` for the downs). The same lane does the same calls in the same order with the same integer and float operations on the same words; the entries of a group are independent, so the staging and the rows-per-lane / rows-per-block choices only change which lane computes what. `expdec_bench`: `identical` on all 1,280 cells (10 pairs x T 1-4 x 2 window shapes x 8 choices x 2 cards). In the engine the greedy text is unchanged. ## Measurements Hardware: RTX 2080 Ti 22 GB (sm_75) + RTX 3060 12 GB (sm_86), Windows 11, CUDA 12.6, layer split (3060: layers 0-15, 2080 Ti: 16-47 + head), `--pipeline-windows 2`, Qwen3.8-Flash-Next IQ3_XXS pack, 131k context. Micro-bench, gate/up phase, us per call with 20 distinct experts (the shape of a real window), `def` -> ratio with EXPDEC: | card | format | T=1 | T=2 | T=3 | T=4 | |---|---|---|---|---|---| | RTX 2080 Ti | IQ3_S | 72 / 0.73x | 112 / 0.79x | 124 / 0.75x | 138 / 0.74x | | RTX 2080 Ti | IQ3_XXS | 61 / 0.79x | 95 / 0.84x | 107 / 0.79x | 119 / 0.77x | | RTX 2080 Ti | IQ2_XXS / IQ2_XS / IQ2_S | 42-62 / 0.75-0.98x | 74-83 / 0.92x | 86-96 / 0.86-0.87x | 97-109 / 0.84x | | RTX 3060 | IQ3_S | 74 / 1.09x | 144 / 0.99x | 154 / 0.89x | 174 / 0.87x | | RTX 3060 | IQ3_XXS | 78 / 0.89x | 132 / 0.92x | 143 / 0.90x | 162 / 0.82x | | RTX 3060 | IQ2_XXS / IQ2_XS / IQ2_S | 60-68 / 0.90-1.04x | 112-118 / 0.89-1.00x | 119-136 / 0.89-0.95x | 132-149 / 0.87-0.91x | Per-expert / per-entry cost (us per layer call, RTX 2080 Ti, IQ3_S): 5.20 / 1.26 default -> 3.73 / 0.64. nsys in the engine (96 verify windows, median kernel time per window at T = 1 / 2 / 3 / 4): RTX 2080 Ti gate/up 1.60 / 2.58 / 3.51 / 4.16 ms -> **1.42 / 2.07 / 2.75 / 3.20 ms (-11 / -20 / -22 / -23%)**, slope per row +0.86 -> +0.60 ms; RTX 3060 gate/up 0.77 / 1.43 / 1.97 / 2.23 -> 0.80 / 1.21 / 1.83 / 1.84 ms. End to end, deterministic 131k bench (forced tokens and window sizes, 1024 decode tokens, three interleaved runs lever on / off / on): **+1.1%** decode tokens/s (+0.0 and +2.1 in the two B arms, no B arm below A), text identical. With the companion QSA scores lever (`STRATA_QSA_ROWS=1`, separate PR) on top: +2.9% together. The expert kernels are ~20% of the critical card's window and EXPDEC takes ~20% of that, so the end-to-end number is small by construction. Ported and validated on v0.1.43 (micro-bench bit-exact on both cards; our full stack on 0.1.43 matches our 0.1.42 daily on the forced 131k bench). ## Risks / limits - Off by default; a format or geometry outside the list keeps the default kernel for that role; `STRATA_OLD_IQ_MMVQ=1` turns it off as it turns the multi kernels off. Captured graphs keep the kernels they captured. - Shared memory per block: 11.5 KB of activations + 1-8 KB of codebook (IQ2_S: 19.7 KB); registers 64 with the default `MINB` 4 (116-128 with `2`). Checked with `cuobjdump -res-usage` on sm_75 and sm_86. - Not measured on sm_80/89/120 or on HIP (CUDA only). The weight side (the 8-9 two-byte loads per call and row) is untouched and is the next step; T=1 gains less (the flat cost is the weight side). **Scope.** The gain comes from VRAM-resident routed experts / decode QSA scores on Turing/Ampere setups where these kernels are a large share of the verify window. On rigs where most routed experts are off-GPU (e.g. 4x RTX 5080 16 GB, see the test in #2058) the window is dominated by stage hand-offs and host experts, and the levers measure neutral.
関連リンク
インストール・モデル・リリースへの站内リンク。