diff --git a/c/fused_simd.h b/c/fused_simd.h new file mode 100644 index 000000000..594bcf6d3 --- /dev/null +++ b/c/fused_simd.h @@ -0,0 +1,156 @@ +/* fused_simd.h — FUSED3: faster AVX2 kernels for the int8 (Q8_0-per-row) + * routed-expert path of olmoe.c. Three cooperating pieces: + * + * 1) dot_i8idot_avx2_v2 / matmul_q_idot_v2: re-vectorized per-row IDOT matmul + * (2 independent accumulator chains instead of matmul_q's 1 per row). + * 2) quant_x_q8_avx2: AVX2 activation quantization, BIT-IDENTICAL to the + * scalar loop in matmul_q (see the proof comment above it). + * 3) matmul_q_idot_v3 / matmul_q_idot_pair_v3: v2 + vectorized quant, plus a + * gate/up pair that quantizes the shared input ONCE instead of twice. + * + * Numerics: all integer dots are exact (sign-extend + madd, as dot_i8_16); + * activation quantization is byte-exact vs the scalar path, so v3/pair-v3 + * outputs are bit-identical to the stock matmul_q calls — verified by + * memcmp in tests/bench_fused3.c. The flag only changes instruction + * scheduling, never values. + * + * Header-only, all static, no dependencies beyond immintrin — mirrors quant.h + * conventions. Guarded by __AVX2__; callers keep the stock fallback. + */ +#ifndef COLI_FUSED_SIMD_H +#define COLI_FUSED_SIMD_H + +#include +#include + +#ifdef __AVX2__ +#include + +/* One Q8_0-style row dot: y = sum_b xs[b] * dot16(xi_b, w_b), 16 int8 per + * block — the same activation block contract as olmoe.c matmul_q's IDOT + * branch. Integer dots are exact (cvtepi8_epi16 + madd_epi16), so the only + * numeric delta vs the stock path is fp32 summation order of the per-block + * fold, exactly as in matmul_q. Requires nb%2==0 for the paired fast loop + * (activation blocks come in pairs); tail block scalar with v1 semantics. */ +static inline float dot_i8idot_avx2_v2(const int8_t *xi, const int8_t *w, + const float *xs, int nb){ + __m256 tot0=_mm256_setzero_ps(), tot1=_mm256_setzero_ps(); + int b=0; + for(; b+2<=nb; b+=2){ + __m256i pa=_mm256_madd_epi16(_mm256_cvtepi8_epi16(_mm_loadu_si128((const __m128i*)(xi+b*16))), + _mm256_cvtepi8_epi16(_mm_loadu_si128((const __m128i*)(w +b*16)))); + tot0=_mm256_fmadd_ps(_mm256_cvtepi32_ps(pa),_mm256_set1_ps(xs[b]),tot0); + __m256i pb=_mm256_madd_epi16(_mm256_cvtepi8_epi16(_mm_loadu_si128((const __m128i*)(xi+b*16+16))), + _mm256_cvtepi8_epi16(_mm_loadu_si128((const __m128i*)(w +b*16+16)))); + tot1=_mm256_fmadd_ps(_mm256_cvtepi32_ps(pb),_mm256_set1_ps(xs[b+1]),tot1); + } + __m256 tot=_mm256_add_ps(tot0,tot1); + __m128 lo4=_mm256_castps256_ps128(tot), hi4=_mm256_extractf128_ps(tot,1); + lo4=_mm_add_ps(lo4,hi4); __m128 sh=_mm_movehl_ps(lo4,lo4); lo4=_mm_add_ps(lo4,sh); + sh=_mm_shuffle_ps(lo4,lo4,1); lo4=_mm_add_ss(lo4,sh); + float acc=_mm_cvtss_f32(lo4); + for(; bam) am=a; } + float s=am/127.f; if(s<1e-12f) s=1e-12f; + xs[b]=s; float inv=1.f/s; + for(int i=0;i<16;i++) xi[b*16+i]=(int8_t)lrintf(xb[i]*inv); + } + #pragma omp parallel for schedule(static) + for(int o=0;o permute4x64) + * only reorder, they do not round. + * Verified byte-for-byte (xi memcmp + xs bit-compare) in bench_fused3. + * Requires I%16==0, I<=4096 — the same contract as the v2 IDOT path. */ +static inline void quant_x_q8_avx2(const float *x, int8_t *xi, float *xs, int I){ + const __m256 sgn=_mm256_castsi256_ps(_mm256_set1_epi32(0x80000000)); + int nb=I/16; + for(int b=0;b output bit-identical to v2. */ +static void matmul_q_idot_v3(float *y, const float *x, const int8_t *q, + const float *scale, int I, int O){ + int nb=I/16; int8_t xi[4096]; float xs[256]; + quant_x_q8_avx2(x,xi,xs,I); + #pragma omp parallel for schedule(static) + for(int o=0;o last) ? (clock - last) : 0; @@ -202,6 +206,9 @@ static inline int32_t dot_i8_16(const int8_t *a, const int8_t *b) { } #define HAVE_FAST_DOT_I8 1 #endif +#if defined(__AVX2__) +#include "fused_simd.h" /* FUSED3=1: quant_x_q8_avx2 + matmul_q_idot{,_pair}_v3 (bit-exact) */ +#endif static void matmul_q(float *y, const float *x, const int8_t *q, const float *scale, int I, int O) { #if defined(HAVE_FAST_DOT_I8) static int idot = -1; @@ -687,10 +694,25 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out) { const float *xs = x + (int64_t)s*D; for (int kk = 0; kk < K; kk++) { Slot *e; expert_get(m, layer, idx[kk], &e); +#if defined(__AVX2__) + /* FUSED3: same contract as matmul_q's IDOT fast branch (IDOT env, + * dims %16==0, <=4096) — outside it the stock calls below run + * unchanged. Exact integer arithmetic only: bit-identical output + * (verified by memcmp in tests/bench_fused3.c). OFF by default. */ + static int idot_moe = -1; + if (idot_moe < 0) { const char *ie = getenv("IDOT"); idot_moe = !(ie && *ie == '0'); } + if (g_fused3 && idot_moe && D % 16 == 0 && D <= 4096 && I % 16 == 0 && I <= 4096) { + matmul_q_idot_pair_v3(g, u, xs, e->g, e->gs, e->u, e->us, D, I); /* gate+up share one quant of xs */ + for (int i = 0; i < I; i++) { float gv = g[i]; g[i] = (gv / (1.f + expf(-gv))) * u[i]; } + matmul_q_idot_v3(hh, g, e->d, e->ds, I, D); /* down_proj [D,I] */ + } else +#endif + { matmul_q(g, xs, e->g, e->gs, D, I); /* gate_proj [I,D] */ matmul_q(u, xs, e->u, e->us, D, I); /* up_proj [I,D] */ for (int i = 0; i < I; i++) { float gv = g[i]; g[i] = (gv / (1.f + expf(-gv))) * u[i]; } matmul_q(hh, g, e->d, e->ds, I, D); /* down_proj [D,I] */ + } float w = val[kk]; float *os = out + (int64_t)s*D; for (int d = 0; d < D; d++) os[d] += w * hh[d]; @@ -1301,6 +1323,7 @@ int main(int argc, char **argv) { g_wide = getenv("WIDE") ? atoi(getenv("WIDE")) : 1; g_pilot_evict_guard = getenv("PILOT_EVICT_GUARD") ? atoi(getenv("PILOT_EVICT_GUARD")) : 1; g_expert_drop = getenv("EXPERT_DROP") ? atoi(getenv("EXPERT_DROP")) : 0; + g_fused3 = getenv("FUSED3") ? atoi(getenv("FUSED3")) : 0; if (g_wide < 1) g_wide = 1; if (g_wide > 4) g_wide = 4; int hot_n = getenv("HOT") ? atoi(getenv("HOT")) : 0; @@ -1369,8 +1392,8 @@ int main(int argc, char **argv) { float smooth = getenv("SMOOTH") ? (float)atof(getenv("SMOOTH")) : 0.3f; float conf = getenv("CONF_LIMIT") ? (float)atof(getenv("CONF_LIMIT")) : 0.92f; - printf("== Streaming C engine v2.2 | cache=%d/layer bits=%d pilot=%d wide=%d guard=%d hot=%d smooth=%.2f conf=%.2f ==\n", - cap, bits, g_pilot, g_wide, g_pilot_evict_guard, hot_n, smooth, conf); + printf("== Streaming C engine v2.2 | cache=%d/layer bits=%d pilot=%d wide=%d guard=%d hot=%d smooth=%.2f conf=%.2f fused3=%d ==\n", + cap, bits, g_pilot, g_wide, g_pilot_evict_guard, hot_n, smooth, conf, g_fused3); FILE *f = fopen(refpath, "rb"); if (!f) { perror(refpath); return 1; } fseek(f,0,SEEK_END); long n=ftell(f); fseek(f,0,SEEK_SET); diff --git a/c/tests/bench_fused3.c b/c/tests/bench_fused3.c new file mode 100644 index 000000000..28bf3c6f2 --- /dev/null +++ b/c/tests/bench_fused3.c @@ -0,0 +1,183 @@ +/* bench_fused3.c — microbenchmark + correctness harness for the FUSED3 + * kernels (c/fused_simd.h) on OLMoE expert shapes: stock matmul_q IDOT branch + * (v1, copied verbatim below) vs v2 (re-vectorized row dots) vs v3 (+ + * vectorized activation quant, gate/up pair). + * Build: gcc -O3 -mavx2 -mfma -fopenmp -I.. bench_fused3.c -o bench_fused3.exe -lm + * Prints KEY=VALUE lines. Correctness gates: + * (a) quant_x_q8_avx2 output byte-identical to the scalar quant + * (xi memcmp, xs bit-compare), incl. the all-zero clamp path; + * (b) v3/pair-v3 row outputs bit-identical to matmul_q_idot_v2. */ +#include +#include +#include +#include +#include +#include +#ifdef _OPENMP +#include +#endif + +#include "fused_simd.h" + +/* dot_i8_16: verbatim AVX2 copy from olmoe.c (static there) — v1 baseline dot */ +static inline int32_t dot_i8_16(const int8_t *a, const int8_t *b) { + __m256i va16 = _mm256_cvtepi8_epi16(_mm_loadu_si128((const __m128i*)a)); + __m256i vb16 = _mm256_cvtepi8_epi16(_mm_loadu_si128((const __m128i*)b)); + __m256i prod = _mm256_madd_epi16(va16, vb16); + __m128i sum128 = _mm_add_epi32(_mm256_castsi256_si128(prod), _mm256_extractf128_si256(prod, 1)); + __m128i hi64 = _mm_unpackhi_epi64(sum128, sum128); + __m128i sum64 = _mm_add_epi32(sum128, hi64); + __m128i hi32 = _mm_shuffle_epi32(sum64, _MM_SHUFFLE(2, 3, 0, 1)); + __m128i sum32 = _mm_add_epi32(sum64, hi32); + return _mm_cvtsi128_si32(sum32); +} + +static double now_s(void){ + struct timespec ts; clock_gettime(CLOCK_MONOTONIC,&ts); + return ts.tv_sec + ts.tv_nsec*1e-9; +} + +static unsigned rng_s=12345; +static unsigned rnd(void){ rng_s=rng_s*1664525u+1013904223u; return rng_s>>8; } + +int main(int argc, char **argv){ + int reps = argc>1 ? atoi(argv[1]) : 200; + printf("# bench_fused3 (S=1 GEMV, int8 Q8_0-per-row, olmoe expert path)\n"); +#ifdef _OPENMP + printf("omp_threads=%d\n", omp_get_max_threads()); +#else + printf("omp_threads=1\n"); +#endif + + /* ---- int8 per-row IDOT (olmoe.c expert path) -------------------------- */ + /* v1 baseline: verbatim copy of olmoe.c matmul_q's HAVE_FAST_DOT_I8 branch */ + int i8shapes[2][2] = { {1024,2048}, {2048,1024} }; + for(int sh=0; sh<2; sh++){ + int O=i8shapes[sh][0], I=i8shapes[sh][1]; + int8_t *q8=malloc((int64_t)O*I); + float *scl=malloc((int64_t)O*sizeof(float)); + float *x=malloc(I*sizeof(float)); + float *y1=malloc((int64_t)O*sizeof(float)); + float *y2=malloc((int64_t)O*sizeof(float)); + if(!q8||!scl||!x||!y1||!y2){ fprintf(stderr,"OOM\n"); return 1; } + for(int64_t i=0;i<(int64_t)O*I;i++) q8[i]=(int8_t)(rnd()%255-127); + for(int o=0;oam) am=a; } \ + float s=am/127.f; if(s<1e-12f) s=1e-12f; \ + xs[b]=s; float inv=1.f/s; \ + for(int i=0;i<16;i++) xi[b*16+i]=(int8_t)lrintf(xb[i]*inv); } \ + _Pragma("omp parallel for schedule(static)") \ + for(int o=0;oam) am=a; } + float s=am/127.f; if(s<1e-12f) s=1e-12f; + xs[b]=s; float inv=1.f/s; + for(int i=0;i<16;i++) xi[b*16+i]=(int8_t)lrintf(xb[i]*inv); } + double r=0; const int8_t *w=q8+(int64_t)o*I; + for(int b=0;be1)e1=d1; if(d2>e2)e2=d2; + } + printf("i8_O%d_I%d: relerr_v1=%.3g relerr_v2=%.3g\n",O,I,e1,e2); + + double t1=1e30,t2=1e30; + for(int k=0;k<3;k++){ MATMUL_Q_IDOT_V1(y1); matmul_q_idot_v2(y2,x,q8,scl,I,O); } + for(int r=0;ram) am=a; } + float s=am/127.f; if(s<1e-12f) s=1e-12f; + xs_s[b]=s; float inv=1.f/s; + for(int i=0;i<16;i++) xi_s[b*16+i]=(int8_t)lrintf(xb[i]*inv); } + quant_x_q8_avx2(x,xi_v,xs_v,I); + int qbad=memcmp(xi_s,xi_v,I)!=0 || memcmp(xs_s,xs_v,nb*sizeof(float))!=0; + /* all-zero block: clamp path */ + { float xz[64]={0}; int8_t xz_s[64],xz_v[64]; float xs2_s[4],xs2_v[4]; + for(int b=0;b<4;b++){ float am=0.f; float s=am/127.f; if(s<1e-12f) s=1e-12f; + xs2_s[b]=s; float inv=1.f/s; + for(int i=0;i<16;i++) xz_s[b*16+i]=(int8_t)lrintf(xz[b*16+i]*inv); } + quant_x_q8_avx2(xz,xz_v,xs2_v,64); + if(memcmp(xz_s,xz_v,64)!=0 || memcmp(xs2_s,xs2_v,4*sizeof(float))!=0) qbad=1; } + printf("v3_quant_bitexact=%s\n", qbad?"FAIL":"yes"); + + /* (b) v3 / pair-v3 outputs bit-identical to v2 */ + matmul_q_idot_v2(yg,x,qg,sg,I,O); matmul_q_idot_v3(yg2,x,qg,sg,I,O); + int bid = memcmp(yg,yg2,O*sizeof(float))==0; + matmul_q_idot_v2(yg,x,qg,sg,I,O); matmul_q_idot_v2(yu,x,qu,su,I,O); + matmul_q_idot_pair_v3(yg2,yu2,x,qg,sg,qu,su,I,O); + bid = bid && memcmp(yg,yg2,O*sizeof(float))==0 && memcmp(yu,yu2,O*sizeof(float))==0; + matmul_q_idot_v2(yd,xd,qd,sd,Id,Od); matmul_q_idot_v3(yd2,xd,qd,sd,Id,Od); + bid = bid && memcmp(yd,yd2,Od*sizeof(float))==0; + printf("v3_output_bitidentical_v2=%s\n", bid?"yes":"FAIL"); + + /* (c) per-expert timing: v2 (3 separate calls) vs v3 (pair + down) */ + double t2e=1e30,t3e=1e30; + for(int k=0;k<3;k++){ + matmul_q_idot_v2(yg,x,qg,sg,I,O); matmul_q_idot_v2(yu,x,qu,su,I,O); matmul_q_idot_v2(yd,xd,qd,sd,Id,Od); + matmul_q_idot_pair_v3(yg2,yu2,x,qg,sg,qu,su,I,O); matmul_q_idot_v3(yd,xd,qd,sd,Id,Od); } + for(int r=0;r