diff --git a/zig/kernels/cuda/flashnext_batch.cu b/zig/kernels/cuda/flashnext_batch.cu index 8b66e200aa..3a0726e4c9 100644 --- a/zig/kernels/cuda/flashnext_batch.cu +++ b/zig/kernels/cuda/flashnext_batch.cu @@ -16,8 +16,8 @@ extern "C" __global__ void fn_sum_slices(const float* __restrict__ part, void* _ } // Shifts the last-`taps`-rows window [taps, c] in place past `rows`, n rows `stride` apart (conv history, n-gram tail). -extern "C" __global__ void fn_shift_window(__nv_bfloat16* __restrict__ win, const __nv_bfloat16* __restrict__ rows, - long long stride, int n, int taps, int c) { +static __device__ __forceinline__ void shift_window(__nv_bfloat16* __restrict__ win, const __nv_bfloat16* __restrict__ rows, + long long stride, int n, int taps, int c) { const int ch = blockIdx.x * blockDim.x + threadIdx.x; if (ch >= c) return; __nv_bfloat16 next[16]; @@ -29,6 +29,11 @@ extern "C" __global__ void fn_shift_window(__nv_bfloat16* __restrict__ win, cons for (int j = 0; j < taps; ++j) win[static_cast(j) * c + ch] = next[j]; } +extern "C" __global__ void fn_shift_window(__nv_bfloat16* __restrict__ win, const __nv_bfloat16* __restrict__ rows, + long long stride, int n, int taps, int c) { + shift_window(win, rows, stride, n, taps, c); +} + // ---- the DeltaNet chain over a chunk: flashnext_gdn.cu's chain cut into front, rows and norm, with its bits ---- namespace fn_gdn { @@ -58,13 +63,22 @@ static __device__ __forceinline__ float conv_act(const __nv_bfloat16* __restrict } } // namespace fn_gdn +// A layer's pending-row log: log_rows rows of k [NK, DK], v [NV, DV], decay and beta [NV, 2], fp32; ctl first | n << 16 +namespace fn_late { +using namespace fn_gdn; +static __device__ __forceinline__ const float* vs(const float* log, int log_rows) { return log + (size_t)log_rows * NK * DK; } +static __device__ __forceinline__ const float* gbs(const float* log, int log_rows) { return vs(log, log_rows) + (size_t)log_rows * NV * DV; } +} // namespace fn_late + // Program (row, unit), 256 threads: key head u's normed q and k (u < NK), else a value head's v, decay g and beta. -extern "C" __global__ void __launch_bounds__(256) fn_gdn_front( +template // LOG: k, v and the gates go into the log's row after the pending ones +static __device__ __forceinline__ void gdn_front( 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, float* __restrict__ qn, - float* __restrict__ kn, float* __restrict__ v, float* __restrict__ gb) { + float* __restrict__ kn, float* __restrict__ v, float* __restrict__ gb, const unsigned* __restrict__ ctl, int log_rows) { using namespace fn_gdn; const int r = blockIdx.x, u = blockIdx.y, t = threadIdx.x, warp = t >> 5, lane = t & 31; + const int kr = LOG ? (int)(((*ctl & 0xffffu) + (*ctl >> 16) + (unsigned)r) % (unsigned)log_rows) : r; if (u < NK) { __shared__ float x[2][DK]; const int hk = u; @@ -80,22 +94,39 @@ extern "C" __global__ void __launch_bounds__(256) fn_gdn_front( ss = warp_sum(ss); float inv = 1.0f / sqrtf(ss + 1e-6f); if (warp == 0) inv = inv * (1.0f / sqrtf((float)DK)); - float* dst = (warp == 0 ? qn : kn) + ((size_t)r * NK + hk) * DK; + float* dst = warp == 0 ? qn + ((size_t)r * NK + hk) * DK : kn + ((size_t)kr * NK + hk) * DK; #pragma unroll for (int i = 0; i < 4; ++i) dst[lane * 4 + i] = v4[i] * inv; } } else { const int hv = u - NK; - if (t < DV) v[(size_t)r * NV * DV + hv * DV + t] = conv_act(P, cs, cw, r, 2 * NK * DK + hv * DV + t); + if (t < DV) v[(size_t)kr * NV * DV + hv * DV + t] = conv_act(P, cs, cw, r, 2 * NK * DK + hv * DV + t); else if (t == DV) { 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]); - gb[((size_t)r * NV + hv) * 2] = expf(-expf(a_log[hv]) * softplusf_(a + dt_bias[hv])); - gb[((size_t)r * NV + hv) * 2 + 1] = bf(sigmoidf_(b)); + gb[((size_t)kr * NV + hv) * 2] = expf(-expf(a_log[hv]) * softplusf_(a + dt_bias[hv])); + gb[((size_t)kr * NV + hv) * 2 + 1] = bf(sigmoidf_(b)); } } } +extern "C" __global__ void __launch_bounds__(256) fn_gdn_front( + 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, float* __restrict__ qn, + float* __restrict__ kn, float* __restrict__ v, float* __restrict__ gb) { + gdn_front(P, cs, cw, a_log, dt_bias, qn, kn, v, gb, nullptr, 1); +} + +// fn_gdn_front with the window's k, v and gates into the stream's log after its pending rows (q as fn_gdn_front's). +extern "C" __global__ void __launch_bounds__(256) fn_gdn_front_log( + 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, float* __restrict__ qn, + float* __restrict__ log, const unsigned* __restrict__ ctl, int log_rows) { + using namespace fn_gdn; + float* vl = log + (size_t)log_rows * NK * DK; + gdn_front(P, cs, cw, a_log, dt_bias, qn, log, vl, vl + (size_t)log_rows * NV * DV, ctl, log_rows); +} + // The chain through the chunk from state_in to state_out (in place, or apart for a rollback): a warp 4 state rows. static __device__ __forceinline__ void gdn_rows( const float* __restrict__ qn, const float* __restrict__ kn, const float* __restrict__ v, @@ -175,6 +206,100 @@ extern "C" __global__ void __launch_bounds__(128) fn_gdn_rows_to( gdn_rows(qn, kn, v, gb, state_in, state_out, rows, ys_out); } +// A window's chain after its stream's pending rows: those first (state only), then its rows read out; fn_gdn_rows' sums +static __device__ __forceinline__ void gdn_rows_late( + const float* __restrict__ qn, const float* __restrict__ log, const unsigned* __restrict__ ctl, int log_rows, + int store_above, float* state_in, float* state_out, int rows, float* __restrict__ ys_out) { + using namespace fn_gdn; + const int hv = blockIdx.x, hk = hv / (NV / NK), warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int row0 = blockIdx.y * 16 + warp * 4; + const size_t sbase = (size_t)hv * DV * DK; + const float* vl = fn_late::vs(log, log_rows); + const float* gl = fn_late::gbs(log, log_rows); + const unsigned first = *ctl & 0xffffu, np = *ctl >> 16; + float s[4][4]; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) s[j][i] = state_in[sbase + (size_t)(row0 + j) * DK + lane * 4 + i]; + const unsigned total = np + (unsigned)rows; + for (unsigned p = 0; p < total; ++p) { + if (p == np && (int)np > store_above) { +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) state_in[sbase + (size_t)(row0 + j) * DK + lane * 4 + i] = s[j][i]; + } + const size_t lr = (first + p) % (unsigned)log_rows; + const float4 k4 = *reinterpret_cast(log + (lr * NK + hk) * DK + lane * 4); + const float4 v4 = *reinterpret_cast(vl + lr * NV * DV + hv * DV + row0); + const float kk[4] = {k4.x, k4.y, k4.z, k4.w}, vv[4] = {v4.x, v4.y, v4.z, v4.w}; + const float g = gl[(lr * NV + hv) * 2], beta = gl[(lr * NV + hv) * 2 + 1]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float kv = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + s[j][i] = s[j][i] * g; + kv = kv + s[j][i] * kk[i]; + } + kv = warp_sum(kv); + const float delta = (vv[j] - kv) * beta; +#pragma unroll + for (int i = 0; i < 4; ++i) s[j][i] = s[j][i] + kk[i] * delta; + } + if (p < np) continue; + const int r = (int)(p - np); + const float4 q4 = *reinterpret_cast(qn + ((size_t)r * NK + hk) * DK + lane * 4); + const float qq[4] = {q4.x, q4.y, q4.z, q4.w}; + float o4[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float o = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) o = o + s[j][i] * qq[i]; + o4[j] = bf(warp_sum(o)); + } + if (lane == 0) *reinterpret_cast(ys_out + ((size_t)r * NV + hv) * DV + row0) = make_float4(o4[0], o4[1], o4[2], o4[3]); + } +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) state_out[sbase + (size_t)(row0 + j) * DK + lane * 4 + i] = s[j][i]; +} + +extern "C" __global__ void __launch_bounds__(128) fn_gdn_rows_late( + const float* __restrict__ qn, const float* __restrict__ log, const unsigned* __restrict__ ctl, int log_rows, + int store_above, float* state_in, float* state_out, int rows, float* __restrict__ ys_out) { + gdn_rows_late(qn, log, ctl, log_rows, store_above, state_in, state_out, rows, ys_out); +} + +// A rolled-back lane's conv windows and n-gram tail in one launch: each window is its copy moved on by `keep` rows. +namespace fn_roll { +constexpr int MAXJ = 40; +struct Job { + __nv_bfloat16* win; + const __nv_bfloat16* hist; + const __nv_bfloat16* rows; + long long stride; + int taps, c; +}; +struct Jobs { + Job j[MAXJ]; + int n; +}; +} // namespace fn_roll + +extern "C" __global__ void fn_late_roll(const fn_roll::Jobs t, int keep) { + const fn_roll::Job& j = t.j[blockIdx.y]; + const int ch = blockIdx.x * blockDim.x + threadIdx.x; + if (ch >= j.c) return; + for (int r = 0; r < j.taps; ++r) { + const int src = r + keep; // row src of [window (taps rows); the pass's rows] + j.win[(long long)r * j.c + ch] = src < j.taps ? j.hist[(long long)src * j.c + ch] : j.rows[(long long)(src - j.taps) * j.stride + ch]; + } +} + extern "C" __global__ void __launch_bounds__(128) fn_gdn_norm( const float* __restrict__ ys_in, const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ norm_w, float eps, __nv_bfloat16* __restrict__ out, float* __restrict__ xs) { @@ -237,6 +362,45 @@ __device__ __forceinline__ int fn_block_scan(int v, int* warp_sums, int* total) return v; } +// A radix pass's cut by one warp (8 bins a lane): the highest bin whose suffix sum reaches need, and how many it gives. +__device__ __forceinline__ void fn_qsa_bin_cut(const int* hist, int need, uint32_t cut, int shift, uint32_t* cut_s, int* need_s) { + const int lane = threadIdx.x & 31; + int c[8], local = 0; +#pragma unroll + for (int j = 0; j < 8; ++j) { + c[j] = hist[lane * 8 + j]; + local += c[j]; + } + int above = local; // the bins of this lane and every lane above it +#pragma unroll + for (int o = 1; o < 32; o <<= 1) { + const int v = __shfl_down_sync(0xffffffffu, above, o); + if (lane + o < 32) above += v; + } + above -= local; + int best = -1, rest = 0, s = above; // s = S[d + 1] as d walks down the lane's bins +#pragma unroll + for (int j = 7; j >= 0; --j) { + const int d = lane * 8 + j; + if (best < 0 && d >= 1 && s + c[j] >= need) { + best = d; + rest = need - s; + } + s += c[j]; + } + const int top_d = __reduce_max_sync(0xffffffffu, best); + if (top_d < 0) { // no bin from 1 up reaches it: bin 0, need - S[1] + const int s1 = __shfl_sync(0xffffffffu, s - c[0], 0); + if (lane == 0) { + *need_s = need - s1; + *cut_s = cut; + } + } else if (best == top_d) { + *need_s = rest; + *cut_s = cut | (static_cast(top_d) << shift); + } +} + // attention._select and _select_tiles, a row a block: a row past `top` complete blocks lists the best, then its tail. extern "C" __global__ void __launch_bounds__(1024) fn_qsa_select(const float* __restrict__ sc, const int* __restrict__ pos0, int* __restrict__ ids, @@ -260,48 +424,63 @@ extern "C" __global__ void __launch_bounds__(1024) fn_qsa_select(const float* __ int* out = ids + static_cast(r) * idw; uint32_t cut = 0; int need = top; + const int lane = threadIdx.x & 31; for (int p = 0; p < 4; ++p) { const int shift = 24 - 8 * p; for (int i = threadIdx.x; i < 256; i += blockDim.x) hist[i] = 0; __syncthreads(); - for (int b = threadIdx.x; b < complete; b += blockDim.x) { - const uint32_t k = fn_qsa_key(row[b]); - if (p == 0 || (k >> (shift + 8)) == (cut >> (shift + 8))) atomicAdd(&hist[(k >> shift) & 0xFF], 1); - } - __syncthreads(); - if (threadIdx.x == 0) { - int above = 0; // keys under the prefix with a larger byte - int d = 255; - for (; d > 0; --d) { - if (above + hist[d] >= need) break; - above += hist[d]; + constexpr int U = 8; // a thread's keys in flight before their adds (one L2 wait a key was the pass's time) + for (int b0 = 0; b0 < complete; b0 += U * blockDim.x) { + float v[U]; +#pragma unroll + for (int u = 0; u < U; ++u) { + const int b = b0 + u * blockDim.x + threadIdx.x; + v[u] = b < complete ? row[b] : 0.0f; + } +#pragma unroll + for (int u = 0; u < U; ++u) { + const int b = b0 + u * blockDim.x + threadIdx.x; + const uint32_t k = fn_qsa_key(v[u]); + if (b < complete && (p == 0 || (k >> (shift + 8)) == (cut >> (shift + 8)))) atomicAdd(&hist[(k >> shift) & 0xFF], 1); } - need_s = need - above; - cut_s = cut | (static_cast(d) << shift); } __syncthreads(); + if (threadIdx.x < 32) fn_qsa_bin_cut(hist, need, cut, shift, &cut_s, &need_s); + __syncthreads(); cut = cut_s; need = need_s; __syncthreads(); } - // every key above the cut, then the first `need` keys equal to it, placed in block order - int seen = 0; - int placed = 0; - for (int t0 = 0; t0 < complete && placed < top; t0 += blockDim.x) { - const int b = t0 + threadIdx.x; - const bool ok = b < complete; - const uint32_t k = ok ? fn_qsa_key(row[b]) : 0u; - const bool above = ok && k > cut; - const bool equal = ok && k == cut; - int eq_total = 0; - const int rank = seen + fn_block_scan(equal ? 1 : 0, warp_sums, &eq_total); - const bool chosen = above || (equal && rank <= need); - int ch_total = 0; - const int place = placed + fn_block_scan(chosen ? 1 : 0, warp_sums, &ch_total) - 1; - if (chosen) - for (int j = 0; j < ratio; ++j) out[place * ratio + j] = b * ratio + j; - seen += eq_total; - placed += ch_total; + // keys above the cut, then the first `need` equal ones in block order: a warp a run, ballots and two block scans + const int warps = static_cast(blockDim.x) >> 5, warp = threadIdx.x >> 5; + const int steps = (complete + 32 * warps - 1) / (32 * warps); + const int w0 = warp * steps * 32; + const unsigned below = (1u << lane) - 1u; + int na = 0, ne = 0; + for (int st = 0; st < steps; ++st) { + const int b = w0 + st * 32 + lane; + const uint32_t k = b < complete ? fn_qsa_key(row[b]) : 0u; + na += __popc(__ballot_sync(0xffffffffu, b < complete && k > cut)); + ne += __popc(__ballot_sync(0xffffffffu, b < complete && k == cut)); + } + int total = 0; + const int eq_before = __shfl_sync(0xffffffffu, fn_block_scan(lane == 0 ? ne : 0, warp_sums, &total), 0) - ne; + const int take = max(0, min(ne, need - eq_before)); + const int mine = na + take; + int place = __shfl_sync(0xffffffffu, fn_block_scan(lane == 0 ? mine : 0, warp_sums, &total), 0) - mine; + int eq_seen = 0; + for (int st = 0; st < steps; ++st) { + const int b = w0 + st * 32 + lane; + const uint32_t k = b < complete ? fn_qsa_key(row[b]) : 0u; + const unsigned eq = __ballot_sync(0xffffffffu, b < complete && k == cut); + const bool chosen = b < complete && (k > cut || (k == cut && eq_seen + __popc(eq & below) < take)); + const unsigned ch = __ballot_sync(0xffffffffu, chosen); + if (chosen) { + const int at = place + __popc(ch & below); + for (int j = 0; j < ratio; ++j) out[at * ratio + j] = b * ratio + j; + } + place += __popc(ch); + eq_seen += __popc(eq); } if (threadIdx.x < ratio) { const int tail = ratio * complete + threadIdx.x; diff --git a/zig/kernels/cuda/flashnext_gdn_fused.cuh b/zig/kernels/cuda/flashnext_gdn_fused.cuh index 3b15c092e0..50b541611a 100644 --- a/zig/kernels/cuda/flashnext_gdn_fused.cuh +++ b/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(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(base + j * ROW + 2 * DK * 4)[p] = o; // v + } + } + if (p < n) { // row p's g and beta float* gbp = reinterpret_cast(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 diff --git a/zig/kernels/cuda/flashnext_hc.cu b/zig/kernels/cuda/flashnext_hc.cu index 925506cb7a..947802fea8 100644 --- a/zig/kernels/cuda/flashnext_hc.cu +++ b/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 then fn_hc_normalize in one block a row, slices summed and added in their order: the same bits. +template +__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(r) * D + d]; + for (int k = 1; k < WORLD; ++k) acc = __fadd_rn(acc, BR[k * RS + static_cast(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(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(r) * S * D + static_cast(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(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(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(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, diff --git a/zig/kernels/cuda/flashnext_mtp.cu b/zig/kernels/cuda/flashnext_mtp.cu index 1ac2e6fa4b..f90fcd1925 100644 --- a/zig/kernels/cuda/flashnext_mtp.cu +++ b/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(L + static_cast(r) * ld); + unsigned short* ov = out_v + static_cast(r) * K; + int* oc = out_c + static_cast(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(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(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; + } + } +} diff --git a/zig/kernels/cuda/prompt16.cu b/zig/kernels/cuda/prompt16.cu index 8323fb6330..545058b43c 100644 --- a/zig/kernels/cuda/prompt16.cu +++ b/zig/kernels/cuda/prompt16.cu @@ -11,7 +11,7 @@ namespace tf_prompt16 { using namespace qmm_frag; -enum Mode : int { FP4 = 0, FP8 = 1, MXFP8 = 2 }; +enum Mode : int { FP4 = 0, FP8 = 1, MXFP8 = 2, BF16 = 4 }; constexpr int GS = 64; @@ -41,7 +41,7 @@ struct Tile { static constexpr int ROW = GS * 2; static constexpr int CHUNKS = ROW / 16; static constexpr int X = BM * ROW; - static constexpr int W = MODE == FP4 ? BN * GS / 2 : BN * GS; + static constexpr int W = MODE == FP4 ? BN * GS / 2 : MODE == BF16 ? BN * GS * 2 : BN * GS; static constexpr int S = MODE == FP4 ? BN * 4 : MODE == MXFP8 ? BN * 2 : 0; static constexpr int STAGE = (X + W + S + 127) / 128 * 128; static constexpr int SMEM = STAGES * STAGE; @@ -422,4 +422,132 @@ template __global__ void prompt_lean( template __global__ void prompt_lean( const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, int, int, int, int, int, int); + +// BF16 prompt rows on qmmf's BF16 words with qmmf_kernel's sums (K slices from zero, added in order): the lane's bits. +template +__device__ __forceinline__ void prompt_bf16(const __nv_bfloat16* __restrict__ x, const unsigned char* __restrict__ w, + float scale, void* __restrict__ out, int M, int N, int K, int SK, int npad, + int ldx, int group) { + constexpr int BM = 128, BN = 128, WM = 2, WN = 4, STAGES = 3; + using T = Tile; + extern __shared__ __align__(128) unsigned char buf[]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int wm = warp / WN, wn = warp % WN; + const int KG = K / GS, slice_groups = KG / SK; + const int2 at = tile_of(blockIdx.x, M, N, BM, BN, group); + const int m0 = at.x, n0 = at.y; + + auto stage = [&](int s) { return buf + s * T::STAGE; }; + auto load = [&](int s, int g) { + unsigned char* p = stage(s); + for (int c = tid; c < BM * T::CHUNKS; c += T::THREADS) { + const int r = c / T::CHUNKS, ch = c % T::CHUNKS; + cp16z(p + r * T::ROW + swz(r, ch) * 16, + x + static_cast(min(m0 + r, M - 1)) * ldx + g * GS + ch * 8, m0 + r < M); + } + unsigned char* pw = p + T::X; + constexpr int TILE_BYTES = 64 * GS * 2; + for (int c = tid; c < T::W / 16; c += T::THREADS) { + const int t = c / (TILE_BYTES / 16), off = c % (TILE_BYTES / 16); + if (n0 + t * 64 < npad) + cp16(pw + c * 16, w + (static_cast(n0 / 64 + t) * KG + g) * TILE_BYTES + off * 16); + } + }; + + float acc[T::MT][T::NT][4]; + float tot[T::MT][T::NT][4]; +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) acc[i][j][e] = tot[i][j][e] = 0.0f; +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < KG) load(s, s); + commit(); + } + for (int g = 0; g < KG; ++g) { + if (SK > 1 && g > 0 && g % slice_groups == 0) { +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) { + tot[i][j][e] = g == slice_groups ? acc[i][j][e] : tot[i][j][e] + acc[i][j][e]; + acc[i][j][e] = 0.0f; + } + } + wait(); + __syncthreads(); + if (g + STAGES - 1 < KG) load((g + STAGES - 1) % STAGES, g + STAGES - 1); + commit(); + const unsigned char* p = stage(g % STAGES); + const uint2* pw = reinterpret_cast(p + T::X); +#pragma unroll + for (int kt = 0; kt < GS / 16; ++kt) { + uint32_t a[T::MT][4]; +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + const int r = wm * (BM / WM) + i * 16 + (lane & 7) + ((lane >> 3) & 1) * 8; + const int ch = kt * 2 + (lane >> 4); + ldmatrix4(a[i], p + r * T::ROW + swz(r, ch) * 16); + } +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + // n8 tile jj of the block's two 64-column tiles, laid out back to back: (jj, kt, lane) as qmmf reads + const uint2 v = pw[((wn * T::NT + j) * 4 + kt) * 32 + lane]; +#pragma unroll + for (int i = 0; i < T::MT; ++i) mma(acc[i][j], a[i], v.x, v.y); + } + } + } + wait<0>(); + if (SK > 1) { +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) acc[i][j][e] = tot[i][j][e] + acc[i][j][e]; + } +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + const int col = n0 + wn * (BN / WN) + j * 8 + (lane & 3) * 2; +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int row = m0 + wm * (BM / WM) + i * 16 + (lane >> 2) + h * 8; + if (row >= M) continue; + const float v0 = acc[i][j][2 * h] * scale, v1 = acc[i][j][2 * h + 1] * scale; + if (F32) { + float* dst = reinterpret_cast(out) + static_cast(row) * N + col; + if (col < N) dst[0] = v0; + if (col + 1 < N) dst[1] = v1; + } else { + __nv_bfloat16* dst = reinterpret_cast<__nv_bfloat16*>(out) + static_cast(row) * N + col; + if (col + 1 < N && (N & 1) == 0) + *reinterpret_cast<__nv_bfloat162*>(dst) = __floats2bfloat162_rn(v0, v1); + else { + if (col < N) dst[0] = __float2bfloat16_rn(v0); + if (col + 1 < N) dst[1] = __float2bfloat16_rn(v1); + } + } + } + } +} + +extern "C" __global__ void __launch_bounds__(256) tf_prompt_bf16(const __nv_bfloat16* __restrict__ x, const unsigned char* __restrict__ w, + float scale, void* __restrict__ out, int M, int N, int K, int SK, + int npad, int ldx, int group) { + prompt_bf16(x, w, scale, out, M, N, K, SK, npad, ldx, group); +} + +extern "C" __global__ void __launch_bounds__(256) tf_prompt_bf16_f32(const __nv_bfloat16* __restrict__ x, const unsigned char* __restrict__ w, + float scale, void* __restrict__ out, int M, int N, int K, int SK, + int npad, int ldx, int group) { + prompt_bf16(x, w, scale, out, M, N, K, SK, npad, ldx, group); +} } diff --git a/zig/kernels/cuda/route_rows.cu b/zig/kernels/cuda/route_rows.cu index 94779d3985..f11d876e37 100644 --- a/zig/kernels/cuda/route_rows.cu +++ b/zig/kernels/cuda/route_rows.cu @@ -46,13 +46,13 @@ __device__ __forceinline__ void block_best(float& v, int& i, float (*sv)[NT / 32 using namespace tf_route; // OUT[m, e] = x[m] . w[e], 8 experts a warp, RT tiles of 16 rows, one mma chain over K, STAGES of KC (row sums fixed). -template +template // EFAST: x walks the expert blocks, rows read once __device__ __forceinline__ void route_logits(const __nv_bfloat16* __restrict__ X, int x_ld, const __nv_bfloat16* __restrict__ W, float* __restrict__ OUT, int M, int D, int E) { constexpr int LD = KC + 8, EB = 8 * WARPS, XR = 16 * RT, ROWS = XR + EB; extern __shared__ __align__(16) unsigned char smem_raw[]; auto tile = reinterpret_cast<__nv_bfloat16 (*)[ROWS][LD]>(smem_raw); // [STAGES] XR input rows, then EB expert rows - const int m0 = blockIdx.x * XR, e0 = blockIdx.y * EB; + const int m0 = (EFAST ? blockIdx.y : blockIdx.x) * XR, e0 = (EFAST ? blockIdx.x : blockIdx.y) * EB; const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; const int stages = (D + KC - 1) / KC; const int xr = min(XR, M - m0); // input rows past M stay zero: written once, never staged @@ -116,11 +116,11 @@ extern "C" __global__ void __launch_bounds__(128) tf_route_logits(const __nv_bfl route_logits<4, 2, 128>(X, x_ld, W, OUT, M, D, E); } -// route_logits over 128 rows and 64 experts a block, eight warps (from 64 rows), the same sums. +// route_logits over 128 rows and 64 experts a block, eight warps (from 64 rows), the same sums; EFAST block order. extern "C" __global__ void __launch_bounds__(256) tf_route_logits_wide(const __nv_bfloat16* __restrict__ X, int x_ld, const __nv_bfloat16* __restrict__ W, float* __restrict__ OUT, int M, int D, int E) { - route_logits<8, 2, 64, 8>(X, x_ld, W, OUT, M, D, E); + route_logits<8, 2, 64, 8, true>(X, x_ld, W, OUT, M, D, E); } // route_logits with three stages of 256 in flight (up to 16 rows, where the router rows' reads set its time). diff --git a/zig/src/core/contract/flashnext_tile_records.zig b/zig/src/core/contract/flashnext_tile_records.zig index dd72ba9e68..d3d11dbd8b 100644 --- a/zig/src/core/contract/flashnext_tile_records.zig +++ b/zig/src/core/contract/flashnext_tile_records.zig @@ -17,6 +17,15 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "flashnext_hc fn_hc_rinv", .symbol = "fn_hc_rinv", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "flashnext_hc fn_hc_mix_h", .symbol = "fn_hc_mix_h", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 27, .static = 1024, .bound = 256 }, .{ .regs = 27, .static = 1024, .bound = 256 }, .{ .regs = 27, .static = 1024, .bound = 256 }, .{ .regs = 40, .static = 1024, .bound = 256 }, .{ .regs = 40, .static = 1024, .bound = 256 } } }, .{ .launch = .{ .name = "route_rows tf_route_logits_wide", .symbol = "tf_route_logits_wide", .image = "route_rows", .threads = 256, .dynamic = 55296, .opt_in = true }, .usage = .{ .{ .regs = 47, .static = 0, .bound = 256 }, .{ .regs = 47, .static = 0, .bound = 256 }, .{ .regs = 47, .static = 0, .bound = 256 }, .{ .regs = 61, .static = 0, .bound = 256 }, .{ .regs = 61, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_batch fn_gdn_front_log", .symbol = "fn_gdn_front_log", .image = "flashnext_batch", .threads = 256 }, .usage = .{ .{ .regs = 26, .static = 1024, .bound = 256 }, .{ .regs = 26, .static = 1024, .bound = 256 }, .{ .regs = 26, .static = 1024, .bound = 256 }, .{ .regs = 32, .static = 1024, .bound = 256 }, .{ .regs = 32, .static = 1024, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_batch fn_gdn_rows_late", .symbol = "fn_gdn_rows_late", .image = "flashnext_batch", .threads = 128 }, .usage = .{ .{ .regs = 44, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_batch fn_late_roll", .symbol = "fn_late_roll", .image = "flashnext_batch", .threads = 256 }, .usage = .{ .{ .regs = 22, .static = 0, .bound = 0 }, .{ .regs = 22, .static = 0, .bound = 0 }, .{ .regs = 22, .static = 0, .bound = 0 }, .{ .regs = 20, .static = 0, .bound = 0 }, .{ .regs = 20, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_writeback_norm0", .symbol = "fn_hc_writeback_norm0", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 59, .static = 1296, .bound = 256 }, .{ .regs = 60, .static = 1296, .bound = 256 }, .{ .regs = 60, .static = 1296, .bound = 256 }, .{ .regs = 62, .static = 1296, .bound = 256 }, .{ .regs = 62, .static = 1296, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_writeback_norm_moe", .symbol = "fn_hc_writeback_norm_moe", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 64, .static = 1296, .bound = 256 }, .{ .regs = 62, .static = 1296, .bound = 256 }, .{ .regs = 62, .static = 1296, .bound = 256 }, .{ .regs = 62, .static = 1296, .bound = 256 }, .{ .regs = 62, .static = 1296, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_writeback_norm_slices", .symbol = "fn_hc_writeback_norm_slices", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 64, .static = 1296, .bound = 256 }, .{ .regs = 64, .static = 1296, .bound = 256 }, .{ .regs = 64, .static = 1296, .bound = 256 }, .{ .regs = 64, .static = 1296, .bound = 256 }, .{ .regs = 64, .static = 1296, .bound = 256 } } }, + .{ .launch = .{ .name = "prompt16 tf_prompt_bf16", .symbol = "tf_prompt_bf16", .image = "prompt16", .threads = 256, .dynamic = 98304, .opt_in = true }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 256 }, .{ .regs = 255, .static = 0, .bound = 256 }, .{ .regs = 255, .static = 0, .bound = 256 }, .{ .regs = 255, .static = 0, .bound = 256 }, .{ .regs = 255, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "prompt16 tf_prompt_bf16_f32", .symbol = "tf_prompt_bf16_f32", .image = "prompt16", .threads = 256, .dynamic = 98304, .opt_in = true }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 256 }, .{ .regs = 255, .static = 0, .bound = 256 }, .{ .regs = 255, .static = 0, .bound = 256 }, .{ .regs = 255, .static = 0, .bound = 256 }, .{ .regs = 255, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_mtp fn_row_topk", .symbol = "fn_row_topk", .image = "flashnext_mtp", .threads = 1024 }, .usage = .{ .{ .regs = 32, .static = 9228, .bound = 1024 }, .{ .regs = 59, .static = 9228, .bound = 1024 }, .{ .regs = 59, .static = 9228, .bound = 1024 }, .{ .regs = 31, .static = 9228, .bound = 1024 }, .{ .regs = 31, .static = 9228, .bound = 1024 } } }, }; /// nvcc 13.2's records where it compiles these tiles differently from 13.0. diff --git a/zig/src/core/contract/nvcc_133_records.zig b/zig/src/core/contract/nvcc_133_records.zig index 5f157a01b2..aefa642f8f 100644 --- a/zig/src/core/contract/nvcc_133_records.zig +++ b/zig/src/core/contract/nvcc_133_records.zig @@ -95,7 +95,7 @@ pub const cuda_133 = [_]registry.Override{ .{ .name = "flashnext_attn fn_attn_prep_media_paged", .usage = .{ .{ .regs = 32, .static = 128, .bound = 256 }, .{ .regs = 36, .static = 128, .bound = 256 }, .{ .regs = 36, .static = 128, .bound = 256 }, .{ .regs = 32, .static = 128, .bound = 256 }, .{ .regs = 32, .static = 128, .bound = 256 } } }, .{ .name = "flashnext_attn fn_qsa_scores_paged", .usage = .{ .{ .regs = 220, .static = 16384, .bound = 256 }, .{ .regs = 219, .static = 16384, .bound = 256 }, .{ .regs = 219, .static = 16384, .bound = 256 }, .{ .regs = 219, .static = 16384, .bound = 256 }, .{ .regs = 219, .static = 16384, .bound = 256 } } }, .{ .name = "flashnext_attn qsa_tile scores", .usage = .{ .{ .regs = 222, .static = 16384, .bound = 256 }, .{ .regs = 222, .static = 16384, .bound = 256 }, .{ .regs = 222, .static = 16384, .bound = 256 }, .{ .regs = 223, .static = 16384, .bound = 256 }, .{ .regs = 223, .static = 16384, .bound = 256 } } }, - .{ .name = "flashnext_batch fn_gdn_fused", .usage = .{ .{ .regs = 138, .static = 0, .bound = 384 }, .{ .regs = 142, .static = 0, .bound = 384 }, .{ .regs = 142, .static = 0, .bound = 384 }, .{ .regs = 164, .static = 0, .bound = 384 }, .{ .regs = 164, .static = 0, .bound = 384 } } }, + .{ .name = "flashnext_batch fn_gdn_fused", .usage = .{ .{ .regs = 168, .static = 0, .bound = 384 }, .{ .regs = 166, .static = 0, .bound = 384 }, .{ .regs = 166, .static = 0, .bound = 384 }, .{ .regs = 168, .static = 0, .bound = 384 }, .{ .regs = 168, .static = 0, .bound = 384 } } }, .{ .name = "flashnext_gdn flashnext_gdn_chain", .usage = .{ .{ .regs = 51, .static = 2060, .bound = 1024 }, .{ .regs = 56, .static = 2060, .bound = 1024 }, .{ .regs = 56, .static = 2060, .bound = 1024 }, .{ .regs = 61, .static = 2060, .bound = 1024 }, .{ .regs = 61, .static = 2060, .bound = 1024 } } }, .{ .name = "flashnext_hc fn_hc_dense_mix8", .usage = .{ .{ .regs = 32, .static = 1024, .bound = 256 }, .{ .regs = 34, .static = 1024, .bound = 256 }, .{ .regs = 34, .static = 1024, .bound = 256 }, .{ .regs = 38, .static = 1024, .bound = 256 }, .{ .regs = 38, .static = 1024, .bound = 256 } } }, .{ .name = "flashnext_hc fn_hc_normalize_cols", .usage = .{ .{ .regs = 32, .static = 16, .bound = 256 }, .{ .regs = 40, .static = 16, .bound = 256 }, .{ .regs = 40, .static = 16, .bound = 256 }, .{ .regs = 40, .static = 16, .bound = 256 }, .{ .regs = 40, .static = 16, .bound = 256 } } }, diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index ac87ea207c..28138574d2 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -317,7 +317,7 @@ pub const cuda_sets = [_]CudaSet{ .{ .name = "mamba tf_nemo_scan", .usage = .{ .{ .regs = 138, .static = 41088, .bound = 128 }, .{ .regs = 144, .static = 41088, .bound = 128 }, .{ .regs = 144, .static = 41088, .bound = 128 }, .{ .regs = 118, .static = 41088, .bound = 128 }, .{ .regs = 118, .static = 41088, .bound = 128 } } }, } ++ learning.cuda_132 ++ @import("paging_pending.zig").cuda_132 ++ @import("nemotron_prompt_records.zig").cuda_132 ++ @import("glm_cuda_records.zig").cuda_132 ++ @import("qwen_served_records.zig").cuda_132 ++ @import("flashnext_tile_records.zig").cuda_132 ++ @import("nvfp4_fact_records.zig").cuda_132 ++ @import("pocket_records.zig").cuda_132 ++ @import("minimax_cuda_records.zig").cuda_132) }, .{ .release = .{ .major = 13, .minor = 3 }, .usage = &([_]Override{ - .{ .name = "flashnext_batch fn_qsa_select", .usage = .{ .{ .regs = 32, .static = 1168, .bound = 1024 }, .{ .regs = 50, .static = 1168, .bound = 1024 }, .{ .regs = 50, .static = 1168, .bound = 1024 }, .{ .regs = 44, .static = 1168, .bound = 1024 }, .{ .regs = 44, .static = 1168, .bound = 1024 } } }, + .{ .name = "flashnext_batch fn_qsa_select", .usage = .{ .{ .regs = 32, .static = 1168, .bound = 1024 }, .{ .regs = 52, .static = 1168, .bound = 1024 }, .{ .regs = 52, .static = 1168, .bound = 1024 }, .{ .regs = 63, .static = 1168, .bound = 1024 }, .{ .regs = 63, .static = 1168, .bound = 1024 } } }, } ++ @import("flashnext_tile_records.zig").cuda_133 ++ @import("glm_cuda_records.zig").cuda_133 ++ @import("nvcc_133_records.zig").cuda_133 ++ @import("minimax_cuda_records.zig").cuda_133) }, }; pub fn newer(x: limits.Nvcc, y: limits.Nvcc) bool { diff --git a/zig/src/cuda/nvfp4.zig b/zig/src/cuda/nvfp4.zig index 0403c9c5fa..ce9689b0d5 100644 --- a/zig/src/cuda/nvfp4.zig +++ b/zig/src/cuda/nvfp4.zig @@ -66,6 +66,8 @@ pub const Prompt = struct { fp4: Function, fp8: Function, fp4_f32: Function, // FP4 writing the fp32 sums unrounded + bf16: Function, // BF16 on 2 x 4 warps, 3 stages, the lane's K slices (bf16Shared): the lane's bits + bf16_f32: Function, lean: Function, // prompt_lean FP4, bf16 out lean_f32: Function, // and fp32 out sms: usize, // the current device's multiprocessors: prompt_lean's blocks (one a multiprocessor, tiles in turn) @@ -82,6 +84,13 @@ pub const Prompt = struct { std.fmt.comptimePrint(name, .{ 1, 0 }), std.fmt.comptimePrint(name, .{ 0, 1 }), }; + const bf16_threads = 256; + + /// prompt16.cu Tile::SMEM: x and the bf16 words of 64 inputs, three stages. + pub fn bf16Shared() u32 { + return 3 * ((bm * group * 2 + bn * group * 2 + 127) / 128 * 128); + } + const lean_bm = 128; const lean_bn = 256; const lean_threads = 256; @@ -113,6 +122,10 @@ pub const Prompt = struct { try fp8.allowDynamicShared(sharedBytes(.fp8)); const fp4_f32 = try mod.function(symbols[2]); try fp4_f32.allowDynamicShared(sharedBytes(.fp4)); + const bf16 = try mod.function("tf_prompt_bf16"); + try bf16.allowDynamicShared(bf16Shared()); + const bf16_f32 = try mod.function("tf_prompt_bf16_f32"); + try bf16_f32.allowDynamicShared(bf16Shared()); const lean = try mod.function(lean_symbols[0]); try lean.allowDynamicShared(leanSharedBytes()); const lean_f32 = try mod.function(lean_symbols[1]); @@ -121,7 +134,7 @@ pub const Prompt = struct { try d.check(d.api.cuCtxGetDevice(&dev), "cuCtxGetDevice"); var sms: c_int = 0; try d.check(d.api.cuDeviceGetAttribute(&sms, .multiprocessor_count, dev), "cuDeviceGetAttribute"); - return .{ .mod = mod, .fp4 = fp4, .fp8 = fp8, .fp4_f32 = fp4_f32, .lean = lean, .lean_f32 = lean_f32, .sms = @intCast(@max(sms, 1)) }; + return .{ .mod = mod, .fp4 = fp4, .fp8 = fp8, .fp4_f32 = fp4_f32, .bf16 = bf16, .bf16_f32 = bf16_f32, .lean = lean, .lean_f32 = lean_f32, .sms = @intCast(@max(sms, 1)) }; } pub fn unload(p: *Prompt) void { @@ -133,8 +146,9 @@ pub const Prompt = struct { return p.matmulOut(s, x, ldx, m, w, y, .bf16); } - /// matmul() writing y as `out` says: bf16, or fp32 (FP4 only). + /// matmul() writing y as `out` says: bf16, or fp32 (FP4 and BF16 only). pub fn matmulOut(p: Prompt, s: Stream, x: u64, ldx: usize, m: usize, w: qmmf.Weight, y: u64, out: qmmf.Out) !void { + if (w.mode == .bf16) return p.matmulBf16(s, x, ldx, m, w, y, out); return p.matmulTile(s, x, ldx, m, w, y, out, .lean); } @@ -178,6 +192,29 @@ pub const Prompt = struct { .shared = t.shared, }, s, &args); } + + /// A BF16 weight in the lane's word order: the K slices the lane would take (qmmf.splitK), folded in order. + fn matmulBf16(p: Prompt, s: Stream, x: u64, ldx: usize, m: usize, w: qmmf.Weight, y: u64, out: qmmf.Out) !void { + const rows_t = (m + bm - 1) / bm; + const band = @max(1, @min(rows_t, (12 << 20) / (@as(usize, bm) * w.k * 2))); + var args: launch_.Args = .{}; + args.add(x); + args.add(w.codes); + args.add(w.scale); + args.add(y); + args.add(@as(i32, @intCast(m))); + args.add(@as(i32, @intCast(w.n))); + args.add(@as(i32, @intCast(w.k))); + args.add(@as(i32, @intCast(qmmf.splitK(w.n, w.k)))); + args.add(@as(i32, @intCast(w.npad))); + args.add(@as(i32, @intCast(if (m == 1) w.k else ldx))); + args.add(@as(i32, @intCast(band))); + try launch_.launch(if (out == .f32) p.bf16_f32 else p.bf16, .{ + .grid = .{ .x = @intCast(rows_t * ((w.n + bn - 1) / bn)), .y = 1, .z = 1 }, + .block = .{ .x = bf16_threads, .y = 1, .z = 1 }, + .shared = bf16Shared(), + }, s, &args); + } }; test "the prompt tile's shared bytes and instance names" { diff --git a/zig/src/cuda/route_rows.zig b/zig/src/cuda/route_rows.zig index 4aeb28202a..43cacaeb3e 100644 --- a/zig/src/cuda/route_rows.zig +++ b/zig/src/cuda/route_rows.zig @@ -54,8 +54,8 @@ pub const RouteRows = struct { // up to 16 rows, deeper staging (the same sums: the router rows' reads set the time) const few = rows <= 16; if (rows >= 64) { - // 128 rows and 64 experts a block: each router row read once per 128 input rows (the same sums) - try launch_.launch(g.wide_fn, .{ .grid = .{ .x = @intCast((rows + 127) / 128), .y = @intCast((e + 63) / 64), .z = 1 }, .block = .{ .x = 256 }, .shared = wide_shared }, s, &a); + // 128 rows and 64 experts a block (the same sums); x walks the expert blocks: a row tile's run together + try launch_.launch(g.wide_fn, .{ .grid = .{ .x = @intCast((e + 63) / 64), .y = @intCast((rows + 127) / 128), .z = 1 }, .block = .{ .x = 256 }, .shared = wide_shared }, s, &a); return; } try launch_.launch(if (few) g.few_fn else g.logits_fn, .{ .grid = .{ .x = @intCast((rows + 15) / 16), .y = @intCast((e + 31) / 32), .z = 1 }, .block = .{ .x = 128 }, .shared = if (few) few_shared else many_shared }, s, &a); diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 6066b1ea63..78765513bf 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -8,6 +8,9 @@ const ple = @import("cuda_ple.zig"); const wt = @import("cuda_prompt_wt.zig"); const load = @import("cuda_prompt_load.zig"); const parts = @import("cuda_batch_state.zig"); +const late = @import("cuda_batch_late.zig"); +const shared = @import("cuda_batch_side.zig"); +const hc = @import("cuda_batch_hc.zig"); const batch_lanes = @import("cuda_batch_lanes.zig"); const tiles = @import("cuda_tiles.zig"); const tk = @import("cuda_tile_kernels.zig"); @@ -38,7 +41,6 @@ const out_k = wt.out_k; const nch = wt.nch; const ple_layer = wt.ple_layer; const ple_tail_n = wt.ple_tail_n; -const state_n = wt.state_n; const hc_sk = wt.hc_sk; const out_sk = wt.out_sk; const Q4 = wt.Q4; @@ -77,6 +79,9 @@ pub const Span = parts.Span; pub const MoeAdd = parts.MoeAdd; pub const pocket_stride = parts.pocket_stride; pub const Hooks = parts.Hooks; +pub const late_rows = parts.late_rows; +pub const late_store = parts.late_store; +pub const late_row_bytes = parts.late_row_bytes; pub const Kernels = parts.Kernels; pub const Scratch = parts.Scratch; pub const SeqState = parts.SeqState; @@ -118,6 +123,8 @@ pub fn Pass(comptime Tri: type) type { pending_f32: bool = false, /// This layer's out projection leaves part's other slices alone; its writeback adds the zero slice once. zero_slices: bool = false, + /// stageWindow queued the pass's uploads (ids, positions, n-gram rows): it launches kernels alone (graphs). + staged: bool = false, // ---- launches ---------------------------------------------------------------------------------------- @@ -183,11 +190,6 @@ pub fn Pass(comptime Tri: type) type { try self.tri.hcWriteback(h, h, pss, dummy, "*bf16", inj, y, "*fp32", wts, "*fp32", rs, n, dims, streams, 2, 10, 11, 1); } - /// Prompt rows' read-out on flashnext_hcf.cu's tiles, fn_hc_down's and fn_hc_upmix's bits (K slices folded). - fn readoutTiles(self: *Self, hc_w: HcW, r0: usize, n: usize) !void { - return tiles.readout(self.tri, self.k, self.stream, self.s, hc_w, r0, n); - } - /// A read-out over rows [r0, r0 + n) of the streams: mixed and its group sums, and the inject gates. pub fn readout(self: *Self, hc_w: HcW, r0: usize, n: usize) !void { return self.readoutKeep(hc_w, r0, n, .{}); @@ -202,18 +204,13 @@ pub fn Pass(comptime Tri: type) type { const norm = s.norm + r0 * wide * 2; const act = s.act + r0 * low * 2; const xs_act = s.xs_act + r0 * (low / 32) * 4; - if (hc_w.down.native) |down| { - const up = hc_w.up.native orelse return error.UnexpectedTensor; + if (hc_w.down.native != null) { try self.tri.hcNormalize(h, pss, hc_w.scale.ptr, norm, eps, n, dims, streams); - try self.k.native_lane.matmulF32(self.stream, norm, wide, n, down, s.hc_part, s.native_part); - const inj_ptr = if (hc_w.has_inj == 1) inj else act; - try self.tri.hcReduceAct(s.hc_part, act, xs_act, inj_ptr, 1, n, streams, low, hc_w.n_down, hc_w.has_inj); - try self.k.native_lane.matmul(self.stream, act, low, n, up, s.hc_up, s.native_part); - try self.tri.hcDenseMix(s.hc_up, norm, s.mixed + r0 * dims * 2, s.xsm + r0 * (dims / 32) * 4, n, dims, streams); - return; + return hc.readoutNative(self, hc_w, r0, n); } - if (n >= 64 and self.tri.readoutTilesOwned() and std.meta.eql(keep, HcKeep{})) return self.readoutTiles(hc_w, r0, n); + if (n >= 64 and self.tri.readoutTilesOwned() and std.meta.eql(keep, HcKeep{})) return tiles.readout(self.tri, self.k, self.stream, self.s, hc_w, r0, n); try self.copy(keep.h, h, n * wide * 2); + if (hc_sk * n * hc_w.n_down * 4 > s.hc_part_bytes) return error.UnexpectedTensor; try self.tri.hcDown(h, pss, hc_w.scale.ptr, norm, hc_w.down.w.ptr, hc_w.down.s.ptr, hc_w.down.b.ptr, s.dn, s.hc_part, eps, n, hc_w.n_down, wide, dims, dims / 256, streams, hc_sk); if (keep.v != 0) try self.sumSlices(s.hc_part, keep.v, false, n * hc_w.n_down, hc_sk); const inj_ptr = if (hc_w.has_inj == 1) inj else act; @@ -262,11 +259,16 @@ pub fn Pass(comptime Tri: type) type { } // every row's conv, norms and gates, the recurrence, then the gated norm: the one-block chain's bits const fused = self.lanes.len == 0 and n >= 64 and gdn_fused.on(); + if (fused and keep.log != 0) return error.UnexpectedTensor; if (fused) { // the three in one kernel a value head: the same per-row op sequences, no fp32 rows through memory const st = self.seq.state[li].ptr; const in: gdn_fused.In = .{ .proj = s.proj, .conv = self.seq.conv[li].ptr, .cw = lw.cw.ptr, .a_log = lw.a_log.ptr, .dt = lw.dt.ptr, .state = st, .state_out = if (keep.state_out != 0) keep.state_out else st, .norm = lw.norm.ptr }; try gdn_fused.launch(self.k.gdn_fused, self.stream, in, n, s.gout, s.gxs); + } else if (self.lanes.len == 0 and keep.log != 0) { + // pending rows: the window goes into the log after them, the chain runs from before them + try late.front(self.k, self.stream, lw, s.proj, self.seq.conv[li].ptr, n, s.gq, keep); + try late.chain(self.k, self.stream, s.gq, keep, self.seq.state[li].ptr, n, s.ys); } else if (self.lanes.len == 0) { try self.gdnFront(lw, s.proj, self.seq.conv[li].ptr, n, s.gq, s.gk, s.gv, s.ggb); try self.gdnChain(s.gq, s.gk, s.gv, s.ggb, self.seq.state[li].ptr, keep.state_out, n, s.ys); @@ -281,8 +283,13 @@ pub fn Pass(comptime Tri: type) type { const gb = row0 * gdn.nv * 2 * 4; try self.copy(lk.proj, proj, l.rows * gdn.proj_width * 2); try self.copy(lk.hist, l.seq.conv[li].ptr, (gdn.taps - 1) * gdn.conv_dim * 2); - try self.gdnFront(lw, proj, l.seq.conv[li].ptr, l.rows, s.gq + qk, s.gk + qk, s.gv + v, s.ggb + gb); - try self.gdnChain(s.gq + qk, s.gk + qk, s.gv + v, s.ggb + gb, l.seq.state[li].ptr, lk.state_out, l.rows, s.ys + v); + if (lk.log != 0) { + try late.front(self.k, self.stream, lw, proj, l.seq.conv[li].ptr, l.rows, s.gq + qk, lk); + try late.chain(self.k, self.stream, s.gq + qk, lk, l.seq.state[li].ptr, l.rows, s.ys + v); + } else { + try self.gdnFront(lw, proj, l.seq.conv[li].ptr, l.rows, s.gq + qk, s.gk + qk, s.gv + v, s.ggb + gb); + try self.gdnChain(s.gq + qk, s.gk + qk, s.gv + v, s.ggb + gb, l.seq.state[li].ptr, lk.state_out, l.rows, s.ys + v); + } // the lane's conv window moves on now: the gated norm below reads none of it try self.shiftWindow(l.seq.conv[li].ptr, proj, gdn.proj_width, l.rows, gdn.taps - 1, gdn.conv_dim); row0 += l.rows; @@ -440,16 +447,10 @@ pub fn Pass(comptime Tri: type) type { const pt = nvfp4_tile.size(n); try self.k.router.route(self.stream, s.picks, pairs, e, pt, plan); const items = cuda.grouped.maxItems(pairs, e, pt); + const side = try shared.fork(self, native, n); try nvfp4_tile.run(self.k, pt, self.stream, .up, .{ .x = s.mixed, .stride = dims, .slots = slots }, owned.view(), plan, s.moe_act, items, experts_n); try nvfp4_tile.run(self.k, pt, self.stream, .down_f32, .{ .x = s.moe_act, .stride = owned.width }, owned.view(), plan, s.y, items, experts_n); - if (native.rank == 0) { - const glue = self.tri.own orelse return error.NativeGlueRequired; - try self.k.native_lane.matmul(self.stream, s.mixed, dims, n, native.shared[0].native.?, s.shared_gate, s.native_part); - try self.k.native_lane.matmul(self.stream, s.mixed, dims, n, native.shared[1].native.?, s.shared_up, s.native_part); - try glue.sharedSwiGLU(self.stream, s.shared_gate, s.shared_up, s.shared_act, n, moe_w); - try self.k.native_lane.matmulF32(self.stream, s.shared_act, moe_w, n, native.shared[2].native.?, s.shared_out, s.native_part); - try glue.sharedScatter(self.stream, s.shared_out, s.y, n, dims, slots); - } + try shared.join(self, side, native, n); if (native.world != 1) { const collective = self.m.collective orelse return error.TensorParallelUnavailable; try collective.sum(self.stream, s.y, pairs * dims); @@ -509,13 +510,13 @@ pub fn Pass(comptime Tri: type) type { self.pending = null; self.pending_f32 = false; // the gather writes the scratch's pinned rows, so the last pass's uploads must have read them first - if (!self.prepared) try s.ready(); + if (!self.prepared and !self.staged) try s.ready(); const emb = s.hostEmb(n); const xs = s.hostXs(n); var gather: Gather = .{ .ngram = self.m.ngram, .history = self.seq.history[0..], .ids = ids, .emb = emb, .xs = xs, .io = self.io, .lanes = self.lanes }; - var thread: ?std.Thread = if (self.prepared) null else try std.Thread.spawn(.{}, Gather.main, .{&gather}); + var thread: ?std.Thread = if (self.prepared or self.staged) null else try std.Thread.spawn(.{}, Gather.main, .{&gather}); defer if (thread) |th| th.join(); - try self.stage(ids); + if (!self.staged) try self.stage(ids); if (self.m.table_dense) { try self.tri.embedDense(s.ids, self.m.table_w, s.h, dims, streams, n); } else { @@ -529,14 +530,21 @@ pub fn Pass(comptime Tri: type) type { thread = null; if (gather.failed) |err| return err; if (profile()) std.debug.print("batch: {d} rows' n-gram rows gathered on the host in {d:.1} ms\n", .{ n, @as(f64, @floatFromInt(gather.ns)) / 1e6 }); - try self.upload(s.ple_emb, std.mem.sliceAsBytes(emb)); - try self.upload(s.ple_xs, std.mem.sliceAsBytes(xs)); - if (!self.prepared) try s.uploaded(); + if (!self.staged) { + try self.upload(s.ple_emb, std.mem.sliceAsBytes(emb)); + try self.upload(s.ple_xs, std.mem.sliceAsBytes(xs)); + if (!self.prepared) try s.uploaded(); + } if (self.pending != null) try self.writeback(0, n); try self.pleBlock(lw, ids); } - try self.writeback(0, n); - try self.readoutKeep(lw.attn_hc, 0, n, keep.attn_hc); + if (hc.fuses(self, lw.attn_hc, n, keep.attn_hc)) { + try hc.writebackNorm(self, lw.attn_hc, 0, n); + try hc.readoutNative(self, lw.attn_hc, 0, n); + } else { + try self.writeback(0, n); + try self.readoutKeep(lw.attn_hc, 0, n, keep.attn_hc); + } self.zero_slices = zs.usable(self.s.part_zero, keep.mix_branch, self.m.collective != null, lw.out.native != null); if (lw.linear) { try self.gdnBlock(lw, layer, wt.linearIndex(layer), keep); @@ -546,8 +554,13 @@ pub fn Pass(comptime Tri: type) type { if (keep.mix_branch != 0) try self.sumSlices(s.part, keep.mix_branch, true, n * dims, self.out_slices); const sl = zs.slices(self.zero_slices, s.cap, n, self.out_slices); self.zero_slices = false; - try self.tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", sl.stride, n, dims, streams, 4, 1, 1, sl.count); - try self.readoutKeep(lw.mlp_hc, 0, n, keep.mlp_hc); + if (hc.fuses(self, lw.mlp_hc, n, keep.mlp_hc)) { + try self.tri.hcWritebackNorm(s.h, s.h, s.pss, s.part, s.inj, s.dummy, s.dummy, lw.mlp_hc.scale.ptr, s.norm, eps, sl.stride, n, dims, streams, 4, 1, 1, sl.count); + try hc.readoutNative(self, lw.mlp_hc, 0, n); + } else { + try self.tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", sl.stride, n, dims, streams, 4, 1, 1, sl.count); + try self.readoutKeep(lw.mlp_hc, 0, n, keep.mlp_hc); + } try self.moeBlockKeep(lw, keep); try self.moeAdds(layer, keep); self.pending = layer; @@ -578,6 +591,7 @@ pub fn Pass(comptime Tri: type) type { test "a pass's scratch at the dense cap stays near a gigabyte and a half" { _ = @import("cuda_paging_batch_test.zig"); - try std.testing.expect(Scratch.bytes(default_chunk) < 2 << 30); - try std.testing.expect(Scratch.bytes(1) < Scratch.bytes(2)); + try std.testing.expect(Scratch.bytes(default_chunk, .{ .head_rows = default_chunk }) < 2 << 30); + try std.testing.expect(Scratch.bytes(1, .{}) < Scratch.bytes(2, .{})); + try std.testing.expect(Scratch.bytes(4096, .{ .lane_reduce = false }) < Scratch.bytes(4096, .{})); } diff --git a/zig/src/families/flashnext/cuda_batch_hc.zig b/zig/src/families/flashnext/cuda_batch_hc.zig new file mode 100644 index 0000000000..a4970fbd15 --- /dev/null +++ b/zig/src/families/flashnext/cuda_batch_hc.zig @@ -0,0 +1,62 @@ +//! A layer's writeback with its read-out's normalize in one launch, and the native read-out (TF_FLASHNEXT_HC_FUSE). +const std = @import("std"); +const wt = @import("cuda_prompt_wt.zig"); +const parts = @import("cuda_batch_state.zig"); + +const dims = wt.dims; +const streams = wt.streams; +const wide = wt.wide; +const low = wt.low; +const slots = wt.slots; +const eps = wt.eps; +const pss_n = (dims / 256) * streams; + +/// On by default (the same bits either way); TF_FLASHNEXT_HC_FUSE=0: the writeback and the normalize as two launches. +pub fn on() bool { + const v = std.c.getenv("TF_FLASHNEXT_HC_FUSE") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// Whether hc_w's read-out takes its normed rows from the writeback: native, nothing kept, 64 rows or more, our glue. +pub fn fuses(pass: anytype, hc_w: wt.HcW, n: usize, keep: parts.HcKeep) bool { + return n >= 64 and hc_w.down.native != null and std.meta.eql(keep, parts.HcKeep{}) and dims == 2560 and + streams == 4 and pass.tri.writebackOwned() and on(); +} + +/// The pass's pending writeback with hc_w's normalize in the same launch: the normed rows readoutNative reads. +pub fn writebackNorm(pass: anytype, hc_w: wt.HcW, r0: usize, n: usize) !void { + const s = pass.s; + const dummy = s.dummy; + const h = s.h + r0 * wide * 2; + const pss = s.pss + r0 * pss_n * 4; + const inj = s.inj + r0 * streams * 2; + const norm = s.norm + r0 * wide * 2; + const rs = pass.rows * dims; + const scale = hc_w.scale.ptr; + if (pass.pending == null) return pass.tri.hcWritebackNorm(h, h, pss, dummy, dummy, dummy, dummy, scale, norm, eps, rs, n, dims, streams, 0, 1, 1, 1); + pass.pending = null; + if (pass.pending_f32) { + pass.pending_f32 = false; + return pass.tri.hcWritebackNorm(h, h, pss, s.moe_f32 + r0 * dims * 4, inj, dummy, dummy, scale, norm, eps, rs, n, dims, streams, 4, 1, 1, 1); + } + const y = s.y + r0 * slots * dims * 4; + const wts = s.wts + r0 * slots * 4; + try pass.tri.hcWritebackNorm(h, h, pss, dummy, inj, y, wts, scale, norm, eps, rs, n, dims, streams, 2, 10, 11, 1); +} + +/// A native read-out over rows [r0, r0 + n) whose normed rows are in: down, act and gates, up, the mix. +pub fn readoutNative(pass: anytype, hc_w: wt.HcW, r0: usize, n: usize) !void { + const s = pass.s; + const inj = s.inj + r0 * streams * 2; + const norm = s.norm + r0 * wide * 2; + const act = s.act + r0 * low * 2; + const xs_act = s.xs_act + r0 * (low / 32) * 4; + const down = hc_w.down.native orelse return error.UnexpectedTensor; + const up = hc_w.up.native orelse return error.UnexpectedTensor; + if (n * hc_w.n_down * 4 > s.hc_part_bytes) return error.UnexpectedTensor; // one slice within hc_part (hcPartRows) + try pass.k.nativeMatmul(pass.stream, norm, wide, n, down, s.hc_part, s.native_part, .f32); + const inj_ptr = if (hc_w.has_inj == 1) inj else act; + try pass.tri.hcReduceAct(s.hc_part, act, xs_act, inj_ptr, 1, n, streams, low, hc_w.n_down, hc_w.has_inj); + try pass.k.nativeMatmul(pass.stream, act, low, n, up, s.hc_up, s.native_part, .bf16); + try pass.tri.hcDenseMix(s.hc_up, norm, s.mixed + r0 * dims * 2, s.xsm + r0 * (dims / 32) * 4, n, dims, streams); +} diff --git a/zig/src/families/flashnext/cuda_batch_head.zig b/zig/src/families/flashnext/cuda_batch_head.zig index c33fe9bc1a..0057b548df 100644 --- a/zig/src/families/flashnext/cuda_batch_head.zig +++ b/zig/src/families/flashnext/cuda_batch_head.zig @@ -81,7 +81,7 @@ pub fn projectHead(self: anytype, out: u64, n: usize) !void { _ = try self.lane(s.mixed, s.xsm, self.m.head, out, s.dummy, n, 1); return; } - if (self.m.head.n != vocab / 2 or n > s.cap) return error.UnexpectedTensor; + if (self.m.head.n != vocab / 2 or n > s.cap or n > s.head_rows) return error.UnexpectedTensor; const c = self.m.collective orelse return error.TensorParallelUnavailable; _ = try self.lane(s.mixed, s.xsm, self.m.head, s.head_local, s.dummy, n, 1); for (0..n) |r| try c.gather(self.stream, s.head_local + r * vocab, out + r * vocab * 2, vocab / 2); diff --git a/zig/src/families/flashnext/cuda_batch_kernels.zig b/zig/src/families/flashnext/cuda_batch_kernels.zig index 1443cf1398..75c6620d95 100644 --- a/zig/src/families/flashnext/cuda_batch_kernels.zig +++ b/zig/src/families/flashnext/cuda_batch_kernels.zig @@ -19,11 +19,19 @@ pub const Kernels = struct { /// The prompt tiles that keep #952's bits (cuda_tile_kernels.zig). tiles: tk.TileKernels, native_lane: cuda.qmmf.Lane, + /// BF16 weights from prompt_rows rows on prompt16.cu's tile, the lane's bits (TF_FLASHNEXT_PROMPT_BF16=0: lane). + prompt: ?cuda.nvfp4.Prompt, + prompt_rows: usize, nvfp4: cuda.experts.Nvfp4Experts, expert_tile: ?cuda.experts_tile.Tile = null, front: cuda.Function, chain: cuda.Function, chain_to: cuda.Function, + /// TF_FLASHNEXT_KEEP_LATE: the front into a stream's pending-row log, and the chain with the pending rows. + front_log: cuda.Function, + chain_late: cuda.Function, + /// A rolled-back lane's conv windows and n-gram tail in one launch (pending rows). + late_roll: cuda.Function, chain_tpr: cuda.Function, /// front + chain + gated norm in one kernel a value head (prompt passes of 64 rows or more). gdn_fused: cuda.Function, @@ -56,6 +64,9 @@ pub const Kernels = struct { k.moe = try cuda.experts.RowsExperts.resolve(k.rows); k.native_lane = try cuda.qmmf.Lane.load(d, (try d.capability(0)) / 10); errdefer k.native_lane.unload(); + k.prompt_rows = envRows("TF_FLASHNEXT_PROMPT_ROWS", 1024); + k.prompt = if (envOn("TF_FLASHNEXT_PROMPT_BF16")) try cuda.nvfp4.Prompt.load(d) else null; + errdefer if (k.prompt) |*p| p.unload(); var device: c_int = 0; var sms: c_int = 0; try d.check(d.api.cuDeviceGet(&device, 0), "cuDeviceGet"); @@ -72,6 +83,9 @@ pub const Kernels = struct { k.front = try k.batch.function("fn_gdn_front"); k.chain = try k.batch.function("fn_gdn_rows"); k.chain_to = try k.batch.function("fn_gdn_rows_to"); + k.front_log = try k.batch.function("fn_gdn_front_log"); + k.chain_late = try k.batch.function("fn_gdn_rows_late"); + k.late_roll = try k.batch.function("fn_late_roll"); k.chain_tpr = try k.batch.function("fn_gdn_chain_tpr"); try k.chain_tpr.allowDynamicShared(chain_tpr_smem); k.gdn_fused = try k.batch.function("fn_gdn_fused"); @@ -80,9 +94,16 @@ pub const Kernels = struct { return k; } + /// A native projection over m rows: a BF16 weight takes the prompt tile from prompt_rows rows, the rest the lane. + pub fn nativeMatmul(k: *const Kernels, s: cuda.Stream, x: u64, ldx: usize, m: usize, w: cuda.qmmf.Weight, y: u64, part: u64, out: cuda.qmmf.Out) !void { + if (k.prompt) |p| if (m >= k.prompt_rows and w.mode == .bf16) return p.matmulOut(s, x, ldx, m, w, y, out); + return k.native_lane.matmulOut(s, x, ldx, m, w, y, out, part); + } + pub fn deinit(k: *Kernels) void { if (k.expert_tile) |*tile| tile.unload(); k.nvfp4.unload(); + if (k.prompt) |*p| p.unload(); k.native_lane.unload(); k.batch.unload(); k.rows.unload(); @@ -91,3 +112,15 @@ pub const Kernels = struct { k.qmm.unload(); } }; + +/// A switch on unless its variable is "0". +fn envOn(name: [*:0]const u8) bool { + const v = std.c.getenv(name) orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// A row count from the environment, or `default`. +fn envRows(name: [*:0]const u8, default: usize) usize { + const v = std.c.getenv(name) orelse return default; + return std.fmt.parseInt(usize, std.mem.span(v), 10) catch default; +} diff --git a/zig/src/families/flashnext/cuda_batch_lanes.zig b/zig/src/families/flashnext/cuda_batch_lanes.zig index 7970b02e3b..1b1eeef4b4 100644 --- a/zig/src/families/flashnext/cuda_batch_lanes.zig +++ b/zig/src/families/flashnext/cuda_batch_lanes.zig @@ -27,6 +27,8 @@ pub fn rollbackOnly(k: *const Keep) bool { rest.proj = 0; rest.hist = 0; rest.state_out = 0; + rest.log = 0; + rest.log_ctl = 0; if (!std.meta.eql(rest, LayerKeep{})) return false; } return true; diff --git a/zig/src/families/flashnext/cuda_batch_late.zig b/zig/src/families/flashnext/cuda_batch_late.zig new file mode 100644 index 0000000000..dc557e733f --- /dev/null +++ b/zig/src/families/flashnext/cuda_batch_late.zig @@ -0,0 +1,40 @@ +//! Pending DeltaNet rows (TF_FLASHNEXT_KEEP_LATE): a window's rows into its stream's log, its chain after them. +const cuda = @import("cuda"); +const gdn = @import("cuda_gdn.zig"); +const parts = @import("cuda_batch_state.zig"); +const Layer = @import("cuda_prompt_load.zig").Layer; + +fn cint(v: usize) c_int { + return @intCast(v); +} + +/// fn_gdn_front_log: q as fn_gdn_front's, the rows' k, v and gates into `keep`'s log after its pending rows. +pub fn front(k: anytype, stream: cuda.Stream, lw: *const Layer, proj: u64, conv: u64, rows: usize, gq: u64, keep: parts.LayerKeep) !void { + var f: cuda.Args = .{}; + f.add(proj); + f.add(conv); + f.add(lw.cw.ptr); + f.add(lw.a_log.ptr); + f.add(lw.dt.ptr); + f.add(gq); + f.add(keep.log); + f.add(keep.log_ctl); + f.add(cint(parts.late_rows)); + try cuda.launch.launch(k.front_log, .{ .grid = .{ .x = @intCast(rows), .y = gdn.nk + gdn.nv }, .block = .{ .x = 256 } }, stream, &f); +} + +/// fn_gdn_rows_late: the pending rows on `state`, then the window's rows into ys, the end to keep.state_out. +pub fn chain(k: anytype, stream: cuda.Stream, gq: u64, keep: parts.LayerKeep, state: u64, rows: usize, ys: u64) !void { + if (keep.state_out == 0) return error.UnexpectedTensor; + var a: cuda.Args = .{}; + a.add(gq); + a.add(keep.log); + a.add(keep.log_ctl); + a.add(cint(parts.late_rows)); + a.add(cint(parts.late_store)); + a.add(state); + a.add(keep.state_out); + a.add(cint(rows)); + a.add(ys); + try cuda.launch.launch(k.chain_late, .{ .grid = .{ .x = gdn.nv, .y = gdn.dv / 16 }, .block = .{ .x = 128 } }, stream, &a); +} diff --git a/zig/src/families/flashnext/cuda_batch_side.zig b/zig/src/families/flashnext/cuda_batch_side.zig new file mode 100644 index 0000000000..c15df1e439 --- /dev/null +++ b/zig/src/families/flashnext/cuda_batch_side.zig @@ -0,0 +1,63 @@ +//! A decode or verify pass's shared expert on its own stream beside the routed experts (TF_FLASHNEXT_SHARED_SIDE). +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); + +const dims = wt.dims; +const slots = wt.slots; +const moe_w = wt.moe_w; + +/// On by default: its slot rows of y are apart from the routed experts', so the bits match; =0 runs it after them. +pub fn on() bool { + const v = std.c.getenv("TF_FLASHNEXT_SHARED_SIDE") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// The shared expert's stream and the two points it forks and joins at. +pub const Side = struct { + stream: cuda.Stream, + fork: cuda.Event, + join: cuda.Event, + + pub fn init(d: *const cuda.Driver) !Side { + var st = try cuda.Stream.init(d, true); + errdefer st.deinit(); + var f = try cuda.Event.init(d, false); + errdefer f.deinit(); + return .{ .stream = st, .fork = f, .join = try cuda.Event.init(d, false) }; + } + + pub fn deinit(x: *Side) void { + x.join.deinit(); + x.fork.deinit(); + x.stream.deinit(); + } +}; + +/// The shared expert over the pass's `n` rows on `st`: gate and up from the MoE input, SwiGLU, down, into y's slot. +pub fn run(pass: anytype, st: cuda.Stream, native: anytype, n: usize) !void { + const s = pass.s; + const glue = pass.tri.own orelse return error.NativeGlueRequired; + try pass.k.nativeMatmul(st, s.mixed, dims, n, native.shared[0].native.?, s.shared_gate, s.native_part, .bf16); + try pass.k.nativeMatmul(st, s.mixed, dims, n, native.shared[1].native.?, s.shared_up, s.native_part, .bf16); + try glue.sharedSwiGLU(st, s.shared_gate, s.shared_up, s.shared_act, n, moe_w); + try pass.k.nativeMatmul(st, s.shared_act, moe_w, n, native.shared[2].native.?, s.shared_out, s.native_part, .f32); + try glue.sharedScatter(st, s.shared_out, s.y, n, dims, slots); +} + +/// After y is cleared: the shared expert forked onto the side stream (one GPU, not a prompt fill), or null. +pub fn fork(pass: anytype, native: anytype, n: usize) !?Side { + const x = pass.s.side orelse return null; + if (native.rank != 0 or native.world != 1 or pass.hooks.prompt) return null; + try x.fork.record(pass.stream); + try x.stream.wait(x.fork); + try run(pass, x.stream, native, n); + try x.join.record(x.stream); + return x; +} + +/// After the routed experts: the forked shared expert joined, or run here on rank 0 as before. +pub fn join(pass: anytype, side: ?Side, native: anytype, n: usize) !void { + if (side) |x| return pass.stream.wait(x.join); + if (native.rank == 0) try run(pass, pass.stream, native, n); +} diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index 98faec84be..ae4c2cb6ec 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -7,6 +7,7 @@ const ple = @import("cuda_ple.zig"); const wt = @import("cuda_prompt_wt.zig"); const tk = @import("cuda_tile_kernels.zig"); const load = @import("cuda_prompt_load.zig"); +const shared = @import("cuda_batch_side.zig"); const layers_n = wt.layers_n; const linear_n = wt.linear_n; const dims = wt.dims; @@ -87,6 +88,9 @@ pub const LayerKeep = struct { hist: u64 = 0, /// The chain's end state goes here, the sequence's own left as before the chunk; 0: in place. state_out: u64 = 0, + /// This layer's pending-row log and its control word (fn_late in flashnext_batch.cu); 0: none. + log: u64 = 0, + log_ctl: u64 = 0, attn_hc: HcKeep = .{}, mlp_hc: HcKeep = .{}, /// The mixer's and the MoE's branches [rows, dims] bf16 as the write-backs add them (the MoE's with its pocket). @@ -138,8 +142,33 @@ pub const Hooks = struct { streams: u64 = 0, /// A training pass's adds to each layer's MoE output, in place of the model's own (Model.pocket). moe_add: ?MoeAdd = null, + /// The rows are prompt tokens (a fill), not a decode or verify window. + prompt: bool = false, + /// The last row's scores also in seq.logits (a sync, a vocab-row copy); drafted windows' draws read device rows. + host_logits: bool = true, }; +/// Rows a layer's pending-row log holds, and the pending count past which a chain stores its state in place. +pub const late_rows: usize = 48; +pub const late_store: usize = 12; +/// One log row of a layer: k [nk, dk], v [nv, dv], decay and beta [nv, 2], fp32. +pub const late_row_bytes: usize = (gdn.nk * gdn.dk + gdn.nv * gdn.dv + gdn.nv * 2) * 4; + +/// A rolled-back lane's windows for fn_late_roll (fn_roll::Job, by value): its conv windows and the n-gram tail. +pub const RollJob = extern struct { win: u64 = 0, hist: u64 = 0, rows: u64 = 0, stride: c_longlong = 0, taps: c_int = 0, c: c_int = 0 }; +pub const roll_jobs_max = 40; +pub const RollJobs = extern struct { j: [roll_jobs_max]RollJob = @splat(.{}), n: c_int = 0 }; + +comptime { + std.debug.assert(@sizeOf(RollJob) == 40 and @sizeOf(RollJobs) == 1608); +} + +/// TF_FLASHNEXT_VERIFY_HOST_LOGITS=1: drafted windows also copy their last row's scores to the host, as before. +pub fn verifyHostLogits() bool { + const v = std.c.getenv("TF_FLASHNEXT_VERIFY_HOST_LOGITS") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); +} + /// The kernels the pass launches besides the captured Triton ones, loaded once. pub const Kernels = @import("cuda_batch_kernels.zig").Kernels; @@ -191,9 +220,17 @@ pub fn Staging(comptime Q: type) type { }; } +const sizes = @import("cuda_scratch_sizes.zig"); +pub const head_rows_one = sizes.head_rows_one; +pub const Sizes = sizes.Sizes; +const hcPartRows = sizes.hcPartRows; +const partRows = sizes.partRows; + /// Device scratch for chunks of up to `cap` rows. pub const Scratch = struct { cap: usize, + /// The shared expert's stream (cuda_batch_side.zig); null: it runs after the routed experts on the pass's stream. + side: ?shared.Side = null, bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, ids: u64 = 0, pos: u64 = 0, @@ -260,7 +297,11 @@ pub const Scratch = struct { ple_pss: u64 = 0, ple_nrow: u64 = 0, logits: u64 = 0, + /// A split head's local columns before the gather, or the grouped draft head's scores: head_rows * vocab bytes. head_local: u64 = 0, + head_rows: usize = 0, + /// hc_part's bytes (hcPartRows): a read-out that would write more refuses (Pass.readoutKeep / readoutNative). + hc_part_bytes: usize = 0, inv: u64 = 0, dummy: u64 = 0, /// Pinned host copies of what a pass uploads (ids, positions, n-gram rows and sums), owned until the copies land. @@ -313,16 +354,17 @@ pub const Scratch = struct { } /// Bytes `alloc(cap)` takes on the device. - pub fn bytes(cap: usize) usize { + pub fn bytes(cap: usize, z: Sizes) usize { + const head_rows = z.head_rows; const pairs = cap * slots; const nblk = cdiv(pairs, plan_small); const items = cdiv(pairs, tile) + experts_n + 1; var n: usize = 0; n += cap * 4 + (cap + 1) * 16; - n += cap * streams * 4 + cap * wide * 2 * 2 + cap * pss_n * 4 + cap * (low + 4) * 2 + hc_sk * cap * (low + 4) * 4; - n += cap * wide * 2 + 8 * cap * dims * 4; + n += cap * streams * 4 + cap * wide * 2 * 2 + cap * pss_n * 4 + cap * (low + 4) * 2 + hcPartRows(cap) * (low + 4) * 4; + n += cap * wide * 2 + (if (z.lane_reduce) 8 * cap * dims * 4 else 0); n += cap * low * 2 + cap * (low / 32) * 4 + cap * streams * 2 + cap * dims * 2 + cap * (dims / 32) * 4; - n += cap * proj_row * 2 + out_sk * cap * dims * 4 + cap * gdn.value_dim * (2 + 4 + 4) + cap * (gdn.value_dim / 32) * 4; + n += cap * proj_row * 2 + partRows(cap) * dims * 4 + cap * gdn.value_dim * (2 + 4 + 4) + cap * (gdn.value_dim / 32) * 4; n += cap * (2 * gdn.nk * gdn.dk + 2 * gdn.nv) * 4; n += cap * q_heads * head_dim * 2 + cap * index_heads * index_dim * 2; n += att_rows * nch * q_heads * (head_dim + 2) * 4 + cap * out_k * 2 * 2 + cap * (out_k / 32) * 4; @@ -331,13 +373,14 @@ pub const Scratch = struct { n += cap * moe_w * 6 + cap * dims * 4; n += pairs * moe_w * 2 + pairs * dims * 4 + cap * dims * 2; n += cap * dims * 2 + cap * (dims / 32) * 4 + cap * wide * 2 + cap * dims * 2 + 4 * cap * dims * 4; - n += cap * wide * 2 * 2 + cap * streams * 4 + vocab * 2 + cap * vocab + half * 4 + 4096; + n += cap * wide * 2 * 2 + cap * streams * 4 + vocab * 2 + head_rows * vocab + half * 4 + 4096; n += cap * dims * 4 + cap * pocket_stride * 4; return n; } - pub fn alloc(gpa: std.mem.Allocator, d: *const cuda.Driver, stream: cuda.Stream, cap: usize) !Scratch { - var s: Scratch = .{ .cap = cap }; + pub fn alloc(gpa: std.mem.Allocator, d: *const cuda.Driver, stream: cuda.Stream, cap: usize, z: Sizes) !Scratch { + const head_rows = z.head_rows; + var s: Scratch = .{ .cap = cap, .head_rows = head_rows }; errdefer { cuda.completion.drain(d, stream, "Flash Next scratch work did not drain before rollback"); s.free(gpa); @@ -353,16 +396,17 @@ pub const Scratch = struct { s.norm = try s.add(gpa, d, cap * wide * 2); s.rinv = try s.add(gpa, d, cap * streams * 4); s.dn = try s.add(gpa, d, cap * (low + 4) * 2); - s.hc_part = try s.add(gpa, d, hc_sk * cap * (low + 4) * 4); + s.hc_part_bytes = hcPartRows(cap) * (low + 4) * 4; + s.hc_part = try s.add(gpa, d, s.hc_part_bytes); s.hc_up = try s.add(gpa, d, cap * wide * 2); - s.native_part = try s.add(gpa, d, 8 * cap * dims * 4); + if (z.lane_reduce) s.native_part = try s.add(gpa, d, 8 * cap * dims * 4); s.act = try s.add(gpa, d, cap * low * 2); s.xs_act = try s.add(gpa, d, cap * (low / 32) * 4); s.inj = try s.add(gpa, d, cap * streams * 2); s.mixed = try s.add(gpa, d, cap * dims * 2); s.xsm = try s.add(gpa, d, cap * (dims / 32) * 4); s.proj = try s.add(gpa, d, cap * proj_row * 2); - s.part = try s.add(gpa, d, out_sk * cap * dims * 4); + s.part = try s.add(gpa, d, partRows(cap) * dims * 4); try s.bufs.items[s.bufs.items.len - 1].fill8(0, stream.handle); s.part_zero = true; s.gout = try s.add(gpa, d, cap * gdn.value_dim * 2); @@ -409,7 +453,7 @@ pub const Scratch = struct { s.ple_pss = try s.add(gpa, d, cap * streams * 4); s.ple_nrow = try s.add(gpa, d, cap * wide * 2); s.logits = try s.add(gpa, d, vocab * 2); - s.head_local = try s.add(gpa, d, cap * vocab); + s.head_local = try s.add(gpa, d, head_rows * vocab); var inv: [half]f32 = undefined; for (&inv, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / @as(f64, half))); s.inv = try s.add(gpa, d, half * 4); @@ -419,6 +463,7 @@ pub const Scratch = struct { try stream.synchronize(); s.host = try cuda.HostBuffer.alloc(d, hostAt(cap)[4]); s.staging = .{ .q = .{ .stream = stream, .event = try cuda.Event.init(d, false) } }; + if (shared.on()) s.side = try shared.Side.init(d); // last: a refused allocation above never makes a stream return s; } @@ -429,6 +474,7 @@ pub const Scratch = struct { g.q.event.deinit(); } if (s.host) |*h| h.free(); + if (s.side) |*x| x.deinit(); for (s.bufs.items) |*b| b.free(); s.bufs.deinit(gpa); s.* = undefined; @@ -437,6 +483,8 @@ pub const Scratch = struct { /// A sequence's state the pass reads and leaves as the row-by-row path would. pub const SeqState = struct { + /// The sequence's generation (Seq.gen), for graph keys; 0 for the draft head's own views. + gen: u64 = 0, media: ?@import("cuda_media.zig").View = null, conv: []cuda.DeviceBuffer, state: []cuda.DeviceBuffer, diff --git a/zig/src/families/flashnext/cuda_glue.zig b/zig/src/families/flashnext/cuda_glue.zig index 14a17dc451..5bd83a5f3f 100644 --- a/zig/src/families/flashnext/cuda_glue.zig +++ b/zig/src/families/flashnext/cuda_glue.zig @@ -34,6 +34,10 @@ const Fns = struct { wb0: cuda.Function, wb_moe: cuda.Function, wb_slices: cuda.Function, + // each writeback with fn_hc_normalize in one block a row (2560 columns, 4 streams): the same bits + wbn0: cuda.Function, + wbn_moe: cuda.Function, + wbn_slices: cuda.Function, normalize: cuda.Function, normalize_cols: cuda.Function, // fn_hc_normalize's bits over 256-column slices (dims a multiple of 256) dense_mix8: cuda.Function, // fn_hc_dense_mix's bits, a warp a 32-column group, eight a block @@ -121,6 +125,9 @@ pub const Glue = struct { .wb0 = try g.hc.function("fn_hc_writeback0"), .wb_moe = try g.hc.function("fn_hc_writeback_moe"), .wb_slices = try g.hc.function("fn_hc_writeback_slices"), + .wbn0 = try g.hc.function("fn_hc_writeback_norm0"), + .wbn_moe = try g.hc.function("fn_hc_writeback_norm_moe"), + .wbn_slices = try g.hc.function("fn_hc_writeback_norm_slices"), .normalize = try g.hc.function("fn_hc_normalize"), .normalize_cols = try g.hc.function("fn_hc_normalize_cols"), .dense_mix8 = try g.hc.function("fn_hc_dense_mix8"), @@ -389,6 +396,35 @@ pub const Glue = struct { try go(g.f.down, s, .{ cdiv(rows, 16), cdiv(n, 64), sk }, 128, shared, &a); } + /// hcWriteback then hcNormalize with `scale` into `norm`, one block a row: Flash Next's 2560 columns and 4 streams. + pub fn hcWritebackNorm(g: *const Glue, s: cuda.Stream, h: u64, hout: u64, pss: u64, branch: u64, inj: u64, y: u64, wts: u64, scale: u64, norm: u64, eps: f32, rs: usize, rows: usize, dims: usize, streams: usize, mode: usize, topk: usize, slots: usize, world: usize) !void { + if (dims != 2560 or streams != 4) return error.RowShape; + var a: cuda.Args = .{}; + const grid: [3]usize = .{ rows, 1, 1 }; + switch (mode) { + 0 => { + ptrs(&a, &.{ h, pss, scale, norm }); + a.add(eps); + try go(g.f.wbn0, s, grid, 256, 0, &a); + }, + 2 => { + ptrs(&a, &.{ h, hout, pss, inj, y, wts, scale, norm }); + a.add(eps); + ints(&a, &.{ topk, slots }); + try go(g.f.wbn_moe, s, grid, 256, 0, &a); + }, + 4 => { + ptrs(&a, &.{ h, hout, pss, branch, inj }); + a.add(@as(i64, @intCast(rs))); + ptrs(&a, &.{ scale, norm }); + a.add(eps); + ints(&a, &.{world}); + try go(g.f.wbn_slices, s, grid, 256, 0, &a); + }, + else => return error.RowShape, + } + } + /// Mode 0: the streams' sums; 2: the MoE slots (fp32 y and weights) written back; 4: `world` fp32 K slices. pub fn hcWriteback(g: *const Glue, s: cuda.Stream, h: u64, hout: u64, pss: u64, branch: u64, inj: u64, y: u64, wts: u64, rs: usize, rows: usize, dims: usize, streams: usize, mode: usize, topk: usize, slots: usize, world: usize) !void { if (dims % 256 != 0 or streams > 8) return error.RowShape; diff --git a/zig/src/families/flashnext/cuda_images_test.zig b/zig/src/families/flashnext/cuda_images_test.zig index a4a5b4a823..b4edd77ff1 100644 --- a/zig/src/families/flashnext/cuda_images_test.zig +++ b/zig/src/families/flashnext/cuda_images_test.zig @@ -62,6 +62,7 @@ fn run(comptime Cuda: type, comptime open: anytype, comptime drop: anytype, draf self.head = if (drafting) &head else null; self.vision = &v; self.lanes = .empty; + self.parts = .{}; defer { var it = self.lanes.valueIterator(); while (it.next()) |l| drop(&self, l); @@ -149,6 +150,7 @@ fn release(comptime Cuda: type, comptime open: anytype) !void { self.head = null; self.vision = null; self.lanes = .empty; + self.parts = .{}; defer self.lanes.deinit(a); var s = try lanes.Stream.init(a, .{ .id = "release", .prompt = &.{ 1, 2, 3 }, .max_new = 1 }); defer s.deinit(a); diff --git a/zig/src/families/flashnext/cuda_kv_lanes.zig b/zig/src/families/flashnext/cuda_kv_lanes.zig index 296a396be0..0d25dacad9 100644 --- a/zig/src/families/flashnext/cuda_kv_lanes.zig +++ b/zig/src/families/flashnext/cuda_kv_lanes.zig @@ -200,6 +200,8 @@ pub const Pool = struct { x.run = run; try p.mapStates(&x.seq); try p.e.clearSeq(&x.seq, cap); + // a new generation each open: every lane's caches share the arena's address, so the graph keys need it + x.seq.gen = @import("cuda_window_graph.zig").nextGen(); x.used = true; return x.seq; } diff --git a/zig/src/families/flashnext/cuda_lane_parts.zig b/zig/src/families/flashnext/cuda_lane_parts.zig new file mode 100644 index 0000000000..408d9309d8 --- /dev/null +++ b/zig/src/families/flashnext/cuda_lane_parts.zig @@ -0,0 +1,33 @@ +//! A contiguous lane's fixed-size buffers reused across streams (TF_FLASHNEXT_LANE_PARTS=0: allocated each stream). +const std = @import("std"); +const cuda = @import("cuda"); + +/// Fixed-size lane buffers reused across streams: at most `kept` idle blocks of sequence and draft parts each. +pub const Parts = struct { + on: bool = false, + seqs: std.ArrayList(cuda.DeviceBuffer) = .empty, + drafts: std.ArrayList(cuda.DeviceBuffer) = .empty, + const kept = 4; + + pub fn take(list: *std.ArrayList(cuda.DeviceBuffer), d: *cuda.Driver, len: usize) !cuda.DeviceBuffer { + while (list.pop()) |b| { + if (b.len >= len) return b; + var x = b; + x.free(); + } + return cuda.DeviceBuffer.alloc(d, len); + } + + pub fn give(list: *std.ArrayList(cuda.DeviceBuffer), gpa: std.mem.Allocator, b: cuda.DeviceBuffer) void { + var x = b; + if (list.items.len >= kept) return x.free(); + list.append(gpa, b) catch x.free(); + } + + pub fn deinit(p: *Parts, gpa: std.mem.Allocator) void { + for (p.seqs.items) |*b| b.free(); + for (p.drafts.items) |*b| b.free(); + p.seqs.deinit(gpa); + p.drafts.deinit(gpa); + } +}; diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 69ac0f595d..e9ae3406cf 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -34,11 +34,11 @@ pub fn Namespace(comptime Tri: type) type { pub const max_segments: u32 = 1; pub const prompt_rows: u32 = @intCast(prompt.batch.default_chunk); /// The prompt cache's copies of a stream's state (cuda_snap.zig), called with `Loaded.target`. - pub const snap_bytes = snapBytes; + pub const snap_bytes = snaps.bytes; pub const snap_repeat_min: u32 = 16; - pub const snap_save = snapSave; - pub const snap_restore = snapRestore; - pub const snap_drop = snapDrop; + pub const snap_save = snaps.save; + pub const snap_restore = snaps.restore; + pub const snap_drop = snaps.drop; pub const kept_disk = @import("cuda_disk.zig").Kept(@This()); // their disk image, as qwen3_5 declares its own /// TF_FLASHNEXT_KV_PAGES: the pool's lanes, the pages of a request's rows and a drafted window (kvl). pub const window_rows = kvl.window_rows; @@ -131,9 +131,16 @@ pub fn Namespace(comptime Tri: type) type { filled: ?usize = 0, /// Where its caches came from: the pool's lanes, or its own (the one a pool leaves beside it). from: kvl.From = .none, + /// Its sequence's and draft's fixed-size buffers' blocks (Parts), given back when it drops. + part_seq: ?cuda.DeviceBuffer = null, + part_draft: ?cuda.DeviceBuffer = null, }; + const Parts = @import("cuda_lane_parts.zig").Parts; + const seq_parts = @import("cuda_seq_parts.zig"); + const Ops = mtp.Ops(Tri); + const snaps = @import("cuda_lanes_snap.zig").Snaps(Cuda); pub const Cuda = struct { gpa: std.mem.Allocator, @@ -143,12 +150,14 @@ pub fn Namespace(comptime Tri: type) type { ranks: ?*ranks.Session = null, head: ?*mtp.Head = null, depth: usize = mtp.default_depth, + depth_greedy: usize = mtp.default_depth_greedy, confidence: f64 = mtp.default_confidence, /// A captured Triton set (TENSORFOLD_CUDA_KERNELS) for comparison, else null: our own glue. set: ?cuda.aot.Set = null, glue: Glue, stream: cuda.Stream, lanes: std.AutoHashMapUnmanaged(*lanes.Stream, Lane) = .empty, + parts: Parts = .{}, /// A round's scores, every row [round_streams, vocab] bf16. logits: cuda.DeviceBuffer, drawn: [ring]u32 = undefined, @@ -203,7 +212,13 @@ pub fn Namespace(comptime Tri: type) type { .first_masked = firstMaskedFn, }; const ranked_table = @import("cuda_score.zig").withoutScore(scoring_table); - return .{ .ptr = self, .proposal_head = if (self.head) |h| &h.service else null, .vtable = if (self.ranks == null) &scoring_table else &ranked_table }; + return .{ .ptr = self, .proposal_head = if (self.head) |h| &h.service else null, .first_before_draft = firstEarly(), .vtable = if (self.ranks == null) &scoring_table else &ranked_table }; + } + + /// A drafted stream's first token is sent before its first drafts are asked (core/lanes/open_first.zig). + fn firstEarly() bool { + const v = std.c.getenv("TF_FLASHNEXT_FIRST_EARLY") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); } /// With the head: chains of up to 16 rows cut by its confidence, copies of up to 32; else one-row rounds. @@ -216,7 +231,7 @@ pub fn Namespace(comptime Tri: type) type { .mtp = drafting, .speculate = drafting, .speculate_early = false, - .drafts = if (drafting) @intCast(self.depth) else 0, + .drafts = if (drafting) @intCast(@max(self.depth, self.depth_greedy)) else 0, .draft_probabilities = drafting, .streams_exact = true, .hidden_rows = true, @@ -270,10 +285,14 @@ pub fn Namespace(comptime Tri: type) type { // a window's rows past the reply's last token (its drafts) still need cache slots const cap = @min(ids.len + s.max_new + (if (drafting) mtp.max_rows else 0), prompt.max_tokens); const from = self.kv.choose(); - gop.value_ptr.* = .{ .seq = (if (self.ranks) |r| r.create(cap) else if (from == .pool) self.kv.pool.?.open(cap) else self.engine.newSeq(cap)) catch |err| { + // a contiguous lane's fixed-size buffers from the parts kept from an earlier stream + var part_seq: ?cuda.DeviceBuffer = null; + if (self.parts.on and self.ranks == null and from != .pool) part_seq = Parts.take(&self.parts.seqs, self.engine.driver, seq_parts.seqPartsBytes()) catch null; + gop.value_ptr.* = .{ .seq = (if (self.ranks) |r| r.create(cap) else if (from == .pool) self.kv.pool.?.open(cap) else if (part_seq) |b| seq_parts.newSeqIn(self.engine, cap, b) else self.engine.newSeq(cap)) catch |err| { + if (part_seq) |b| Parts.give(&self.parts.seqs, self.gpa, b); self.lanes.removeByPtr(gop.key_ptr); return err; - }, .from = from }; + }, .from = from, .part_seq = part_seq }; self.kv.opened(from); const l = gop.value_ptr; l.seq.untaught = s.untaught; @@ -281,7 +300,10 @@ pub fn Namespace(comptime Tri: type) type { dropLane(self, l); self.lanes.removeByPtr(gop.key_ptr); } - if (drafting) l.draft = if (from == .pool) try self.kv.pool.?.draft(&l.seq, cap) else try mtp.Draft.init(self.gpa, self.engine.driver, self.stream, cap, self.head.?.ids.len, self.engine.chunk); + if (drafting and from != .pool and l.part_seq != null) { + l.part_draft = try Parts.take(&self.parts.drafts, self.engine.driver, mtp.Draft.partsBytes(self.head.?.ids.len, self.engine.chunk)); + l.draft = try mtp.Draft.initParts(self.gpa, self.engine.driver, self.stream, cap, self.head.?.ids.len, self.engine.chunk, l.part_draft.?); + } else if (drafting) l.draft = if (from == .pool) try self.kv.pool.?.draft(&l.seq, cap) else try mtp.Draft.init(self.gpa, self.engine.driver, self.stream, cap, self.head.?.ids.len, self.engine.chunk); if (s.images.len != 0) { const v = self.vision orelse return error.VisionNotEnabled; l.seq.media = try v.state(ids, s.images); @@ -324,8 +346,8 @@ pub fn Namespace(comptime Tri: type) type { if (l.draft) |*d| { try self.ops().prefillChunk(&l.seq, d, ids, at, stop); } else if (self.ranks) |r| { - try r.forward(Tri, self.tri(), &l.seq, ids[at..stop], .{ .logits = if (last) .last else .none }); - } else try self.engine.forward(Tri, self.tri(), &l.seq, ids[at..stop], .{ .logits = if (last) .last else .none }); + try r.forward(Tri, self.tri(), &l.seq, ids[at..stop], .{ .logits = if (last) .last else .none, .prompt = true }); + } else try self.engine.forward(Tri, self.tri(), &l.seq, ids[at..stop], .{ .logits = if (last) .last else .none, .prompt = true }); l.filled = if (last) null else stop; if (!last and std.mem.indexOfScalar(u32, s.reuse.marks, @intCast(stop)) != null) { if (s.reuse.hook) |k| k.at(k.ptr, s, @intCast(stop)); @@ -357,6 +379,7 @@ pub fn Namespace(comptime Tri: type) type { var it = own.lanes.valueIterator(); while (it.next()) |l| dropLane(own, l); own.lanes.deinit(own.gpa); + own.parts.deinit(own.gpa); own.engine.round_graph.clear(); own.logits.free(); if (own.head) |h| h.deinit(); @@ -374,6 +397,11 @@ pub fn Namespace(comptime Tri: type) type { cuda.completion.drain(self.engine.driver, self.stream, "flashnext: a failed or closed lane did not drain"); if (l.draft) |*d| d.deinit(); if (self.ranks) |r| r.drop(&l.seq) else if (!self.kv.close(l.from, &l.seq)) l.seq.deinit(self.gpa); + // the blocks after both: a state / spare swap may have left views of each in the other + if (l.part_draft) |b| Parts.give(&self.parts.drafts, self.gpa, b); + if (l.part_seq) |b| Parts.give(&self.parts.seqs, self.gpa, b); + l.part_draft = null; + l.part_seq = null; } fn of(ptr: *anyopaque) *Cuda { @@ -503,7 +531,7 @@ pub fn Namespace(comptime Tri: type) type { if (r.position != l.seq.fed + 1) return error.PositionMismatch; // the window's rows past the cache's room cannot run: fewer drafts there const room = l.seq.cap - l.seq.fed - 1; - const count = @min(@min(r.depth, self.depth), room); + const count = @min(@min(r.depth, mtp.depthFor(r.stream.sampling, self.depth, self.depth_greedy)), room); if (r.rows) |path| { for (path, 0..) |row, i| if (row != i) return error.TreesNotBuilt; if (r.follow.len < path.len) return error.PositionMismatch; @@ -523,12 +551,16 @@ pub fn Namespace(comptime Tri: type) type { return .{ .count = @intCast(l.chain.n), .tokens = try gpa.dupe(u32, l.chain.slice()) }; } - /// Each held draft's temperature-1 probability under the head. + /// Each held draft's chain chance under the head (windows.nodeChances): p1 * ... * pj at temperature 1. fn probabilitiesFn(ptr: *anyopaque, s: *lanes.Stream, out: []f64) anyerror!bool { const self = of(ptr); const l = try self.lane(s); if (l.draft == null or out.len > l.chain.n) return false; - @memcpy(out, l.chain.probs[0..out.len]); + var chance: f64 = 1; + for (out, l.chain.probs[0..out.len]) |*o, p| { + chance *= p; + o.* = chance; + } return true; } @@ -544,7 +576,7 @@ pub fn Namespace(comptime Tri: type) type { test { _ = @import("cuda_images_test.zig").Cases(Cuda, Cuda.open, dropLane); - _ = @import("cuda_lanes_restore_test.zig").Cases(Cuda, Cuda.open, snapRestore, dropLane); + _ = @import("cuda_lanes_restore_test.zig").Cases(Cuda, Cuda.open, snaps.restore, dropLane); _ = @import("cuda_lanes_admission_test.zig").Cases(Cuda, Cuda.open, dropLane); _ = @import("cuda_score.zig").CapabilityCases(Cuda); } @@ -560,41 +592,5 @@ pub fn Namespace(comptime Tri: type) type { try std.testing.expectError(bad[2], open(std.testing.allocator, std.testing.io, &ctx, "no-such-checkpoint", null, .{ .context = 4096, .drafts = true })); } } - - // -- the prompt cache's copies ------------------------------------------------------------------------ - - fn snapBytes(ptr: *anyopaque, at: u32) u64 { - return snap.bytes(at, of(ptr).head != null); - } - - /// The owner stream's state, its prompt pass standing at `at` (a drafting stream's with the head's rows). - fn snapSave(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { - const self = of(ptr); - const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); - const l = try self.lane(s); - if (l.seq.fed != at) return error.NotAtMark; - if (self.ranks) |r| { - if (l.draft != null) return error.TensorRankDraftSnapshotUnavailable; - return try r.save(&l.seq); - } - return try snap.save(self.gpa, self.engine.driver, self.stream, &l.seq, if (l.draft) |*d| d else null); - } - - fn snapRestore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror!void { - const self = of(ptr); - const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); - const l = try self.lane(s); - if (self.ranks) |r| { - if (l.draft != null) return error.NoHeadState; - return r.restore(&l.seq, @ptrCast(@alignCast(saved))); - } - try snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, @ptrCast(@alignCast(saved))); - } - - fn snapDrop(ptr: *anyopaque, saved: *anyopaque) void { - const self = of(ptr); - if (self.ranks) |r| return r.dropSaved(@ptrCast(@alignCast(saved))); - snap.drop(self.gpa, self.stream, @ptrCast(@alignCast(saved))); - } }; } diff --git a/zig/src/families/flashnext/cuda_lanes_admission_test.zig b/zig/src/families/flashnext/cuda_lanes_admission_test.zig index 898c0bd9c6..7666ad741c 100644 --- a/zig/src/families/flashnext/cuda_lanes_admission_test.zig +++ b/zig/src/families/flashnext/cuda_lanes_admission_test.zig @@ -31,6 +31,7 @@ fn run(comptime Cuda: type, comptime open: anytype, comptime drop: anytype, coun self.ranks = null; self.head = &head; self.lanes = .empty; + self.parts = .{}; self.kv = .{}; defer { var it = self.lanes.valueIterator(); diff --git a/zig/src/families/flashnext/cuda_lanes_draft.zig b/zig/src/families/flashnext/cuda_lanes_draft.zig index 6f0474d22b..b964cfe22f 100644 --- a/zig/src/families/flashnext/cuda_lanes_draft.zig +++ b/zig/src/families/flashnext/cuda_lanes_draft.zig @@ -49,7 +49,7 @@ pub fn batched(self: anytype, requests: []const be.DraftRequest, done: []bool) ! } if (r.position != l.seq.fed + 1) return error.PositionMismatch; const room = l.seq.cap - l.seq.fed - 1; - counts[n] = @min(@min(r.depth, self.depth), room); + counts[n] = @min(@min(r.depth, mtp.depthFor(r.stream.sampling, self.depth, self.depth_greedy)), room); if (r.rows) |path| { nexts[n] = r.follow[0..path.len]; rows[n] = d.streams.ptr; diff --git a/zig/src/families/flashnext/cuda_lanes_load.zig b/zig/src/families/flashnext/cuda_lanes_load.zig index 5830e8705f..d979a6b3c0 100644 --- a/zig/src/families/flashnext/cuda_lanes_load.zig +++ b/zig/src/families/flashnext/cuda_lanes_load.zig @@ -24,10 +24,17 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu if (o.slide and captured != null) return error.TrainingOnCapturedGlue; const page_tokens = try @import("cuda_kv_lanes.zig").paging(); if (page_tokens > 0 and (o.tp != 1 or o.slide or captured != null)) return error.FlashNextPagesUnsupported; + // pending DeltaNet rows run on one GPU only: TF_FLASHNEXT_KEEP_LATE=1 with tensor parallel is refused + const Draft = @import("cuda_mtp_draft.zig").Draft; + const late_fits = o.tp == 1; + if (Draft.lateWanted() == true and !late_fits) return error.FlashNextKeepLateUnsupported; + Draft.late_on = (Draft.lateWanted() orelse true) and late_fits; const driver: *cuda.Driver = @constCast(ctx.d); const own = try gpa.create(Cuda); errdefer gpa.destroy(own); own.* = .{ .gpa = gpa, .context = o.context, .engine = undefined, .glue = undefined, .stream = undefined, .logits = undefined }; + // fixed-size lane buffers reused across streams (cuda_lane_parts.zig; TF_FLASHNEXT_LANE_PARTS=0: allocate each) + own.parts.on = if (std.c.getenv("TF_FLASHNEXT_LANE_PARTS")) |v| !std.mem.eql(u8, std.mem.span(v), "0") else true; own.stream = try cuda.Stream.init(driver, true); errdefer own.stream.deinit(); own.glue = try Glue.load(driver); @@ -69,6 +76,7 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu if (o.drafts and (o.tp == 1 or o.rank == 0)) { own.head = try mtp.Head.init(gpa, io, driver, &own.stream, dir); own.depth = @max(1, @min(mtp.max_chain, F.envUsize("TF_FLASHNEXT_DRAFTS", mtp.default_depth))); + own.depth_greedy = @max(1, @min(mtp.max_chain, F.envUsize("TF_FLASHNEXT_DRAFTS_GREEDY", mtp.default_depth_greedy))); if (std.c.getenv("TF_FLASHNEXT_CONFIDENCE")) |v| own.confidence = std.fmt.parseFloat(f64, std.mem.span(v)) catch mtp.default_confidence; } errdefer if (own.head) |h| h.deinit(); @@ -93,6 +101,8 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu const head_bytes = if (own.head) |h| mtp.Draft.bytes(window, h.ids.len, own.engine.chunk) else 0; own.kv.layout = @import("cuda_kv_lanes.zig").layout(o.context, page_tokens, own.kv.head, stream_bytes - own_seq_bytes - head_bytes); } + // the weights are in: the n-gram table into the page cache behind them, so the first prompt finds its rows in RAM + own.engine.ngram.warm(); return .{ .layout = own.kv.layout, .vision = own.vision, diff --git a/zig/src/families/flashnext/cuda_lanes_restore_test.zig b/zig/src/families/flashnext/cuda_lanes_restore_test.zig index aa441f6166..92176e0dc7 100644 --- a/zig/src/families/flashnext/cuda_lanes_restore_test.zig +++ b/zig/src/families/flashnext/cuda_lanes_restore_test.zig @@ -88,6 +88,7 @@ pub fn run(comptime Cuda: type, comptime open: anytype, comptime restore: anytyp self.ranks = &rank; self.head = if (drafting) &head else null; self.lanes = .empty; + self.parts = .{}; defer self.lanes.deinit(a); var kept: snap.Snap = .{ .rank_id = 99, .buf = undefined, .at = 2, .history = .{ 8, 9 }, .head = false }; kept.buf.ptr = 0x100; diff --git a/zig/src/families/flashnext/cuda_lanes_snap.zig b/zig/src/families/flashnext/cuda_lanes_snap.zig new file mode 100644 index 0000000000..9ce4e6d0ca --- /dev/null +++ b/zig/src/families/flashnext/cuda_lanes_snap.zig @@ -0,0 +1,47 @@ +//! The prompt cache's copies of a Flash Next lane's state (cuda_snap.zig), called with the lane namespace's target. +const lanes = @import("lanes"); +const snap = @import("cuda_snap.zig"); + +pub fn Snaps(comptime Cuda: type) type { + return struct { + fn of(ptr: *anyopaque) *Cuda { + return @ptrCast(@alignCast(ptr)); + } + + pub fn bytes(ptr: *anyopaque, at: u32) u64 { + return snap.bytes(at, of(ptr).head != null); + } + + /// The owner stream's state, its prompt pass standing at `at` (a drafting stream's with the head's rows). + pub fn save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { + const self = of(ptr); + const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); + const l = try self.lane(s); + if (l.seq.fed != at) return error.NotAtMark; + if (self.ranks) |r| { + if (l.draft != null) return error.TensorRankDraftSnapshotUnavailable; + return try r.save(&l.seq); + } + // pending DeltaNet rows into the states first: a snapshot holds them as they stand + if (l.draft) |*d| try self.ops().flushLate(&l.seq, d); + return try snap.save(self.gpa, self.engine.driver, self.stream, &l.seq, if (l.draft) |*d| d else null); + } + + pub fn restore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror!void { + const self = of(ptr); + const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); + const l = try self.lane(s); + if (self.ranks) |r| { + if (l.draft != null) return error.NoHeadState; + return r.restore(&l.seq, @ptrCast(@alignCast(saved))); + } + try snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, @ptrCast(@alignCast(saved))); + } + + pub fn drop(ptr: *anyopaque, saved: *anyopaque) void { + const self = of(ptr); + if (self.ranks) |r| return r.dropSaved(@ptrCast(@alignCast(saved))); + snap.drop(self.gpa, self.stream, @ptrCast(@alignCast(saved))); + } + }; +} diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 27d78bba79..ea972c1170 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -86,6 +86,11 @@ pub const Tri = struct { return if (g.mixed & glue_.entryBit(name) != 0) g else null; } + /// The writeback runs on our glue (hcWritebackNorm fuses it with the normalize). + pub fn writebackOwned(t: Tri) bool { + return t.glue("_hc_writeback") != null; + } + /// The read-out tiles replace down, reduction and upmix together; mixed operation ownership stays on its route. pub fn readoutTilesOwned(t: Tri) bool { return t.glue("_qmm_hcdown") != null and t.glue("_hc_reduce_act") != null and t.glue("_qmm_upmix") != null; @@ -335,6 +340,12 @@ pub const Tri = struct { }, &.{ ci("D", dims), ci("S", streams), ci("MODE", mode), ci("TOPK", topk), ci("SLOTS", slots), ci("BLOCK", 256), ci("WORLD", world) }); } + /// hcWriteback and hcNormalize (`scale` into `norm`) in one launch on the native glue: the same bits as the two. + pub fn hcWritebackNorm(t: Tri, h: u64, hout: u64, pss: u64, branch: u64, inj: u64, y: u64, wts: u64, scale: u64, norm: u64, eps: f32, rs: usize, rows: usize, dims: usize, streams: usize, mode: usize, topk: usize, slots: usize, world: usize) !void { + const g = t.glue("_hc_writeback") orelse return error.NativeGlueRequired; + return g.hcWritebackNorm(t.s, h, hout, pss, branch, inj, y, wts, scale, norm, eps, rs, rows, dims, streams, mode, topk, slots, world); + } + /// glue.ple_embed: gathered n-gram rows, MLX 4-bit group 32. `scale` is a constexpr. pub fn pleEmbed(t: Tri, weight: u64, scales: u64, biases: u64, out: u64, xs: u64, rows: usize, heads: usize, dh: usize, scale: f32) !void { if (t.set == null) return error.CapturedKernelsRequired; // the row path's n-gram check only diff --git a/zig/src/families/flashnext/cuda_mtp.zig b/zig/src/families/flashnext/cuda_mtp.zig index 8fe177ec9a..13be6cfad2 100644 --- a/zig/src/families/flashnext/cuda_mtp.zig +++ b/zig/src/families/flashnext/cuda_mtp.zig @@ -31,8 +31,49 @@ pub const Sampling = lanes.sampling.Sampling; pub const max_chain: usize = 15; /// Preserve the existing recipe unless an explicit depth is selected. pub const default_depth: usize = 6; -/// The live recipe's --mtp-confidence (Python's CONFIDENCE default is 0.7). -pub const default_confidence: f64 = 0.6; +/// A greedy stream's chain at most (TF_FLASHNEXT_DRAFTS_GREEDY): greedy drafts land more often, so a deeper chain pays. +pub const default_depth_greedy: usize = 8; + +/// The chain a stream drafts at most: `greedy` for greedy streams (no sampling or temperature 0), else `depth`. +pub fn depthFor(sm: ?Sampling, depth: usize, greedy: usize) usize { + return if (isSampled(sm)) depth else greedy; +} +/// The chain's stop rule (Stop): the running product of the drafts' probabilities at 0.4; 0.6 per draft before it. +pub const default_confidence: f64 = -0.4; + +/// TF_FLASHNEXT_CONFIDENCE c > 0: a later draft only at or above c; c < 0: while p1 * ... * pj >= -c; 0: every draft. +pub const Stop = struct { + c: f64, + product: f64 = 1.0, + + /// Whether draft j (0-based) with probability `p` is proposed, and whether the chain may go on. + pub fn take(s: *Stop, j: usize, p: f64) struct { propose: bool, more: bool } { + if (s.c > 0) { + const low = p < s.c; + if (low and j > 0) return .{ .propose = false, .more = false }; + return .{ .propose = true, .more = !low }; + } + if (s.c < 0) { + s.product *= p; + // below the line it ends the chain: no later product can reach it again + if (s.product < -s.c) return .{ .propose = j == 0, .more = false }; + } + return .{ .propose = true, .more = true }; + } +}; + +test "the stop rule: per draft above zero, the running product below it, every draft at zero" { + var per: Stop = .{ .c = 0.6 }; + try std.testing.expect(per.take(0, 0.5).propose and !per.take(0, 0.5).more); + try std.testing.expect(!per.take(1, 0.5).propose); + var prod: Stop = .{ .c = -0.4 }; + try std.testing.expect(prod.take(0, 0.8).more); // 0.8 + try std.testing.expect(prod.take(1, 0.6).more); // 0.48 + const t = prod.take(2, 0.8); // 0.384: below 0.4 + try std.testing.expect(!t.propose and !t.more); + var all: Stop = .{ .c = 0 }; + try std.testing.expect(all.take(5, 0.01).propose and all.take(6, 0.01).more); +} /// A verify window's rows at most: the pending token and a chain, or a copied continuation. pub const max_rows: usize = 32; /// A drafting round's rows at most: a shared verify's (TF_FLASHNEXT_WINDOW_ROWS lowers it; 0: each alone). @@ -165,6 +206,7 @@ pub const Head = struct { service: lanes.proposal_head.Service, module: cuda.Module, best_fn: cuda.Function, + topk_fn: cuda.Function, // fn_row_topk: a row's top_k candidates as drawRow ranks them (cuda_mtp_draw.zig) pub fn init(gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, dir: []const u8) !*Head { if (!cuda.kernels.available) return error.BuiltWithoutKernels; @@ -212,6 +254,7 @@ pub const Head = struct { h.module = try cuda.Module.load(driver, cuda.kernels.flashnext_mtp); errdefer h.module.unload(); h.best_fn = try h.module.function("fn_row_best"); + h.topk_fn = try h.module.function("fn_row_topk"); return h; } @@ -247,6 +290,18 @@ pub const Head = struct { _ = try pass.lane(input, sums, h.draft, out, pass.s.dummy, rows, 1); } + /// Each of `rows` rows' `k` best of its first `v` values (fn_row_topk): bf16 bits to out_v, columns to out_c. + pub fn topk(h: *const Head, stream: cuda.Stream, scores: u64, v: usize, ld: usize, rows: usize, k: usize, out_v: u64, out_c: u64) !void { + var a: cuda.Args = .{}; + a.add(scores); + a.add(cint(v)); + a.add(@as(c_longlong, @intCast(ld))); + a.add(cint(k)); + a.add(out_v); + a.add(out_c); + try cuda.launch.launch(h.topk_fn, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 1024 } }, stream, &a); + } + /// fn_row_best over `rows` rows of `v` scores `ld` apart: [rows, 4] fp32 at `out` (best, log-sum-exp, id bits). pub fn best(h: *const Head, stream: cuda.Stream, scores: u64, v: usize, ld: usize, rows: usize, out: u64) !void { var a: cuda.Args = .{}; diff --git a/zig/src/families/flashnext/cuda_mtp_draft.zig b/zig/src/families/flashnext/cuda_mtp_draft.zig index c2f23d0598..49cb74958e 100644 --- a/zig/src/families/flashnext/cuda_mtp_draft.zig +++ b/zig/src/families/flashnext/cuda_mtp_draft.zig @@ -26,6 +26,8 @@ const Q4 = wt.Q4; const HcW = wt.HcW; const max_chain = mtp.max_chain; const max_rows = mtp.max_rows; +/// best: [max_rows, 4] fp32 best values, then [max_rows, topk_max] bf16 candidates and i32 columns (cuda_mtp_draw.zig). +pub const best_bytes: usize = max_rows * 16 + max_rows * @import("cuda_mtp_draw.zig").topk_max * 6; const conv_bytes = mtp.conv_bytes; const proj_bytes = mtp.proj_bytes; @@ -45,6 +47,8 @@ pub const Chain = struct { /// A sequence's head cache and the buffers a window needs: its rows' scores and streams, and what it rolls back to. pub const Draft = struct { + /// Its graphs' keys (cuda_window_graph.nextGen). + gen: u64 = 0, stream: ?cuda.Stream = null, media: ?@import("cuda_media.zig").View = null, gpa: std.mem.Allocator, @@ -76,6 +80,10 @@ pub const Draft = struct { spare: [linear_n]cuda.DeviceBuffer = undefined, /// A window ran with the states written to `spare` and is not yet kept or rolled back. open: bool = false, + /// Pending DeltaNet rows: a control word, then each layer's log of batch.late_rows rows; null: rollbacks replay. + log: ?cuda.DeviceBuffer = null, + log_first: u32 = 0, + pending: u32 = 0, keep: batch.Keep = .{}, history: [2]i64 = .{ 0, 0 }, unused_history: [2]i64 = .{ 0, 0 }, @@ -95,12 +103,42 @@ pub const Draft = struct { /// Device bytes `init(cap, draft_n, rows)` allocates. pub fn bytes(cap: usize, draft_n: usize, rows: usize) usize { - return cap * (2 * prompt.kv_row_bytes + index_dim * 2) + prompt.pooledBytes(cap) + @max(rows, max_rows) * wide * 2 + wide * 2 + max_rows * vocab * 2 + draft_n * 2 + max_rows * 16 + saved_bytes + linear_n * state_n * 4; + const log = if (late()) logBytes() else 0; + return cap * (2 * prompt.kv_row_bytes + index_dim * 2) + prompt.pooledBytes(cap) + @max(rows, max_rows) * wide * 2 + wide * 2 + max_rows * vocab * 2 + draft_n * 2 + best_bytes + saved_bytes + linear_n * state_n * 4 + log; + } + + /// Drafting streams keep a partly kept window's DeltaNet rows pending (set at load; TF_FLASHNEXT_KEEP_LATE=0: off). + pub var late_on: bool = false; + + pub fn late() bool { + return late_on; + } + + /// TF_FLASHNEXT_KEEP_LATE: "1" asked for, "0" refused, null: not set. + pub fn lateWanted() ?bool { + const v = std.c.getenv("TF_FLASHNEXT_KEEP_LATE") orelse return null; + return !std.mem.eql(u8, std.mem.span(v), "0"); + } + + const log_ctl_bytes: usize = 256; + const layer_log_bytes: usize = batch.late_rows * batch.late_row_bytes; + + /// The pending-row log's bytes: its first and pending rows, then every DeltaNet layer's rows. + pub fn logBytes() usize { + return log_ctl_bytes + linear_n * layer_log_bytes; + } + + /// A partly kept window's rows join the pending ones; past late_store the chain stored its state, so they restart. + pub fn keptLate(d: *Draft, kept: usize) void { + if (d.pending > batch.late_store) { + d.log_first = @intCast((d.log_first + d.pending) % batch.late_rows); + d.pending = @intCast(kept); + } else d.pending += @intCast(kept); } /// A head cache of `cap` entries; `rows` the widest prompt chunk whose streams it keeps. pub fn init(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: cuda.Stream, cap: usize, draft_n: usize, rows: usize) !Draft { - var d: Draft = .{ .stream = stream, .gpa = gpa, .cap = cap, .rows_cap = @max(rows, max_rows), .kc = undefined, .vc = undefined, .ikc = undefined, .pooled = undefined, .streams = undefined, .out = undefined, .logits = undefined, .dlogits = undefined, .best = undefined, .saved = undefined, .host = undefined }; + var d: Draft = .{ .gen = @import("cuda_window_graph.zig").nextGen(), .stream = stream, .gpa = gpa, .cap = cap, .rows_cap = @max(rows, max_rows), .kc = undefined, .vc = undefined, .ikc = undefined, .pooled = undefined, .streams = undefined, .out = undefined, .logits = undefined, .dlogits = undefined, .best = undefined, .saved = undefined, .host = undefined }; d.kc = try cuda.DeviceBuffer.alloc(driver, cap * kv_heads * head_dim * 2); errdefer d.kc.free(); d.vc = try cuda.DeviceBuffer.alloc(driver, cap * kv_heads * head_dim * 2); @@ -117,10 +155,12 @@ pub const Draft = struct { errdefer d.logits.free(); d.dlogits = try cuda.DeviceBuffer.alloc(driver, draft_n * 2); errdefer d.dlogits.free(); - d.best = try cuda.DeviceBuffer.alloc(driver, max_rows * 16); + d.best = try cuda.DeviceBuffer.alloc(driver, best_bytes); errdefer d.best.free(); d.saved = try cuda.DeviceBuffer.alloc(driver, saved_bytes); errdefer d.saved.free(); + if (late()) d.log = try cuda.DeviceBuffer.alloc(driver, logBytes()); + errdefer if (d.log) |*b| b.free(); var made: usize = 0; errdefer for (d.spare[0..made]) |*b| b.free(); while (made < linear_n) : (made += 1) d.spare[made] = try cuda.DeviceBuffer.alloc(driver, state_n * 4); @@ -133,6 +173,54 @@ pub const Draft = struct { return d; } + /// The fixed-size buffers of a draft in one block (initParts): everything but the head's `cap`-entry caches. + pub fn partsBytes(draft_n: usize, rows: usize) usize { + const a = std.mem.alignForward; + const rows_cap = @max(rows, max_rows); + var n: usize = 0; + for ([_]usize{ rows_cap * wide * 2, wide * 2, max_rows * vocab * 2, draft_n * 2, best_bytes, saved_bytes, if (late()) logBytes() else 0 }) |x| n += a(usize, x, 256); + return n + linear_n * a(usize, state_n * 4, 256); + } + + /// init's draft whose fixed-size buffers (partsBytes) are zeroed views of `block`; its caches allocated as init's. + pub fn initParts(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: cuda.Stream, cap: usize, draft_n: usize, rows: usize, block: cuda.DeviceBuffer) !Draft { + const bytes_in = partsBytes(draft_n, rows); + if (block.len < bytes_in) return error.FlashNextRunTooShort; + var d: Draft = .{ .gen = @import("cuda_window_graph.zig").nextGen(), .stream = stream, .gpa = gpa, .cap = cap, .rows_cap = @max(rows, max_rows), .kc = undefined, .vc = undefined, .ikc = undefined, .pooled = undefined, .streams = undefined, .out = undefined, .logits = undefined, .dlogits = undefined, .best = undefined, .saved = undefined, .spare = undefined, .host = undefined }; + d.kc = try cuda.DeviceBuffer.alloc(driver, cap * kv_heads * head_dim * 2); + errdefer d.kc.free(); + d.vc = try cuda.DeviceBuffer.alloc(driver, cap * kv_heads * head_dim * 2); + errdefer d.vc.free(); + d.ikc = try cuda.DeviceBuffer.alloc(driver, cap * index_dim * 2); + errdefer d.ikc.free(); + d.pooled = try cuda.DeviceBuffer.alloc(driver, prompt.pooledBytes(cap)); + errdefer d.pooled.free(); + const Carve = struct { + fn take(dr: *const cuda.Driver, p: *u64, n: usize) cuda.DeviceBuffer { + const b: cuda.DeviceBuffer = .{ .d = dr, .ptr = p.*, .len = n, .borrowed = true }; + p.* += std.mem.alignForward(usize, n, 256); + return b; + } + }; + var at = block.ptr; + d.streams = Carve.take(driver, &at, d.rows_cap * wide * 2); + d.out = Carve.take(driver, &at, wide * 2); + d.logits = Carve.take(driver, &at, max_rows * vocab * 2); + d.dlogits = Carve.take(driver, &at, draft_n * 2); + d.best = Carve.take(driver, &at, best_bytes); + d.saved = Carve.take(driver, &at, saved_bytes); + if (late()) d.log = Carve.take(driver, &at, logBytes()); + for (&d.spare) |*b| b.* = Carve.take(driver, &at, state_n * 4); + std.debug.assert(at - block.ptr == bytes_in); + d.host = try gpa.alloc(u16, @max(max_rows * vocab, draft_n)); + errdefer gpa.free(d.host); + errdefer cuda.completion.drain(driver, stream, "Flash Next draft fills did not drain before rollback"); + for ([_]cuda.DeviceBuffer{ d.kc, d.vc, d.ikc, d.pooled }) |b| try b.fill8(0, stream.handle); + try stream.d.check(stream.d.api.cuMemsetD8Async(block.ptr, 0, bytes_in, stream.handle), "cuMemsetD8Async"); + d.point(); + return d; + } + /// The keep's history and projection rows in `saved`, and the end states into the spares. fn point(d: *Draft) void { var li: usize = 0; @@ -140,6 +228,10 @@ pub const Draft = struct { if (!wt.linearLayer(layer)) continue; lk.hist = d.saved.ptr + hist_at + li * conv_bytes; lk.proj = d.saved.ptr + proj_at + li * max_rows * proj_bytes; + if (d.log) |b| { + lk.log = b.ptr + log_ctl_bytes + li * layer_log_bytes; + lk.log_ctl = b.ptr; + } li += 1; } d.aim(); @@ -147,7 +239,7 @@ pub const Draft = struct { /// Device bytes `initIn` takes from a run: every buffer but the head's cache, each on a 256-byte boundary. pub fn runBytes(draft_n: usize, rows: usize) usize { - return bytes(0, draft_n, rows) + (6 + linear_n) * 256; + return bytes(0, draft_n, rows) + (7 + linear_n) * 256; } /// A pool lane's head: its cache on `paged`'s pages after the model's layers, its buffers from `base`'s run. @@ -165,13 +257,14 @@ pub const Draft = struct { } }; const head = wt.attn_n; - var d: Draft = .{ .gpa = gpa, .cap = cap, .rows_cap = @max(rows, max_rows), .kc = Carve.page(driver, paged.at(head, .key)), .vc = Carve.page(driver, paged.at(head, .value)), .ikc = Carve.page(driver, paged.at(head, .index)), .pooled = Carve.page(driver, paged.at(head, .pooled)), .pg = paged.pg(), .streams = undefined, .out = undefined, .logits = undefined, .dlogits = undefined, .best = undefined, .saved = undefined, .host = undefined }; + var d: Draft = .{ .gen = @import("cuda_window_graph.zig").nextGen(), .gpa = gpa, .cap = cap, .rows_cap = @max(rows, max_rows), .kc = Carve.page(driver, paged.at(head, .key)), .vc = Carve.page(driver, paged.at(head, .value)), .ikc = Carve.page(driver, paged.at(head, .index)), .pooled = Carve.page(driver, paged.at(head, .pooled)), .pg = paged.pg(), .streams = undefined, .out = undefined, .logits = undefined, .dlogits = undefined, .best = undefined, .saved = undefined, .host = undefined }; d.streams = Carve.take(driver, &at, d.rows_cap * wide * 2); d.out = Carve.take(driver, &at, wide * 2); d.logits = Carve.take(driver, &at, max_rows * vocab * 2); d.dlogits = Carve.take(driver, &at, draft_n * 2); - d.best = Carve.take(driver, &at, max_rows * 16); + d.best = Carve.take(driver, &at, best_bytes); d.saved = Carve.take(driver, &at, saved_bytes); + if (late()) d.log = Carve.take(driver, &at, logBytes()); for (&d.spare) |*b| b.* = Carve.take(driver, &at, state_n * 4); std.debug.assert(at - base <= len); d.host = try gpa.alloc(u16, @max(max_rows * vocab, draft_n)); @@ -189,18 +282,21 @@ pub const Draft = struct { } } - /// An open window kept whole: its end states become the sequence's (buffers swapped, nothing copied). + /// An open window kept whole: its end states become the sequence's (buffers swapped), no rows left pending. pub fn settle(d: *Draft, seq: *prompt.Seq) void { if (!d.open) return; for (seq.state, &d.spare) |*a, *b| std.mem.swap(cuda.DeviceBuffer, a, b); d.aim(); d.open = false; + d.log_first = 0; + d.pending = 0; } pub fn deinit(d: *Draft) void { if (d.stream) |stream| cuda.completion.drain(stream.d, stream, "Flash Next draft work did not drain before release"); d.gpa.free(d.host); for (&d.spare) |*b| b.free(); + if (d.log) |*b| b.free(); d.saved.free(); d.best.free(); d.dlogits.free(); @@ -218,5 +314,7 @@ pub const Draft = struct { d.len = 0; d.drafted = 0; d.rows = 0; + d.log_first = 0; + d.pending = 0; } }; diff --git a/zig/src/families/flashnext/cuda_mtp_draw.zig b/zig/src/families/flashnext/cuda_mtp_draw.zig new file mode 100644 index 0000000000..4a52c094ce --- /dev/null +++ b/zig/src/families/flashnext/cuda_mtp_draw.zig @@ -0,0 +1,87 @@ +//! Draws from score rows: a sampled draw reads fn_row_topk's top_k candidates (TF_FLASHNEXT_GPU_TOPK), not the row. +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const mtp = @import("cuda_mtp.zig"); +const lanes = @import("lanes"); + +const vocab = wt.vocab; +const max_rows = mtp.max_rows; +const Sampling = mtp.Sampling; +const Pick = mtp.Pick; + +/// A sampled draw with top_k up to this reads candidates; a wider one (or top_k 0) downloads its whole row. +pub const topk_max: usize = 64; + +/// On by default (the same draws); TF_FLASHNEXT_GPU_TOPK=0 downloads whole rows. +pub fn on() bool { + const v = std.c.getenv("TF_FLASHNEXT_GPU_TOPK") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// Whether a draw under `s` can read fn_row_topk's candidates instead of its whole row. +pub fn usable(s: Sampling) bool { + return s.temperature > 0 and s.top_k > 0 and s.top_k <= topk_max and on(); +} + +/// drawRow over a row's candidates (its top_k bf16 values and their columns, any order): choose draws from these alone. +pub fn drawTopk(gpa: std.mem.Allocator, vals: []const u16, cols: []const i32, ids: ?[]const u32, position: u64, s: Sampling) !u32 { + var v: [topk_max]f64 = undefined; + var keys: [topk_max]u64 = undefined; + for (vals, cols, 0..) |b, c, i| { + v[i] = mtp.promote(b); + keys[i] = if (ids) |m| m[@intCast(c)] else @intCast(c); + } + return @intCast(try lanes.sampling.choose(gpa, v[0..vals.len], keys[0..vals.len], position, s)); +} + +/// The candidates of `rows` rows of `v` scores `ld` apart, k a row, into d.best past its best values, then to d.host. +fn candidates(h: *const mtp.Head, st: cuda.Stream, d: *mtp.Draft, scores: u64, v: usize, ld: usize, rows: usize, k: usize) !struct { vals: []u16, cols: []i32 } { + const tv = d.best.ptr + max_rows * 16; + const tc = tv + max_rows * topk_max * 2; + try h.topk(st, scores, v, ld, rows, k, tv, tc); + const vals = d.host[0 .. rows * k]; + const cols: []i32 = @alignCast(std.mem.bytesAsSlice(i32, std.mem.sliceAsBytes(d.host[max_rows * topk_max ..][0 .. 2 * rows * k]))); + try download(st, std.mem.sliceAsBytes(vals), tv); + try download(st, std.mem.sliceAsBytes(cols), tc); + return .{ .vals = vals, .cols = cols }; +} + +/// The head's sampled draft from its candidates: drawRow's draw, and the drawn token's probability under `lse`. +pub fn pickHead(h: *const mtp.Head, st: cuda.Stream, gpa: std.mem.Allocator, d: *mtp.Draft, lse: f32, position: u64, s: Sampling) !Pick { + const n = h.ids.len; + const c = try candidates(h, st, d, d.dlogits.ptr, n, n, 1, s.top_k); + const tok = try drawTopk(gpa, c.vals, c.cols, h.ids, position, s); + for (c.vals, c.cols) |v, col| if (h.ids[@intCast(col)] == tok) return .{ .token = tok, .prob = @exp(@as(f64, mtp.promote(v)) - @as(f64, lse)) }; + return .{ .token = tok, .prob = 0 }; +} + +/// The window's rows drawn at their keyed positions (fed0 + 1 + r), each within its `masks` row if any. +pub fn draws(h: *const mtp.Head, st: cuda.Stream, gpa: std.mem.Allocator, d: *mtp.Draft, sm: ?Sampling, masks: []const u32, out: []u32) !void { + const n = d.rows; + if (!mtp.isSampled(sm) and masks.len == 0) { + try h.best(st, d.logits.ptr, vocab, vocab, n, d.best.ptr); + try download(st, std.mem.sliceAsBytes(d.best_host[0 .. 4 * n]), d.best.ptr); + for (out[0..n], 0..) |*t, r| t.* = @intCast(@as(i32, @bitCast(d.best_host[4 * r + 2]))); + return; + } + if (masks.len == 0) if (sm) |s| if (usable(s)) { + const k: usize = s.top_k; + const c = try candidates(h, st, d, d.logits.ptr, vocab, vocab, n, k); + for (out[0..n], 0..) |*t, r| t.* = try drawTopk(gpa, c.vals[r * k ..][0..k], c.cols[r * k ..][0..k], null, d.fed0 + 1 + r, s); + return; + }; + const rows = d.host[0 .. n * vocab]; + try download(st, std.mem.sliceAsBytes(rows), d.logits.ptr); + const words = masks.len / n; + for (out[0..n], 0..) |*t, r| { + const row = rows[r * vocab ..][0..vocab]; + if (words > 0) mtp.maskRow(row, masks[r * words ..][0..words]); + t.* = try mtp.drawRow(gpa, row, null, d.fed0 + 1 + r, sm); + } +} + +fn download(st: cuda.Stream, dst: []u8, src: u64) !void { + try st.synchronize(); + try st.d.check(st.d.api.cuMemcpyDtoH_v2(dst.ptr, src, dst.len), "cuMemcpyDtoH"); +} diff --git a/zig/src/families/flashnext/cuda_mtp_ops.zig b/zig/src/families/flashnext/cuda_mtp_ops.zig index 993e5328ce..1e34128f3d 100644 --- a/zig/src/families/flashnext/cuda_mtp_ops.zig +++ b/zig/src/families/flashnext/cuda_mtp_ops.zig @@ -8,6 +8,9 @@ const prompt = @import("cuda_prompt.zig"); const gdn = @import("cuda_gdn.zig"); const mtp = @import("cuda_mtp.zig"); const plan = @import("cuda_mtp_plan.zig"); +const roll = @import("cuda_mtp_roll.zig"); +const pieces = @import("cuda_mtp_pieces.zig"); +const draw = @import("cuda_mtp_draw.zig"); /// One group of lanes in the batched head's pass: the lanes' drafts, their rows and whether the head scores. const Group = struct { @@ -46,8 +49,6 @@ const Draft = mtp.Draft; const Chain = mtp.Chain; const Pick = mtp.Pick; const Sampling = mtp.Sampling; -const drawRow = mtp.drawRow; -const maskRow = mtp.maskRow; const promote = mtp.promote; fn cint(v: usize) c_int { @@ -96,10 +97,10 @@ pub fn Ops(comptime Tri: type) type { /// Prompt rows [at, stop): the pass with its streams kept, then the head over each row with a known next token. pub fn prefillChunk(o: Self, seq: *prompt.Seq, d: *Draft, ids: []const u32, at: usize, stop: usize) !void { if (stop > d.cap or stop <= at or stop - at > d.rows_cap) return error.PromptTooLong; - if (seq.fed != at or d.len != at or d.drafted != 0) return error.PositionMismatch; + if (seq.fed != at or d.len != at or d.drafted != 0 or d.pending != 0) return error.PositionMismatch; d.media = if (seq.media) |*m| m.view() else null; const last = stop == ids.len; - try o.forwardMain(seq, ids[at..stop], .{ .logits = if (last) .last else .none, .streams = d.streams.ptr }); + try o.forwardMain(seq, ids[at..stop], .{ .logits = if (last) .last else .none, .streams = d.streams.ptr, .prompt = true }); const head_to = @min(stop, ids.len - 1); if (head_to > at) try o.forwardHead(d, ids[at + 1 .. head_to + 1], d.streams.ptr, false); if (last) d.prompt_last = d.streams.ptr + (stop - 1 - at) * wide * 2; @@ -144,44 +145,11 @@ pub fn Ops(comptime Tri: type) type { const m = @min(piece, next.len - at); p.rows = m; p.pos0 = d.len; - p.pending = null; try p.stage(next[at..][0..m]); - // the embedding branch: the token's row (stream 0 of the four copies), normed, fc_embedding - if (e.table.dense) { - try tri.embedDense(s.ids, e.table.w.ptr, s.ple_keys, dims, streams, m); - } else { - try tri.embed(s.ids, e.table.w.ptr, e.table.s.ptr, e.table.b.ptr, s.ple_keys, dims, streams, m); - } - try @import("cuda_media.zig").injectNext(d.media, o.stream(), s.ple_keys, d.len, m); - try tri.rmsnorm(s.ple_keys, h.norm_e.ptr, s.mixed, s.xsm, eps, wide, m, dims, dims); - const esk = try p.lane(s.mixed, s.xsm, h.fc_e, s.dummy, s.part, m, fc_sk); - try p.sumSlices(s.part, s.branch, true, m * dims, esk); - // the hidden branch: the four streams normed as one row, fc_hidden on each stream - try tri.rmsnorm(rows + at * wide * 2, h.norm_h.ptr, s.norm, s.gq, eps, wide, m, wide, wide); - const hsk = try p.lane(s.norm, s.gq, h.fc_h, s.dummy, s.part, streams * m, fc_sk); - try p.sumSlices(s.part, s.ple_gated, true, streams * m * dims, hsk); - try tri.addStreams(s.branch, s.ple_gated, s.h, m, dims, streams); - // the layer: attention over the head's own cache, its experts - try p.writeback(0, m); - try p.readout(h.layer.attn_hc, 0, m); - try p.attnBlock(&h.layer, 0); - try tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", m * dims, m, dims, streams, 4, 1, 1, p.out_slices); - try p.readout(h.layer.mlp_hc, 0, m); - try p.moeBlock(&h.layer); - p.pending = 0; - try p.writeback(0, m); - try p.copy(d.out.ptr, s.h + (m - 1) * wide * 2, wide * 2); + const one: pieces.Piece(Self, batch.Pass(Tri)) = .{ .o = o, .p = &p, .d = d, .input = rows + at * wide * 2, .m = m, .score = at + m == next.len and scores }; + try pieces.runPiece(one, score_rows); d.len += m; at += m; - if (at == next.len and scores) { - // the last row's streams and sums at row 0, where the captured kernels' pointers stay aligned - if (m > 1) { - try p.copy(s.h, s.h + (m - 1) * wide * 2, wide * 2); - try p.copy(s.pss, s.pss + (m - 1) * pss_n * 4, pss_n * 4); - } - try p.readout(h.mixer, 0, 1); - try h.project(&p, s.mixed, s.xsm, d.dlogits.ptr, 1); - } } } @@ -208,24 +176,26 @@ pub fn Ops(comptime Tri: type) type { const col: usize = @intCast(@as(i32, @bitCast(d.best_host[2]))); return .{ .token = try view_.global(@intCast(col)), .prob = @exp(@as(f64, top) - @as(f64, lse)) }; } + if (sm) |x| if (draw.usable(x)) return draw.pickHead(h, o.stream(), o.e.gpa, d, lse, position, x); const row = d.host[0..n]; try download(o.stream(), std.mem.sliceAsBytes(row), d.dlogits.ptr); return mtp.drawPick(o.e.gpa, row, h.ids, lse, position, sm); } - /// Python decode.draft: absorb the kept rows, then chain up to `count` drafts, the first always kept. + /// Python decode.draft: absorb the kept rows, then chain up to `count` drafts under mtp.Stop, the first kept. pub fn draft(o: Self, d: *Draft, next: []const u32, rows: u64, position: u64, count: usize, sm: ?Sampling, confidence: f64, out: *Chain) !void { out.n = 0; try o.absorb(d, next, rows, count > 0); + var stop: mtp.Stop = .{ .c = confidence }; var j: usize = 0; while (j < @min(count, max_chain)) : (j += 1) { const got = try o.pick(d, position + j, sm); - const low = confidence > 0 and got.prob < confidence; - if (low and j > 0) break; + const t = stop.take(j, got.prob); + if (!t.propose) break; out.tokens[out.n] = got.token; out.probs[out.n] = got.prob; out.n += 1; - if (low) break; + if (!t.more) break; if (j + 1 < count) { try o.forwardHead(d, &.{got.token}, d.out.ptr, true); d.drafted += 1; @@ -299,49 +269,8 @@ pub fn Ops(comptime Tri: type) type { p.pos0 = 0; p.pending = null; try p.stage(ids_buf[0..total]); - if (e.table.dense) { - try tri.embedDense(s.ids, e.table.w.ptr, s.ple_keys, dims, streams, total); - } else { - try tri.embed(s.ids, e.table.w.ptr, e.table.s.ptr, e.table.b.ptr, s.ple_keys, dims, streams, total); - } - try tri.rmsnorm(s.ple_keys, h.norm_e.ptr, s.mixed, s.xsm, eps, wide, total, dims, dims); - const esk = try p.lane(s.mixed, s.xsm, h.fc_e, s.dummy, s.part, total, fc_sk); - try p.sumSlices(s.part, s.branch, true, total * dims, esk); - // each lane's rows of main-model streams, side by side (the head reads them as one block of rows) - at = 0; - for (group.rows, nexts) |r, nx| { - try p.copy(s.ple_nrow + at * wide * 2, r, nx.len * wide * 2); - at += nx.len; - } - try tri.rmsnorm(s.ple_nrow, h.norm_h.ptr, s.norm, s.gq, eps, wide, total, wide, wide); - const hsk = try p.lane(s.norm, s.gq, h.fc_h, s.dummy, s.part, streams * total, fc_sk); - try p.sumSlices(s.part, s.ple_gated, true, streams * total * dims, hsk); - try tri.addStreams(s.branch, s.ple_gated, s.h, total, dims, streams); - try p.writeback(0, total); - try p.readout(h.layer.attn_hc, 0, total); - try p.attnBlock(&h.layer, 0); - try tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", total * dims, total, dims, streams, 4, 1, 1, p.out_slices); - try p.readout(h.layer.mlp_hc, 0, total); - try p.moeBlock(&h.layer); - p.pending = 0; - try p.writeback(0, total); - // each lane's last row: its d.out, and gathered to row i for the draft head (row i <= its last row) - at = 0; - for (ds, nexts, 0..) |d, nx, i| { - const last = at + nx.len - 1; - try p.copy(d.out.ptr, s.h + last * wide * 2, wide * 2); - if (group.scores and last != i) { - try p.copy(s.h + i * wide * 2, s.h + last * wide * 2, wide * 2); - try p.copy(s.pss + i * pss_n * 4, s.pss + last * pss_n * 4, pss_n * 4); - } - d.len += nx.len; - at += nx.len; - } - if (!group.scores) return; - try p.readout(h.mixer, 0, n_l); - const n_ids = h.ids.len; - _ = try p.lane(s.mixed, s.xsm, h.draft, s.head_local, s.dummy, n_l, 1); - for (ds, 0..) |d, i| try p.copy(d.dlogits.ptr, s.head_local + i * n_ids * 2, n_ids * 2); + try pieces.runGroup(pieces.Group(Self, batch.Pass(Tri), Group){ .o = o, .p = &p, .group = group, .total = total }, score_rows); + for (ds, nexts) |d, nx| d.len += nx.len; } /// draft over several drafts, greedy or sampled: kept rows in one pass, then one a step, each as it would alone @@ -362,6 +291,8 @@ pub fn Ops(comptime Tri: type) type { act[n_act] = i; n_act += 1; }; + var stops: [max_lanes]mtp.Stop = undefined; + for (stops[0..ds.len]) |*st| st.* = .{ .c = confidence }; var j: usize = 0; var best_host: [max_lanes * 4]f32 = undefined; var step_ds: [max_lanes]*Draft = undefined; @@ -380,17 +311,18 @@ pub fn Ops(comptime Tri: type) type { const col: usize = @intCast(@as(i32, @bitCast(best_host[r * 4 + 2]))); // a sampled lane draws from its own row, as its lone pick does; a greedy one takes its best column const got: Pick = if (mtp.isSampled(sms[i])) sampled: { + if (draw.usable(sms[i].?)) break :sampled try draw.pickHead(h, o.stream(), o.e.gpa, ds[i], lse, positions[i] + j, sms[i].?); const row = ds[i].host[0..h.view().logical]; try download(o.stream(), std.mem.sliceAsBytes(row), ds[i].dlogits.ptr); break :sampled try mtp.drawPick(o.e.gpa, row, h.ids, lse, positions[i] + j, sms[i]); } else .{ .token = h.ids[col], .prob = @exp(@as(f64, top) - @as(f64, lse)) }; const out = outs[i]; - const low = confidence > 0 and got.prob < confidence; - if (low and j > 0) continue; + const t = stops[i].take(j, got.prob); + if (!t.propose) continue; out.tokens[out.n] = got.token; out.probs[out.n] = got.prob; out.n += 1; - if (low) continue; + if (!t.more) continue; if (j + 1 < @min(counts[i], max_chain)) { step_ds[n_step] = ds[i]; step_next[n_step] = .{got.token}; @@ -418,13 +350,13 @@ pub fn Ops(comptime Tri: type) type { if (n == 0 or n > max_rows) return error.WindowTooWide; if (seq.fed + n > seq.cap) return error.PromptTooLong; d.settle(seq); + const keeps = try o.lateBefore(seq, d, n); d.fed0 = seq.fed; d.rows = n; d.history = seq.history; @memcpy(d.ids[0..n], ids); - const keeps = n > 1; if (keeps) try copyAsync(o.stream(), d.saved.ptr + Draft.tail_at, seq.tail.ptr, ple_tail_n * wide * 2); - try o.e.forward(Tri, o.tri, seq, ids, .{ .logits = .all, .logits_out = d.logits.ptr, .keep = if (keeps) &d.keep else null, .streams = d.streams.ptr }); + try o.e.forward(Tri, o.tri, seq, ids, .{ .logits = .all, .logits_out = d.logits.ptr, .keep = if (keeps) &d.keep else null, .streams = d.streams.ptr, .host_logits = @import("cuda_batch_state.zig").verifyHostLogits() }); if (keeps) try copyAsync(o.stream(), d.saved.ptr + Draft.nrow_at, o.e.batch_s.?.ple_nrow, n * wide * 2); d.open = keeps; } @@ -443,42 +375,69 @@ pub fn Ops(comptime Tri: type) type { for (seqs, ds, idss, wins) |seq, d, ids, *w| { const n = ids.len; d.settle(seq); + const keeps = try o.lateBefore(seq, d, n); d.fed0 = seq.fed; d.rows = n; d.history = seq.history; @memcpy(d.ids[0..n], ids); - const keeps = n > 1; if (keeps) try copyAsync(o.stream(), d.saved.ptr + Draft.tail_at, seq.tail.ptr, ple_tail_n * wide * 2); w.* = .{ .seq = seq, .ids = ids, .keep = if (keeps) &d.keep else null, .streams = d.streams.ptr, .logits = d.logits.ptr }; } try o.e.windows(Tri, o.tri, wins, out); var row0: usize = 0; - for (ds, idss) |d, ids| { - const keeps = ids.len > 1; + for (ds, idss, wins) |d, ids, w| { + const keeps = w.keep != null; if (keeps) try copyAsync(o.stream(), d.saved.ptr + Draft.nrow_at, o.e.batch_s.?.ple_nrow + row0 * wide * 2, ids.len * wide * 2); d.open = keeps; row0 += ids.len; } } + /// Whether an n-row window keeps a rollback; with a log its control word goes to the device (full: flushed). + fn lateBefore(o: Self, seq: *prompt.Seq, d: *Draft, n: usize) !bool { + if (d.log == null) return n > 1; + if (d.pending + n > batch.late_rows) try o.flushLate(seq, d); + if (n == 1 and d.pending == 0) return false; + try o.lateCtl(d); + return true; + } + + /// The log's first pending row | its pending rows << 16 on the device (fn_late's ctl), in stream order. + fn lateCtl(o: Self, d: *const Draft) !void { + const st = o.stream(); + try st.d.check(st.d.api.cuMemsetD32Async(d.log.?.ptr, d.log_first | d.pending << 16, 1, st.handle), "cuMemsetD32Async"); + } + + /// The pending rows applied to the DeltaNet states in place: they stand at seq.fed as without a log. + pub fn flushLate(o: Self, seq: *prompt.Seq, d: *Draft) !void { + if (d.log == null or d.pending == 0) return; + if (d.open) return error.CannotRewind; + const st = o.stream(); + try o.lateCtl(d); + var li: usize = 0; + for (o.e.layers[0..o.e.n_layers], 0..) |*lw, layer| { + if (!lw.linear) continue; + defer li += 1; + const lk = d.keep.layers[layer]; + var a: cuda.Args = .{}; + a.add(@as(u64, 0)); + a.add(lk.log); + a.add(lk.log_ctl); + a.add(cint(batch.late_rows)); + a.add(cint(batch.late_rows)); + a.add(seq.state[li].ptr); + a.add(seq.state[li].ptr); + a.add(cint(0)); + a.add(@as(u64, 0)); + try cuda.launch.launch(o.e.batch_k.chain_late, .{ .grid = .{ .x = gdn.nv, .y = gdn.dv / 16 }, .block = .{ .x = 128 } }, st, &a); + } + d.log_first = 0; + d.pending = 0; + } + /// The window's rows drawn at their keyed positions (fed0 + 1 + r), each within its `masks` row if any. pub fn draws(o: Self, d: *Draft, sm: ?Sampling, masks: []const u32, out: []u32) !void { - const n = d.rows; - const greedy = if (sm) |s| s.temperature <= 0 else true; - if (greedy and masks.len == 0) { - try o.h.best(o.stream(), d.logits.ptr, vocab, vocab, n, d.best.ptr); - try download(o.stream(), std.mem.sliceAsBytes(d.best_host[0 .. 4 * n]), d.best.ptr); - for (out[0..n], 0..) |*t, r| t.* = @intCast(@as(i32, @bitCast(d.best_host[4 * r + 2]))); - return; - } - const rows = d.host[0 .. n * vocab]; - try download(o.stream(), std.mem.sliceAsBytes(rows), d.logits.ptr); - const words = masks.len / n; - for (out[0..n], 0..) |*t, r| { - const row = rows[r * vocab ..][0..vocab]; - if (words > 0) maskRow(row, masks[r * words ..][0..words]); - t.* = try drawRow(o.e.gpa, row, null, d.fed0 + 1 + r, sm); - } + return draw.draws(o.h, o.stream(), o.e.gpa, d, sm, masks, out); } /// Keep the window's first `keep` rows (Python forward.commit) from its saved rows; a repeat is a no-op. @@ -489,57 +448,22 @@ pub fn Ops(comptime Tri: type) type { if (keep == d.rows) return d.settle(seq); if (!d.open) return error.CannotRewind; const e = o.e; - const st = o.stream(); const s = if (e.batch_s) |*bs| bs else return error.CannotRewind; - const k = &e.batch_k; seq.fed = d.fed0 + keep; var hist = d.history; for (d.ids[0..keep]) |t| hist = .{ hist[1], t }; seq.history = hist; - try copyAsync(st, seq.tail.ptr, d.saved.ptr + Draft.tail_at, ple_tail_n * wide * 2); - try shift(k, st, seq.tail.ptr, d.saved.ptr + Draft.nrow_at, wide, keep, ple_tail_n, wide); - var li: usize = 0; - for (e.layers[0..e.n_layers], 0..) |*lw, layer| { - if (!lw.linear) continue; - defer li += 1; - // the sequence's state is still the window's start (the pass wrote its end state to the spare) - const lk = d.keep.layers[layer]; - try copyAsync(st, seq.conv[li].ptr, lk.hist, conv_bytes); - var f: cuda.Args = .{}; - f.add(lk.proj); - f.add(seq.conv[li].ptr); - f.add(lw.cw.ptr); - f.add(lw.a_log.ptr); - f.add(lw.dt.ptr); - f.add(s.gq); - f.add(s.gk); - f.add(s.gv); - f.add(s.ggb); - try cuda.launch.launch(k.front, .{ .grid = .{ .x = @intCast(keep), .y = gdn.nk + gdn.nv }, .block = .{ .x = 256 } }, st, &f); - var a: cuda.Args = .{}; - a.add(s.gq); - a.add(s.gk); - a.add(s.gv); - a.add(s.ggb); - a.add(seq.state[li].ptr); - a.add(cint(keep)); - a.add(s.ys); - try cuda.launch.launch(k.chain, .{ .grid = .{ .x = gdn.nv, .y = gdn.dv / 16 }, .block = .{ .x = 128 } }, st, &a); - try shift(k, st, seq.conv[li].ptr, lk.proj, gdn.proj_width, keep, gdn.taps - 1, gdn.conv_dim); + if (d.log != null) { + // pending rows: the kept rows stay in the log; the conv windows and n-gram tail move on in one launch + try roll.lateRoll(o, seq, d, keep); + d.keptLate(keep); + d.rows = keep; + d.open = false; + return; } + try roll.replay(e, s, seq, d, keep); d.rows = keep; d.open = false; } - - fn shift(k: *const batch.Kernels, st: cuda.Stream, win: u64, rows_ptr: u64, stride: usize, n: usize, taps: usize, c: usize) !void { - var a: cuda.Args = .{}; - a.add(win); - a.add(rows_ptr); - a.add(@as(c_longlong, @intCast(stride))); - a.add(cint(n)); - a.add(cint(taps)); - a.add(cint(c)); - try cuda.launch.launch(k.shift, .{ .grid = .{ .x = @intCast((c + 255) / 256) }, .block = .{ .x = 256 } }, st, &a); - } }; } diff --git a/zig/src/families/flashnext/cuda_mtp_pieces.zig b/zig/src/families/flashnext/cuda_mtp_pieces.zig new file mode 100644 index 0000000000..f8382ea50f --- /dev/null +++ b/zig/src/families/flashnext/cuda_mtp_pieces.zig @@ -0,0 +1,154 @@ +//! The draft head's passes as kernels alone (their ids staged first), so a graph can hold them on one GPU. +const std = @import("std"); +const wt = @import("cuda_prompt_wt.zig"); +const mtp = @import("cuda_mtp.zig"); +const window_graph = @import("cuda_window_graph.zig"); + +const dims = wt.dims; +const streams = wt.streams; +const wide = wt.wide; +const eps = wt.eps; +const fc_sk = mtp.fc_sk; +const pss_n = mtp.pss_n; + +/// One piece of forwardHead: the head's layer over `m` rows from `input` (main-model streams), scored when `score`. +pub fn Piece(comptime O: type, comptime P: type) type { + return struct { + o: O, + p: *P, + d: *mtp.Draft, + input: u64, + m: usize, + score: bool, + + pub fn run(r: @This()) !void { + const pp = r.p; + const ps = pp.s; + const ee = r.o.e; + const t = r.o.tri; + const hh = r.o.h; + const m = r.m; + pp.pending = null; + // the embedding branch: the token's row (stream 0 of the four copies), normed, fc_embedding + if (ee.table.dense) { + try t.embedDense(ps.ids, ee.table.w.ptr, ps.ple_keys, dims, streams, m); + } else { + try t.embed(ps.ids, ee.table.w.ptr, ee.table.s.ptr, ee.table.b.ptr, ps.ple_keys, dims, streams, m); + } + try @import("cuda_media.zig").injectNext(r.d.media, ee.stream.*, ps.ple_keys, r.d.len, m); + try t.rmsnorm(ps.ple_keys, hh.norm_e.ptr, ps.mixed, ps.xsm, eps, wide, m, dims, dims); + const esk = try pp.lane(ps.mixed, ps.xsm, hh.fc_e, ps.dummy, ps.part, m, fc_sk); + try pp.sumSlices(ps.part, ps.branch, true, m * dims, esk); + // the hidden branch: the four streams normed as one row, fc_hidden on each stream + try t.rmsnorm(r.input, hh.norm_h.ptr, ps.norm, ps.gq, eps, wide, m, wide, wide); + const hsk = try pp.lane(ps.norm, ps.gq, hh.fc_h, ps.dummy, ps.part, streams * m, fc_sk); + try pp.sumSlices(ps.part, ps.ple_gated, true, streams * m * dims, hsk); + try t.addStreams(ps.branch, ps.ple_gated, ps.h, m, dims, streams); + // the layer: attention over the head's own cache, its experts + try pp.writeback(0, m); + try pp.readout(hh.layer.attn_hc, 0, m); + try pp.attnBlock(&hh.layer, 0); + try t.hcWriteback(ps.h, ps.h, ps.pss, ps.part, "*fp32", ps.inj, ps.dummy, "*bf16", ps.dummy, "*bf16", m * dims, m, dims, streams, 4, 1, 1, pp.out_slices); + try pp.readout(hh.layer.mlp_hc, 0, m); + try pp.moeBlock(&hh.layer); + pp.pending = 0; + try pp.writeback(0, m); + try pp.copy(r.d.out.ptr, ps.h + (m - 1) * wide * 2, wide * 2); + if (!r.score) return; + // the last row's streams and sums at row 0, where the captured kernels' pointers stay aligned + if (m > 1) { + try pp.copy(ps.h, ps.h + (m - 1) * wide * 2, wide * 2); + try pp.copy(ps.pss, ps.pss + (m - 1) * pss_n * 4, pss_n * 4); + } + try pp.readout(hh.mixer, 0, 1); + try hh.project(pp, ps.mixed, ps.xsm, r.d.dlogits.ptr, 1); + } + }; +} + +/// A piece on one GPU as its graph, keyed by what its kernels read (eagerly with media or tensor parallelism). +pub fn runPiece(one: anytype, score_rows: u64) !void { + const e = one.o.e; + const d = one.d; + const s = one.p.s; + if (!window_graph.enabled() or d.media != null or e.collective != null) return one.run(); + const key: window_graph.HeadKey = .{ .gen = d.gen, .rows = @intCast(one.m), .geometry = .at(d.len + one.m, d.cap), .scores = one.score, .cap = d.cap, .kc = d.kc.ptr, .input = one.input, .out = d.out.ptr, .dlogits = d.dlogits.ptr, .scratch = s.h, .scratch_cap = s.cap, .score_rows = score_rows }; + try e.graphs.head.run(e.gpa, e.stream.*, key, one); +} + +/// A group of lanes' rows in one head pass (headGroup), `total` rows staged: each lane's last row its own. +pub fn Group(comptime O: type, comptime P: type, comptime G: type) type { + return struct { + o: O, + p: *P, + group: G, + total: usize, + + pub fn run(r: @This()) !void { + const pp = r.p; + const ps = pp.s; + const ee = r.o.e; + const t = r.o.tri; + const hh = r.o.h; + const tot = r.total; + pp.pending = null; + if (ee.table.dense) { + try t.embedDense(ps.ids, ee.table.w.ptr, ps.ple_keys, dims, streams, tot); + } else { + try t.embed(ps.ids, ee.table.w.ptr, ee.table.s.ptr, ee.table.b.ptr, ps.ple_keys, dims, streams, tot); + } + try t.rmsnorm(ps.ple_keys, hh.norm_e.ptr, ps.mixed, ps.xsm, eps, wide, tot, dims, dims); + const esk = try pp.lane(ps.mixed, ps.xsm, hh.fc_e, ps.dummy, ps.part, tot, fc_sk); + try pp.sumSlices(ps.part, ps.branch, true, tot * dims, esk); + // each lane's rows of main-model streams, side by side (the head reads them as one block of rows) + var row0: usize = 0; + for (r.group.rows, r.group.nexts) |rw, nx| { + try pp.copy(ps.ple_nrow + row0 * wide * 2, rw, nx.len * wide * 2); + row0 += nx.len; + } + try t.rmsnorm(ps.ple_nrow, hh.norm_h.ptr, ps.norm, ps.gq, eps, wide, tot, wide, wide); + const hsk = try pp.lane(ps.norm, ps.gq, hh.fc_h, ps.dummy, ps.part, streams * tot, fc_sk); + try pp.sumSlices(ps.part, ps.ple_gated, true, streams * tot * dims, hsk); + try t.addStreams(ps.branch, ps.ple_gated, ps.h, tot, dims, streams); + try pp.writeback(0, tot); + try pp.readout(hh.layer.attn_hc, 0, tot); + try pp.attnBlock(&hh.layer, 0); + try t.hcWriteback(ps.h, ps.h, ps.pss, ps.part, "*fp32", ps.inj, ps.dummy, "*bf16", ps.dummy, "*bf16", tot * dims, tot, dims, streams, 4, 1, 1, pp.out_slices); + try pp.readout(hh.layer.mlp_hc, 0, tot); + try pp.moeBlock(&hh.layer); + pp.pending = 0; + try pp.writeback(0, tot); + // each lane's last row: its d.out, and gathered to row i for the draft head (row i <= its last row) + row0 = 0; + for (r.group.ds, r.group.nexts, 0..) |d, nx, i| { + const last = row0 + nx.len - 1; + try pp.copy(d.out.ptr, ps.h + last * wide * 2, wide * 2); + if (r.group.scores and last != i) { + try pp.copy(ps.h + i * wide * 2, ps.h + last * wide * 2, wide * 2); + try pp.copy(ps.pss + i * pss_n * 4, ps.pss + last * pss_n * 4, pss_n * 4); + } + row0 += nx.len; + } + if (!r.group.scores) return; + const lanes_n = r.group.ds.len; + try pp.readout(hh.mixer, 0, lanes_n); + const n_ids = hh.ids.len; + _ = try pp.lane(ps.mixed, ps.xsm, hh.draft, ps.head_local, ps.dummy, lanes_n, 1); + for (r.group.ds, 0..) |d, i| try pp.copy(d.dlogits.ptr, ps.head_local + i * n_ids * 2, n_ids * 2); + } + }; +} + +/// A group on one GPU as its graph, keyed by each lane's draft buffers, rows and geometry, then the scratch. +pub fn runGroup(run: anytype, score_rows: u64) !void { + const e = run.o.e; + const s = run.p.s; + if (!window_graph.enabled() or e.collective != null) return run.run(); + var h = std.hash.Wyhash.init(0); + for (run.group.ds, run.group.nexts, run.group.rows) |d, nx, rw| { + const g: window_graph.Geometry = .at(d.len + nx.len, d.cap); + for ([_]u64{ d.gen, d.kc.ptr, d.cap, nx.len, g.chunks, g.score_grid, rw, d.out.ptr, d.dlogits.ptr }) |v| h.update(std.mem.asBytes(&v)); + } + const key: window_graph.HeadManyKey = .{ .lanes = @intCast(run.group.ds.len), .rows = @intCast(run.total), .scores = run.group.scores, .scratch = s.h, .scratch_cap = s.cap, .score_rows = score_rows, .each = h.final() }; + try e.graphs.head_many.run(e.gpa, e.stream.*, key, run); +} diff --git a/zig/src/families/flashnext/cuda_mtp_roll.zig b/zig/src/families/flashnext/cuda_mtp_roll.zig new file mode 100644 index 0000000000..a4c512587e --- /dev/null +++ b/zig/src/families/flashnext/cuda_mtp_roll.zig @@ -0,0 +1,97 @@ +//! A partial keep: the kept rows replayed over the DeltaNet layers (a graph on one GPU), or moved on with pending rows. +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const parts = @import("cuda_batch_state.zig"); +const prompt = @import("cuda_prompt.zig"); +const gdn = @import("cuda_gdn.zig"); +const Draft = @import("cuda_mtp_draft.zig").Draft; +const window_graph = @import("cuda_window_graph.zig"); +const std = @import("std"); +const conv_bytes = (gdn.taps - 1) * gdn.conv_dim * 2; + +fn cint(v: usize) c_int { + return @intCast(v); +} + +/// fn_late_roll: each DeltaNet layer's conv window and the n-gram tail from their saved copies, `keep` rows on. +pub fn lateRoll(o: anytype, seq: *prompt.Seq, d: *Draft, keep: usize) !void { + var t: parts.RollJobs = .{}; + var n: usize = 0; + for (o.e.layers[0..o.e.n_layers], 0..) |*lw, layer| { + if (!lw.linear) continue; + const lk = d.keep.layers[layer]; + t.j[n] = .{ .win = seq.conv[n].ptr, .hist = lk.hist, .rows = lk.proj, .stride = @intCast(gdn.proj_width), .taps = @intCast(gdn.taps - 1), .c = @intCast(gdn.conv_dim) }; + n += 1; + } + t.j[n] = .{ .win = seq.tail.ptr, .hist = d.saved.ptr + Draft.tail_at, .rows = d.saved.ptr + Draft.nrow_at, .stride = @intCast(wt.wide), .taps = @intCast(wt.ple_tail_n), .c = @intCast(wt.wide) }; + n += 1; + t.n = @intCast(n); + var a: cuda.Args = .{}; + a.add(t); + a.add(cint(keep)); + const widest = @max(gdn.conv_dim, wt.wide); + try cuda.launch.launch(o.e.batch_k.late_roll, .{ .grid = .{ .x = @intCast((widest + 255) / 256), .y = @intCast(n) }, .block = .{ .x = 256 } }, o.e.stream.*, &a); +} + +/// The kept rows replayed on the window's start state (device work alone: a graph holds it). +const Replay = struct { + e: *prompt.Engine, + st: cuda.Stream, + s: *parts.Scratch, + seq: *prompt.Seq, + d: *Draft, + keep: usize, + + pub fn run(r: @This()) !void { + const k = &r.e.batch_k; + const rs = r.s; + try copyAsync(r.st, r.seq.tail.ptr, r.d.saved.ptr + Draft.tail_at, wt.ple_tail_n * wt.wide * 2); + try shift(k, r.st, r.seq.tail.ptr, r.d.saved.ptr + Draft.nrow_at, wt.wide, r.keep, wt.ple_tail_n, wt.wide); + var li: usize = 0; + for (r.e.layers[0..r.e.n_layers], 0..) |*lw, layer| { + if (!lw.linear) continue; + defer li += 1; + // the sequence's state is still the window's start (the pass wrote its end state to the spare) + const lk = r.d.keep.layers[layer]; + try copyAsync(r.st, r.seq.conv[li].ptr, lk.hist, conv_bytes); + var f: cuda.Args = .{}; + for ([_]u64{ lk.proj, r.seq.conv[li].ptr, lw.cw.ptr, lw.a_log.ptr, lw.dt.ptr, rs.gq, rs.gk, rs.gv, rs.ggb }) |v| f.add(v); + try cuda.launch.launch(k.front, .{ .grid = .{ .x = @intCast(r.keep), .y = gdn.nk + gdn.nv }, .block = .{ .x = 256 } }, r.st, &f); + var a: cuda.Args = .{}; + for ([_]u64{ rs.gq, rs.gk, rs.gv, rs.ggb, r.seq.state[li].ptr }) |v| a.add(v); + a.add(cint(r.keep)); + a.add(rs.ys); + try cuda.launch.launch(k.chain, .{ .grid = .{ .x = gdn.nv, .y = gdn.dv / 16 }, .block = .{ .x = 128 } }, r.st, &a); + try shift(k, r.st, r.seq.conv[li].ptr, lk.proj, gdn.proj_width, r.keep, gdn.taps - 1, gdn.conv_dim); + } + } +}; + +/// The replay of `keep` kept rows; on one GPU its ~146 launches as one graph, keyed by every buffer they touch. +pub fn replay(e: *prompt.Engine, s: *parts.Scratch, seq: *prompt.Seq, d: *Draft, keep: usize) !void { + const r: Replay = .{ .e = e, .st = e.stream.*, .s = s, .seq = seq, .d = d, .keep = keep }; + if (!window_graph.enabled() or e.collective != null) return r.run(); + var h = std.hash.Wyhash.init(0); + for ([_]u64{ seq.gen, d.gen, seq.tail.ptr, d.saved.ptr }) |v| h.update(std.mem.asBytes(&v)); + for (seq.conv) |b| h.update(std.mem.asBytes(&b.ptr)); + for (seq.state) |b| h.update(std.mem.asBytes(&b.ptr)); + h.update(std.mem.asBytes(&d.keep)); + const key: window_graph.RollbackKey = .{ .keep = @intCast(keep), .scratch = s.h, .scratch_cap = s.cap, .bufs = h.final() }; + try e.graphs.rollback.run(e.gpa, e.stream.*, key, r); +} + +fn copyAsync(stream: cuda.Stream, dst: u64, src: u64, n: usize) !void { + if (n == 0) return; + try stream.d.check(stream.d.api.cuMemcpyDtoDAsync_v2(dst, src, n, stream.handle), "cuMemcpyDtoDAsync"); +} + +fn shift(k: *const parts.Kernels, st: cuda.Stream, win: u64, rows_ptr: u64, stride: usize, n: usize, taps: usize, c: usize) !void { + var a: cuda.Args = .{}; + a.add(win); + a.add(rows_ptr); + a.add(@as(c_longlong, @intCast(stride))); + a.add(cint(n)); + a.add(cint(taps)); + a.add(cint(c)); + try cuda.launch.launch(k.shift, .{ .grid = .{ .x = @intCast((c + 255) / 256) }, .block = .{ .x = 256 } }, st, &a); +} diff --git a/zig/src/families/flashnext/cuda_ple.zig b/zig/src/families/flashnext/cuda_ple.zig index e4ba3ca438..f1b01a8664 100644 --- a/zig/src/families/flashnext/cuda_ple.zig +++ b/zig/src/families/flashnext/cuda_ple.zig @@ -349,6 +349,7 @@ pub const Ngram = struct { read_mode: ple_reader.Mode = .mmap, warm_ahead: bool = false, pool: ?*ple_pool.Pool = null, // the gather's parked workers (null: every row on the caller's thread) + warmer: ?std.Thread = null, // warm's read-ahead thread until joined (waitWarm, deinit) pub fn init(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8) !Ngram { const cfg = try readSpec(gpa, io, model_dir); @@ -424,6 +425,7 @@ pub const Ngram = struct { } pub fn deinit(g: *Ngram) void { + g.waitWarm(); if (g.pool) |p| p.destroy(); for (g.files.items) |*o| o.close(g.gpa); g.files.deinit(g.gpa); @@ -434,6 +436,40 @@ pub const Ngram = struct { return g.hash.heads * g.hash.dh; } + /// The table read ahead into the page cache after load, so a cold prompt reads RAM. TF_FLASHNEXT_TABLE_WARM=0: off. + pub fn warm(g: *Ngram) void { + if (@import("builtin").os.tag != .linux) return; + if (std.c.getenv("TF_FLASHNEXT_TABLE_WARM")) |v| if (std.mem.eql(u8, std.mem.span(v), "0")) return; + g.warmer = std.Thread.spawn(.{}, warmMain, .{g.*}) catch null; + } + + /// Returns once warm's read-ahead is queued. + pub fn waitWarm(g: *Ngram) void { + if (g.warmer) |t| t.join(); + g.warmer = null; + } + + fn warmMain(g: Ngram) void { + var bytes: usize = 0; + for (0..g.hash.shards) |shard| switch (g.format) { + .nvfp4 => bytes += advise(g.tables[shard].codes) + advise(g.tables[shard].scales), + .mlx => bytes += advise(g.w[shard]) + advise(g.s[shard]) + advise(g.b[shard]), + }; + std.debug.print("n-gram table: {d:.1} GiB read ahead into the page cache\n", .{@as(f64, @floatFromInt(bytes)) / (1 << 30)}); + } + + /// MADV_WILLNEED over the pages of `bytes`, a 128 KiB readahead window a call (the most the kernel reads a call). + fn advise(bytes: []const u8) usize { + if (bytes.len == 0) return 0; + const page = std.heap.pageSize(); + const step = 128 << 10; + const lo = std.mem.alignBackward(usize, @intFromPtr(bytes.ptr), page); + const hi = std.mem.alignForward(usize, @intFromPtr(bytes.ptr) + bytes.len, page); + var at = lo; + while (at < hi) : (at += step) _ = std.os.linux.madvise(@ptrFromInt(at), @min(step, hi - at), std.os.linux.MADV.WILLNEED); + return bytes.len; + } + /// Rows [lo, hi) of a chunk's token-major row ids `ids`, each dequantized into its head's slice of `out`. pub fn fill(g: *const Ngram, ids: []const i64, out: []u16, lo: usize, hi: usize) !void { const dh = g.hash.dh; diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index fc70793092..da785f8b10 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -59,7 +59,7 @@ pub fn argmax(raw: []const u8) u32 { return @intCast(at); } /// `n` zeroed buffers, the fill on the compute stream: a legacy-stream memset returns before it lands. -fn zeros(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: cuda.Stream, n: usize, len: usize) ![]cuda.DeviceBuffer { +pub fn zeros(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: cuda.Stream, n: usize, len: usize) ![]cuda.DeviceBuffer { const Zeroed = struct { driver: *cuda.Driver, stream: cuda.Stream, @@ -143,6 +143,7 @@ pub const Engine = struct { /// Rows a lane's prompt pass takes while another stream decodes: TF_FLASHNEXT_CHUNK. busy_chunk: usize = batch.default_chunk, round_graph: round_graph.Cache = .{}, + graphs: @import("cuda_window_graph.zig").Graphs = .{}, // drafted windows, draft-head pieces, rollbacks on one GPU image_factory: ?*@import("cuda_images.zig").Vision = null, collective: ?*@import("cuda_collective.zig").Collective = null, /// Each layer's Living Weights pocket, from a taught folder's shards or the learner: every pass adds it. @@ -168,6 +169,7 @@ pub const Engine = struct { e.collective = null; e.image_factory = null; e.round_graph = .{}; + e.graphs = .{}; errdefer for (e.layers[0..e.n_layers]) |*l| l.free(gpa); e.model_dir = try gpa.dupe(u8, model_dir); errdefer gpa.free(e.model_dir); @@ -234,6 +236,7 @@ pub const Engine = struct { pub fn deinit(e: *Engine) void { e.round_graph.clear(); + e.graphs.clear(e.gpa); e.pockets.deinit(); if (e.batch_s) |*bs| bs.free(e.gpa); if (e.scores) |*b| b.free(); @@ -304,6 +307,7 @@ pub const Engine = struct { if (cap == 0 or cap > max_tokens) return error.UnexpectedTensor; _ = try seqBytesPaged(cap, if (paged) |p| p.tokens else 0); var seq: Seq = .{ + .gen = @import("cuda_window_graph.zig").nextGen(), .stream = e.stream.*, .paged = paged, .conv = &.{}, @@ -362,6 +366,7 @@ pub const Engine = struct { if (b.len >= need) return b.ptr; try e.complete(); e.round_graph.clear(); + e.graphs.clear(e.gpa); b.free(); e.scores = null; } @@ -380,10 +385,14 @@ pub const Engine = struct { if (bs.cap >= rows) return bs; try e.complete(); e.round_graph.clear(); + e.graphs.clear(e.gpa); bs.free(e.gpa); e.batch_s = null; } - e.batch_s = try batch.Scratch.alloc(e.gpa, e.driver, e.stream.*, rows); + // head_local: the draft head's lanes on one GPU, a pass's rows split; native_part only where the lane reduces + const parts = @import("cuda_batch_state.zig"); + const z: parts.Sizes = .{ .head_rows = if (e.collective != null) @max(rows, parts.head_rows_one) else parts.head_rows_one, .lane_reduce = !e.batch_k.native_lane.clusters }; + e.batch_s = try batch.Scratch.alloc(e.gpa, e.driver, e.stream.*, rows, z); return &e.batch_s.?; } @@ -404,6 +413,7 @@ pub const Engine = struct { .seq = seq.view(scores), .hooks = hooks, }; + if (try @import("cuda_prompt_windows.zig").forwardGraph(e, Tri, &p, seq, ids, hooks)) return; var at: usize = 0; while (at < ids.len) { const n = @min(s.cap, ids.len - at); @@ -411,7 +421,7 @@ pub const Engine = struct { e.ngram.warmAhead(ids, at + n, @min(ids.len, at + n + s.cap)); p.hooks.logits = if (last) hooks.logits else .none; p.hooks.streams = if (hooks.streams != 0) hooks.streams + at * wide * 2 else 0; - try p.chunk(ids[at..][0..n], seq.fed + at, if (last and hooks.logits != .none) seq.logits else null); + try p.chunk(ids[at..][0..n], seq.fed + at, if (last and hooks.logits != .none and hooks.host_logits) seq.logits else null); at += n; } if (hooks.logits == .none) try e.complete(); @@ -478,6 +488,7 @@ pub const Seq = struct { stream: ?cuda.Stream = null, media: ?@import("cuda_media.zig").State = null, rank_id: u64 = 0, + gen: u64 = 0, // its graphs' keys (cuda_window_graph.nextGen): unique among the engine's sequences conv: []cuda.DeviceBuffer, state: []cuda.DeviceBuffer, kc: []cuda.DeviceBuffer, @@ -496,7 +507,7 @@ pub const Seq = struct { /// The caches a batched pass reads and leaves; `scores` the engine's block scores (Engine.reserveScores). pub fn view(s: *Seq, scores: u64) batch.SeqState { - return .{ .media = if (s.media) |*m| m.view() else null, .conv = s.conv, .state = s.state, .kc = s.kc, .vc = s.vc, .ikc = s.ikc, .pooled = s.pooled, .tail = s.tail, .history = &s.history, .cap = s.cap, .scores = scores, .pg = s.pg() }; + return .{ .gen = s.gen, .media = if (s.media) |*m| m.view() else null, .conv = s.conv, .state = s.state, .kc = s.kc, .vc = s.vc, .ikc = s.ikc, .pooled = s.pooled, .tail = s.tail, .history = &s.history, .cap = s.cap, .scores = scores, .pg = s.pg() }; } /// The attention kernels' page argument (zero: contiguous caches). diff --git a/zig/src/families/flashnext/cuda_prompt_windows.zig b/zig/src/families/flashnext/cuda_prompt_windows.zig index 0f04a3f590..0fc3e3d951 100644 --- a/zig/src/families/flashnext/cuda_prompt_windows.zig +++ b/zig/src/families/flashnext/cuda_prompt_windows.zig @@ -4,9 +4,21 @@ const prompt = @import("cuda_prompt.zig"); const batch = prompt.batch; const pocket = @import("cuda_pocket.zig"); const vocab = @import("cuda_prompt_wt.zig").vocab; +const window_graph = @import("cuda_window_graph.zig"); const Engine = prompt.Engine; const Seq = prompt.Seq; +/// A drafted window on one GPU as a graph (one chunk, scores left on the device, no media or adds); true once run. +pub fn forwardGraph(e: *Engine, comptime Tri: type, p: *batch.Pass(Tri), seq: *Seq, ids: []const u32, hooks: batch.Hooks) !bool { + if (!window_graph.enabled() or hooks.logits != .all or hooks.host_logits or hooks.moe_add != null) return false; + if (ids.len > @import("cuda_mtp.zig").max_rows or ids.len > p.s.cap or e.collective != null or seq.media != null or p.m.pocket != null) return false; + const key = window_graph.Key.init(seq, seq.fed, ids.len, p.s, p.seq.scores, @intFromPtr(p.m.layers.ptr), hooks); + try window_graph.stageWindow(p, ids, seq.fed); + try e.graphs.window.run(e.gpa, e.stream.*, key, window_graph.Chunk(batch.Pass(Tri)){ .p = p, .ids = ids, .pos0 = seq.fed }); + seq.fed += ids.len; + return true; +} + /// A stream's window in a shared pass: its rows, and where its rollback rows, hidden rows and scores go (0: none). pub const Window = struct { seq: *Seq, ids: []const u32, keep: ?*const batch.Keep = null, streams: u64 = 0, logits: u64 = 0 }; @@ -54,14 +66,23 @@ pub fn windows(e: *Engine, comptime Tri: type, tri: Tri, wins: []const Window, o .lanes = lanes, .hooks = .{ .logits = .all, .logits_out = out }, }; - try p.chunk(ids, 0, null); + const media = for (wins) |w| { + if (w.seq.media != null) break true; + } else false; + if (window_graph.enabled() and reads == .none and p.m.pocket == null and e.collective == null and !media) { + // on one GPU: staged, then the pass as a graph keyed by every lane's buffers + const key = window_graph.SharedKey.init(lanes, s, scores, @intFromPtr(p.m.layers.ptr), out); + try window_graph.stageWindow(&p, ids, 0); + try e.graphs.shared.run(e.gpa, e.stream.*, key, window_graph.Chunk(batch.Pass(Tri)){ .p = &p, .ids = ids, .pos0 = 0 }); + } else try p.chunk(ids, 0, null); try e.complete(); - // each stream's last row in its host scores, as its lone pass leaves them + // each stream's last row in its host scores only when asked: the drafted windows' draws read `out` const d = e.stream.d; + const host = @import("cuda_batch_state.zig").verifyHostLogits(); at = 0; for (wins) |w| { at += w.ids.len; - try d.check(d.api.cuMemcpyDtoH_v2(w.seq.logits.ptr, out + (at - 1) * vocab * 2, vocab * 2), "cuMemcpyDtoH"); + if (host) try d.check(d.api.cuMemcpyDtoH_v2(w.seq.logits.ptr, out + (at - 1) * vocab * 2, vocab * 2), "cuMemcpyDtoH"); w.seq.fed += w.ids.len; } } diff --git a/zig/src/families/flashnext/cuda_prompt_wt.zig b/zig/src/families/flashnext/cuda_prompt_wt.zig index 732e0ce77e..cce6d97d36 100644 --- a/zig/src/families/flashnext/cuda_prompt_wt.zig +++ b/zig/src/families/flashnext/cuda_prompt_wt.zig @@ -78,7 +78,10 @@ const Shard = struct { return .{ .f = try core.safetensors.File.open(gpa, io, path), .io = io }; } + /// Its pages out of the page cache: on unified memory they starve the n-gram table (TF_FLASHNEXT_DROP_LOADED=0). fn close(self: *Shard) void { + const keep = if (std.c.getenv("TF_FLASHNEXT_DROP_LOADED")) |v| std.mem.eql(u8, std.mem.span(v), "0") else false; + if (!keep) self.f.release(); self.f.close(self.io); } diff --git a/zig/src/families/flashnext/cuda_scratch_sizes.zig b/zig/src/families/flashnext/cuda_scratch_sizes.zig new file mode 100644 index 0000000000..2646140991 --- /dev/null +++ b/zig/src/families/flashnext/cuda_scratch_sizes.zig @@ -0,0 +1,21 @@ +//! The prompt scratch's buffers sized by use (Scratch.bytes / alloc): head_local, native_part, hc_part and part. +const wt = @import("cuda_prompt_wt.zig"); + +const hc_sk = wt.hc_sk; +const out_sk = wt.out_sk; + +/// head_local's rows on one GPU: the draft head's lanes (<= 16) times its ids; a split head needs a pass's rows. +pub const head_rows_one: usize = 32; + +/// head_local's rows, and whether the lane reduces K slices through native_part (with clusters it never does). +pub const Sizes = struct { head_rows: usize = head_rows_one, lane_reduce: bool = true }; + +/// hc_part's rows of (low + 4) floats: hc_sk slices below 64 rows (legacy kernel), <= 8 below 256 (lane), one beyond. +pub fn hcPartRows(cap: usize) usize { + return @max(cap, @max(hc_sk * @min(cap, 63), @as(usize, 8) * @min(cap, 255))); +} + +/// part's rows of dims: an out projection's slice plus the zero slice, or a legacy lane tile's out_sk below 64 rows. +pub fn partRows(cap: usize) usize { + return @max(2 * cap, out_sk * @min(cap, 63)); +} diff --git a/zig/src/families/flashnext/cuda_seq_parts.zig b/zig/src/families/flashnext/cuda_seq_parts.zig new file mode 100644 index 0000000000..551e517918 --- /dev/null +++ b/zig/src/families/flashnext/cuda_seq_parts.zig @@ -0,0 +1,73 @@ +//! A contiguous sequence whose conv windows, states and tail are views of one reused block (cuda_lane_parts.zig). +const std = @import("std"); +const cuda = @import("cuda"); +const gdn = @import("cuda_gdn.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const prompt = @import("cuda_prompt.zig"); +const Engine = prompt.Engine; +const Seq = prompt.Seq; +const linear_n = wt.linear_n; +const attn_n = wt.attn_n; +const state_n = wt.state_n; +const ple_tail_n = wt.ple_tail_n; +const wide = wt.wide; +const vocab = wt.vocab; +const kv_heads = wt.kv_heads; +const head_dim = wt.head_dim; +const index_dim = wt.index_dim; + +const conv_part = (gdn.taps - 1) * gdn.conv_dim * 2; + +/// A sequence's fixed-size buffers in one block (newSeqIn): its conv windows, DeltaNet states and n-gram tail. +pub fn seqPartsBytes() usize { + const a = std.mem.alignForward; + return linear_n * a(usize, conv_part, 256) + linear_n * a(usize, state_n * 4, 256) + a(usize, ple_tail_n * wide * 2, 256); +} + +/// `n` borrowed views of `len` bytes from `p` on, 256-byte aligned (the block stays its owner's). +fn viewsAt(gpa: std.mem.Allocator, d: *const cuda.Driver, p: *u64, n: usize, len: usize) ![]cuda.DeviceBuffer { + const bs = try gpa.alloc(cuda.DeviceBuffer, n); + for (bs) |*b| b.* = viewAt(d, p, len); + return bs; +} + +fn viewAt(d: *const cuda.Driver, p: *u64, len: usize) cuda.DeviceBuffer { + const b: cuda.DeviceBuffer = .{ .d = d, .ptr = p.*, .len = len, .borrowed = true }; + p.* += std.mem.alignForward(usize, len, 256); + return b; +} + +/// A contiguous sequence whose conv windows, states and tail are views of `block` (zeroed here); its caches its own. +pub fn newSeqIn(e: *Engine, cap: usize, block: cuda.DeviceBuffer) !Seq { + if (block.len < seqPartsBytes()) return error.UnexpectedTensor; + e.round_graph.clear(); + if (cap == 0 or cap > prompt.max_tokens) return error.UnexpectedTensor; + _ = try prompt.seqBytesPaged(cap, 0); + var seq: Seq = .{ + .gen = @import("cuda_window_graph.zig").nextGen(), + .stream = e.stream.*, + .conv = &.{}, + .state = &.{}, + .kc = &.{}, + .vc = &.{}, + .ikc = &.{}, + .pooled = &.{}, + .tail = .{ .d = e.driver, .ptr = 0, .len = 0 }, + .history = .{ e.eos, e.eos }, + .fed = 0, + .logits = &.{}, + .cap = cap, + }; + errdefer seq.deinit(e.gpa); + seq.logits = try e.gpa.alloc(u8, vocab * 2); + try e.stream.d.check(e.stream.d.api.cuMemsetD8Async(block.ptr, 0, seqPartsBytes(), e.stream.handle), "cuMemsetD8Async"); + var at = block.ptr; + seq.conv = try viewsAt(e.gpa, e.driver, &at, linear_n, conv_part); + seq.state = try viewsAt(e.gpa, e.driver, &at, linear_n, state_n * 4); + seq.tail = viewAt(e.driver, &at, ple_tail_n * wide * 2); + seq.kc = try prompt.zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * kv_heads * head_dim * 2); + seq.vc = try prompt.zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * kv_heads * head_dim * 2); + seq.ikc = try prompt.zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * index_dim * 2); + seq.pooled = try prompt.zeros(e.gpa, e.driver, e.stream.*, attn_n, prompt.pooledBytes(cap)); + return seq; +} diff --git a/zig/src/families/flashnext/cuda_sequence_completion_test.zig b/zig/src/families/flashnext/cuda_sequence_completion_test.zig index ddf4e69bfa..2a90cf5f6f 100644 --- a/zig/src/families/flashnext/cuda_sequence_completion_test.zig +++ b/zig/src/families/flashnext/cuda_sequence_completion_test.zig @@ -171,7 +171,7 @@ test "Flash Next completion: scratch growth refusal after queued inverse-frequen d.api.cuEventDestroy_v2 = Fake.destroy; d.api.cuMemcpyHtoDAsync_v2 = Fake.copyAndRefuseNext; const stream: cuda.Stream = .{ .d = &d, .handle = @ptrFromInt(330) }; - try t.expectError(error.OutOfDeviceMemory, prompt.batch.Scratch.alloc(t.allocator, &d, stream, 1)); + try t.expectError(error.OutOfDeviceMemory, prompt.batch.Scratch.alloc(t.allocator, &d, stream, 1, .{})); const cleanup_syncs = Fake.syncs; _ = try cuda.memory.setAllocationLimit(null); try stream.synchronize(); // also clean the old-main control's deliberately unfinished copy source diff --git a/zig/src/families/flashnext/cuda_tiles.zig b/zig/src/families/flashnext/cuda_tiles.zig index a623d411af..b62e0e79fd 100644 --- a/zig/src/families/flashnext/cuda_tiles.zig +++ b/zig/src/families/flashnext/cuda_tiles.zig @@ -35,14 +35,14 @@ fn sweepOf(m: usize, bm: usize, k: usize) usize { pub fn matmul(k: *const Kernels, stream: cuda.Stream, x: u64, xs: u64, w: Q4, out: u64, part: u64, native_part: u64, m: usize, sk: usize) !usize { if (w.native) |native| { if (sk == 1) { - try k.native_lane.matmul(stream, x, w.k, m, native, out, native_part); + try k.nativeMatmul(stream, x, w.k, m, native, out, native_part, .bf16); return 1; } - // Legacy callers consume SK slices; the native lane has already reduced its own shape-driven slices. - try k.native_lane.matmulF32(stream, x, w.k, m, native, part, native_part); + // the native lane has reduced its own slices: that slice and one zero slice give the old sums' bits (-0 to +0) + try k.nativeMatmul(stream, x, w.k, m, native, part, native_part, .f32); const bytes = m * w.n * 4; - if (sk > 1) try stream.d.check(stream.d.api.cuMemsetD8Async(part + bytes, 0, bytes * (sk - 1), stream.handle), "cuMemsetD8Async"); - return sk; + try stream.d.check(stream.d.api.cuMemsetD8Async(part + bytes, 0, bytes, stream.handle), "cuMemsetD8Async"); + return 2; } // from 64 rows the prompt tile (flashnext_fact.cu), bit for bit the 16-row lane tile const big = m >= 64; diff --git a/zig/src/families/flashnext/cuda_window_graph.zig b/zig/src/families/flashnext/cuda_window_graph.zig new file mode 100644 index 0000000000..3612ecdcc7 --- /dev/null +++ b/zig/src/families/flashnext/cuda_window_graph.zig @@ -0,0 +1,248 @@ +//! Drafted windows' passes, the draft head's pieces and rollback replays as CUDA graphs on one GPU, one key a graph. +const std = @import("std"); +const cuda = @import("cuda"); +const parts = @import("cuda_batch_state.zig"); +const wt = @import("cuda_prompt_wt.zig"); + +// A key holds every host value its kernels take; uploads stay outside the graph (stageWindow); caches clear on realloc. + +/// Graphs a cache holds at most; past it the cache starts over (finished sequences' graphs stay until then). +const max_graphs = 128; + +var gens = std.atomic.Value(u64).init(1); + +/// A sequence's or a draft's own number for the keys: a freed one's graphs never match a new one at its addresses. +pub fn nextGen() u64 { + return gens.fetchAdd(1, .monotonic); +} + +/// On by default (the same kernels, launched from a graph); TF_FLASHNEXT_WINDOW_GRAPH=0: every pass eagerly. +pub fn enabled() bool { + const v = std.c.getenv("TF_FLASHNEXT_WINDOW_GRAPH") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// attnRows' host values for rows ending at `ends` in a `cap`-token cache: dense chunks, the score grid (0: none). +pub const Geometry = struct { + chunks: u32, + score_grid: u32, + + pub fn at(ends: usize, cap: usize) Geometry { + const ratio = wt.index_ratio; + const nb = parts.scoreStride(cap); + const scored = ends / ratio > wt.index_top; + return .{ + .chunks = @intCast((@min(ends, wt.dense_keys) + 511) / 512), + .score_grid = if (scored) @intCast((@min(nb, (ends + ratio - 1) / ratio) + 63) / 64) else 0, + }; + } +}; + +/// A verify window's pass (Engine.forward); `hooks` hashes the scores out, the streams and every pointer of its keep. +pub const Key = struct { + gen: u64, + rows: u32, + geometry: Geometry, + cap: u64, + kc: u64, + state: u64, + scratch: u64, + scratch_cap: u64, + scores: u64, + model: u64, + hooks: u64, + + pub fn init(seq: anytype, pos0: usize, rows: usize, s: *const parts.Scratch, scores: u64, model: u64, hooks: parts.Hooks) Key { + var h = std.hash.Wyhash.init(0); + h.update(std.mem.asBytes(&hooks.logits_out)); + h.update(std.mem.asBytes(&hooks.streams)); + if (hooks.keep) |k| h.update(std.mem.asBytes(k)); + return .{ + .gen = seq.gen, + .rows = @intCast(rows), + .geometry = .at(pos0 + rows, seq.cap), + .cap = seq.cap, + .kc = seq.kc[0].ptr, + .state = seq.state[0].ptr, + .scratch = s.h, + .scratch_cap = s.cap, + .scores = scores, + .model = model, + .hooks = h.final(), + }; + } +}; + +/// One piece of the draft head's pass (forwardHead): its rows, where its input rows come from, whether it scores. +pub const HeadKey = struct { + gen: u64, + rows: u32, + geometry: Geometry, + scores: bool, + cap: u64, + kc: u64, + input: u64, + out: u64, + dlogits: u64, + scratch: u64, + scratch_cap: u64, + score_rows: u64, +}; + +/// Several streams' drafted windows in one pass (Engine.windows): each lane's buffers, rows and geometry, hashed. +pub const SharedKey = struct { + lanes: u32, + rows: u32, + scratch: u64, + scratch_cap: u64, + scores: u64, + model: u64, + out: u64, + each: u64, + + pub fn init(lanes: []const parts.Lane, s: *const parts.Scratch, scores: u64, model: u64, out: u64) SharedKey { + var h = std.hash.Wyhash.init(0); + var rows: usize = 0; + for (lanes) |l| { + const g: Geometry = .at(l.pos + l.rows, l.seq.cap); + for ([_]u64{ l.seq.gen, l.seq.kc[0].ptr, l.seq.state[0].ptr, l.seq.cap, l.rows, g.chunks, g.score_grid, l.streams, l.logits }) |v| h.update(std.mem.asBytes(&v)); + if (l.keep) |k| h.update(std.mem.asBytes(k)); + rows += l.rows; + } + return .{ .lanes = @intCast(lanes.len), .rows = @intCast(rows), .scratch = s.h, .scratch_cap = s.cap, .scores = scores, .model = model, .out = out, .each = h.final() }; + } +}; + +/// A group of streams' draft-head pieces in one pass (Ops.headGroup): each lane's buffers, rows and geometry hashed. +pub const HeadManyKey = struct { + lanes: u32, + rows: u32, + scores: bool, + scratch: u64, + scratch_cap: u64, + score_rows: u64, + each: u64, +}; + +/// A partial keep's replay over the DeltaNet layers (Ops.rollback): rows kept, scratch, every buffer it touches hashed. +pub const RollbackKey = struct { + keep: u32, + scratch: u64, + scratch_cap: u64, + bufs: u64, +}; + +pub fn Cache(comptime K: type) type { + return struct { + const Self = @This(); + map: std.AutoHashMapUnmanaged(K, cuda.graph.Exec) = .empty, + /// Keys run once eagerly and are captured the second time (a one-off key would never pay its capture back). + seen: std.AutoHashMapUnmanaged(K, void) = .empty, + captures: u64 = 0, + replays: u64 = 0, + /// A capture failed: every later pass runs eagerly (logged once). + off: bool = false, + + pub fn clear(c: *Self, gpa: std.mem.Allocator) void { + var it = c.map.valueIterator(); + while (it.next()) |e| e.deinit(); + c.map.clearAndFree(gpa); + c.seen.clearAndFree(gpa); + } + + /// `pass.run()` (kernels only: its uploads already queued) as `key`'s graph, captured first when new. + pub fn run(c: *Self, gpa: std.mem.Allocator, stream: cuda.Stream, key: K, pass: anytype) !void { + if (c.off) return pass.run(); + if (c.map.get(key)) |exec| { + c.replays += 1; + return exec.launchOn(stream); + } + if (c.map.count() >= max_graphs or c.seen.count() >= 4 * max_graphs) c.clear(gpa); + if (!c.seen.contains(key)) { + try c.seen.put(gpa, key, {}); + return pass.run(); + } + var exec = capture(stream, pass) catch |err| { + // the stream is out of capture again: this pass eagerly, and the ones after it + std.log.warn("{s} graphs off: capture failed ({s}), eager from here", .{ @typeName(K), @errorName(err) }); + c.off = true; + return pass.run(); + }; + errdefer exec.deinit(); + try c.map.put(gpa, key, exec); + c.captures += 1; + if (c.captures % 64 == 1) std.log.info("{s} graphs: {d} captured, {d} replayed, {d} held", .{ @typeName(K), c.captures, c.replays, c.map.count() }); + try exec.launchOn(stream); + } + }; +} + +/// The engine's graph caches, one per kind of pass. +pub const Graphs = struct { + window: Cache(Key) = .{}, + head: Cache(HeadKey) = .{}, + rollback: Cache(RollbackKey) = .{}, + shared: Cache(SharedKey) = .{}, + head_many: Cache(HeadManyKey) = .{}, + + pub fn clear(g: *Graphs, gpa: std.mem.Allocator) void { + g.window.clear(gpa); + g.head.clear(gpa); + g.rollback.clear(gpa); + g.shared.clear(gpa); + g.head_many.clear(gpa); + } +}; + +/// The pass's kernels captured and instantiated (uploaded to the device), nothing launched; the stream left uncaptured. +fn capture(stream: cuda.Stream, pass: anytype) !cuda.graph.Exec { + try cuda.graph.beginCapture(stream, .thread_local); + pass.run() catch |err| { + if (cuda.graph.endCapture(stream)) |g| { + var discarded = g; + discarded.deinit(); + } else |_| {} + return err; + }; + var graph = try cuda.graph.endCapture(stream); + defer graph.deinit(); + var exec = try graph.instantiate(); + errdefer exec.deinit(); + try exec.upload(stream); + return exec; +} + +/// A window's uploads ahead of its pass: ids and positions, then its n-gram rows gathered and queued behind one fence. +pub fn stageWindow(p: anytype, ids: []const u32, pos0: usize) !void { + const s = p.s; + const n = ids.len; + if (n == 0 or n > s.cap) return error.UnexpectedTensor; + try s.ready(); + p.pos0 = pos0; + const was = p.prepared; + p.prepared = true; + defer p.prepared = was; + try p.stage(ids); + const emb = s.hostEmb(n); + const xs = s.hostXs(n); + var gather: parts.Gather = .{ .ngram = p.m.ngram, .history = p.seq.history[0..], .ids = ids, .emb = emb, .xs = xs, .io = p.io, .lanes = p.lanes }; + gather.main(); + if (gather.failed) |err| return err; + try p.upload(s.ple_emb, std.mem.sliceAsBytes(emb)); + try p.upload(s.ple_xs, std.mem.sliceAsBytes(xs)); + try s.uploaded(); + p.staged = true; +} + +/// A pass over a window whose uploads stageWindow queued: Pass.chunk launches its kernels alone. +pub fn Chunk(comptime P: type) type { + return struct { + p: *P, + ids: []const u32, + pos0: usize, + + pub fn run(r: @This()) !void { + return r.p.chunk(r.ids, r.pos0, null); + } + }; +} diff --git a/zig/src/families/flashnext/cuda_zero_slice.zig b/zig/src/families/flashnext/cuda_zero_slice.zig index 0a082e6306..17e91dc66f 100644 --- a/zig/src/families/flashnext/cuda_zero_slice.zig +++ b/zig/src/families/flashnext/cuda_zero_slice.zig @@ -16,7 +16,7 @@ pub fn lane(k: *const parts.Kernels, stream: cuda.Stream, s: *parts.Scratch, zer if (spends(w.native != null, sk, m, w.n, s.cap)) s.part_zero = false; const native = w.native orelse return null; if (!zero or sk == 1) return null; - try k.native_lane.matmulF32(stream, x, w.k, m, native, part, s.native_part); + try k.nativeMatmul(stream, x, w.k, m, native, part, s.native_part, .f32); return 1; } diff --git a/zig/src/families/flashnext/cuda_zero_slice_test.zig b/zig/src/families/flashnext/cuda_zero_slice_test.zig index 7df165bc86..adb7a28f75 100644 --- a/zig/src/families/flashnext/cuda_zero_slice_test.zig +++ b/zig/src/families/flashnext/cuda_zero_slice_test.zig @@ -91,6 +91,7 @@ const Rig = struct { r.k.native_lane.f32_fns = @splat(@splat(f(&r.d, qmmf_id))); r.k.native_lane.reduce = @splat(f(&r.d, reduce_id)); r.k.native_lane.clusters = false; + r.k.prompt = null; // no BF16 prompt tile: the lane's handles alone r.s = undefined; r.s.cap = cap; r.s.part = base; diff --git a/zig/src/native/draft_sources_cuda_test.zig b/zig/src/native/draft_sources_cuda_test.zig index 2f652631a6..6bd1fc3e20 100644 --- a/zig/src/native/draft_sources_cuda_test.zig +++ b/zig/src/native/draft_sources_cuda_test.zig @@ -97,6 +97,7 @@ test "draft sources: CUDA Flash Next follows the loaded checkpoint head" { var b: @import("flashnext").native.Cuda = undefined; b.head = null; b.depth = 8; + b.depth_greedy = 8; try check(b.facts(), null, true, false); b.head = present(@TypeOf(b.head.?)); try check(b.facts(), null, true, true); diff --git a/zig/tests/cuda/contract_family_sets.zig b/zig/tests/cuda/contract_family_sets.zig index a602175cb3..4c420d0ce5 100644 --- a/zig/tests/cuda/contract_family_sets.zig +++ b/zig/tests/cuda/contract_family_sets.zig @@ -1,4 +1,4 @@ -//! The family sets in the full kernel contract: Nemotron's pocket, MiniMax-M3's, MOSS-Image's, GB10's grouped affine-4. +//! The contract's family sets: Nemotron's pocket, MiniMax-M3's, MOSS-Image's, GB10's grouped affine-4, Flash Next's. const std = @import("std"); const cuda = @import("cuda"); const check = @import("check.zig"); @@ -12,6 +12,7 @@ pub const Sets = struct { minimax: *@import("contract_minimax.zig").Set, moss: *@import("contract_moss_image.zig").Set, grouped: *Grouped, + flashnext: *@import("contract_flashnext_more.zig").Set, /// Each set's kernels appended to `list`, in this order. pub fn init(gpu: check.Gpu, s: cuda.Stream, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel), affine4: *cuda.qlinear.Affine4Kernels, cap: u32) !Sets { @@ -26,10 +27,12 @@ pub const Sets = struct { grouped.* = Grouped.init(gpu, s, affine4); errdefer grouped.deinit(); try grouped.append(cap, a, list); - return .{ .gpa = gpu.gpa, .pocket = pocket, .minimax = minimax, .moss = moss, .grouped = grouped }; + const flashnext = try @import("contract_flashnext_more.zig").Set.init(gpu, s, a, list); + return .{ .gpa = gpu.gpa, .pocket = pocket, .minimax = minimax, .moss = moss, .grouped = grouped, .flashnext = flashnext }; } pub fn deinit(sets: Sets) void { + sets.flashnext.deinit(); sets.grouped.deinit(); sets.gpa.destroy(sets.grouped); sets.moss.deinit(); diff --git a/zig/tests/cuda/contract_flashnext_late.zig b/zig/tests/cuda/contract_flashnext_late.zig new file mode 100644 index 0000000000..2e6c053141 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_late.zig @@ -0,0 +1,302 @@ +//! Pending DeltaNet rows under the contract: fn_gdn_front_log and fn_gdn_rows_late, the replaced kernels' bytes. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const parts = flash.batch_state; + +pub const Kind = enum { front, chain }; + +const nk = 16; +const nv = 48; +const dk = 128; +const dv = 128; +const taps = 4; +const conv_c = 2 * nk * dk + nv * dv; // a projection row's conv channels (q, k, v) +const proj_w = conv_c + nv * dv + 2 * nv; // a projection row: conv channels, z, b, a +const lr = parts.late_rows; +const k_row = nk * dk * 4; // a log row's bytes in each of its three parts +const v_row = nv * dv * 4; +const g_row = nv * 2 * 4; +const state_len = nv * dv * dk * 4; + +fn noSums(_: c.Shape) usize { + return 0; +} + +/// A verify window's rows: one, a greedy window and the widest a window may be (mtp.max_rows). +const served = [_]c.Shape{ + .{ .rows = 1, .n = nv * dv, .k = dk }, + .{ .rows = 9, .n = nv * dv, .k = dk }, + .{ .rows = 32, .n = nv * dv, .k = dk }, +}; + +pub fn spec(kind: Kind) c.Spec { + return switch (kind) { + .front => .{ .name = "cuda flashnext DeltaNet window rows into the pending log fn_gdn_front_log", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = "q, and k, v and the gates in the log's slots after the pending rows, bit-equal to fn_gdn_front's; other slots untouched" }, .replaces = "fn_gdn_front", .served = &served }, + .chain => .{ .name = "cuda flashnext DeltaNet chain after pending rows fn_gdn_rows_late", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = "ys, the end state and the stored state bit-equal to fn_gdn_rows_to over the pending rows, then over the window's" }, .replaces = "fn_gdn_rows_to", .in_place = true, .served = &served }, + }; +} + +pub const Fns = struct { front: cuda.Function, front_log: cuda.Function, rows_to: cuda.Function, rows_late: cuda.Function }; + +pub const Late = struct { + rig: rig_.Rig, + kind: Kind, + f: *const Fns, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + ins: [6]u64 = @splat(0), + outs: [3]u64 = @splat(0), + lens: [3]usize = @splat(0), + nout: usize = 0, + base: u64 = 0, + first: usize = 0, + pending: usize = 0, + /// The chain stores its state in place: every launch restores these bytes first. + state: u64 = 0, + fresh: ?[]u8 = null, + seen: std.ArrayList(u8) = .empty, + msg: [160]u8 = undefined, + + pub fn deinit(self: *Late) void { + if (self.fresh) |f| self.rig.gpu.gpa.free(f); + self.seen.deinit(self.rig.gpu.gpa); + self.rig.deinit(); + } + + pub fn kernel(self: *Late) g.Kernel { + return .{ .spec = spec(self.kind), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Late); + + fn f32s(self: *Late, n: usize, r: std.Random, lo: f32, hi: f32) ![]f32 { + const v = try self.rig.gpu.gpa.alloc(f32, n); + for (v) |*x| x.* = lo + r.float(f32) * (hi - lo); + return v; + } + + fn bfs(self: *Late, n: usize, r: std.Random, lo: f32, hi: f32) ![]u16 { + const v = try self.rig.gpu.gpa.alloc(u16, n); + for (v) |*x| x.* = c.toBf16(lo + r.float(f32) * (hi - lo)); + return v; + } + + fn out(self: *Late, len: usize) !u64 { + const p = try self.rig.output(len); + self.outs[self.nout] = p; + self.lens[self.nout] = len; + self.nout += 1; + return p; + } + + /// The log's first pending slot and pending count: none, at the store threshold, past it; the window may wrap. + fn place(self: *Late, rows: usize, seed: u64) void { + const want = ([_]usize{ 0, parts.late_store, parts.late_store + 5 })[seed % 3]; + self.pending = @min(want, lr - rows); + self.first = (lr - 3 + seed / 3) % lr; + } + + pub fn prepare(self: *Late, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + if (self.fresh) |f| self.rig.gpu.gpa.free(f); + self.fresh = null; + self.state = 0; + self.nout = 0; + self.shape = shape; + const m = shape.rows; + if (m == 0 or m > 32 or shape.n != nv * dv or shape.k != dk) return false; + var prng = std.Random.DefaultPrng.init(seed ^ 0x6c617465); + const r = prng.random(); + const a = self.rig.gpu.gpa; + const bad = fill == .nonfinite; + self.place(m, seed); + const ctl = [4]u32{ @intCast(self.first | self.pending << 16), 0, 0, 0 }; + switch (self.kind) { + .front => { + const proj = try self.bfs(m * proj_w, r, -1, 1); + defer a.free(proj); + const cs = try self.bfs((taps - 1) * conv_c, r, -1, 1); + defer a.free(cs); + const cw = try self.bfs(conv_c * taps, r, -0.5, 0.5); + defer a.free(cw); + const a_log = try self.f32s(nv, r, -1, 1); + defer a.free(a_log); + const dt = try self.f32s(nv, r, -1, 1); + defer a.free(dt); + if (bad) { + proj[3] = c.toBf16(std.math.nan(f32)); + proj[(m / 2) * proj_w + 2 * nk * dk + 7] = c.toBf16(std.math.inf(f32)); + proj[(m - 1) * proj_w + conv_c + nv * dv + 1] = c.toBf16(-std.math.inf(f32)); + } + for ([_][]const u8{ std.mem.sliceAsBytes(proj), std.mem.sliceAsBytes(cs), std.mem.sliceAsBytes(cw), std.mem.sliceAsBytes(a_log), std.mem.sliceAsBytes(dt), std.mem.sliceAsBytes(&ctl) }, 0..) |b, i| self.ins[i] = try self.rig.input(b); + _ = try self.out(m * k_row); // q, as fn_gdn_front's + _ = try self.out(lr * (k_row + v_row + g_row)); // the log: only the window's slots are written + }, + .chain => { + const qn = try self.f32s(m * nk * dk, r, -0.1, 0.1); + defer a.free(qn); + const ks = try self.f32s(lr * nk * dk, r, -0.1, 0.1); + defer a.free(ks); + const vs = try self.f32s(lr * nv * dv, r, -1, 1); + defer a.free(vs); + const gs = try self.f32s(lr * nv * 2, r, 0.5, 1); // a decay, then beta + defer a.free(gs); + const st = try self.f32s(nv * dv * dk, r, -0.5, 0.5); + defer a.free(st); + if (bad) { + qn[3] = std.math.nan(f32); + vs[((self.first + self.pending) % lr) * nv * dv + 5] = std.math.inf(f32); + st[dv * dk + 7] = -std.math.inf(f32); + } + const log = try std.mem.concat(a, u8, &.{ std.mem.sliceAsBytes(ks), std.mem.sliceAsBytes(vs), std.mem.sliceAsBytes(gs) }); + defer a.free(log); + for ([_][]const u8{ std.mem.sliceAsBytes(qn), log, std.mem.sliceAsBytes(&ctl), std.mem.sliceAsBytes(st) }, 0..) |b, i| self.ins[i] = try self.rig.input(b); + _ = try self.out(m * v_row); // ys + self.state = try self.out(state_len); // the state read, and stored in place past the threshold + _ = try self.out(state_len); // the state after the window + self.fresh = try a.dupe(u8, std.mem.sliceAsBytes(st)); + }, + } + var total: usize = 0; + for (self.lens[0..self.nout]) |l| total += l; + self.base = try self.rig.scratch(total); + return true; + } + + fn launchFront(self: *Late, fun: cuda.Function, q: u64, tail: []const u64, log: bool) !void { + var a: cuda.Args = .{}; + for ([_]u64{ self.ins[0], self.ins[1], self.ins[2], self.ins[3], self.ins[4], q }) |v| a.add(v); + for (tail) |v| a.add(v); + if (log) a.add(@as(c_int, @intCast(lr))); + try cuda.launch.launch(fun, .{ .grid = .{ .x = @intCast(self.shape.rows), .y = nk + nv }, .block = .{ .x = 256 } }, self.rig.s, &a); + } + + fn rowsTo(self: *Late, q: u64, k: u64, v: u64, gb: u64, in: u64, to: u64, rows: usize, ys: u64) !void { + var a: cuda.Args = .{}; + for ([_]u64{ q, k, v, gb, in, to }) |x| a.add(x); + a.add(@as(c_int, @intCast(rows))); + a.add(ys); + try cuda.launch.launch(self.f.rows_to, .{ .grid = .{ .x = nv, .y = dv / 16 }, .block = .{ .x = 128 } }, self.rig.s, &a); + } + + fn copy(self: *Late, dst: u64, src: u64, len: usize) !void { + try (cuda.DeviceBuffer{ .d = self.rig.gpu.d, .ptr = dst, .len = len }).copyFrom(0, src, len, self.rig.s.handle); + } + + /// `n` rows from log slot `from` on (wrapping) to or from contiguous rows: the log's three parts. + fn slots(self: *Late, log: u64, rows: [3]u64, from: usize, n: usize, into_log: bool) !void { + const sizes = [3]usize{ k_row, v_row, g_row }; + var at = log; + for (sizes, rows) |size, flat| { + for (0..n) |i| { + const slot = at + ((from + i) % lr) * size; + if (into_log) try self.copy(slot, flat + i * size, size) else try self.copy(flat + i * size, slot, size); + } + at += lr * size; + } + } + + fn ours(self: *Late, o: []const u64) !void { + switch (self.kind) { + .front => try self.launchFront(self.f.front_log, o[0], &.{ o[1], self.ins[5] }, true), + .chain => { + var a: cuda.Args = .{}; + for ([_]u64{ self.ins[0], self.ins[1], self.ins[2] }) |v| a.add(v); + a.add(@as(c_int, @intCast(lr))); + a.add(@as(c_int, @intCast(parts.late_store))); + a.add(o[1]); + a.add(o[2]); + a.add(@as(c_int, @intCast(self.shape.rows))); + a.add(o[0]); + try cuda.launch.launch(self.f.rows_late, .{ .grid = .{ .x = nv, .y = dv / 16 }, .block = .{ .x = 128 } }, self.rig.s, &a); + }, + } + } + + /// The replaced kernels: fn_gdn_front, rows copied to slots; fn_gdn_rows_to over the pending rows, then the window. + fn theirs(self: *Late, o: []const u64) !void { + const m = self.shape.rows; + const after = self.first + self.pending; + switch (self.kind) { + .front => { + const flat = [3]u64{ try self.rig.scratch(m * k_row), try self.rig.scratch(m * v_row), try self.rig.scratch(m * g_row) }; + try self.launchFront(self.f.front, o[0], &flat, false); + try self.slots(o[1], flat, after, m, true); + }, + .chain => { + const mid = try self.rig.scratch(state_len); + if (self.pending > 0) { + const np = self.pending; + const pend = [3]u64{ try self.rig.scratch(np * k_row), try self.rig.scratch(np * v_row), try self.rig.scratch(np * g_row) }; + try self.slots(self.ins[1], pend, self.first, np, false); + // the pending rows' read-outs are dropped: their q is zero, their ys go to scratch + try self.rowsTo(try self.rig.scratch(np * k_row), pend[0], pend[1], pend[2], self.ins[3], mid, np, try self.rig.scratch(np * v_row)); + } else try self.copy(mid, self.ins[3], state_len); + const win = [3]u64{ try self.rig.scratch(m * k_row), try self.rig.scratch(m * v_row), try self.rig.scratch(m * g_row) }; + try self.slots(self.ins[1], win, after, m, false); + try self.rowsTo(self.ins[0], win[0], win[1], win[2], mid, o[2], m, o[0]); + try self.copy(o[1], if (self.pending > parts.late_store) mid else self.ins[3], state_len); + }, + } + } + + fn restoreState(self: *Late) !void { + const f = self.fresh orelse return; + try (cuda.DeviceBuffer{ .d = self.rig.gpu.d, .ptr = self.state, .len = f.len }).upload(0, f); + } + + pub fn run(self: *Late) !void { + try self.restoreState(); + try self.ours(self.outs[0..self.nout]); + try self.rig.s.synchronize(); + } + + /// The scratch the replaced kernels write: the outputs' layout back to back, poisoned like them (0xff). + fn baseline(self: *Late) ![]const u8 { + var at = self.base; + var o: [3]u64 = undefined; + var total: usize = 0; + for (self.lens[0..self.nout], 0..) |l, i| { + o[i] = at; + at += l; + total += l; + } + const ff = try self.rig.gpu.gpa.alloc(u8, total); + defer self.rig.gpu.gpa.free(ff); + @memset(ff, 0xff); + try (cuda.DeviceBuffer{ .d = self.rig.gpu.d, .ptr = self.base, .len = total }).upload(0, ff); + try self.theirs(o[0..self.nout]); + try self.rig.s.synchronize(); + return self.rig.bytes(self.base, total); + } + + pub fn replaced(self: *Late) ![]const u8 { + return self.baseline(); + } + + /// Gate 1's stated rule: the outputs on the planted inputs equal the replaced kernels' bytes. + pub fn check(self: *Late, gate: g.Gate) !?[]const u8 { + _ = gate; + self.seen.clearRetainingCapacity(); + try self.seen.appendSlice(self.rig.gpu.gpa, try self.rig.outputBytes()); + const want = try self.baseline(); + if (std.mem.eql(u8, self.seen.items, want)) return null; + return try std.fmt.bufPrint(&self.msg, "bytes differ from the replaced kernels' from byte {d}", .{std.mem.indexOfDiff(u8, self.seen.items, want).?}); + } + + pub fn count(self: *Late) usize { + var total: usize = 0; + for (self.lens[0..self.nout]) |l| total += l; + return total; + } + + pub fn reference(_: *Late, _: []f64, _: []f64) !void { + return error.NoFloat64Reference; + } +}; diff --git a/zig/tests/cuda/contract_flashnext_more.zig b/zig/tests/cuda/contract_flashnext_more.zig new file mode 100644 index 0000000000..71a974c0bb --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_more.zig @@ -0,0 +1,61 @@ +//! Flash Next's newer CUDA kernels under the contract, registered as one group. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const tc = @import("check.zig"); +const late_k = @import("contract_flashnext_late.zig"); +const roll_k = @import("contract_flashnext_roll.zig"); +const wbnorm_k = @import("contract_flashnext_wbnorm.zig"); +const topk_k = @import("contract_flashnext_topk.zig"); +const flash = @import("flashnext"); + +const g = core.contract.gates; + +/// The kernels and their modules, registered in the contract's list; destroyed after the gates have run. +pub const Set = struct { + a: std.mem.Allocator, + batch: cuda.Module, + late_fns: late_k.Fns, + late: [2]late_k.Late, + roll_fns: roll_k.Fns, + roll: roll_k.Roll, + glue: flash.Glue, + wbnorm: [3]wbnorm_k.WbNorm, + mtp: cuda.Module, + topk: topk_k.TopK, + + pub fn init(gpu: tc.Gpu, s: cuda.Stream, a: std.mem.Allocator, kernels: *std.ArrayList(g.Kernel)) !*Set { + const self = try a.create(Set); + errdefer a.destroy(self); + self.a = a; + self.batch = try cuda.Module.load(gpu.d, cuda.kernels.flashnext_batch); + errdefer self.batch.unload(); + const b = &self.batch; + self.late_fns = .{ .front = try b.function("fn_gdn_front"), .front_log = try b.function("fn_gdn_front_log"), .rows_to = try b.function("fn_gdn_rows_to"), .rows_late = try b.function("fn_gdn_rows_late") }; + for (std.enums.values(late_k.Kind), &self.late) |kind, *x| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .kind = kind, .f = &self.late_fns }; + for (&self.late) |*x| try kernels.append(a, x.kernel()); + self.roll_fns = .{ .shift = try b.function("fn_shift_window"), .roll = try b.function("fn_late_roll") }; + self.roll = .{ .rig = .{ .gpu = gpu, .s = s }, .f = &self.roll_fns }; + try kernels.append(a, self.roll.kernel()); + self.glue = try flash.Glue.load(gpu.d); + errdefer self.glue.deinit(); + for (std.enums.values(wbnorm_k.Mode), &self.wbnorm) |mode, *x| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .mode = mode, .glue = &self.glue }; + for (&self.wbnorm) |*x| try kernels.append(a, x.kernel()); + self.mtp = try cuda.Module.load(gpu.d, cuda.kernels.flashnext_mtp); + errdefer self.mtp.unload(); + self.topk = .{ .rig = .{ .gpu = gpu, .s = s }, .f = try self.mtp.function("fn_row_topk") }; + try kernels.append(a, self.topk.kernel()); + return self; + } + + pub fn deinit(self: *Set) void { + for (&self.late) |*x| x.deinit(); + self.roll.deinit(); + for (&self.wbnorm) |*x| x.deinit(); + self.topk.deinit(); + self.mtp.unload(); + self.glue.deinit(); + self.batch.unload(); + self.a.destroy(self); + } +}; diff --git a/zig/tests/cuda/contract_flashnext_roll.zig b/zig/tests/cuda/contract_flashnext_roll.zig new file mode 100644 index 0000000000..c58b2c1175 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_roll.zig @@ -0,0 +1,154 @@ +//! A rolled-back stream's windows in one launch (fn_late_roll) against a copy and fn_shift_window a window. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; + +const nk = 16; +const nv = 48; +const dk = 128; +const dv = 128; +const taps = 4; +const conv_c = 2 * nk * dk + nv * dv; +const proj_w = conv_c + nv * dv + 2 * nv; +const hist_len = (taps - 1) * conv_c * 2; + +fn noSums(_: c.Shape) usize { + return 0; +} + +/// rows: the rollback's jobs (a stream's DeltaNet layers' windows); n and k fixed by the model. +pub fn spec() c.Spec { + return .{ .name = "cuda flashnext rollback windows fn_late_roll", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = "every window bit-equal to its saved copy then fn_shift_window by the kept rows" }, .replaces = "a copy and fn_shift_window a window", .served = &.{ .{ .rows = 2, .n = conv_c, .k = taps - 1 }, .{ .rows = 37, .n = conv_c, .k = taps - 1 } } }; +} + +pub const Fns = struct { shift: cuda.Function, roll: cuda.Function }; + +pub const Roll = struct { + rig: rig_.Rig, + f: *const Fns, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + keep: usize = 0, + ins: [2]u64 = @splat(0), + region: u64 = 0, + base: u64 = 0, + len: usize = 0, + fresh: ?[]u8 = null, + seen: std.ArrayList(u8) = .empty, + msg: [160]u8 = undefined, + + pub fn deinit(self: *Roll) void { + if (self.fresh) |f| self.rig.gpu.gpa.free(f); + self.seen.deinit(self.rig.gpu.gpa); + self.rig.deinit(); + } + + const vt = g.hooks(Roll); + + pub fn kernel(self: *Roll) g.Kernel { + return .{ .spec = spec(), .ctx = self, .vt = &vt }; + } + + /// Jobs of conv windows (taps - 1 rows of conv_c), each `keep` rows on from its saved copy over the pass's rows. + pub fn prepare(self: *Roll, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + if (self.fresh) |f| self.rig.gpu.gpa.free(f); + self.fresh = null; + self.shape = shape; + if (shape.rows < 2 or shape.rows > 40 or shape.n != conv_c or shape.k != taps - 1) return false; + const a = self.rig.gpu.gpa; + var prng = std.Random.DefaultPrng.init(seed ^ 0x726f6c6c); + const r = prng.random(); + self.keep = 1 + r.uintLessThan(usize, 32); + const rows_n = self.keep + taps; + const hist = try a.alloc(u16, shape.rows * hist_len / 2); + defer a.free(hist); + for (hist) |*x| x.* = c.toBf16(r.float(f32) * 2 - 1); + const rows = try a.alloc(u16, shape.rows * rows_n * proj_w); + defer a.free(rows); + for (rows) |*x| x.* = c.toBf16(r.float(f32) * 2 - 1); + if (fill == .nonfinite) hist[5] = c.toBf16(std.math.inf(f32)); + self.ins[0] = try self.rig.input(std.mem.sliceAsBytes(hist)); + self.ins[1] = try self.rig.input(std.mem.sliceAsBytes(rows)); + self.len = shape.rows * hist_len; + const fresh = try a.alloc(u8, self.len); + self.fresh = fresh; + @memset(fresh, 0xff); + self.region = try self.rig.output(self.len); + self.base = try self.rig.scratch(self.len); + return true; + } + + fn cint(v: usize) c_int { + return @intCast(v); + } + + /// The rollback's launch as cuda_mtp_roll.lateRoll makes it. + fn ours(self: *Roll, region: u64) !void { + var t: @import("flashnext").batch_state.RollJobs = .{ .n = @intCast(self.shape.rows) }; + for (t.j[0..self.shape.rows], 0..) |*j, i| j.* = .{ .win = region + i * hist_len, .hist = self.ins[0] + i * hist_len, .rows = self.ins[1] + i * (self.keep + taps) * proj_w * 2, .stride = proj_w, .taps = taps - 1, .c = conv_c }; + var a: cuda.Args = .{}; + a.add(t); + a.add(cint(self.keep)); + try cuda.launch.launch(self.f.roll, .{ .grid = .{ .x = @intCast((conv_c + 255) / 256), .y = @intCast(self.shape.rows) }, .block = .{ .x = 256 } }, self.rig.s, &a); + } + + /// Each window's saved copy, then fn_shift_window by the kept rows: the launches the one launch replaces. + fn theirs(self: *Roll, region: u64) !void { + for (0..self.shape.rows) |i| { + const win = region + i * hist_len; + try (cuda.DeviceBuffer{ .d = self.rig.gpu.d, .ptr = win, .len = hist_len }).copyFrom(0, self.ins[0] + i * hist_len, hist_len, self.rig.s.handle); + var a: cuda.Args = .{}; + a.add(win); + a.add(self.ins[1] + i * (self.keep + taps) * proj_w * 2); + a.add(@as(c_longlong, proj_w)); + a.add(cint(self.keep)); + a.add(cint(taps - 1)); + a.add(cint(conv_c)); + try cuda.launch.launch(self.f.shift, .{ .grid = .{ .x = @intCast((conv_c + 255) / 256) }, .block = .{ .x = 256 } }, self.rig.s, &a); + } + } + + fn restore(self: *Roll, at: u64) !void { + try (cuda.DeviceBuffer{ .d = self.rig.gpu.d, .ptr = at, .len = self.len }).upload(0, self.fresh.?); + } + + pub fn run(self: *Roll) !void { + try self.restore(self.region); + try self.ours(self.region); + try self.rig.s.synchronize(); + } + + fn baseline(self: *Roll) ![]const u8 { + try self.restore(self.base); + try self.theirs(self.base); + try self.rig.s.synchronize(); + return self.rig.bytes(self.base, self.len); + } + + pub fn replaced(self: *Roll) ![]const u8 { + return self.baseline(); + } + + /// Gate 1's stated rule: every window on the planted inputs equals the copy and shift's bytes. + pub fn check(self: *Roll, gate: g.Gate) !?[]const u8 { + _ = gate; + self.seen.clearRetainingCapacity(); + try self.seen.appendSlice(self.rig.gpu.gpa, try self.rig.outputBytes()); + const want = try self.baseline(); + if (std.mem.eql(u8, self.seen.items, want)) return null; + return try std.fmt.bufPrint(&self.msg, "bytes differ from the copy and shift's from byte {d}", .{std.mem.indexOfDiff(u8, self.seen.items, want).?}); + } + + pub fn count(self: *Roll) usize { + return self.len; + } + + pub fn reference(_: *Roll, _: []f64, _: []f64) !void { + return error.NoFloat64Reference; + } +}; diff --git a/zig/tests/cuda/contract_flashnext_topk.zig b/zig/tests/cuda/contract_flashnext_topk.zig new file mode 100644 index 0000000000..2c1064cb65 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_topk.zig @@ -0,0 +1,151 @@ +//! fn_row_topk under the contract: each row's K best by drawRow's order against a host oracle, every byte. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); + +const c = core.contract.spec; +const g = core.contract.gates; + +const vocab = 248320; // Flash Next's score rows +const k_max = 64; // cuda_mtp_draw.topk_max +const pad = 8; // the row stride's extra columns + +fn noSums(_: c.Shape) usize { + return 0; +} + +/// rows = score rows, n = the columns ranked (V), k = the candidates kept (K). +pub fn spec() c.Spec { + return .{ .name = "cuda flashnext fn_row_topk (a row's top_k candidates for sampled draws)", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = "bf16 bits as drawRow orders them (-0 as +0, NaN above +Inf); ties to the lowest columns; the values above the K-th, then the ties, each in column order" }, .tiles = .{ .{}, .{ .size = 1024, .step = 1, .min = 1 }, .{ .size = k_max, .step = 1, .min = 1 } }, .served = &.{ .{ .rows = 1, .n = vocab, .k = 20 }, .{ .rows = 9, .n = vocab, .k = k_max }, .{ .rows = 17, .n = vocab, .k = 1 } } }; +} + +/// drawRow's order as an unsigned key: larger value, larger key; -0 as +0. +fn key(u: u16) u32 { + const v: u32 = if (u == 0x8000) 0 else u; + return if (v & 0x8000 != 0) ~v & 0xffff else v | 0x8000; +} + +pub const TopK = struct { + rig: rig_.Rig, + f: cuda.Function, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + x: u64 = 0, + ov: u64 = 0, + oc: u64 = 0, + want: []u8 = &.{}, + msg: [160]u8 = undefined, + + pub fn deinit(self: *TopK) void { + self.rig.gpu.gpa.free(self.want); + self.rig.deinit(); + } + + const vt = g.hooks(TopK); + + pub fn kernel(self: *TopK) g.Kernel { + return .{ .spec = spec(), .ctx = self, .vt = &vt }; + } + + /// Random rows, rows of few distinct values (many ties), and planted NaNs, infinities and signed zeros. + pub fn prepare(self: *TopK, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + const a = self.rig.gpu.gpa; + a.free(self.want); + self.want = &.{}; + self.shape = shape; + const m = shape.rows; + const n = shape.n; + const k = shape.k; + if (m == 0 or n == 0 or k == 0 or k > k_max or k > n or n > vocab) return false; + const ld = n + pad; + var prng = std.Random.DefaultPrng.init(seed ^ 0x746f706b); + const r = prng.random(); + const rows = try a.alloc(u16, m * ld); + defer a.free(rows); + for (rows, 0..) |*v, i| { + // the stride's pad holds NaNs, which rank first: any read past a row's n columns shows + v.* = if (i % ld >= n) 0x7fc0 else if (fill == .stress) c.toBf16(@floatFromInt(r.uintLessThan(u32, 4))) else c.toBf16(r.float(f32) * 8 - 4); + } + if (fill == .nonfinite) for (0..m) |row| { + const base = (row * 7919) % n; + for ([_]u16{ 0x7fc1, 0xffc1, 0x7f80, 0xff80, 0x8000, 0x0000 }, 0..) |bits, j| rows[row * ld + (base + j) % n] = bits; + }; + self.want = try oracle(a, rows, m, n, ld, k); + self.x = try self.rig.input(std.mem.sliceAsBytes(rows)); + self.ov = try self.rig.output(m * k * 2); + self.oc = try self.rig.output(m * k * 4); + return true; + } + + /// A row's K best by drawRow's order (ties to lower columns): those above the K-th, then the ties, by column. + fn oracle(a: std.mem.Allocator, rows: []const u16, m: usize, n: usize, ld: usize, k: usize) ![]u8 { + const vals = try a.alloc(u16, m * k); + defer a.free(vals); + const cols = try a.alloc(i32, m * k); + defer a.free(cols); + const order = try a.alloc(u32, n); + defer a.free(order); + for (0..m) |row| { + const x = rows[row * ld ..][0..n]; + for (order, 0..) |*o, i| o.* = @intCast(i); + const Ctx = struct { + x: []const u16, + fn less(ctx: @This(), p: u32, q: u32) bool { + const kp = key(ctx.x[p]); + const kq = key(ctx.x[q]); + return kp > kq or (kp == kq and p < q); + } + }; + std.sort.pdq(u32, order, Ctx{ .x = x }, Ctx.less); + const kth = key(x[order[k - 1]]); + var at: usize = row * k; + for (x, 0..) |v, i| if (key(v) > kth) { + vals[at] = v; + cols[at] = @intCast(i); + at += 1; + }; + var need = k - (at - row * k); + for (x, 0..) |v, i| if (need > 0 and key(v) == kth) { + vals[at] = v; + cols[at] = @intCast(i); + at += 1; + need -= 1; + }; + } + return std.mem.concat(a, u8, &.{ std.mem.sliceAsBytes(vals), std.mem.sliceAsBytes(cols) }); + } + + /// The launch as Head.topk makes it: a row a block of 1,024 threads. + pub fn run(self: *TopK) !void { + var a: cuda.Args = .{}; + a.add(self.x); + a.add(@as(c_int, @intCast(self.shape.n))); + a.add(@as(c_longlong, @intCast(self.shape.n + pad))); + a.add(@as(c_int, @intCast(self.shape.k))); + a.add(self.ov); + a.add(self.oc); + try cuda.launch.launch(self.f, .{ .grid = .{ .x = @intCast(self.shape.rows) }, .block = .{ .x = 1024 } }, self.rig.s, &a); + try self.rig.s.synchronize(); + } + + pub fn count(self: *TopK) usize { + return self.want.len; + } + + pub fn reference(self: *TopK, want: []f64, mag: []f64) !void { + for (want, mag, self.want) |*w, *m, byte| { + w.* = @floatFromInt(byte); + m.* = 0; + } + } + + /// Gates 1 to 3's stated rule: every byte equals the oracle's. + pub fn check(self: *TopK, at: g.Gate) !?[]const u8 { + _ = at; + if (!try self.rig.intact()) return "fn_row_topk changed a buffer guard"; + const got = try self.rig.outputBytes(); + if (std.mem.eql(u8, got, self.want)) return null; + return try std.fmt.bufPrint(&self.msg, "candidates differ from the oracle's from byte {d}", .{std.mem.indexOfDiff(u8, got, self.want).?}); + } +}; diff --git a/zig/tests/cuda/contract_flashnext_wbnorm.zig b/zig/tests/cuda/contract_flashnext_wbnorm.zig new file mode 100644 index 0000000000..2f390f15f2 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_wbnorm.zig @@ -0,0 +1,189 @@ +//! The fused writeback and normalize under the contract: Glue.hcWritebackNorm against hcWriteback then hcNormalize. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const rig_ = @import("contract_rig.zig"); + +const c = core.contract.spec; +const g = core.contract.gates; + +pub const Mode = enum { sums, moe, slices }; + +const dims = 2560; +const streams = 4; +const wide = dims * streams; +const nc = dims / 256; +const slots = 11; // the routed experts' slots and the shared one, as the pass binds them +const eps: f32 = 1e-6; + +fn noSums(_: c.Shape) usize { + return 0; +} + +pub fn spec(mode: Mode) c.Spec { + const rule = "the written rows, their 256-column square sums and the normalized rows bit-equal to the two launches"; + return .{ + .name = switch (mode) { + .sums => "cuda flashnext fn_hc_writeback_norm0 (square sums, then normalize)", + .moe => "cuda flashnext fn_hc_writeback_norm_moe (MoE slots written back, then normalize)", + .slices => "cuda flashnext fn_hc_writeback_norm_slices (K slices written back, then normalize)", + }, + .accum = .none, + .depth = noSums, + .out = .raw, + .nan = .{ .stated = rule }, + .replaces = "fn_hc_writeback0 / _moe / _slices, then fn_hc_normalize", + .served = &.{ .{ .rows = 1, .n = dims, .k = 1 }, .{ .rows = 17, .n = dims, .k = 2 }, .{ .rows = 300, .n = dims, .k = 3 } }, + }; +} + +pub const WbNorm = struct { + rig: rig_.Rig, + mode: Mode, + glue: *const flash.Glue, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + ins: [6]u64 = @splat(0), + lens: [3]usize = @splat(0), + outs: [3]u64 = @splat(0), + nout: usize = 0, + base: u64 = 0, + seen: std.ArrayList(u8) = .empty, + msg: [160]u8 = undefined, + + pub fn deinit(self: *WbNorm) void { + self.seen.deinit(self.rig.gpu.gpa); + self.rig.deinit(); + } + + const vt = g.hooks(WbNorm); + + pub fn kernel(self: *WbNorm) g.Kernel { + return .{ .spec = spec(self.mode), .ctx = self, .vt = &vt }; + } + + fn modeNumber(self: *const WbNorm) usize { + return switch (self.mode) { + .sums => 0, + .moe => 2, + .slices => 4, + }; + } + + pub fn prepare(self: *WbNorm, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + self.nout = 0; + self.shape = shape; + const m = shape.rows; + if (m == 0 or m > 300 or shape.n != dims or shape.k == 0 or shape.k > 4) return false; + const a = self.rig.gpu.gpa; + var prng = std.Random.DefaultPrng.init(seed ^ 0x77626e); + const r = prng.random(); + const h = try a.alloc(u16, m * wide); + defer a.free(h); + for (h) |*x| x.* = c.toBf16(r.float(f32) * 2 - 1); + const inj = try a.alloc(u16, m * streams); + defer a.free(inj); + for (inj) |*x| x.* = c.toBf16(r.float(f32) * 4 - 2); + const y = try a.alloc(f32, m * slots * dims); + defer a.free(y); + for (y) |*x| x.* = r.float(f32) * 2 - 1; + const wts = try a.alloc(f32, m * slots); + defer a.free(wts); + for (wts) |*x| x.* = r.float(f32) * 0.2; + const br = try a.alloc(f32, shape.k * m * dims); + defer a.free(br); + for (br) |*x| x.* = r.float(f32) - 0.5; + const scale = try a.alloc(f32, wide); + defer a.free(scale); + for (scale) |*x| x.* = 0.5 + r.float(f32); + if (fill == .nonfinite) for (0..m) |row| { + h[row * wide + (row % streams) * dims + 3] = switch (row % 4) { + 0 => 0x7fc1, + 1 => 0x7f80, + 2 => 0xff80, + else => 0x8000, + }; + }; + for ([_][]const u8{ std.mem.sliceAsBytes(h), std.mem.sliceAsBytes(inj), std.mem.sliceAsBytes(y), std.mem.sliceAsBytes(wts), std.mem.sliceAsBytes(br), std.mem.sliceAsBytes(scale) }, 0..) |b, i| self.ins[i] = try self.rig.input(b); + if (self.mode != .sums) try self.out(m * wide * 2); // the written-back rows + try self.out(m * nc * streams * 4); // square sums a 256-column slice and stream + try self.out(m * wide * 2); // the normalized rows + var total: usize = 0; + for (self.lens[0..self.nout]) |l| total += l; + self.base = try self.rig.scratch(total); + return true; + } + + fn out(self: *WbNorm, len: usize) !void { + self.outs[self.nout] = try self.rig.output(len); + self.lens[self.nout] = len; + self.nout += 1; + } + + /// The written rows, the sums and the normalized rows in `o`: the first is absent for plain sums. + fn split(self: *const WbNorm, o: []const u64) struct { hout: u64, pss: u64, norm: u64 } { + return if (self.mode == .sums) .{ .hout = 0, .pss = o[0], .norm = o[1] } else .{ .hout = o[0], .pss = o[1], .norm = o[2] }; + } + + fn ours(self: *WbNorm, o: []const u64) !void { + const p = self.split(o); + const m = self.shape.rows; + try self.glue.hcWritebackNorm(self.rig.s, self.ins[0], p.hout, p.pss, self.ins[4], self.ins[1], self.ins[2], self.ins[3], self.ins[5], p.norm, eps, m * dims, m, dims, streams, self.modeNumber(), slots - 1, slots, self.shape.k); + } + + fn theirs(self: *WbNorm, o: []const u64) !void { + const p = self.split(o); + const m = self.shape.rows; + try self.glue.hcWriteback(self.rig.s, self.ins[0], p.hout, p.pss, self.ins[4], self.ins[1], self.ins[2], self.ins[3], m * dims, m, dims, streams, self.modeNumber(), slots - 1, slots, self.shape.k); + try self.glue.hcNormalize(self.rig.s, if (self.mode == .sums) self.ins[0] else p.hout, p.pss, self.ins[5], p.norm, eps, m, dims, streams); + } + + pub fn run(self: *WbNorm) !void { + try self.ours(self.outs[0..self.nout]); + try self.rig.s.synchronize(); + } + + /// The scratch the two launches write: the outputs' layout back to back, poisoned like them (0xff). + fn baseline(self: *WbNorm) ![]const u8 { + var at = self.base; + var o: [3]u64 = undefined; + var total: usize = 0; + for (self.lens[0..self.nout], 0..) |l, i| { + o[i] = at; + at += l; + total += l; + } + const ff = try self.rig.gpu.gpa.alloc(u8, total); + defer self.rig.gpu.gpa.free(ff); + @memset(ff, 0xff); + try (cuda.DeviceBuffer{ .d = self.rig.gpu.d, .ptr = self.base, .len = total }).upload(0, ff); + try self.theirs(o[0..self.nout]); + try self.rig.s.synchronize(); + return self.rig.bytes(self.base, total); + } + + pub fn replaced(self: *WbNorm) ![]const u8 { + return self.baseline(); + } + + /// Gate 1's stated rule: the outputs on the planted inputs equal the two launches' bytes. + pub fn check(self: *WbNorm, gate: g.Gate) !?[]const u8 { + _ = gate; + self.seen.clearRetainingCapacity(); + try self.seen.appendSlice(self.rig.gpu.gpa, try self.rig.outputBytes()); + const want = try self.baseline(); + if (std.mem.eql(u8, self.seen.items, want)) return null; + return try std.fmt.bufPrint(&self.msg, "bytes differ from the two launches' from byte {d}", .{std.mem.indexOfDiff(u8, self.seen.items, want).?}); + } + + pub fn count(self: *WbNorm) usize { + var total: usize = 0; + for (self.lens[0..self.nout]) |l| total += l; + return total; + } + + pub fn reference(_: *WbNorm, _: []f64, _: []f64) !void { + return error.NoFloat64Reference; + } +}; diff --git a/zig/tests/cuda/contract_projection_f32.zig b/zig/tests/cuda/contract_projection_f32.zig index ec1666e6ed..4f9483058d 100644 --- a/zig/tests/cuda/contract_projection_f32.zig +++ b/zig/tests/cuda/contract_projection_f32.zig @@ -52,7 +52,7 @@ pub const Product = struct { mode: cuda.qmmf.Mode, clusters: bool, out_kind: cuda.qmmf.Out = .f32, - prompt: ?*const cuda.nvfp4.Prompt = null, // set: prompt16's FP8 tile instead of the lane + prompt: ?*const cuda.nvfp4.Prompt = null, // set: prompt16's FP8 or BF16 tile instead of the lane d: ?Dense = null, weight: cuda.qmmf.Weight = undefined, x: u64 = 0, @@ -65,15 +65,19 @@ pub const Product = struct { } pub fn kernel(self: *Product) g.Kernel { - if (self.prompt != null) return .{ .spec = .{ - .name = "cuda prompt16 FP8 tile 4", - .accum = .f32, - .depth = promptDepth, - .ops = 1, - .out = .bf16, - .tiles = .{ .{ .size = 128 }, .{ .size = 128, .step = 64, .min = 64 }, .{ .size = 128, .step = 64, .min = 64 } }, - .served = &served_prompt, - }, .ctx = self, .vt = &vt }; + if (self.prompt != null) return .{ + .spec = .{ + .name = if (self.mode != .bf16) "cuda prompt16 FP8 tile 4" else if (self.out_kind == .f32) "cuda prompt16 BF16 tile, fp32 out" else "cuda prompt16 BF16 tile", + .accum = .f32, + .depth = if (self.mode == .bf16) depth else promptDepth, // BF16: the lane's K slices, folded in order + .ops = if (self.mode == .bf16) 0 else 1, + .out = if (self.out_kind == .f32) .f32 else .bf16, + .tiles = .{ .{ .size = 128 }, .{ .size = 128, .step = 64, .min = 64 }, .{ .size = 128, .step = 64, .min = 64 } }, + .served = &served_prompt, + }, + .ctx = self, + .vt = &vt, + }; const wide = self.out_kind == .f32; const name = switch (self.mode) { .bf16 => if (!wide) (if (self.clusters) "cuda BF16 bf16 clustered" else "cuda BF16 bf16 reduce") else if (self.clusters) "cuda BF16 fp32 clustered" else "cuda BF16 fp32 reduce", @@ -112,7 +116,7 @@ pub const Product = struct { self.x = try self.rig.input(std.mem.sliceAsBytes(d.x)); const bytes = if (self.prompt != null) 0 else self.lane.partBytes(shape.rows, self.weight); if (bytes > 0) self.part = try self.rig.scratch(bytes); - self.out = try self.rig.output(shape.rows * shape.n * @as(usize, if (self.out_kind == .f32 and self.prompt == null) 4 else 2)); + self.out = try self.rig.output(shape.rows * shape.n * @as(usize, if (self.out_kind == .f32) 4 else 2)); return true; } @@ -209,7 +213,7 @@ pub const Product = struct { pub fn run(self: *Product) !void { const shape = self.d.?.shape; if (self.prompt) |p| { - try p.matmul(self.rig.s, self.x, shape.k, shape.rows, self.weight, self.out); + if (self.mode == .bf16) try p.matmulOut(self.rig.s, self.x, shape.k, shape.rows, self.weight, self.out, self.out_kind) else try p.matmul(self.rig.s, self.x, shape.k, shape.rows, self.weight, self.out); } else try self.lane.matmulOut(self.rig.s, self.x, shape.k, shape.rows, self.weight, self.out, self.out_kind, self.part); try self.rig.s.synchronize(); } diff --git a/zig/tests/cuda/contract_route_rows.zig b/zig/tests/cuda/contract_route_rows.zig index e73f8fbc17..f2189666f7 100644 --- a/zig/tests/cuda/contract_route_rows.zig +++ b/zig/tests/cuda/contract_route_rows.zig @@ -93,7 +93,7 @@ pub const Logits = struct { for ([_]usize{ sh.rows, sh.k, sh.n }) |v| a.add(@as(i32, @intCast(v))); const r = self.route; switch (kind) { - .wide => try cuda.launch.launch(r.wide_fn, .{ .grid = .{ .x = @intCast((sh.rows + 127) / 128), .y = @intCast((sh.n + 63) / 64) }, .block = .{ .x = 256 }, .shared = rr.wide_shared }, self.rig.s, &a), + .wide => try cuda.launch.launch(r.wide_fn, .{ .grid = .{ .x = @intCast((sh.n + 63) / 64), .y = @intCast((sh.rows + 127) / 128) }, .block = .{ .x = 256 }, .shared = rr.wide_shared }, self.rig.s, &a), .few => try cuda.launch.launch(r.few_fn, .{ .grid = .{ .x = @intCast((sh.rows + 15) / 16), .y = @intCast((sh.n + 31) / 32) }, .block = .{ .x = 128 }, .shared = rr.few_shared }, self.rig.s, &a), .logits => try cuda.launch.launch(r.logits_fn, .{ .grid = .{ .x = @intCast((sh.rows + 15) / 16), .y = @intCast((sh.n + 31) / 32) }, .block = .{ .x = 128 }, .shared = rr.many_shared }, self.rig.s, &a), } diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index b9158471f4..bb43b84823 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -229,12 +229,12 @@ pub fn run(gpu: check.Gpu) !void { } defer for (&wide) |*m| m.deinit(); for (&wide) |*m| try kernels.append(a, m.kernel()); - var narrow: [wide_k.bf16_modes.len * 2 + 1]wide_k.Product = undefined; // bf16 out (the 27B's ModelOpt layers) + var narrow: [wide_k.bf16_modes.len * 2 + 3]wide_k.Product = undefined; // bf16 out (the 27B's ModelOpt layers) for (wide_k.bf16_modes, 0..) |mode, i| { narrow[i * 2] = .{ .rig = .{ .gpu = gpu, .s = s }, .lane = &lane, .mode = mode, .clusters = true, .out_kind = .bf16 }; narrow[i * 2 + 1] = .{ .rig = .{ .gpu = gpu, .s = s }, .lane = &lane_reduce, .mode = mode, .clusters = false, .out_kind = .bf16 }; } - narrow[narrow.len - 1] = .{ .rig = .{ .gpu = gpu, .s = s }, .lane = &lane, .mode = .fp8, .clusters = true, .out_kind = .bf16, .prompt = &prompt }; + for ([_]cuda.qmmf.Mode{ .fp8, .bf16, .bf16 }, [_]cuda.qmmf.Out{ .bf16, .bf16, .f32 }, narrow[narrow.len - 3 ..]) |mode, out, *m| m.* = .{ .rig = .{ .gpu = gpu, .s = s }, .lane = &lane, .mode = mode, .clusters = true, .out_kind = out, .prompt = &prompt }; defer for (&narrow) |*m| m.deinit(); for (&narrow) |*m| try kernels.append(a, m.kernel()); var rank = rank_k.Rank{ .rig = .{ .gpu = gpu, .s = s }, .sum = &rank_sum };