flashnext cuda: drafted rounds: first token early, n-gram table resident, deeper greedy drafts, CUDA graphs

zig/kernels/cuda/flashnext_gdn_fused.cuh

@@ -1,4 +1,4 @@
-// Flash Next prompt GDN in one kernel a value head: front, chain and gated norm, rows kept in shared memory.
+// Flash Next prompt GDN, one kernel a value head: 4 producer warps front the next batch, 8 chain warps step and norm.
 namespace gdn_fused {
 using namespace fn_gdn;
 constexpr int STEPS = 16, TPR = 2, CH = 256, PR = 128, THREADS = CH + PR;
@@ -8,49 +8,52 @@ constexpr int SMEM = 2 * STAGE + 2 * STEPS * DV * 4 + STEPS * 2 * DK * 4 + STEPS
 
 __device__ __forceinline__ void pbar() { asm volatile("bar.sync 1, 128;\n" ::: "memory"); }
 
-// conv_act for rows r0 .. r0 + n - 1 of channel c, each P / cs row loaded once (conv_act's taps, order and rounding)
-__device__ __forceinline__ void conv_rows(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs,
-                                          const __nv_bfloat16* __restrict__ cw, int r0, int n, int c, float (&out)[STEPS]) {
-    float w[TAPS];
-#pragma unroll
-    for (int tap = 0; tap < TAPS; ++tap) w[tap] = __bfloat162float(cw[c * TAPS + tap]);
-    float xw[STEPS + TAPS - 1];
-#pragma unroll
-    for (int k = 0; k < STEPS + TAPS - 1; ++k) {
-        const int at = r0 + k;  // conv_act's `at` for row r0 + j, tap k - j
-        xw[k] = k < n + TAPS - 1 ? (at < TAPS - 1 ? __bfloat162float(cs[at * C + c]) : __bfloat162float(P[(size_t)(at - (TAPS - 1)) * PW + c])) : 0.0f;
-    }
+__device__ __forceinline__ void cbar() { asm volatile("bar.sync 2, 256;\n" ::: "memory"); }
+
+constexpr int XW = STEPS + TAPS - 1;  // a batch's conv window rows
+
+// a batch's loads, issued together: rows r0 .. r0 + n - 1's q, k and v conv windows and row r0 + p's b and a
+struct Rows {
+    float xw[3][XW];
+    float b, a;
+};
+
+__device__ __forceinline__ void load_rows(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, int r0,
+                                          int n, const int (&ch)[3], int hv, int p, Rows& in) {
 #pragma unroll
-    for (int j = 0; j < STEPS; ++j) {
-        float acc = 0.0f;
+    for (int c3 = 0; c3 < 3; ++c3)
 #pragma unroll
-        for (int tap = 0; tap < TAPS; ++tap) acc = acc + w[tap] * xw[j + tap];
-        out[j] = bf(acc / (1.0f + expf(-acc)));
+        for (int k = 0; k < XW; ++k) {
+            const int at = r0 + k, c = ch[c3];  // conv_act's `at` for row r0 + j, tap k - j
+            in.xw[c3][k] = k < n + TAPS - 1 ? (at < TAPS - 1 ? __bfloat162float(cs[at * C + c]) : __bfloat162float(P[(size_t)(at - (TAPS - 1)) * PW + c])) : 0.0f;
+        }
+    if (p < n) {
+        const int r = r0 + p;
+        in.b = __bfloat162float(P[(size_t)r * PW + C + NV * DV + hv]);
+        in.a = __bfloat162float(P[(size_t)r * PW + C + NV * DV + NV + hv]);
     }
 }
 
-// producers: rows r0 .. r0 + n - 1's q, k (normed as fn_gdn_front), v, g, beta into the stage
-__device__ __forceinline__ void front_batch(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs,
-                                            const __nv_bfloat16* __restrict__ cw, const float* __restrict__ a_log,
-                                            const float* __restrict__ dt_bias, int r0, int n, int hk, int hv, int p,
-                                            unsigned char* base, float (*x)[2][DK]) {
-    float o[STEPS];
-    conv_rows(P, cs, cw, r0, n, hk * DK + p, o);                       // fn_gdn_front unit hk, t = p: q's channel
-#pragma unroll
-    for (int j = 0; j < STEPS; ++j) if (j < n) x[j][0][p] = o[j];
-    conv_rows(P, cs, cw, r0, n, NK * DK + hk * DK + p, o);             // t = p + DK: k's channel
-#pragma unroll
-    for (int j = 0; j < STEPS; ++j) if (j < n) x[j][1][p] = o[j];
-    conv_rows(P, cs, cw, r0, n, 2 * NK * DK + hv * DV + p, o);         // unit NK + hv, t = p: v
-#pragma unroll
-    for (int j = 0; j < STEPS; ++j) if (j < n) reinterpret_cast<float*>(base + j * ROW + 2 * DK * 4)[p] = o[j];
-    if (p < n) {                                                         // unit NK + hv, t == DV: row r0 + p's g and beta
-        const int r = r0 + p;
-        const float b = __bfloat162float(P[(size_t)r * PW + C + NV * DV + hv]);
-        const float a = __bfloat162float(P[(size_t)r * PW + C + NV * DV + NV + hv]);
+// producers: the batch's q, k, v, g, beta into the stage as fn_gdn_front makes them (nea = -exp(a_log), dtb dt bias)
+__device__ __forceinline__ void front_rows(const Rows& in, const float (&w)[3][TAPS], float nea, float dtb, int n, int p,
+                                           unsigned char* base, float (*x)[2][DK]) {
+#pragma unroll
+    for (int c3 = 0; c3 < 3; ++c3)
+#pragma unroll
+        for (int j = 0; j < STEPS; ++j) {
+            float acc = 0.0f;
+#pragma unroll
+            for (int tap = 0; tap < TAPS; ++tap) acc = acc + w[c3][tap] * in.xw[c3][j + tap];
+            const float o = bf(acc / (1.0f + expf(-acc)));
+            if (j < n) {
+                if (c3 < 2) x[j][c3][p] = o;                                              // q, k: normed below
+                else reinterpret_cast<float*>(base + j * ROW + 2 * DK * 4)[p] = o;       // v
+            }
+        }
+    if (p < n) {                                                                          // row p's g and beta
         float* gbp = reinterpret_cast<float*>(base + p * ROW + (2 * DK + DV) * 4);
-        gbp[0] = expf(-expf(a_log[hv]) * softplusf_(a + dt_bias[hv]));
-        gbp[1] = bf(sigmoidf_(b));
+        gbp[0] = expf(nea * softplusf_(in.a + dtb));
+        gbp[1] = bf(sigmoidf_(in.b));
     }
     pbar();
     const int warp = p >> 5, lane = p & 31;
@@ -71,12 +74,13 @@ __device__ __forceinline__ void front_batch(const __nv_bfloat16* __restrict__ P,
     }
 }
 
-// producers: fn_gdn_norm for n rows from row r0 (readouts ys[j][dv])
+// the chain's 256 threads after a batch's steps: fn_gdn_norm's arithmetic for its n rows from row r0 (ys[j][dv])
 __device__ __forceinline__ void norm_rows(const float (*ys)[DV], const __nv_bfloat16* __restrict__ P,
                                           const __nv_bfloat16* __restrict__ norm_w, float eps, __nv_bfloat16* __restrict__ out,
-                                          float* __restrict__ xs, int r0, int n, int hv, int p, float* rinv) {
-    const int warp = p >> 5, lane = p & 31;
-    for (int j = warp; j < n; j += PR / 32) {
+                                          float* __restrict__ xs, int r0, int n, int hv, float* rinv) {
+    const int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31;
+    cbar();  // every readout of the batch is in ys
+    for (int j = warp; j < n; j += CH / 32) {
         float ss = 0.0f;
 #pragma unroll
         for (int i = 0; i < 4; ++i) {
@@ -86,23 +90,28 @@ __device__ __forceinline__ void norm_rows(const float (*ys)[DV], const __nv_bflo
         ss = warp_sum(ss);
         if (lane == 0) rinv[j] = 1.0f / sqrtf(ss / (float)DV + eps);
     }
-    pbar();
-    float zs[STEPS];
-#pragma unroll
-    for (int j = 0; j < STEPS; ++j) zs[j] = j < n ? __bfloat162float(P[(size_t)(r0 + j) * PW + C + hv * DV + p]) : 0.0f;
+    cbar();
+    constexpr int HALF = STEPS / 2;
+    const int h = tid >> 7, p = tid & 127;
+    float zs[HALF];
+#pragma unroll
+    for (int k = 0; k < HALF; ++k) {
+        const int j = h * HALF + k;
+        zs[k] = j < n ? __bfloat162float(P[(size_t)(r0 + j) * PW + C + hv * DV + p]) : 0.0f;
+    }
     const float nw = __bfloat162float(norm_w[p]);
 #pragma unroll
-    for (int j = 0; j < STEPS; ++j) {
+    for (int k = 0; k < HALF; ++k) {
+        const int j = h * HALF + k;
         if (j >= n) break;
         const int r = r0 + j;
         const float yn = bf(bf(ys[j][p] * rinv[j]) * nw);
-        const float z = zs[j];
+        const float z = zs[k];
         const float o = bf(yn * sigmoidf_(z));
         out[(size_t)r * NV * DV + hv * DV + p] = __float2bfloat16_rn(o);
         const float gs = warp_sum(o);
-        if (lane == 0) xs[(size_t)r * (NV * DV / 32) + hv * (DV / 32) + warp] = gs;
+        if (lane == 0) xs[(size_t)r * (NV * DV / 32) + hv * (DV / 32) + (p >> 5)] = gs;
     }
-    pbar();
 }
 
 __device__ __forceinline__ void run(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs,
@@ -119,9 +128,15 @@ __device__ __forceinline__ void run(const __nv_bfloat16* __restrict__ P, const _
     const bool producer = threadIdx.x >= CH;
     const int p = threadIdx.x - CH;
     const int batches = (rows + STEPS - 1) / STEPS;
+    // producers: the three channels' taps and the gate constants once, then a batch's loads together, then its front
+    const int ch[3] = {hk * DK + p, NK * DK + hk * DK + p, 2 * NK * DK + hv * DV + p};
+    float w[3][TAPS];
+    float nea = 0.0f, dtb = 0.0f;
     auto fill = [&](int bi) {
         const int r0 = bi * STEPS, n = min(STEPS, rows - r0);
-        front_batch(P, cs, cw, a_log, dt_bias, r0, n, hk, hv, p, stage + (bi & 1) * STAGE, xb);
+        Rows in;
+        load_rows(P, cs, r0, n, ch, hv, p, in);
+        front_rows(in, w, nea, dtb, n, p, stage + (bi & 1) * STAGE, xb);
     };
     // the chain's registers (gdn_tpr::chain<2, 16>'s)
     constexpr int L = 32 / TPR;
@@ -133,13 +148,20 @@ __device__ __forceinline__ void run(const __nv_bfloat16* __restrict__ P, const _
         for (int l = 0; l < L; ++l)
 #pragma unroll
             for (int i = 0; i < 4; ++i) s[l][i] = state_in[sbase + (l * TPR + t) * 4 + i];
-    } else if (batches > 0) fill(0);
+    } else {
+#pragma unroll
+        for (int c3 = 0; c3 < 3; ++c3)
+#pragma unroll
+            for (int tap = 0; tap < TAPS; ++tap) w[c3][tap] = __bfloat162float(cw[ch[c3] * TAPS + tap]);
+        nea = -expf(a_log[hv]);
+        dtb = dt_bias[hv];
+        if (batches > 0) fill(0);
+    }
     __syncthreads();
     for (int bi = 0; bi < batches; ++bi) {
         const int r0 = bi * STEPS, n = min(STEPS, rows - r0);
         if (producer) {
             if (bi + 1 < batches) fill(bi + 1);
-            if (bi >= 1) norm_rows(ysb[(bi - 1) & 1], P, norm_w, eps, out, xs, (bi - 1) * STEPS, STEPS, hv, p, rinv);
         } else {
             const unsigned char* base = stage + (bi & 1) * STAGE;
             for (int rr = 0; rr < n; ++rr) {
@@ -181,12 +203,11 @@ __device__ __forceinline__ void run(const __nv_bfloat16* __restrict__ P, const _
                     if (ys_out) ys_out[((size_t)(r0 + rr) * NV + hv) * DV + row] = o;
                 }
             }
+            norm_rows(ysb[bi & 1], P, norm_w, eps, out, xs, r0, n, hv, rinv);  // on the chain: it would wait at the barrier
         }
         __syncthreads();
     }
-    if (producer) {
-        if (batches > 0) norm_rows(ysb[(batches - 1) & 1], P, norm_w, eps, out, xs, (batches - 1) * STEPS, rows - (batches - 1) * STEPS, hv, p, rinv);
-    } else {
+    if (!producer) {
 #pragma unroll
         for (int l = 0; l < L; ++l)
 #pragma unroll

zig/kernels/cuda/flashnext_hc.cu

@@ -122,6 +122,84 @@ extern "C" __global__ void __launch_bounds__(256) fn_hc_writeback_slices(const _
     hc_writeback<4>(H, HOUT, PSS, BR, INJ, nullptr, nullptr, RS, D, S, 0, 1, WORLD);
 }
 
+// hc_writeback<MODE> then fn_hc_normalize in one block a row, slices summed and added in their order: the same bits.
+template <int MODE, int NC, int S>
+__device__ __forceinline__ void hc_writeback_norm(const __nv_bfloat16* H, __nv_bfloat16* HOUT, float* __restrict__ PSS,
+                                                  const float* __restrict__ BR, const __nv_bfloat16* __restrict__ INJ,
+                                                  const float* __restrict__ Y, const float* __restrict__ WTS,
+                                                  const float* __restrict__ SCALE, __nv_bfloat16* __restrict__ NORM,
+                                                  long long RS, float eps, int TOPK, int SLOTS, int WORLD) {
+    constexpr int D = NC * 256;
+    __shared__ float red[NC][8][S];
+    __shared__ float rinv[S];
+    const int r = blockIdx.x, warp = threadIdx.x >> 5;
+    float hv[NC][S];
+#pragma unroll
+    for (int c = 0; c < NC; ++c) {
+        const int d = c * 256 + threadIdx.x;
+        float branch = 0.0f;
+        if constexpr (MODE == 4) {
+            float acc = BR[static_cast<long long>(r) * D + d];
+            for (int k = 1; k < WORLD; ++k) acc = __fadd_rn(acc, BR[k * RS + static_cast<long long>(r) * D + d]);
+            branch = rbf(acc);
+        } else if constexpr (MODE == 2) {
+            float acc = 0.0f;
+            for (int k = 0; k <= TOPK; ++k) {
+                const long long at = static_cast<long long>(r) * SLOTS + k;
+                acc = __fmaf_rn(Y[at * D + d], WTS[at], acc);
+            }
+            branch = rbf(acc);
+        }
+#pragma unroll
+        for (int s = 0; s < S; ++s) {
+            const long long at = static_cast<long long>(r) * S * D + static_cast<long long>(s) * D + d;
+            float v = bf(H[at]);
+            if constexpr (MODE != 0) {
+                v = rbf(__fmaf_rn(branch, bf(INJ[r * S + s]), v));
+                HOUT[at] = __float2bfloat16_rn(v);
+            }
+            hv[c][s] = v;
+            const float ss = warp_sum(__fmul_rn(v, v));
+            if ((threadIdx.x & 31) == 0) red[c][warp][s] = ss;
+        }
+    }
+    __syncthreads();
+    if (threadIdx.x < S) {
+        float all = 0.0f;
+        for (int c = 0; c < NC; ++c) {
+            float total = 0.0f;
+            for (int w = 0; w < 8; ++w) total = __fadd_rn(total, red[c][w][threadIdx.x]);
+            PSS[(static_cast<long long>(r) * NC + c) * S + threadIdx.x] = total;
+            all = __fadd_rn(all, total);
+        }
+        rinv[threadIdx.x] = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(all, static_cast<float>(D)), eps)));
+    }
+    __syncthreads();
+#pragma unroll
+    for (int c = 0; c < NC; ++c)
+#pragma unroll
+        for (int s = 0; s < S; ++s) {
+            const int col = s * D + c * 256 + threadIdx.x;
+            NORM[static_cast<long long>(r) * S * D + col] =
+                __float2bfloat16_rn(__fmul_rn(__fmul_rn(hv[c][s], rinv[s]), SCALE[col]));
+        }
+}
+
+// Flash Next's shape: 2560 columns (10 slices), 4 streams.
+extern "C" __global__ void __launch_bounds__(256) fn_hc_writeback_norm0(const __nv_bfloat16* H, float* PSS, const float* SCALE, __nv_bfloat16* NORM, float eps) {
+    hc_writeback_norm<0, 10, 4>(H, nullptr, PSS, nullptr, nullptr, nullptr, nullptr, SCALE, NORM, 0, eps, 0, 1, 1);
+}
+
+extern "C" __global__ void __launch_bounds__(256) fn_hc_writeback_norm_moe(const __nv_bfloat16* H, __nv_bfloat16* HOUT, float* PSS, const __nv_bfloat16* INJ,
+    const float* Y, const float* WTS, const float* SCALE, __nv_bfloat16* NORM, float eps, int TOPK, int SLOTS) {
+    hc_writeback_norm<2, 10, 4>(H, HOUT, PSS, nullptr, INJ, Y, WTS, SCALE, NORM, 0, eps, TOPK, SLOTS, 1);
+}
+
+extern "C" __global__ void __launch_bounds__(256) fn_hc_writeback_norm_slices(const __nv_bfloat16* H, __nv_bfloat16* HOUT, float* PSS, const float* BR,
+    const __nv_bfloat16* INJ, long long RS, const float* SCALE, __nv_bfloat16* NORM, float eps, int WORLD) {
+    hc_writeback_norm<4, 10, 4>(H, HOUT, PSS, BR, INJ, nullptr, nullptr, SCALE, NORM, RS, eps, 0, 1, WORLD);
+}
+
 // qmm._qmm_hcdown, block (16 rows, 64 columns, K slice): the slice's normed rows times the down rows, fp32 to PART.
 extern "C" __global__ void __launch_bounds__(128) fn_hc_down(const __nv_bfloat16* __restrict__ H, const float* __restrict__ PSS,
                                                              const float* __restrict__ SCALE, __nv_bfloat16* __restrict__ NORMED,

zig/kernels/cuda/flashnext_mtp.cu

@@ -68,3 +68,96 @@ extern "C" __global__ void __launch_bounds__(1024) fn_row_best(const __nv_bfloat
         }
     }
 }
+
+namespace fn_mtp {
+// A bf16 value's order as an unsigned key (larger value, larger key); -0 as +0, so equal values tie as drawRow's do.
+static __device__ __forceinline__ unsigned int order_key(unsigned short u) {
+    if (u == 0x8000u) u = 0;
+    return (u & 0x8000u) ? (~u & 0xFFFFu) : (u | 0x8000u);
+}
+}  // namespace fn_mtp
+
+// A row's K best of its first V values as drawRow ranks them (ties: lowest columns): those above the K-th, then ties.
+extern "C" __global__ void __launch_bounds__(1024) fn_row_topk(const __nv_bfloat16* __restrict__ L, int V, long long ld,
+                                                               int K, unsigned short* __restrict__ out_v,
+                                                               int* __restrict__ out_c) {
+    using namespace fn_mtp;
+    const int r = blockIdx.x, t = threadIdx.x, nt = blockDim.x;
+    const unsigned short* row = reinterpret_cast<const unsigned short*>(L + static_cast<long long>(r) * ld);
+    unsigned short* ov = out_v + static_cast<long long>(r) * K;
+    int* oc = out_c + static_cast<long long>(r) * K;
+    __shared__ unsigned int hist[256];
+    __shared__ unsigned int scan[1024], scan_eq[1024];
+    __shared__ unsigned int pick_hi, pick_lo, above;
+    for (int i = t; i < 256; i += nt) hist[i] = 0;
+    __syncthreads();
+    for (int i = t; i < V; i += nt) atomicAdd(&hist[order_key(row[i]) >> 8], 1u);
+    __syncthreads();
+    if (t == 0) {
+        unsigned int cum = 0;
+        int b = 255;
+        for (; b > 0; --b) {
+            if (cum + hist[b] >= static_cast<unsigned int>(K)) break;
+            cum += hist[b];
+        }
+        pick_hi = b;
+        above = cum;
+    }
+    __syncthreads();
+    const unsigned int hi = pick_hi;
+    for (int i = t; i < 256; i += nt) hist[i] = 0;
+    __syncthreads();
+    for (int i = t; i < V; i += nt) {
+        const unsigned int k = order_key(row[i]);
+        if ((k >> 8) == hi) atomicAdd(&hist[k & 255u], 1u);
+    }
+    __syncthreads();
+    if (t == 0) {
+        unsigned int cum = above;
+        int b = 255;
+        for (; b > 0; --b) {
+            if (cum + hist[b] >= static_cast<unsigned int>(K)) break;
+            cum += hist[b];
+        }
+        pick_lo = b;
+        above = cum;
+    }
+    __syncthreads();
+    const unsigned int kth = (hi << 8) | pick_lo, high = above, need = K - high;
+    // each thread's contiguous columns: the values above the K-th and those equal to it, counted, scanned in order
+    const int per = (V + nt - 1) / nt, lo = t * per, end = min(V, lo + per);
+    unsigned int gt = 0, eq = 0;
+    for (int i = lo; i < end; ++i) {
+        const unsigned int k = order_key(row[i]);
+        gt += k > kth;
+        eq += k == kth;
+    }
+    scan[t] = gt;
+    scan_eq[t] = eq;
+    __syncthreads();
+    if (t == 0) {
+        unsigned int s = 0, q = 0;
+        for (int i = 0; i < nt; ++i) {
+            const unsigned int x = scan[i], y = scan_eq[i];
+            scan[i] = s;
+            scan_eq[i] = q;
+            s += x;
+            q += y;
+        }
+    }
+    __syncthreads();
+    // in column order: the values above the K-th fill the first `high` slots, the first `need` ties the rest
+    unsigned int at = scan[t], rank = scan_eq[t];
+    for (int i = lo; i < end; ++i) {
+        const unsigned int k = order_key(row[i]);
+        if (k > kth) {
+            ov[at] = row[i];
+            oc[at] = i;
+            ++at;
+        } else if (k == kth && rank < need) {
+            ov[high + rank] = row[i];
+            oc[high + rank] = i;
+            ++rank;
+        }
+    }
+}