Lab 2：MoE 的向量化计算

即权重预打包，重排为 AMX/VNNI 布局。

以 VNNI 为例，VNNI 布局按 64 行分块 [row_block/64][K/4][256]，使每条 vpdpbusd 的权重在一次 512-bit load 中完成。

// VNNI 布局：[row_block/64][K/4][256]
//每条 vpdpbusd 的权重在一次 512b load 中完备
static void pack_weight_vnni(const int8_t *src, int rows, int K, int8_t *out) {
    for (int rb = 0; rb < rows; rb += 64)   // 64 行一块（一个 ZMM）
        for (int kb = 0; kb < K / 4; ++kb)  // 沿 K/4
            for (int r = 0; r < 64; ++r)    // 块内 64 行
                for (int j = 0; j < 4; ++j) // 块内 4 字节连续
                    out[((rb/64)*(K/4) + kb)*256 + r*4 + j] =
                    src[(rb + r)*K + kb*4 + j]; // 线性打包
}

在预打包时变成线性，缓存友好，计算时下面的内联汇编在连续内存中取数据。

asm volatile(
    "vpdpbusd   0(%[bp]), %[av], %[a0]\n\t"
    "vpdpbusd  64(%[bp]), %[av], %[a1]\n\t"
    "vpdpbusd 128(%[bp]), %[av], %[a2]\n\t"
    "vpdpbusd 192(%[bp]), %[av], %[a3]\n\t"
    : [a0] "+v"(a0), [a1] "+v"(a1), [a2] "+v"(a2), [a3] "+v"(a3)
    : [av] "v"(av), [bp] "r"(bp)
    : "memory");
bp += 256;

这句话的前提是在合理范围内，比如下面的节选片段并不是全用满 16 线程。

// S2/S3 各使用 8 个物理核，S4 使用全部 16 个逻辑线程。(S1 是常驻线程池)
g_compute_threads = E > 16 ? OPENMP_NTHREADS_S4
: (D > 256 ? OPENMP_NTHREADS_LARGE : OPENMP_NTHREADS_SMALL);

经过实测，S3 8 线程效果最优，可以解释为线程更多时一个核心内部的两个线程会抢用 AMX 等计算单元导致效率下降。

题外话，S2 10核心的时候有时候能打出很好的分数，但是 oj 波动很大，最后求稳改回了 8 核心。暂时归结为玄学问题。（上图想复现，发现异常了？这个真不知道为什么）

如：用 AMX 加速的 INT8 矩阵乘加内核

// ── AMX int8 GEMM: A[M][K] × B[rows][K]^T → C[M][rows] ──
// tile: tmm0=A(M×64B), tmm1=B(16×64B), tmm2..5=四个M×16的int32输出
// 四条C链交错: load B块i → dpbssd tile i → load B块i+1...隐藏TMUL延迟
static inline void proj_amx(xxx) {
    for (int col = 0; col < rows; col += 64) {
        _tile_zero(2..5);
        for (int k = 0; k < K; k += 64) {
            _tile_loadd(0, a + k, lda);
            const int8_t *bp = pack+(size_t)(k/4)* rows * 4 + col * 4;
            _tile_loadd(1, bp, rows * 4);
            _tile_dpbssd(2, 0, 1);
            ...
        }
        _tile_stored(2, out + col, rows * (int)sizeof(int32_t));
        ...
    }
}

1. FP32->int8

   // 对单个 token 做 per-token 对称量化：float → int8。
   // 每个 token 用独立的 scale，把动态范围充分利用到 [-127, 127]。
   static inline void quantize_token(xxx) {
       // 1) 扫描整个 token 向量，求所有元素的绝对值最大值 vmax
       //    ZMM lane 上逐块取 max，最后归并成一个标量 x_amax，为了算scale
       ...details omitted...
   
       // 2) 算出该 token 的 scale = max/127，随后反量化时乘回它即可还原；
       //    （若全为 0 则 scale 取 1 避免除零。
       *s_x = (x_amax > 0.0f) ? x_amax / INT8_MAX_ABS : 1.0f;
       const __m512 inv_sx = _mm512_set1_ps(1.0f / (*s_x));
   
       // 3) 逐 16 元素块量化：先把 x 乘上 1/scale，再 float→int32 取整，
       //    再由 int32 压窄成 int8 存回 xq。
       for (int d = 0; d < D; d += 16) {
           __m512i q = _mm512_cvtps_epi32
           (_mm512_mul_ps(_mm512_loadu_ps(xt + d), inv_sx));   
           // * (1/scale) → int32
           _mm_storeu_si128((__m128i *)(xq + d),_mm512_cvtepi32_epi8(q));  
           // int32 → int8
       }
       // 结果：xq 为量化的 int8 激活，s_x 为该 token 的量化 scale。
   }
   

2. int8->int32，如 2.3 中的函数签名

   static inline void proj_amx(const int8_t *a, int lda, int m, int K, const int8_t *pack, int rows, int32_t *out) //int8 输入、int32 输出

3. 反量化int32->FP32

这部分优化体现在其他章节里。

看可支配的精度，简化一些函数（如用低阶多项式替代 exp）

不同的简化程度用于不同测试点。

// exp(x) = 2^{x·log2(e)} = 2^{t+f}
// 令 t = round(x·log2(e))，f = x·log2(e)−t ∈ [-0.5,0.5]
// 则 exp(x) = 2^t × 2^f；2^t 直接用整数 bit-field 写入 FP32 指数位，
// 2^f 在小区间上用多项式 p(f) ≈ 2^f 近似。
static inline __m512 exp_simd(__m512 x) {
    x = _mm512_mul_ps(x, _mm512_set1_ps(1.4426950408889634f));
    __m512 t = _mm512_roundscale_ps(x, 0x00);
    __m512 f = _mm512_sub_ps(x, t);
    // 在 f∈[-0.5,0.5] 上近似 2^f。
    __m512 p = _mm512_fmadd_ps(f, _mm512_set1_ps(0.009618129f), _mm512_set1_ps(0.05550411f));
    p = _mm512_fmadd_ps(f, p, _mm512_set1_ps(0.24022650f));
    p = _mm512_fmadd_ps(f, p, _mm512_set1_ps(0.69314718f));
    p = _mm512_fmadd_ps(f, p, _mm512_set1_ps(1.0f));
    // 2^t 直接写入 FP32 exponent。
    ...
}
static inline __m512 exp_simd_linear(__m512 x) {
static inline __m512 exp_simd_quadratic(__m512 x) {

不同阶数的 exp：

对应到激活函数上，各阶近似的偏差：

内联汇编减少无用搬运/寄存器融合

如：下图中的汇编就有无效搬运。

内联汇编可见3.1.1中的代码块。

具体说明在正文"优化记录"部分。

static inline uint16_t to_bf16(float x) {
    uint32_t bits;
    __builtin_memcpy(&bits, &x, sizeof(bits));
    bits += 0x7fffu + ((bits >> 16) & 1u);
    return (uint16_t)(bits >> 16);
}

router 排序优化（仅S4）

用 BF16 AMX 全量粗筛出 3 候选、再仅对 3 个候选做 FP32 精排取前 2。这样实现了排序计算的并行化，同时权衡了精度和结果可信性。如果都用 FP32 排序，无法用 AMX 并行，速度慢。

Shared/Router 重叠调度（仅S1）

S1 用常驻 worker + 主线程共 5 个执行器，把 shared expert 先单独派给 worker 1、与主线程的 Router/Top-K 重叠执行，省掉每轮 OpenMP fork/join 和串行等待。

优化相对收益

最终相对 RMSE（均方根误差） 。

实验记录