From e7aa8129a42036efd44d24e4accc5c1218fcb45e Mon Sep 17 00:00:00 2001 From: imanu Date: Mon, 12 Oct 2026 07:20:38 +0200 Subject: [PATCH] Verify window: the routed experts' gate/up with the group's q8_1 rows in shared memory (STRATA_EXPDEC=1, opt-in, bitwise); RTX 2080 Ti gate/up kernel -11/-23% at T 1-4, decode +1.1% on a 131k layer split native_expert_grouped already decodes each weight part once per block, but the apply re-reads every entry's q8_1 block from global memory with 9 scattered 4-byte loads, so an extra row of the window costs ~60% of a whole expert. expdec_gu_kernel stages the group's entries in shared memory (two conflict-free 16-byte loads per call, the fp16 scale as a float), takes rb rows per block and stages the IQ3 codebooks too; the same lane does the same calls in the same order with the same operations (apply_u = apply on the block's words), so the output is bitwise the default kernels' (expdec_bench: identical on every cell, 2 cards). Gate/up IQ2_XXS/IQ2_XS/IQ3_XXS/IQ3_S/IQ2_S at n_embd 2560; the down (Q2_0, IQ4_NL) behind STRATA_EXPDEC_DOWN=1 (in the engine the Q2_0 down was ~10% slower). Knobs STRATA_EXPDEC_R/_DR/_MINB/_RB/_DRB/_MINBLK/_D20; one log line per (role, format) at the first launch. docs/EXPERTDEC.md. Co-Authored-By: Claude Opus 5.5 --- CMakeLists.txt | 7 + docs/EXPERTDEC.md | 107 +++++ include/strata/kernels/iq_kernels.hpp | 22 + src/kernels/cuda/iq_kernels.cu | 605 +++++++++++++++++++++++++- src/kernels/expdec_bench.cpp | 480 ++++++++++++++++++++ 5 files changed, 1216 insertions(+), 5 deletions(-) create mode 100644 docs/EXPERTDEC.md create mode 100644 src/kernels/expdec_bench.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index de94fcdd85a..d1dc6e249e3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1492,6 +1492,13 @@ if(STRATA_NATIVE_EXPERTS) add_executable(native_expert_bench src/kernels/native_expert_bench.cpp) target_include_directories(native_expert_bench PRIVATE ${_strata_gpu_include_directories}) target_link_libraries(native_expert_bench PRIVATE strata_kernels strata_kernels_cpu strata_artifact ${_strata_gpu_runtime_target}) + if(STRATA_ENABLE_CUDA) + # EXPDEC (docs/EXPERTDEC.md): the grouped experts' default kernels against STRATA_EXPDEC=1 at T = 1..4, bit for bit + # and in microseconds, on synthetic blobs read cold from DRAM + add_executable(expdec_bench EXCLUDE_FROM_ALL src/kernels/expdec_bench.cpp) + target_include_directories(expdec_bench PRIVATE ${_strata_gpu_include_directories}) + target_link_libraries(expdec_bench PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + endif() if(STRATA_BUILD_TESTS AND STRATA_ENABLE_CUDA) # Unsloth UD-Q4_K_XL's three expert pairs (and Q8_0 alone, its embedding) on synthetic ggml-quantized weights # at 2560/640: GPU and CPU experts against the float reference, GPU dequantizers against to_float. diff --git a/docs/EXPERTDEC.md b/docs/EXPERTDEC.md new file mode 100644 index 00000000000..42563abb0db --- /dev/null +++ b/docs/EXPERTDEC.md @@ -0,0 +1,107 @@ +# EXPDEC: the verify window's routed experts with the group's activations in shared memory (`STRATA_EXPDEC=1`, opt-in) + +Opt-in: `STRATA_EXPDEC=1`. Unset (or `0`) the engine launches exactly the kernels it launched before. CUDA only; the +output is bitwise the default kernels' (see "Bit for bit" below; `expdec_bench` checks it). + +## What it changes + +`native_expert_grouped` (`src/kernels/cuda/iq_kernels.cu`) computes a verify window's routed experts that are resident +in VRAM: group g is one expert's blob, its entries are the window's rows that picked it (1 to T, T = the window's +tokens). The default gate/up kernel (`native_gu_multi_kernel`) already decodes each weight part once per block and +applies it to the group's entries (`row_dot_80_sub16`), 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). On an RTX 2080 Ti +at a 131k decode that makes an extra entry cost about 60% of a whole expert (~1.6 us against ~2.6 us per expert and +layer) although it reads nothing from DRAM. + +`expdec_gu_kernel` (gate/up: IQ2_XXS, IQ2_XS, IQ3_XXS, IQ3_S, IQ2_S at n_embd 2560) and `expdec_down_kernel` (down: +Q2_0 at n_ff 640; IQ4_NL with `STRATA_EXPDEC_D20=1`) instead: + +- stage the group's entries (4 at a time, as the default) in shared memory in a layout where a call's 8 words are two + 16-byte loads with no bank conflicts (the block's two halves in separate planes), the fp16 block scale already a + float (`__low2float(ds)`, the value `apply` computes); for IQ4_NL the words stay together and a call reads two + 8-byte loads; +- compute R rows per lane (`STRATA_EXPDEC_R`, 1 or 2 for gate/up, default 1; `STRATA_EXPDEC_DR` 1/2/4 for down, + default 4): the activation words are loaded once per R rows; +- take `rb` rows per block instead of 16 (`STRATA_EXPDEC_RB` / `_DRB`; 0 = the most that keep `STRATA_EXPDEC_MINBLK`, + default 800, blocks in the launch), so the staging is amortised over 35-140 KB of weights; +- stage the IQ3_XXS and IQ3_S codebooks in shared memory too (the IQ2 ones already were); +- gate/up `__launch_bounds__(256, STRATA_EXPDEC_MINB)`, default 4 (up to 64 registers, 4 blocks per SM like the + default kernels; `2` allows 128). + +Formats outside the list keep the default kernel for that role; `STRATA_OLD_IQ_MMVQ=1` turns EXPDEC off as it turns +the multi kernels off. The down is behind `STRATA_EXPDEC_DOWN=1`: in the engine the Q2_0 down with EXPDEC was about +10% slower than the default kernel although the micro-bench alone had it at 0.78x (the default down already keeps the +group's h rows in shared memory), so by default EXPDEC takes the gate/up only. + +The first time `STRATA_EXPDEC=1` meets a (role, format) the engine writes one line, e.g. +`strata expdec: CUDA1 gate/up type 21: expdec_gu_kernel R 1 MINB 4, 16 rows per block, grid 80 x 10 (cap 10 groups), captured into a graph` +(or `... not covered: the default kernel runs`). + +API (`include/strata/kernels/iq_kernels.hpp`): `native_expert_set_expdec`, `native_expert_set_expdec_params` / +`native_expert_expdec_params` (`NativeExpertExpdecParams`), `native_expert_expdec_takes`. The knobs are read from the +environment at startup; captured graphs keep the kernels they captured. + +## Bit for bit + +For every (row, entry) the sum is the default kernel's tree - `row_dot_80_sub16` for gate/up (lane t of a 16-lane +sub-warp: calls t, t+32, t+64 into `s`, t+16, t+48 into `s16`, `s = __fadd_rn(s, s16)`, butterfly xor 8-4-2-1), +`row_dot_40_sub8` for IQ4_NL, `row_dot_20_sub16` for Q2_0: the same lane does the same calls in the same order, and +`apply_u` is `apply` with the same integer and float operations in the same order on the same words (the q8_1 block +read from shared memory instead of global). The group's entries are independent of each other, so how many are staged +per pass does not matter; R rows per lane and `rb` rows per block only change which lane computes what. The FMA +contraction of `s += d * sumi` is the same compiler decision on the same expression in both kernels. The bench +confirms it: `identical` on every cell (10 format pairs x T 1-4 x 2 window shapes x every variant x 2 cards). + +## Micro-bench: `expdec_bench` + +Not in the default build: `cmake --build --target expdec_bench`. + +``` +expdec_bench [reps=300] [device=-1] [groups=20] [pairs=all] [choices=def,dec,dec:r2,dec:b4,dec:r2:b4] [shape=real] +``` + +Pairs: the model's 5 gate/up x 2 down formats or `21/20,16/42`. Choices: `def` (the reference) and +`dec[:r<1|2>][:d<1|2|4>][:b<2|4>][:w<0|1>][:x<0|1>][:g][:h][:m]` (rows per lane for gate/up and down, min blocks +per SM, EXPDEC for the down too, also IQ4_NL, rows per block, min blocks; a missing field takes the engine's value). +Shapes: `real` = every token picks 10 experts out of `groups` distinct ones (T=1: 10 groups x 1 entry; T=2: 20 x 1; +T=4: 20 groups, 40 entries); `single` = 10 T groups of one entry; `both` also prints the per-expert / per-entry +decomposition. One line per cell: `dev T shape gu/d choice all_us gu_us down_us GB/s pct identical|DIFFERENT` (the +whole call, then the gate/up and down phases alone; memcmp of the output against `def`; exit 1 if anything differs). +The blobs rotate over >= 96 MB so every call reads its weights cold from DRAM, as the decode does each layer. + +Gate/up phase, us per call, synthetic blobs, 20 distinct experts (shape `real`); `dec` = the defaults above: + +| card | format | T=1 def / dec | 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 | 42 / 0.98x | 74 / 0.92x | 86 / 0.86x | 97 / 0.84x | +| RTX 2080 Ti | IQ2_XS | 60 / 0.75x | 80 / 0.92x | 92 / 0.87x | 104 / 0.84x | +| RTX 2080 Ti | IQ2_S | 62 / 0.75x | 83 / 0.92x | 96 / 0.86x | 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 / XS / 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 gate/up): 5.20 / 1.26 default, 3.73 / 0.64 with +EXPDEC. Down Q2_0 with `w1` (EXPDEC for the down): 0.65-0.85x in the bench on both cards; down IQ4_NL 0.95-1.07x +(no gain, hence the default kernel). + +## In the engine + +nsys of a 131k decode (layer split: RTX 3060 layers 0-15, RTX 2080 Ti layers 16-47 and the head, `--pipeline-windows +2`, Windows, 96 verify windows), median kernel time per window, T = 1 / 2 / 3 / 4: + +| | default | EXPDEC | +|---|---|---| +| RTX 2080 Ti, gate/up kernel | 1.60 / 2.58 / 3.51 / 4.16 ms | 1.42 / 2.07 / 2.75 / 3.20 ms (-11 / -20 / -22 / -23%) | +| RTX 2080 Ti, gate/up slope per row | +0.86 ms | +0.60 ms | +| RTX 3060, gate/up kernel | 0.77 / 1.43 / 1.97 / 2.23 ms | 0.80 / 1.21 / 1.83 / 1.84 ms | +| RTX 2080 Ti, down Q2_0 with `STRATA_EXPDEC_DOWN=1` | - | ~10% slower than the default (hence off by default) | + +End to end, same machine, a deterministic 131k bench (forced tokens and window sizes, 1024 tokens, B-A-B with the +lever on / off / on): **+1.1%** decode tokens/s (+0.0 and +2.1 in the two B arms, no B below A), the text identical. +The expert kernels are about 20% of the critical stage's time and EXPDEC takes about 20% of that, so the end-to-end +figure is small by construction; the lever is cheap, bitwise and off by default. + +Not done here: the weight side (the 8-9 two-byte loads per call and row of `Split::load`; a weight tile staged in +shared memory with 16-byte loads would be the next step), and the T=1 flat cost. diff --git a/include/strata/kernels/iq_kernels.hpp b/include/strata/kernels/iq_kernels.hpp index bd7599f7bea..bb47ecbb956 100644 --- a/include/strata/kernels/iq_kernels.hpp +++ b/include/strata/kernels/iq_kernels.hpp @@ -79,6 +79,28 @@ void native_grouped_set_v1(bool v1); /// The bench only: the AMD kernel layout (STRATA_EXP_MODE values; -1 = the environment's) and the phase /// (0 all, 1 gate/up + SwiGLU + quantize, 2 down). void native_expert_set_mode(int mode, int phase); + +/// EXPDEC (STRATA_EXPDEC=1, docs/EXPERTDEC.md): `native_expert_grouped` stages the group's q8_1 activations in shared +/// memory and computes several rows per lane, the same sums per (row, entry) as the default kernels (bitwise), for +/// gate/up IQ2_XXS / IQ2_XS / IQ3_XXS / IQ3_S / IQ2_S at n_embd 2560 and down IQ4_NL / Q2_0 at n_ff 640; other layouts +/// keep the default kernels. Set before graph capture; captured graphs keep the kernels they captured. +void native_expert_set_expdec(int on); +int native_expert_expdec(); +/// Its knobs, read from the environment at startup (STRATA_EXPDEC_R 1 | 2, _DR 1 | 2 | 4, _MINB 2 | 4, _D20 0 | 1, +/// _RB, _DRB, _MINBLK): gate/up and down rows per lane, the gate/up kernel's __launch_bounds__ minimum blocks per SM +/// (2 = up to 128 registers, 4 = up to 64), EXPDEC for the IQ4_NL down too, gate/up and down rows per block (0 = auto +/// from the launch's groups), the auto rule's minimum blocks. In the setter a field <= 0 (< 0 for the two rows per +/// block and down_iq4_nl) keeps the current value. +/// `down` (STRATA_EXPDEC_DOWN, default 0): EXPDEC for the down too; by default only the gate/up (in the engine the +/// Q2_0 down was ~10% slower with EXPDEC, docs/EXPERTDEC.md par. 7). +struct NativeExpertExpdecParams { + int gu_rows_per_lane = 0, down_rows_per_lane = 0, gu_min_blocks_per_sm = 0, down_iq4_nl = -1, down = -1; + int gu_rows_per_block = -1, down_rows_per_block = -1, min_blocks = 0; +}; +void native_expert_set_expdec_params(const NativeExpertExpdecParams& p); +NativeExpertExpdecParams native_expert_expdec_params(); +/// Whether EXPDEC has kernels for this layout's gate/up and down (the bench's check). +bool native_expert_expdec_takes(const NativeExpertLayout& L, bool& gu, bool& down); /// `iq_mmvq` and `native_expert_grouped` decode each weight part once and apply it to every column / entry; /// true selects the older kernels that decode it again per column (STRATA_OLD_IQ_MMVQ=1 at startup). Both give /// bitwise the same results. Set before graph capture; captured graphs keep the kernels they captured. diff --git a/src/kernels/cuda/iq_kernels.cu b/src/kernels/cuda/iq_kernels.cu index c2e44270f5a..6ca7cb552d7 100644 --- a/src/kernels/cuda/iq_kernels.cu +++ b/src/kernels/cuda/iq_kernels.cu @@ -697,6 +697,16 @@ template<> struct Split<16> { // IQ2_XXS const float d = r.dw * __low2float(bq8_1[iqs / 2].ds); return d * sumi; } + // `apply` with the q8_1 block already in registers: u = its 8 words, d8 = __low2float(ds). Same operations in + // the same order (the EXPDEC kernels, see expdec below). + __device__ static float apply_u(const W& r, const int (&u)[8], float d8) { + int sumi = 0; +#pragma unroll + for (int j = 0; j < 8; ++j) sumi = ggml_cuda_dp4a(r.g[j], u[j], sumi); + sumi = sumi * r.ls / 8; + const float d = r.dw * d8; + return d * sumi; + } }; // IQ2_XS and IQ2_S share the apply: two half sums, two 4-bit scales struct SplitLs2 { @@ -711,6 +721,16 @@ struct SplitLs2 { const float d = r.dw * __low2float(bq8_1[iqs / 2].ds); return d * sumi; } + __device__ static float apply_u(const W& r, const int (&u)[8], float d8) { + int sumi0 = 0, sumi1 = 0; +#pragma unroll + for (int j = 0; j < 4; ++j) sumi0 = ggml_cuda_dp4a(r.g[j], u[j], sumi0); +#pragma unroll + for (int j = 4; j < 8; ++j) sumi1 = ggml_cuda_dp4a(r.g[j], u[j], sumi1); + const int sumi = (sumi0 * r.ls0 + sumi1 * r.ls1 + (sumi0 + sumi1) / 2) / 4; + const float d = r.dw * d8; + return d * sumi; + } }; template<> inline constexpr bool kSplit<17> = true; template<> struct Split<17> : SplitLs2 { // IQ2_XS @@ -770,15 +790,18 @@ template<> inline constexpr bool kSplit<18> = true; template<> struct Split<18> { // IQ3_XXS struct W { int g[8]; int ls; float dw; }; template - __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs, const uint32_t* __restrict__ = nullptr) { + __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs, const uint32_t* __restrict__ s_grid = nullptr) { const block_iq3_xxs* bq3 = (const block_iq3_xxs*) vbq + kbx; const int2 q3_packed = make_int2(get_int_b2(bq3->qs, iqs), get_int_b2(bq3->qs, iqs + 1)); const uint8_t* q3 = (const uint8_t*) &q3_packed; const uint32_t aux32 = get_int_b2(bq3->qs, QK_K / 16 + iqs / 2); + const uint32_t* grid_lut; // the same 256 words either way (expdec stages them) + if constexpr (STAGE_GRID) grid_lut = s_grid; + else grid_lut = iq3xxs_grid; W r; #pragma unroll for (int l0 = 0; l0 < 8; l0 += 2) { - const int2 grid_pos = make_int2(iq3xxs_grid[q3[l0 + 0]], iq3xxs_grid[q3[l0 + 1]]); + const int2 grid_pos = make_int2(grid_lut[q3[l0 + 0]], grid_lut[q3[l0 + 1]]); const uint32_t signs = unpack_ksigns(aux32 >> (7 * l0 / 2)); const int signs0 = __vcmpne4(signs & 0x08040201, 0); r.g[l0 + 0] = __vsub4(grid_pos.x ^ signs0, signs0); @@ -797,23 +820,34 @@ template<> struct Split<18> { // IQ3_XXS const float d = r.dw * __low2float(bq8_1[iqs / 2].ds); return d * sumi; } + __device__ static float apply_u(const W& r, const int (&u)[8], float d8) { + int sumi = 0; +#pragma unroll + for (int j = 0; j < 8; ++j) sumi = ggml_cuda_dp4a(r.g[j], u[j], sumi); + sumi = (r.ls * sumi + sumi / 2) / 2; + const float d = r.dw * d8; + return d * sumi; + } }; template<> inline constexpr bool kSplit<21> = true; template<> struct Split<21> { // IQ3_S struct W { int g[8]; int ls; float dw; }; template - __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs, const uint32_t* __restrict__ = nullptr) { + __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs, const uint32_t* __restrict__ s_grid = nullptr) { const block_iq3_s* bq3 = (const block_iq3_s*) vbq + kbx; const int2 qs_packed = make_int2(get_int_b2(bq3->qs, iqs + 0), get_int_b2(bq3->qs, iqs + 1)); const uint8_t* qs = (const uint8_t*) &qs_packed; const int qh = bq3->qh[iqs / 2]; const int signs_packed_32 = get_int_b2(bq3->signs, iqs / 2); const uint8_t* signs_packed_8 = (const uint8_t*) &signs_packed_32; + const uint32_t* grid_lut; // the same 512 words either way (expdec stages them) + if constexpr (STAGE_GRID) grid_lut = s_grid; + else grid_lut = iq3s_grid; W r; #pragma unroll for (int l0 = 0; l0 < 8; l0 += 2) { - const int2 grid_pos = make_int2(iq3s_grid[qs[l0 + 0] | ((qh << (8 - l0)) & 0x100)], - iq3s_grid[qs[l0 + 1] | ((qh << (7 - l0)) & 0x100)]); + const int2 grid_pos = make_int2(grid_lut[qs[l0 + 0] | ((qh << (8 - l0)) & 0x100)], + grid_lut[qs[l0 + 1] | ((qh << (7 - l0)) & 0x100)]); const int signs0 = __vcmpne4(((signs_packed_8[l0 / 2] & 0x03) << 7) | ((signs_packed_8[l0 / 2] & 0x0C) << 21), 0x00000000); const int signs1 = __vcmpne4(((signs_packed_8[l0 / 2] & 0x30) << 3) | ((signs_packed_8[l0 / 2] & 0xC0) << 17), 0x00000000); r.g[l0 + 0] = __vsub4(grid_pos.x ^ signs0, signs0); @@ -831,6 +865,14 @@ template<> struct Split<21> { // IQ3_S const float d = r.dw * __low2float(bq8_1[iqs / 2].ds); return d * sumi; } + __device__ static float apply_u(const W& r, const int (&u)[8], float d8) { + int sumi = 0; +#pragma unroll + for (int j = 0; j < 8; ++j) sumi = ggml_cuda_dp4a(r.g[j], u[j], sumi); + sumi *= r.ls; + const float d = r.dw * d8; + return d * sumi; + } }; template<> inline constexpr bool kSplit<29> = true; template<> struct Split<29> { // IQ1_M @@ -902,6 +944,16 @@ template<> struct Split<20> { // IQ4_NL const float d = r.dw * __low2float(bq8_1->ds); return d * sumi; } + // `apply` with the block's words in registers: p0 = words iqs, iqs + 1; p1 = words iqs + 4, iqs + 5; d8 = __low2float(ds) + __device__ static float apply_u(const W& r, int2 p0, int2 p1, float d8) { + int sumi = 0; + sumi = ggml_cuda_dp4a(r.v[0].x, p0.x, sumi); + sumi = ggml_cuda_dp4a(r.v[0].y, p1.x, sumi); + sumi = ggml_cuda_dp4a(r.v[1].x, p0.y, sumi); + sumi = ggml_cuda_dp4a(r.v[1].y, p1.y, sumi); + const float d = r.dw * d8; + return d * sumi; + } }; template<> inline constexpr bool kSplit<23> = true; template<> struct Split<23> { // IQ4_XS @@ -962,6 +1014,16 @@ template<> struct Split<42> { // Q2_0 const float d8 = __low2float(bq8_1_chunk->ds); return r.d2 * d8 * sumi; } + // `apply` with the q8_1 block (2 kbx + iqs) in registers: u = its 8 words, d8 = __low2float(ds) + __device__ static float apply_u(const W& r, const int (&u)[8], float d8) { + int sumi = 0; +#pragma unroll + for (int j = 0; j < 4; ++j) { + sumi = ggml_cuda_dp4a(u[j * 2 + 0], r.qx[j], sumi); + sumi = ggml_cuda_dp4a(u[j * 2 + 1], r.qy[j], sumi); + } + return r.d2 * d8 * sumi; + } }; // One row against the n <= NC activations x + off[0..n) (n >= 1, warp-uniform; offsets in q8_1 blocks, 32-bit to @@ -1514,6 +1576,369 @@ __global__ void __launch_bounds__(256) native_down_multi_kernel(const unsigned l } } +// ---------------------------------------------------------------- EXPDEC (STRATA_EXPDEC=1): the group's activations +// in shared memory, R rows per lane +// +// native_gu_multi_kernel / native_down_multi_kernel decode each weight part once and apply it to every entry of the +// group, but the gate/up apply re-reads the entry's q8_1 block from global memory with 9 scattered 4-byte loads +// (36-byte blocks: a warp's load touches 32 sectors), so an entry costs about 60% of a whole expert (nsys, 2080 Ti, +// L02 at 131k: ~2.6 us per expert and ~1.6 us per entry and layer) although it reads nothing from DRAM. Here a +// block stages its group's entries (ED_NC at a time, as the default) in shared memory laid out so that a call's +// words are two 16-byte loads (gate/up, Q2_0 down: the block's two halves in separate planes, conflict-free) or two +// 8-byte loads (IQ4_NL down), the fp16 block scale already a float; every lane computes R rows at once, so the +// activation words are loaded once per R rows; the IQ3_XXS / IQ3_S codebooks are staged too (IQ2_* already were), +// and a block takes `rb` rows (auto: the most that keep ~STRATA_EXPDEC_MINBLK blocks in flight) instead of 16. +// Each (row, entry) sum is the default kernel's sum tree (row_dot_80_sub16 / row_dot_40_sub8 / row_dot_20_sub16: the +// same lane does the same calls in the same order with the same operations, `apply_u` = `apply` on the block's +// words), so the output is bitwise the same (expdec_bench checks it; docs/EXPERTDEC.md). +#ifndef STRATA_EXPDEC_MINB +#define STRATA_EXPDEC_MINB 4 // gate/up __launch_bounds__ min blocks per SM by default (STRATA_EXPDEC_MINB at runtime): 4 = up to 64 + // registers, as the default kernels (bench: equal on the 2080 Ti, 0.87-0.95 on the 3060); 2 = up to 128 +#endif +constexpr int ED_NC = GRP_NC; // entries staged at once +constexpr int ED_XB = 80; // q8_1 blocks of a gate/up row (n_embd 2560) +constexpr int ED_HB = 20; // q8_1 blocks of a down row (n_ff 640) + +// the codebook of a gate/up format in shared memory: words, and the copy (every thread of the block; sync after) +template struct EdGrid; +template<> struct EdGrid<16> { static constexpr int words = 512; + __device__ static void copy(uint32_t* s, int tid) { const uint2* src = (const uint2*) iq2xxs_grid; uint2* dst = (uint2*) s; for (int i = tid; i < 256; i += 256) dst[i] = src[i]; } }; +template<> struct EdGrid<17> { static constexpr int words = 1024; + __device__ static void copy(uint32_t* s, int tid) { const uint2* src = (const uint2*) iq2xs_grid; uint2* dst = (uint2*) s; for (int i = tid; i < 512; i += 256) dst[i] = src[i]; } }; +template<> struct EdGrid<22> { static constexpr int words = 2048; + __device__ static void copy(uint32_t* s, int tid) { const uint2* src = (const uint2*) iq2s_grid; uint2* dst = (uint2*) s; for (int i = tid; i < 1024; i += 256) dst[i] = src[i]; } }; +template<> struct EdGrid<18> { static constexpr int words = 256; + __device__ static void copy(uint32_t* s, int tid) { for (int i = tid; i < 256; i += 256) s[i] = iq3xxs_grid[i]; } }; +template<> struct EdGrid<21> { static constexpr int words = 512; + __device__ static void copy(uint32_t* s, int tid) { for (int i = tid; i < 512; i += 256) s[i] = iq3s_grid[i]; } }; + +// One entry's q8_1 row (XB blocks of 36 bytes, 4-byte aligned) into shared memory as entry c: the 8 words of block k +// in planes (HALVES: s_u[(2 c + h) XB + k] = words 4h..4h+3, two 16-byte loads per call, conflict-free over 16 +// consecutive k; else s_u[(c XB + k) 2 + h], the block's 32 bytes together, for the IQ4_NL pairs) and the fp16 scale +// as s_ds[c XB + k] = __low2float(ds), the value `apply` computes. All 256 threads. +template +__device__ __forceinline__ void ed_stage_entry(int4* __restrict__ s_u, float* __restrict__ s_ds, int c, + const block_q8_1* __restrict__ src, int tid) { + int* su = (int*) s_u; + const int* w32 = (const int*) src; + for (int w = tid; w < XB * 9; w += 256) { + const int k = w / 9, j = w - k * 9; + const int v = w32[w]; + if (j == 0) { + const half2 h = *reinterpret_cast(&v); + s_ds[c * XB + k] = __low2float(h); + } else { + const int q = j - 1; + if constexpr (HALVES) su[((c * 2 + (q >> 2)) * XB + k) * 4 + (q & 3)] = v; + else su[(c * XB + k) * 8 + q] = v; + } + } +} +template +__device__ __forceinline__ void ed_load_u(const int4* __restrict__ s_u, int c, int k, int (&u)[8]) { + const int4 a = s_u[(c * 2 + 0) * XB + k], b = s_u[(c * 2 + 1) * XB + k]; + u[0] = a.x; u[1] = a.y; u[2] = a.z; u[3] = a.w; + u[4] = b.x; u[5] = b.y; u[6] = b.z; u[7] = b.w; +} + +// R gate/up rows against the NC staged entries, lane t of a 16-lane sub-warp: row_dot_80_sub16's tree (calls t, +// t+32, t+64 into s; t+16, t+48 into s16; s = __fadd_rn(s, s16); xor 8-4-2-1), every call's weight part decoded once +// per row and the entry's block loaded once per NC x R products. +template +__device__ __forceinline__ void ed_gu_rows(const uint8_t* const* wr, const int4* __restrict__ s_u, + const float* __restrict__ s_ds, const uint32_t* __restrict__ s_grid, int t, + float (&s)[R][NC]) { + using F = Fmt; + using S = Split; + float s16[R][NC]; +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int c = 0; c < NC; ++c) { s[r][c] = 0.0f; s16[r][c] = 0.0f; } +#pragma unroll + for (int j = 0; j < 3; ++j) { + const int k = t + 32 * j; + const int kbx = k / F::ipb, iqs = F::step * (k % F::ipb); + typename S::W w[R]; +#pragma unroll + for (int r = 0; r < R; ++r) w[r] = S::template load(wr[r], kbx, iqs, s_grid); +#pragma unroll + for (int c = 0; c < NC; ++c) { + int u[8]; + ed_load_u(s_u, c, k, u); + const float d8 = s_ds[c * ED_XB + k]; +#pragma unroll + for (int r = 0; r < R; ++r) s[r][c] += S::apply_u(w[r], u, d8); + } + } +#pragma unroll + for (int j = 0; j < 2; ++j) { + const int k = t + 16 + 32 * j; + const int kbx = k / F::ipb, iqs = F::step * (k % F::ipb); + typename S::W w[R]; +#pragma unroll + for (int r = 0; r < R; ++r) w[r] = S::template load(wr[r], kbx, iqs, s_grid); +#pragma unroll + for (int c = 0; c < NC; ++c) { + int u[8]; + ed_load_u(s_u, c, k, u); + const float d8 = s_ds[c * ED_XB + k]; +#pragma unroll + for (int r = 0; r < R; ++r) s16[r][c] += S::apply_u(w[r], u, d8); + } + } +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int c = 0; c < NC; ++c) { + s[r][c] = __fadd_rn(s[r][c], s16[r][c]); +#pragma unroll + for (int o = 8; o > 0; o >>= 1) s[r][c] += __shfl_xor_sync(0xffffffffu, s[r][c], o); + } +} + +// gate/up: 16 sub-warps of 16 lanes, R rows each = 16 R rows per pass, `rb` rows per block (a multiple of 16 R; 2 n_ff +// is one too, so a pass is whole or absent for the entire block). Group stride as native_gu_multi_kernel. +template +__global__ void __launch_bounds__(256, MINB) expdec_gu_kernel(const unsigned long long* __restrict__ grp_ptr, + const int32_t* __restrict__ grp_start, + const int32_t* __restrict__ n_groups, + const int32_t* __restrict__ ent_tok, + const block_q8_1* __restrict__ xq, NativeExpertLayout L, + float* __restrict__ gate, float* __restrict__ up, int rb) { + const int ng = *n_groups; + if (blockIdx.y >= ng) return; + __shared__ alignas(16) int4 s_u[ED_NC * 2 * ED_XB]; + __shared__ float s_ds[ED_NC * ED_XB]; + STRATA_SHARED_ALIGN16_U32(s_grid_buf, EdGrid::words); + const int tid = threadIdx.x; + EdGrid::copy(s_grid_buf, tid); // visible after the staging's __syncthreads + const int sw = tid >> 4, t = tid & 15; + const int nrow = 2 * (int) L.n_ff, n_ff = (int) L.n_ff; + const int row_base = blockIdx.x * rb; + for (int g = blockIdx.y; g < ng; g += gridDim.y) { + const uint8_t* blob = (const uint8_t*) grp_ptr[g]; + const int e0 = grp_start[g], e1 = grp_start[g + 1]; + for (int e = e0; e < e1; e += ED_NC) { + const int n = min(ED_NC, e1 - e); + __syncthreads(); + for (int c = 0; c < n; ++c) ed_stage_entry(s_u, s_ds, c, xq + (size_t) ent_tok[e + c] * ED_XB, tid); + __syncthreads(); + for (int p = 0; p * 16 * R < rb; ++p) { + const int row0 = row_base + p * 16 * R + sw * R; + if (row0 >= nrow) break; // uniform over the block + const bool is_up = row0 >= n_ff; + const int r0 = is_up ? row0 - n_ff : row0; + const uint8_t* wr[R]; +#pragma unroll + for (int r = 0; r < R; ++r) wr[r] = blob + (is_up ? L.up_off : 0) + (size_t) (r0 + r) * L.gu_row; + float* dst = is_up ? up : gate; + auto run = [&](auto nc) { + constexpr int NC = decltype(nc)::value; + float s[R][NC]; + ed_gu_rows(wr, s_u, s_ds, s_grid_buf, t, s); +#pragma unroll + for (int c = 0; c < NC; ++c) + if (t == c) +#pragma unroll + for (int r = 0; r < R; ++r) dst[(size_t) (e + c) * n_ff + r0 + r] = s[r][c]; + }; + switch (n) { + case 1: run(std::integral_constant{}); break; + case 2: run(std::integral_constant{}); break; + case 3: run(std::integral_constant{}); break; + default: run(std::integral_constant{}); break; + } + } + } + } +} + +// R IQ4_NL down rows, lane t of an 8-lane sub-warp: row_dot_40_sub8's tree. s_p: the block's 8 words together +// (ed_stage_entry), read as the int2 pairs (iqs, iqs + 1) and (iqs + 4, iqs + 5). +template +__device__ __forceinline__ void ed_d20_rows(const uint8_t* const* wr, const int2* __restrict__ s_p, + const float* __restrict__ s_ds, int t, float (&s)[R][NC]) { + using F = Fmt<20>; + using S = Split<20>; + auto term = [&](int k, int r, int c, const typename S::W& w) -> float { + const int kbx = k / F::ipb, par = k & 1; + const int2 p0 = s_p[(c * ED_HB + kbx) * 4 + par], p1 = s_p[(c * ED_HB + kbx) * 4 + 2 + par]; + return S::apply_u(w, p0, p1, s_ds[c * ED_HB + kbx]); + }; + auto loadw = [&](int k, typename S::W (&w)[R]) { + const int kbx = k / F::ipb, iqs = F::step * (k % F::ipb); +#pragma unroll + for (int r = 0; r < R; ++r) w[r] = S::load(wr[r], kbx, iqs); + }; +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int c = 0; c < NC; ++c) s[r][c] = 0.0f; +#pragma unroll + for (int j = 0; j < 2; ++j) { // k = t, t + 32 + const int k = t + 32 * j; + typename S::W w[R]; + loadw(k, w); +#pragma unroll + for (int c = 0; c < NC; ++c) +#pragma unroll + for (int r = 0; r < R; ++r) s[r][c] += term(k, r, c, w[r]); + } + { + const int k = t + 16; + typename S::W w[R]; + loadw(k, w); +#pragma unroll + for (int c = 0; c < NC; ++c) +#pragma unroll + for (int r = 0; r < R; ++r) { + float s16 = 0.0f; + s16 += term(k, r, c, w[r]); + s[r][c] = __fadd_rn(s[r][c], s16); + } + } + float s8[R][NC]; +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int c = 0; c < NC; ++c) s8[r][c] = 0.0f; + { + const int k = t + 8; + typename S::W w[R]; + loadw(k, w); +#pragma unroll + for (int c = 0; c < NC; ++c) +#pragma unroll + for (int r = 0; r < R; ++r) s8[r][c] += term(k, r, c, w[r]); + } + { + const int k = t + 24; + typename S::W w[R]; + loadw(k, w); +#pragma unroll + for (int c = 0; c < NC; ++c) +#pragma unroll + for (int r = 0; r < R; ++r) { + float s24 = 0.0f; + s24 += term(k, r, c, w[r]); + s[r][c] = __fadd_rn(s[r][c], __fadd_rn(s8[r][c], s24)); + } + } +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int c = 0; c < NC; ++c) +#pragma unroll + for (int o = 4; o > 0; o >>= 1) s[r][c] += __shfl_xor_sync(0xffffffffu, s[r][c], o); +} + +// R Q2_0 down rows, lane t of a 16-lane sub-warp: row_dot_20_sub16's tree (call t; call t + 16 for t < 4 into s16). +template +__device__ __forceinline__ void ed_d42_rows(const uint8_t* const* wr, const int4* __restrict__ s_u, + const float* __restrict__ s_ds, int t, float (&s)[R][NC]) { + using F = Fmt<42>; + using S = Split<42>; + { + const int k = t, kbx = k / F::ipb, iqs = F::step * (k % F::ipb); + typename S::W w[R]; +#pragma unroll + for (int r = 0; r < R; ++r) w[r] = S::load(wr[r], kbx, iqs); +#pragma unroll + for (int c = 0; c < NC; ++c) { + int u[8]; + ed_load_u(s_u, c, k, u); + const float d8 = s_ds[c * ED_HB + k]; +#pragma unroll + for (int r = 0; r < R; ++r) { s[r][c] = 0.0f; s[r][c] += S::apply_u(w[r], u, d8); } + } + } + float s16[R][NC]; +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int c = 0; c < NC; ++c) s16[r][c] = 0.0f; + if (t < 4) { + const int k = t + 16, kbx = k / F::ipb, iqs = F::step * (k % F::ipb); + typename S::W w[R]; +#pragma unroll + for (int r = 0; r < R; ++r) w[r] = S::load(wr[r], kbx, iqs); +#pragma unroll + for (int c = 0; c < NC; ++c) { + int u[8]; + ed_load_u(s_u, c, k, u); + const float d8 = s_ds[c * ED_HB + k]; +#pragma unroll + for (int r = 0; r < R; ++r) s16[r][c] += S::apply_u(w[r], u, d8); + } + } +#pragma unroll + for (int r = 0; r < R; ++r) +#pragma unroll + for (int c = 0; c < NC; ++c) { + s[r][c] = __fadd_rn(s[r][c], s16[r][c]); +#pragma unroll + for (int o = 8; o > 0; o >>= 1) s[r][c] += __shfl_xor_sync(0xffffffffu, s[r][c], o); + } +} + +// down (IQ4_NL: 32 sub-warps of 8 lanes; Q2_0: 16 of 16), R rows per lane, `rb` rows per block (a multiple of the +// rows per pass; n_embd too). The entries' h rows are contiguous in hq (entry e at e * ED_HB). +template +__global__ void __launch_bounds__(256, STRATA_EXPDEC_MINB) expdec_down_kernel(const unsigned long long* __restrict__ grp_ptr, + const int32_t* __restrict__ grp_start, + const int32_t* __restrict__ n_groups, + const int32_t* __restrict__ ent_dst, + const block_q8_1* __restrict__ hq, NativeExpertLayout L, + float* __restrict__ out, int rb) { + const int ng = *n_groups; + if (blockIdx.y >= ng) return; + constexpr int LANES = TD == 20 ? 8 : 16; + constexpr int RPP = (256 / LANES) * R; // rows per pass + __shared__ alignas(16) int4 s_u[ED_NC * 2 * ED_HB]; + __shared__ float s_ds[ED_NC * ED_HB]; + const int tid = threadIdx.x; + const int sw = tid / LANES, t = tid % LANES; + const int nrow = (int) L.n_embd; + const int row_base = blockIdx.x * rb; + for (int g = blockIdx.y; g < ng; g += gridDim.y) { + const uint8_t* blob = (const uint8_t*) grp_ptr[g] + L.down_off; + const int e0 = grp_start[g], e1 = grp_start[g + 1]; + for (int e = e0; e < e1; e += ED_NC) { + const int n = min(ED_NC, e1 - e); + __syncthreads(); + for (int c = 0; c < n; ++c) ed_stage_entry(s_u, s_ds, c, hq + (size_t) (e + c) * ED_HB, tid); + __syncthreads(); + for (int p = 0; p * RPP < rb; ++p) { + const int row0 = row_base + p * RPP + sw * R; + if (row0 >= nrow) break; // uniform over the block + const uint8_t* wr[R]; +#pragma unroll + for (int r = 0; r < R; ++r) wr[r] = blob + (size_t) (row0 + r) * L.d_row; + auto run = [&](auto nc) { + constexpr int NC = decltype(nc)::value; + float s[R][NC]; + if constexpr (TD == 20) ed_d20_rows(wr, (const int2*) s_u, s_ds, t, s); + else ed_d42_rows(wr, s_u, s_ds, t, s); +#pragma unroll + for (int c = 0; c < NC; ++c) + if (t == c) { + const size_t o = (size_t) ent_dst[e + c] * L.n_embd + row0; +#pragma unroll + for (int r = 0; r < R; ++r) out[o + r] = s[r][c]; + } + }; + switch (n) { + case 1: run(std::integral_constant{}); break; + case 2: run(std::integral_constant{}); break; + case 3: run(std::integral_constant{}); break; + default: run(std::integral_constant{}); break; + } + } + } + } +} + // ---------------------------------------------------------------- q8_1 (quantize.cu) // value i of a q8_1 row set (the warp holds block i / 32, lane = i % 32) __device__ __forceinline__ void q8_1_store(const float xi, block_q8_1* __restrict__ y, const long long i) { @@ -2663,6 +3088,144 @@ void s27_launch_ts(bool ts, const NativeExpertLayout& L, int64_t cap_groups, cud else s27_launch(L, cap_groups, s, grp_ptr, grp_start, n_groups, ent_dst, ent_tok, X, gate, up, h, hq, out, nh); } +// ---- EXPDEC launch: knobs, shapes, rows per block +int env_int_or(const char* name, int dflt) { + const char* v = std::getenv(name); + return v != nullptr && v[0] != '\0' ? std::atoi(v) : dflt; +} +int g_expdec = env_on("STRATA_EXPDEC") ? 1 : 0; +int g_expdec_r = env_int_or("STRATA_EXPDEC_R", 1); // gate/up rows per lane: 1, 2 (bench: R=1 wins on the 2080 Ti) +int g_expdec_dr = env_int_or("STRATA_EXPDEC_DR", 4); // down rows per lane: 1, 2, 4 (Q2_0: R=4 wins on both cards) +int g_expdec_minb = env_int_or("STRATA_EXPDEC_MINB", STRATA_EXPDEC_MINB); // gate/up __launch_bounds__ min blocks: 2 (<= 128 regs) or 4 (<= 64) +int g_expdec_d20 = env_int_or("STRATA_EXPDEC_D20", 0); // 1: EXPDEC also for the IQ4_NL down (bench: no gain, default kernel) +// 1: EXPDEC for the down too (Q2_0; IQ4_NL with _D20). Default 0: in the engine (nsys of a 131k decode) the Q2_0 down +// with EXPDEC was ~10% slower than the default kernel although the bench alone had it at 0.78, so by default EXPDEC +// takes the gate/up only (-20/-23% at T >= 2 on an RTX 2080 Ti in the engine). +int g_expdec_down = env_int_or("STRATA_EXPDEC_DOWN", 0); +int g_expdec_rb = env_int_or("STRATA_EXPDEC_RB", 0); // gate/up rows per block (0 = auto) +int g_expdec_drb = env_int_or("STRATA_EXPDEC_DRB", 0); // down rows per block (0 = auto) +int g_expdec_minblk = env_int_or("STRATA_EXPDEC_MINBLK", 800); // auto: the most rows per block that keep this many blocks +int expdec_r() { return g_expdec_r == 2 ? 2 : 1; } +int expdec_dr() { return g_expdec_dr == 1 || g_expdec_dr == 2 ? g_expdec_dr : 4; } + +bool expdec_gu_ok(const NativeExpertLayout& L) { + const bool fmt = L.gu_type == 16 || L.gu_type == 17 || L.gu_type == 18 || L.gu_type == 21 || L.gu_type == 22; + return fmt && L.n_embd == 32 * ED_XB && L.n_ff % 64 == 0 && L.n_ff > 0; +} +bool expdec_down_ok(const NativeExpertLayout& L) { + return g_expdec_down && ((L.d_type == 20 && g_expdec_d20) || L.d_type == 42) && L.n_ff == 32 * ED_HB && L.n_embd % 128 == 0 && + L.n_embd > 0; +} +// rows per block: `knob` rounded up to the rows per pass, or (0) doubled from one pass while ceil(nrow / rb) x +// cap_groups blocks stay >= g_expdec_minblk (cap_groups: the launch's group rows; every block row g >= n_groups returns) +int expdec_rows_per_block(int nrow, int rpp, int64_t cap_groups, int knob) { + if (knob > 0) { + int rb = (knob + rpp - 1) / rpp * rpp; + return rb > nrow ? nrow : rb; + } + int rb = rpp; + const int64_t cap = cap_groups > 0 ? cap_groups : 1; + while (rb * 2 <= nrow && ((int64_t) ((nrow + 2 * rb - 1) / (2 * rb))) * cap >= g_expdec_minblk) rb *= 2; + return rb; +} + +template +void expdec_gu_r(int64_t gy, int64_t cap_groups, cudaStream_t s, const unsigned long long* grp_ptr, const int32_t* grp_start, + const int32_t* n_groups, const int32_t* ent_tok, const block_q8_1* X, const NativeExpertLayout& L, + float* gate, float* up) { + const int nrow = 2 * (int) L.n_ff; + const int rb = expdec_rows_per_block(nrow, 16 * R, cap_groups, g_expdec_rb); + const dim3 grid((unsigned) ((nrow + rb - 1) / rb), (unsigned) gy); + expdec_gu_kernel<<>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up, rb); +} +template +void expdec_gu_t(int64_t gy, int64_t cap_groups, cudaStream_t s, const unsigned long long* grp_ptr, const int32_t* grp_start, + const int32_t* n_groups, const int32_t* ent_tok, const block_q8_1* X, const NativeExpertLayout& L, + float* gate, float* up) { + const bool b4 = g_expdec_minb >= 4; + if (expdec_r() == 2) { + if (b4) expdec_gu_r(gy, cap_groups, s, grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); + else expdec_gu_r(gy, cap_groups, s, grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); + } else { + if (b4) expdec_gu_r(gy, cap_groups, s, grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); + else expdec_gu_r(gy, cap_groups, s, grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); + } +} +// One line on stderr per (role, format, taken or not), the first time STRATA_EXPDEC=1 meets it: the proof in the +// engine's log that the EXPDEC kernels are the ones launched (or captured into the window's graph), with the +// choices in force. +void expdec_note(const NativeExpertLayout& L, bool gu, bool taken, cudaStream_t s, int64_t gy, int64_t cap_groups) { + static bool seen[2][64][2] = {}; + const int ty = gu ? L.gu_type : L.d_type; + if (ty < 0 || ty >= 64 || seen[gu ? 0 : 1][ty][taken ? 1 : 0]) return; + seen[gu ? 0 : 1][ty][taken ? 1 : 0] = true; + cudaStreamCaptureStatus cs = cudaStreamCaptureStatusNone; + cudaStreamIsCapturing(s, &cs); + int dev = -1; + cudaGetDevice(&dev); + if (!taken) { + std::fprintf(stderr, "strata expdec: CUDA%d %s type %d (n_embd %lld, n_ff %lld) not covered: the default kernel runs\n", dev, + gu ? "gate/up" : "down", ty, (long long) L.n_embd, (long long) L.n_ff); + return; + } + if (gu) { + const int R = expdec_r(), rb = expdec_rows_per_block(2 * (int) L.n_ff, 16 * R, cap_groups, g_expdec_rb); + std::fprintf(stderr, "strata expdec: CUDA%d gate/up type %d: expdec_gu_kernel R %d MINB %d, %d rows per block, grid %d x %lld " + "(cap %lld groups)%s\n", dev, ty, R, g_expdec_minb >= 4 ? 4 : 2, rb, (2 * (int) L.n_ff + rb - 1) / rb, + (long long) gy, (long long) cap_groups, cs != cudaStreamCaptureStatusNone ? ", captured into a graph" : ""); + } else { + const int R = expdec_dr(), lanes = ty == 20 ? 8 : 16; + const int rb = expdec_rows_per_block((int) L.n_embd, (256 / lanes) * R, cap_groups, g_expdec_drb); + std::fprintf(stderr, "strata expdec: CUDA%d down type %d: expdec_down_kernel R %d, %d rows per block, grid %d x %lld " + "(cap %lld groups)%s\n", dev, ty, R, rb, ((int) L.n_embd + rb - 1) / rb, (long long) gy, + (long long) cap_groups, cs != cudaStreamCaptureStatusNone ? ", captured into a graph" : ""); + } +} + +// false: EXPDEC is off or does not take this layout (the caller launches the default kernels) +bool expdec_gu(const NativeExpertLayout& L, int64_t gy, int64_t cap_groups, cudaStream_t s, const unsigned long long* grp_ptr, + const int32_t* grp_start, const int32_t* n_groups, const int32_t* ent_tok, const block_q8_1* X, float* gate, + float* up) { + if (!g_expdec || g_old_kernels) return false; + if (!expdec_gu_ok(L)) { expdec_note(L, true, false, s, gy, cap_groups); return false; } + expdec_note(L, true, true, s, gy, cap_groups); + switch (L.gu_type) { + case 16: expdec_gu_t<16>(gy, cap_groups, s, grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); break; + case 17: expdec_gu_t<17>(gy, cap_groups, s, grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); break; + case 18: expdec_gu_t<18>(gy, cap_groups, s, grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); break; + case 21: expdec_gu_t<21>(gy, cap_groups, s, grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); break; + default: expdec_gu_t<22>(gy, cap_groups, s, grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); break; + } + return true; +} +template +void expdec_down_r(int64_t gy, int64_t cap_groups, cudaStream_t s, const unsigned long long* grp_ptr, const int32_t* grp_start, + const int32_t* n_groups, const int32_t* ent_dst, const block_q8_1* hq, const NativeExpertLayout& L, float* out) { + constexpr int LANES = TD == 20 ? 8 : 16; + const int nrow = (int) L.n_embd; + const int rb = expdec_rows_per_block(nrow, (256 / LANES) * R, cap_groups, g_expdec_drb); + const dim3 grid((unsigned) ((nrow + rb - 1) / rb), (unsigned) gy); + expdec_down_kernel<<>>(grp_ptr, grp_start, n_groups, ent_dst, hq, L, out, rb); +} +template +void expdec_down_t(int64_t gy, int64_t cap_groups, cudaStream_t s, const unsigned long long* grp_ptr, const int32_t* grp_start, + const int32_t* n_groups, const int32_t* ent_dst, const block_q8_1* hq, const NativeExpertLayout& L, float* out) { + switch (expdec_dr()) { + case 1: expdec_down_r(gy, cap_groups, s, grp_ptr, grp_start, n_groups, ent_dst, hq, L, out); break; + case 2: expdec_down_r(gy, cap_groups, s, grp_ptr, grp_start, n_groups, ent_dst, hq, L, out); break; + default: expdec_down_r(gy, cap_groups, s, grp_ptr, grp_start, n_groups, ent_dst, hq, L, out); break; + } +} +bool expdec_down(const NativeExpertLayout& L, int64_t gy, int64_t cap_groups, cudaStream_t s, const unsigned long long* grp_ptr, + const int32_t* grp_start, const int32_t* n_groups, const int32_t* ent_dst, const block_q8_1* hq, float* out) { + if (!g_expdec || g_old_kernels) return false; + if (!expdec_down_ok(L)) { expdec_note(L, false, false, s, gy, cap_groups); return false; } + expdec_note(L, false, true, s, gy, cap_groups); + if (L.d_type == 20) expdec_down_t<20>(gy, cap_groups, s, grp_ptr, grp_start, n_groups, ent_dst, hq, L, out); + else expdec_down_t<42>(gy, cap_groups, s, grp_ptr, grp_start, n_groups, ent_dst, hq, L, out); + return true; +} + template void launch_gu(dim3 grid, cudaStream_t s, const unsigned long long* grp_ptr, const int32_t* grp_start, const int32_t* n_groups, const int32_t* ent_tok, const block_q8_1* X, const NativeExpertLayout& L, @@ -3448,6 +4011,36 @@ bool g_grouped_v1 = env_on("STRATA_GROUPED_V1"); void native_grouped_set_v1(bool v1) { g_grouped_v1 = v1; } +void native_expert_set_expdec(int on) { g_expdec = on ? 1 : 0; } +int native_expert_expdec() { return g_expdec; } +void native_expert_set_expdec_params(const NativeExpertExpdecParams& p) { + if (p.gu_rows_per_lane > 0) g_expdec_r = p.gu_rows_per_lane; + if (p.down_rows_per_lane > 0) g_expdec_dr = p.down_rows_per_lane; + if (p.gu_min_blocks_per_sm > 0) g_expdec_minb = p.gu_min_blocks_per_sm; + if (p.down_iq4_nl >= 0) g_expdec_d20 = p.down_iq4_nl; + if (p.down >= 0) g_expdec_down = p.down; + if (p.gu_rows_per_block >= 0) g_expdec_rb = p.gu_rows_per_block; + if (p.down_rows_per_block >= 0) g_expdec_drb = p.down_rows_per_block; + if (p.min_blocks > 0) g_expdec_minblk = p.min_blocks; +} +NativeExpertExpdecParams native_expert_expdec_params() { + NativeExpertExpdecParams p; + p.gu_rows_per_lane = expdec_r(); + p.down_rows_per_lane = expdec_dr(); + p.gu_min_blocks_per_sm = g_expdec_minb >= 4 ? 4 : 2; + p.down_iq4_nl = g_expdec_d20 ? 1 : 0; + p.down = g_expdec_down ? 1 : 0; + p.gu_rows_per_block = g_expdec_rb; + p.down_rows_per_block = g_expdec_drb; + p.min_blocks = g_expdec_minblk; + return p; +} +bool native_expert_expdec_takes(const NativeExpertLayout& L, bool& gu, bool& down) { + gu = expdec_gu_ok(L); + down = expdec_down_ok(L); + return gu || down; +} + void native_expert_grouped(const NativeExpertLayout& L, const unsigned long long* grp_ptr, const int32_t* grp_start, const int32_t* n_groups, const int32_t* ent_dst, const int32_t* ent_tok, int64_t cap_groups, int64_t cap_entries, const void* x_q8_1, void* scratch, float* out, void* stream, @@ -3551,6 +4144,7 @@ void native_expert_grouped(const NativeExpertLayout& L, const unsigned long long } else #undef STRATA_GU_AMD #endif + if (!expdec_gu(L, gy, cap_groups, s, grp_ptr, grp_start, n_groups, ent_tok, X, gate, up)) // STRATA_EXPDEC=1 switch (L.gu_type) { #define STRATA_GU(T) case T: launch_gu(ggu, s, grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up); break; STRATA_GU_FMTS(STRATA_GU) @@ -3606,6 +4200,7 @@ void native_expert_grouped(const NativeExpertLayout& L, const unsigned long long } else #undef STRATA_D_AMD #endif + if (!expdec_down(L, gy, cap_groups, s, grp_ptr, grp_start, n_groups, ent_dst, hq, out)) // STRATA_EXPDEC=1 switch (L.d_type) { #define STRATA_DOWN(T) case T: launch_down(gd, s, grp_ptr, grp_start, n_groups, ent_dst, hq, L, out); break; STRATA_D_FMTS(STRATA_DOWN) diff --git a/src/kernels/expdec_bench.cpp b/src/kernels/expdec_bench.cpp new file mode 100644 index 00000000000..a3edf0c3901 --- /dev/null +++ b/src/kernels/expdec_bench.cpp @@ -0,0 +1,480 @@ +// src/kernels/expdec_bench.cpp - the VRAM experts of a decode verify window: the default grouped kernels against +// EXPDEC (STRATA_EXPDEC=1, docs/EXPERTDEC.md), bit for bit and in microseconds, per card, per T, per format pair. +// +// build/expdec_bench [reps = 300] [device = -1 (all)] [groups = 20] [pairs = all] [choices = def,dec,dec:r2,dec:b4,dec:r2:b4] [shape = real] +// +// The model's routed expert formats at its geometry (n_embd 2560, n_ff 640): gate/up IQ2_XXS (16), IQ2_XS (17), IQ3_XXS +// (18), IQ3_S (21), IQ2_S (22); down IQ4_NL (20), Q2_0 (42). `pairs` is a comma list of gu/down type ids ("21/20,16/42") +// or "all" (the 10 pairs). `choices` is a comma list of kernel choices: +// def the default kernels (STRATA_EXPDEC off); the reference for the bitwise check +// dec[:r][:d][:b][:w][:x][:g][:h][:m] +// EXPDEC with r = gate/up rows per lane (1, 2), d = down rows per lane (1, 2, 4), b = the +// gate/up kernel's min blocks per SM (2 | 4), w = 1 EXPDEC for the down too (the engine's +// default is the gate/up only), x = 1 also for the IQ4_NL down, g / h = gate/up / down rows +// per block (0 = auto), m = the auto rule's minimum blocks; a missing field takes the +// engine's value (the environment's STRATA_EXPDEC_* or its default) +// `shape`: real = every token picks 10 experts out of `groups` distinct ones (as fp_expert_bench: T=1 10 groups x 1 +// entry, T=2 20 x 1, T=4 20 groups with 40 entries); single = 10 T groups of one entry each (the same entries as `real` +// spread over distinct experts); both = the two, plus the decomposition at the end: per-expert cost a = (single T=4 - +// real T=2) / 20 and per-entry cost b = (real T=4 - real T=2) / 20, in us, for every choice (groups = 20 only). +// +// Per (device, pair, T, shape, choice): the output buffer against `def`'s, bit for bit (memcmp), and the median of +// 5 rounds of `reps` back-to-back `native_expert_grouped` calls with CUDA events: the whole call (gate/up + SwiGLU + +// q8_1 + down), then the gate/up phase and the down phase alone (native_expert_set_mode's phase, 3 rounds each). +// The experts rotate over a pool of >= 96 MB of distinct synthetic blobs (random bytes, sane fp16 block scales), so +// every call reads its weights cold from DRAM as the decode does each layer. +// +// One line per cell: dev T shape gu/d choice all_us gu_us down_us GB/s pct identical|DIFFERENT (GB/s = groups x blob +// bytes / all_us; pct = of the peak DRAM bandwidth 2 x memory clock x bus width / 8). Lines starting with '#' are +// context. Exit code 1 if any choice differs or a CUDA error occurred. GPU time only. +#include "strata/kernels/iq_kernels.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace K = strata::kernels; + +namespace { + +constexpr int kEmbd = 2560, kFf = 640, kTop = 10, kMaxT = 4; +constexpr int kMaxEntries = kMaxT * kTop; + +struct Pair { int gu, d; }; +const Pair kModelPairs[] = { {16, 20}, {17, 20}, {18, 20}, {21, 20}, {22, 20}, {16, 42}, {17, 42}, {18, 42}, {21, 42}, {22, 42} }; + +const char* type_name(int t) { + switch (t) { + case 16: return "IQ2_XXS"; case 17: return "IQ2_XS"; case 18: return "IQ3_XXS"; case 21: return "IQ3_S"; + case 22: return "IQ2_S"; case 20: return "IQ4_NL"; case 42: return "Q2_0"; default: return "?"; + } +} +int qk_of(int t) { return t == 20 ? 32 : t == 42 ? 64 : 256; } + +struct Choice { + std::string name; + bool dec = false; + K::NativeExpertExpdecParams p; // every field set: the engine's startup values, then the choice's fields +}; +K::NativeExpertExpdecParams g_engine_params; // the engine's values at startup (environment or defaults) + +// "def" | "dec[:r][:d][:b][:x][:g][:h][:m]"; false on anything else +bool parse_choice(const std::string& s, Choice& c) { + c = Choice{}; + c.name = s; + c.p = g_engine_params; + if (s == "def") return true; + if (s.rfind("dec", 0) != 0) return false; + c.dec = true; + size_t i = 3; + while (i < s.size()) { + if (s[i] != ':') return false; + size_t j = s.find(':', i + 1); + if (j == std::string::npos) j = s.size(); + if (j < i + 3) return false; + const char key = s[i + 1]; + const int v = std::atoi(s.substr(i + 2, j - i - 2).c_str()); + switch (key) { + case 'r': if (v != 1 && v != 2) return false; c.p.gu_rows_per_lane = v; break; + case 'd': if (v != 1 && v != 2 && v != 4) return false; c.p.down_rows_per_lane = v; break; + case 'b': if (v != 2 && v != 4) return false; c.p.gu_min_blocks_per_sm = v; break; + case 'x': if (v != 0 && v != 1) return false; c.p.down_iq4_nl = v; break; + case 'w': if (v != 0 && v != 1) return false; c.p.down = v; break; // EXPDEC for the down (engine default 0) + case 'g': if (v < 0) return false; c.p.gu_rows_per_block = v; break; + case 'h': if (v < 0) return false; c.p.down_rows_per_block = v; break; + case 'm': if (v <= 0) return false; c.p.min_blocks = v; break; + default: return false; + } + i = j; + } + return true; +} + +void apply_choice(const Choice& c) { + K::native_expert_set_expdec(c.dec ? 1 : 0); + K::native_expert_set_expdec_params(c.p); +} + +bool g_cuda_error = false; + +bool ck(cudaError_t e, const char* what) { + if (e == cudaSuccess) return true; + std::printf("CUDA: %s: %s\n", what, cudaGetErrorString(e)); + g_cuda_error = true; + return false; +} + +uint16_t sane_half(std::mt19937& rng) { // a finite fp16 of magnitude ~2^-10 .. 2^-6 + const uint32_t r = rng(); + return (uint16_t) (((r >> 31) << 15) | ((5u + (r >> 10) % 5u) << 10) | (r & 0x3ffu)); +} + +// One expert blob [gate rows | up rows | down rows]: random bytes (any qs / qh / signs / scales nibble / grid index / +// IQ4_NL nibble / Q2_0 pair is a valid code) with a sane fp16 `d` in every block (offset 0 of every format's block). +void fill_blob(uint8_t* b, const K::NativeExpertLayout& L, std::mt19937& rng) { + for (size_t i = 0; i + 4 <= L.bytes; i += 4) { + const uint32_t r = rng(); + std::memcpy(b + i, &r, 4); + } + const size_t gu_block = K::iq_row_bytes(L.gu_type, qk_of(L.gu_type)), d_block = K::iq_row_bytes(L.d_type, qk_of(L.d_type)); + const size_t gu_blocks = (size_t) 2 * (size_t) L.n_ff * (size_t) (L.n_embd / qk_of(L.gu_type)); + for (size_t k = 0; k < gu_blocks; ++k) { + const uint16_t h = sane_half(rng); + std::memcpy(b + k * gu_block, &h, 2); + } + const size_t d_blocks = (size_t) L.n_embd * (size_t) (L.n_ff / qk_of(L.d_type)); + for (size_t k = 0; k < d_blocks; ++k) { + const uint16_t h = sane_half(rng); + std::memcpy(b + L.down_off + k * d_block, &h, 2); + } +} + +struct Window { + int groups = 0, entries = 0; + std::vector start, tok, dst; // group g: entries [start[g], start[g + 1]); tok = token, dst = out row +}; + +// real: T tokens, kTop experts each, `groups_req` distinct experts in all (capped at kTop * T, floored at kTop). +// Expert j is forced onto token j % T so that every group has at least one entry; the rest of each token's picks are +// random. +Window build_real(int T, int groups_req, std::mt19937& rng) { + Window w; + const int G = std::min(std::max(groups_req, kTop), kTop * T); + std::vector> picks((size_t) G); + for (int k = 0; k < T; ++k) { + std::vector in((size_t) G, 0); + int cnt = 0; + for (int j = k; j < G; j += T) { in[(size_t) j] = 1; ++cnt; } + std::vector rest; + for (int j = 0; j < G; ++j) + if (!in[(size_t) j]) rest.push_back(j); + std::shuffle(rest.begin(), rest.end(), rng); + for (size_t i = 0; cnt < kTop; ++i, ++cnt) in[(size_t) rest[i]] = 1; + for (int j = 0; j < G; ++j) + if (in[(size_t) j]) picks[(size_t) j].push_back(k); + } + w.groups = G; + w.start.push_back(0); + for (int g = 0; g < G; ++g) { + for (int k : picks[(size_t) g]) { + w.tok.push_back(k); + w.dst.push_back((int32_t) w.dst.size()); + } + w.start.push_back((int32_t) w.dst.size()); + } + w.entries = (int) w.dst.size(); + return w; +} +// single: kTop T groups of one entry each (token g % T) +Window build_single(int T) { + Window w; + w.groups = kTop * T; + w.start.push_back(0); + for (int g = 0; g < w.groups; ++g) { + w.tok.push_back(g % T); + w.dst.push_back(g); + w.start.push_back(g + 1); + } + w.entries = w.groups; + return w; +} + +struct DevMem { + std::vector ptrs; + ~DevMem() { for (void* p : ptrs) cudaFree(p); } + void* alloc(size_t n) { + void* p = nullptr; + if (!ck(cudaMalloc(&p, n), "malloc")) return nullptr; + ptrs.push_back(p); + return p; + } +}; + +struct StreamEvents { + cudaStream_t s = nullptr; + cudaEvent_t a = nullptr, b = nullptr; + ~StreamEvents() { + if (a) cudaEventDestroy(a); + if (b) cudaEventDestroy(b); + if (s) cudaStreamDestroy(s); + } +}; + +// (device, pair gu, pair d, T, shape, choice) -> {all, gu, down} us +using Key = std::tuple; +std::map> g_times; + +bool run_device(int dev, int reps, int groups_req, const std::vector& pool, size_t P, size_t stride, + const std::vector& xh, const K::NativeExpertLayout& L, const Pair& pair, + const std::vector& choices, const std::vector& shapes, long long& bad) { + if (!ck(cudaSetDevice(dev), "set device")) return false; + cudaDeviceProp pr{}; + cudaGetDeviceProperties(&pr, dev); + std::printf("# device %d: %s (sm_%d%d, %d SMs, L2 %d KB), %d reps x 5 rounds, <= %d groups, pool %d blobs x %.2f MB\n", dev, + pr.name, pr.major, pr.minor, pr.multiProcessorCount, pr.l2CacheSize / 1024, reps, groups_req, (int) P, + (double) L.bytes / 1e6); + int mclk_khz = 0, bus_bits = 0; + cudaDeviceGetAttribute(&mclk_khz, cudaDevAttrMemoryClockRate, dev); + cudaDeviceGetAttribute(&bus_bits, cudaDevAttrGlobalMemoryBusWidth, dev); + const double peak = 2.0 * (double) mclk_khz * 1e3 * (double) bus_bits / 8.0; // bytes / s + std::printf("# device %d peak DRAM bandwidth %.1f GB/s (memory clock %d kHz, bus %d bit)\n", dev, peak / 1e9, mclk_khz, bus_bits); + const int gmax = kMaxEntries; + DevMem mem; + auto* dpool = (uint8_t*) mem.alloc(pool.size()); + auto* dx = (float*) mem.alloc(xh.size() * sizeof(float)); + void* dxq = mem.alloc((size_t) kMaxT * kEmbd / 32 * 36); + void* dscr = mem.alloc(K::native_expert_scratch_bytes(kMaxEntries, kFf)); + auto* dout = (float*) mem.alloc((size_t) kMaxEntries * kEmbd * sizeof(float)); + auto* dptr = (unsigned long long*) mem.alloc(P * (size_t) gmax * sizeof(unsigned long long)); + auto* dstart = (int32_t*) mem.alloc((size_t) (gmax + 1) * sizeof(int32_t)); + auto* dn = (int32_t*) mem.alloc(sizeof(int32_t)); + auto* ddst = (int32_t*) mem.alloc((size_t) kMaxEntries * sizeof(int32_t)); + auto* dtok = (int32_t*) mem.alloc((size_t) kMaxEntries * sizeof(int32_t)); + if (!dpool || !dx || !dxq || !dscr || !dout || !dptr || !dstart || !dn || !ddst || !dtok) return false; + StreamEvents se; + if (!ck(cudaStreamCreate(&se.s), "stream") || !ck(cudaEventCreate(&se.a), "event") || + !ck(cudaEventCreate(&se.b), "event")) + return false; + void* stream = (void*) se.s; + if (!ck(cudaMemcpy(dpool, pool.data(), pool.size(), cudaMemcpyHostToDevice), "copy pool") || + !ck(cudaMemcpy(dx, xh.data(), xh.size() * sizeof(float), cudaMemcpyHostToDevice), "copy x")) + return false; + K::quantize_q8_1_rows(dx, kMaxT, kEmbd, dxq, stream); + if (!ck(cudaStreamSynchronize(se.s), "quantize")) return false; + + std::mt19937 rng(20261012u); + bool ok = true; + for (int T = 1; T <= kMaxT && ok; ++T) { + for (const std::string& shape : shapes) { + if (!ok) break; + const Window w = shape == "single" ? build_single(T) : build_real(T, groups_req, rng); + const int G = w.groups, NE = w.entries; + // call i reads G experts starting at blob (i * G) % P: a fresh set until the pool wraps (>= 96 MB later) + std::vector hp(P * (size_t) G); + for (size_t i = 0; i < P; ++i) + for (int g = 0; g < G; ++g) + hp[i * (size_t) G + (size_t) g] = (unsigned long long) (dpool + ((i * (size_t) G + (size_t) g) % P) * stride); + ok = ck(cudaMemcpy(dptr, hp.data(), hp.size() * sizeof(unsigned long long), cudaMemcpyHostToDevice), "copy ptrs") && + ck(cudaMemcpy(dstart, w.start.data(), w.start.size() * sizeof(int32_t), cudaMemcpyHostToDevice), "copy start") && + ck(cudaMemcpy(dn, &G, sizeof(int32_t), cudaMemcpyHostToDevice), "copy n") && + ck(cudaMemcpy(ddst, w.dst.data(), w.dst.size() * sizeof(int32_t), cudaMemcpyHostToDevice), "copy dst") && + ck(cudaMemcpy(dtok, w.tok.data(), w.tok.size() * sizeof(int32_t), cudaMemcpyHostToDevice), "copy tok"); + if (!ok) break; + std::printf("# T=%d %s: %d groups, %d entries, %.1f MB of weights per call\n", T, shape.c_str(), G, NE, + (double) G * (double) L.bytes / 1e6); + auto call = [&](size_t i) { + K::native_expert_grouped(L, dptr + (i % P) * (size_t) G, dstart, dn, ddst, dtok, G, NE, dxq, dscr, dout, stream, 0); + }; + auto timed = [&](int phase, int rounds, double& out_us) { + K::native_expert_set_mode(-1, phase); + for (size_t i = 0; i < (phase == 0 ? 100u : 20u); ++i) call(i); + std::vector rs((size_t) rounds); + for (double& r : rs) { + cudaEventRecord(se.a, se.s); + for (int i = 0; i < reps; ++i) call((size_t) i); + cudaEventRecord(se.b, se.s); + cudaEventSynchronize(se.b); + float ms = 0.0f; + ok = ok && ck(cudaEventElapsedTime(&ms, se.a, se.b), "elapsed") && ck(cudaGetLastError(), "kernel"); + r = 1000.0 * (double) ms / reps; + } + std::sort(rs.begin(), rs.end()); + out_us = rs[(size_t) rounds / 2]; + K::native_expert_set_mode(-1, 0); + }; + std::vector ref, res((size_t) NE * kEmbd); + for (size_t vi = 0; vi < choices.size() && ok; ++vi) { + const Choice& c = choices[vi]; + apply_choice(c); + K::native_expert_set_mode(-1, 0); + ok = ck(cudaMemset(dout, 0, (size_t) NE * kEmbd * sizeof(float)), "memset out"); + call(0); + ok = ok && ck(cudaStreamSynchronize(se.s), "run") && ck(cudaGetLastError(), "kernel") && + ck(cudaMemcpy(res.data(), dout, res.size() * sizeof(float), cudaMemcpyDeviceToHost), "copy out"); + if (!ok) break; + bool same = true; + if (vi == 0) { + ref = res; + long long nonfinite = 0; + for (float f : ref) nonfinite += std::isfinite(f) ? 0 : 1; + if (nonfinite) std::printf("# T=%d: %s output has %lld non-finite values\n", T, c.name.c_str(), nonfinite); + } else { + same = std::memcmp(ref.data(), res.data(), ref.size() * sizeof(float)) == 0; + } + double us[3] = {0.0, 0.0, 0.0}; + timed(0, 5, us[0]); + timed(1, 3, us[1]); + timed(2, 3, us[2]); + const double gbps = (double) G * (double) L.bytes / (us[0] * 1e-6) / 1e9; + std::printf("%d %d %s %d/%d %s %.1f %.1f %.1f %.1f %.1f%% %s\n", dev, T, shape.c_str(), L.gu_type, L.d_type, + c.name.c_str(), us[0], us[1], us[2], gbps, peak > 0.0 ? 100.0 * gbps * 1e9 / peak : 0.0, + same ? "identical" : "DIFFERENT"); + g_times[Key(dev, L.gu_type, L.d_type, T, shape, c.name)] = {us[0], us[1], us[2]}; + if (!same) ++bad; + } + } + } + Choice off; + off.p = g_engine_params; + apply_choice(off); + K::native_expert_set_mode(-1, 0); + return ok; +} + +// "gu/d,gu/d,..." -> pairs; false on a malformed item +bool parse_pairs(const std::string& spec, std::vector& out) { + size_t i = 0; + while (i < spec.size()) { + size_t j = spec.find(',', i); + if (j == std::string::npos) j = spec.size(); + const std::string item = spec.substr(i, j - i); + i = j + 1; + if (item.empty()) continue; + const size_t s = item.find('/'); + if (s == std::string::npos || s == 0 || s + 1 >= item.size()) return false; + Pair p; + p.gu = std::atoi(item.substr(0, s).c_str()); + p.d = std::atoi(item.substr(s + 1).c_str()); + if (p.gu <= 0 || p.d <= 0) return false; + out.push_back(p); + } + return !out.empty(); +} + +std::vector split_commas(const std::string& s) { + std::vector out; + size_t i = 0; + while (i <= s.size()) { + size_t j = s.find(',', i); + if (j == std::string::npos) j = s.size(); + if (j > i) out.push_back(s.substr(i, j - i)); + i = j + 1; + } + return out; +} + +} // namespace + +int main(int argc, char** argv) { + setvbuf(stdout, nullptr, _IONBF, 0); + const int reps = argc > 1 ? std::max(20, std::atoi(argv[1])) : 300; + const int only = argc > 2 ? std::atoi(argv[2]) : -1; + const int groups_req = argc > 3 ? std::min(kMaxEntries, std::max(kTop, std::atoi(argv[3]))) : 20; + std::vector pairs; + if (argc > 4 && std::string(argv[4]) != "all") { + if (!parse_pairs(argv[4], pairs)) { + std::printf("expdec_bench: pairs '%s' is not a comma list of gu/down ggml type ids\n", argv[4]); + return 1; + } + } else { + pairs.assign(std::begin(kModelPairs), std::end(kModelPairs)); + } + g_engine_params = K::native_expert_expdec_params(); + std::vector choices; + for (const std::string& s : split_commas(argc > 5 ? argv[5] : "def,dec,dec:r2,dec:b4,dec:r2:b4")) { + Choice c; + if (!parse_choice(s, c)) { + std::printf("expdec_bench: choice '%s' is not def or dec[:r<1|2>][:d<1|2|4>][:b<2|4>][:x<0|1>][:g][:h][:m]\n", s.c_str()); + return 1; + } + choices.push_back(c); + } + if (choices.empty() || choices[0].dec) { + Choice d; + parse_choice("def", d); + choices.insert(choices.begin(), d); // the reference first + } + std::vector shapes; + { + const std::string sh = argc > 6 ? argv[6] : "real"; + if (sh == "both") shapes = {"real", "single"}; + else if (sh == "real" || sh == "single") shapes = {sh}; + else { std::printf("expdec_bench: shape '%s' is not real, single or both\n", sh.c_str()); return 1; } + } + int n_dev = 0; + if (!ck(cudaGetDeviceCount(&n_dev), "device count") || n_dev == 0) { + std::printf("expdec_bench: no CUDA device\n"); + return 1; + } + if (only >= n_dev) { + std::printf("expdec_bench: device %d, only %d present\n", only, n_dev); + return 1; + } + for (const Pair& p : pairs) { + if (!K::native_expert_supported(p.gu, p.d, kEmbd, kFf)) { + std::printf("expdec_bench: native_expert_supported(%d, %d, %d, %d) is false\n", p.gu, p.d, kEmbd, kFf); + return 1; + } + const K::NativeExpertLayout L = K::native_expert_layout(p.gu, p.d, kEmbd, kFf); + bool gu = false, down = false; + K::native_expert_expdec_takes(L, gu, down); + if (!gu || !down) std::printf("# warning: EXPDEC has no %s kernel for pair %d/%d: that phase runs the default\n", + !gu ? "gate/up" : "down", p.gu, p.d); + } + for (const char* name : {"STRATA_NO_SUB16_GU", "STRATA_EXPERT_V2", "STRATA_OLD_IQ_MMVQ", "STRATA_EXP_MODE"}) { + const char* v = std::getenv(name); + if (v != nullptr && v[0] != '\0' && v[0] != '0') std::printf("# warning: %s is set, the kernel choice is bypassed\n", name); + } + K::iq_set_old_kernels(false); + std::printf("# n_embd %d, n_ff %d, %d pair(s), choices:", kEmbd, kFf, (int) pairs.size()); + for (const Choice& c : choices) { + if (c.dec) std::printf(" %s(gu R %d minb %d, down R %d iq4_nl %d, rb %d drb %d, min blocks %d)", c.name.c_str(), + c.p.gu_rows_per_lane, c.p.gu_min_blocks_per_sm, c.p.down_rows_per_lane, c.p.down_iq4_nl, + c.p.gu_rows_per_block, c.p.down_rows_per_block, c.p.min_blocks); + else std::printf(" %s", c.name.c_str()); + } + std::printf("\n# dev T shape gu/d choice all_us gu_us down_us GB/s pct identical|DIFFERENT\n"); + long long bad = 0; + bool ok = true; + for (const Pair& p : pairs) { + if (!ok) break; + const K::NativeExpertLayout L = K::native_expert_layout(p.gu, p.d, kEmbd, kFf); + const size_t stride = (L.bytes + 255) / 256 * 256; + const size_t P = std::max((size_t) (96u << 20) / stride + 1, (size_t) 2 * (size_t) kMaxEntries); + std::printf("# pair %d/%d: gate/up %s, down %s: blob %.2f MB (gu_row %d, d_row %d), pool %.0f MB\n", p.gu, p.d, + type_name(p.gu), type_name(p.d), (double) L.bytes / 1e6, (int) L.gu_row, (int) L.d_row, + (double) (P * stride) / 1e6); + std::mt19937 rng(77u + (unsigned) (p.gu * 131 + p.d)); + std::vector pool(P * stride, 0); + for (size_t i = 0; i < P; ++i) fill_blob(pool.data() + i * stride, L, rng); + std::vector xh((size_t) kMaxT * kEmbd); + std::normal_distribution nd(0.0f, 1.0f); + for (float& v : xh) v = nd(rng); + for (int dev = 0; dev < n_dev; ++dev) { + if (only >= 0 && dev != only) continue; + ok = run_device(dev, reps, groups_req, pool, P, stride, xh, L, p, choices, shapes, bad) && ok; + } + } + if (shapes.size() == 2 && groups_req == 20) { + std::printf("\n# decomposition (us per layer call): a = per expert = (single T=4 - real T=2) / 20, b = per entry = (real T=4 - real T=2) / 20\n"); + std::printf("# dev gu/d choice a_all a_gu a_down b_all b_gu b_down\n"); + for (int dev = 0; dev < n_dev; ++dev) { + if (only >= 0 && dev != only) continue; + for (const Pair& p : pairs) + for (const Choice& c : choices) { + auto f2 = g_times.find(Key(dev, p.gu, p.d, 2, "real", c.name)); + auto f4 = g_times.find(Key(dev, p.gu, p.d, 4, "real", c.name)); + auto s4 = g_times.find(Key(dev, p.gu, p.d, 4, "single", c.name)); + if (f2 == g_times.end() || f4 == g_times.end() || s4 == g_times.end()) continue; + std::printf("%d %d/%d %s", dev, p.gu, p.d, c.name.c_str()); + for (int k = 0; k < 3; ++k) std::printf(" %.2f", (s4->second[k] - f2->second[k]) / 20.0); + for (int k = 0; k < 3; ++k) std::printf(" %.2f", (f4->second[k] - f2->second[k]) / 20.0); + std::printf("\n"); + } + } + } + std::printf("\nexpdec_bench: %s (%lld choice comparisons differ%s)\n", (bad || !ok || g_cuda_error) ? "FAILED" : "ok", bad, + (!ok || g_cuda_error) ? ", CUDA error" : ""); + return (bad || !ok || g_cuda_error) ? 1 : 0; +}