diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 676bda8c8..f670922c1 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -25,6 +25,7 @@ const kernels = [_]Kernel{ .{ .name = "nemotron_ops", .flags = &.{"-O3"} }, // ours: Nemotron's serial feed, routed plan and rests .{ .name = "affine4_pack", .flags = &.{"-O3"} }, // ours: the MLX affine-4 repack into qlinear's tiles .{ .name = "lane_gemv", .flags = &.{"-O3"} }, // ours: qmm_group's arithmetic, a column tile's K slices in one CTA + .{ .name = "qmm_lane", .flags = &.{"-O3"} }, // pre-sm_90 decode matmuls, K slices in slice order .{ .name = "sample", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // ours: the Metal engine's keyed draws .{ .name = "qwen_ops", .flags = &.{"-O3"} }, // ours: the 27B's batched state copies, taps, drafter norms and rotary .{ .name = "nemotron_norms", .flags = glue }, // ours, each with a host reference (glue_ref.zig): Nemotron's glue diff --git a/zig/kernels/cuda/qmm_lane.cu b/zig/kernels/cuda/qmm_lane.cu new file mode 100644 index 000000000..0b5300a2d --- /dev/null +++ b/zig/kernels/cuda/qmm_lane.cu @@ -0,0 +1,509 @@ +// Pre-sm_90 4-bit decode matmuls (no clusters): a tile's K slices added in slice order, 0.6.6's qmm.cu bits. + +#include +#include +#include +#include +#include + +#include "qmm_frag.cuh" + +namespace { + +using namespace qmm_frag; + +// Tile shapes: BM rows by BN columns a block, WM x WN warps, each warp (BM / WM) x (BN / WN). +template +struct Tile { + static constexpr int THREADS = WM * WN * 32; + static constexpr int MT = BM / WM / 16; // m16 tiles a warp + static constexpr int NT = BN / WN / 8; // n8 tiles a warp + static constexpr int ROW = GS * 2; // bytes of one input row a group + static constexpr int CHUNKS = ROW / 16; + static constexpr int X = BM * ROW; // stage bytes: inputs, + static constexpr int W = BN * GS / 2; // weights, + static constexpr int S = BN * 2; // scales, biases (bf16), + static constexpr int XS = BM * 4; // and input sums (fp32) + static constexpr int STAGE = X + W + 2 * S + XS; + static constexpr int PARTIALS = MT * NT * 4 * THREADS * 4; // a K slice's partial, parked for the cluster sum + static constexpr int SMEM = STAGES * STAGE > PARTIALS ? STAGES * STAGE : PARTIALS; +}; + +// A tile's K slice: acc = fma(xs, b, fma(p, s, acc)) a group, in order from zero; every launch form runs these steps. +template +__device__ __forceinline__ void tile_slice(float (&acc)[Tile::MT] + [Tile::NT][4], + unsigned char* buf, const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, + const uint32_t* __restrict__ w, const __nv_bfloat16* __restrict__ scales, + const __nv_bfloat16* __restrict__ biases, int M, int KG, int per, int g0, int m0, int n0, int npad, int ldx) { + using T = Tile; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int wm = warp / WN, wn = warp % WN; + + auto stage = [&](int s) { return buf + s * T::STAGE; }; + // one group's inputs, weights, scales, biases and input sums into stage s + auto load = [&](int s, int g, int m0) { + 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; + const int row = min(m0 + r, M - 1); + cp16z(p + r * T::ROW + swz(r, ch) * 16, x + static_cast(row) * ldx + g * GS + ch * 8, + m0 + r < M); + } + unsigned char* pw = p + T::X; + constexpr int TILE_BYTES = 64 * GS / 2; // one stored 64-column tile's group block + for (int c = tid; c < T::W / 16; c += T::THREADS) { + const int t = c / (TILE_BYTES / 16), off = c % (TILE_BYTES / 16); + const size_t tile = static_cast(n0 / 64 + t) * KG + g; + cp16(pw + c * 16, reinterpret_cast(w) + tile * TILE_BYTES + off * 16); + } + unsigned char* ps = pw + T::W; + for (int c = tid; c < 2 * (T::S / 16); c += T::THREADS) { + const int which = c / (T::S / 16), off = c % (T::S / 16); + const __nv_bfloat16* src = (which ? biases : scales) + static_cast(g) * npad + n0 + off * 8; + cp16(ps + which * T::S + off * 16, src); + } + float* px = reinterpret_cast(ps + 2 * T::S); + for (int r = tid; r < BM; r += T::THREADS) + cp4(px + r, xs + static_cast(min(m0 + r, M - 1)) * KG + g); + }; + + { +#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] = 0.0f; + // acc = fma(xs, b, fma(p, s, acc)) per output, the order every row gets + auto epilogue = [&](float (&acc_)[T::MT][T::NT][4], const float (&d_)[T::MT][T::NT][4], + const float (&s_)[T::NT][2], const float (&b_)[T::NT][2], const float (&x_)[T::MT][2]) { +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int e = 0; e < 4; ++e) + acc_[i][j][e] = __fmaf_rn(x_[i][e >> 1], b_[j][e & 1], __fmaf_rn(d_[i][j][e], s_[j][e & 1], + acc_[i][j][e])); + }; + float dp[T::MT][T::NT][4], sp[T::NT][2], bp[T::NT][2], xp[T::MT][2]; +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < per) load(s, g0 + s, m0); + commit(); + } + for (int it = 0; it < per; ++it) { + wait(); + __syncthreads(); + const int next = it + STAGES - 1; + if (next < per) load(next % STAGES, g0 + next, m0); + commit(); + const unsigned char* p = stage(it % STAGES); + const uint32_t* pw = reinterpret_cast(p + T::X); + const __nv_bfloat16* ps = reinterpret_cast(p + T::X + T::W); + const float* px = reinterpret_cast(p + T::X + T::W + 2 * T::S); + uint32_t words[T::NT][GS / 32]; +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int v = 0; v < GS / 32; ++v) + words[j][v] = pw[((wn * T::NT + j) * 32 + lane) * (GS / 32) + v]; + float d[T::MT][T::NT][4]; +#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) { + const uint32_t b0 = pair(words[j][kt / 2], (kt & 1) * 8); + const uint32_t b1 = pair(words[j][kt / 2], (kt & 1) * 8 + 4); +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + if (kt == 0) mma0(d[i][j], a[i], b0, b1); + else mma(d[i][j], a[i], b0, b1); + } + } + if (PIPE && kt == 0 && it > 0) epilogue(acc, dp, sp, bp, xp); // the previous group's, between MMAs + } + float sv[T::NT][2], bv[T::NT][2], xv[T::MT][2]; +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + const int col = wn * (BN / WN) + j * 8 + (lane & 3) * 2; + const __nv_bfloat162 s2 = *reinterpret_cast(ps + col); + const __nv_bfloat162 b2 = *reinterpret_cast(ps + BN + col); + sv[j][0] = __low2float(s2); sv[j][1] = __high2float(s2); + bv[j][0] = __low2float(b2); bv[j][1] = __high2float(b2); + } +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + const int row = wm * (BM / WM) + i * 16 + (lane >> 2); + xv[i][0] = px[row]; + xv[i][1] = px[row + 8]; + } + if constexpr (PIPE) { +#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) dp[i][j][e] = d[i][j][e]; + xp[i][0] = xv[i][0]; xp[i][1] = xv[i][1]; + } +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + sp[j][0] = sv[j][0]; sp[j][1] = sv[j][1]; bp[j][0] = bv[j][0]; bp[j][1] = bv[j][1]; + } + } else { + epilogue(acc, d, sv, bv, xv); + } + } + if (PIPE && per > 0) epilogue(acc, dp, sp, bp, xp); + wait<0>(); + __syncthreads(); + } +} + +template +__device__ __forceinline__ void qmm_kernel( + const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, const uint32_t* __restrict__ w, + const __nv_bfloat16* __restrict__ scales, const __nv_bfloat16* __restrict__ biases, + void* __restrict__ out, float* __restrict__ part, int M, int N, int K, int SK, int npad, int ldx, int group) { + 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, per = KG / SK; + const int2 at = tile_of(blockIdx.x, M, N, BM, BN, group); + const int m0 = at.x, n0 = at.y, slice = blockIdx.z, g0 = slice * per; + { + float acc[T::MT][T::NT][4]; + tile_slice(acc, buf, x, xs, w, scales, biases, M, KG, per, g0, m0, n0, npad, + ldx); + if constexpr (CLUSTER) { +#if __CUDA_ARCH__ >= 900 + // K slices of a tile form a cluster: slice 0 adds peers' partials in slice order, as reduce_kernel does + constexpr int E = T::MT * T::NT * 4; + auto cluster = cooperative_groups::this_cluster(); + float* mine = reinterpret_cast(buf); + if (slice != 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) mine[((i * T::NT + j) * 4 + e) * T::THREADS + tid] = acc[i][j][e]; + } + cluster.sync(); + if (slice == 0) { + for (int peer = 1; peer < SK; ++peer) { + const float* theirs = cluster.map_shared_rank(mine, peer); +#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] = acc[i][j][e] + theirs[((i * T::NT + j) * 4 + e) * T::THREADS + tid]; + } + } + cluster.sync(); // peers keep their memory until slice 0 has read it + if (slice != 0) return; +#else + __trap(); // clusters are 9.0+; qmm_cuda sends slices to reduce_kernel +#endif + } +#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], v1 = acc[i][j][2 * h + 1]; + if (SK > 1 && !CLUSTER) { + float* dst = part + (static_cast(slice) * M + row) * N + col; + if (col < N) dst[0] = v0; + if (col + 1 < N) dst[1] = v1; + } else 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); + } + } + } + } + } +} + +// The K slices' partial sums added in slice order, then rounded once. +template +__device__ __forceinline__ void reduce_kernel(const float* __restrict__ part, void* __restrict__ out, long long total, int SK) { + const long long i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= total) return; + float acc = part[i]; + for (int s = 1; s < SK; ++s) acc = acc + part[s * total + i]; + if (F32) reinterpret_cast(out)[i] = acc; + else reinterpret_cast<__nv_bfloat16*>(out)[i] = __float2bfloat16_rn(acc); +} + +// tile_slice with the weights as the mma's A and 8-token input tiles as its B: the same steps, so tile_slice's bits. +template +__device__ __forceinline__ void tile_slice_swap(float (&acc)[TT][4], unsigned char* buf, + const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, const uint32_t* __restrict__ w, + const __nv_bfloat16* __restrict__ scales, const __nv_bfloat16* __restrict__ biases, int M, int KG, int per, + int g0, int n0, int npad, int ldx) { + constexpr int BM = 16; + using T = Tile; + static_assert(T::NT == 2 && T::MT == 1, "a warp: one m16 tile of 16 columns"); + const int tid = threadIdx.x % T::THREADS, lane = tid & 31, wn = tid >> 5, gq = lane >> 2, t = lane & 3; + // G > 1: the block holds G slices of a tile, a group of T::THREADS each with its own stages, synced alone + auto sync = [] { + if constexpr (G == 1) __syncthreads(); + else asm volatile("bar.sync %0, %1;" :: "r"(1 + static_cast(threadIdx.x) / T::THREADS), "n"(T::THREADS)); + }; + + auto stage = [&](int s) { return buf + s * T::STAGE; }; + // a group's stage: each thread's 16-byte copies, its sources found once and stepped a group at a time + static_assert(T::THREADS == BM * T::CHUNKS && T::W == T::THREADS * 16 && 2 * (T::S / 16) <= T::THREADS, + "one copy of each kind a thread"); + const int xr = tid / T::CHUNKS, xc = tid % T::CHUNKS; + const bool xread = xr < M; + const int xoff = xr * T::ROW + swz(xr, xc) * 16; + const __nv_bfloat16* xsrc = x + static_cast(min(xr, M - 1)) * ldx + static_cast(g0) * GS + xc * 8; + const unsigned char* wsrc = reinterpret_cast(w) + + (static_cast(n0 / 64) * KG + g0) * (64 * GS / 2) + tid * 16; + const int sw = tid / (T::S / 16), so = tid % (T::S / 16); + const bool sread = tid < 2 * (T::S / 16); + const __nv_bfloat16* ssrc = (sw ? biases : scales) + static_cast(g0) * npad + n0 + so * 8; + const int soff = T::X + T::W + sw * T::S + so * 16; + const bool xsread = tid < BM; + const float* xssrc = xs + static_cast(min(tid, M - 1)) * KG + g0; + const int xsoff = T::X + T::W + 2 * T::S + tid * 4; + auto load = [&](int s, int i) { // group g0 + i into stage s + unsigned char* p = stage(s); + cp16z(p + xoff, xsrc + i * GS, xread); + cp16(p + T::X + tid * 16, wsrc + static_cast(i) * (64 * GS / 2)); + if (sread) cp16(p + soff, ssrc + static_cast(i) * npad); + if (xsread) cp4(p + xsoff, xssrc + i); + }; + +#pragma unroll + for (int u = 0; u < TT; ++u) acc[u][0] = acc[u][1] = acc[u][2] = acc[u][3] = 0.0f; +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < per) load(s, s); + commit(); + } + const int c0 = wn * 16 + gq, c1 = c0 + 8; // the lane's columns in the block + for (int it = 0; it < per; ++it) { + wait(); + sync(); + const int next = it + STAGES - 1; + if (next < per) load(next % STAGES, next); + commit(); + const unsigned char* p = stage(it % STAGES); + const uint32_t* pw = reinterpret_cast(p + T::X); + const __nv_bfloat16* ps = reinterpret_cast(p + T::X + T::W); + const float* px = reinterpret_cast(p + T::X + T::W + 2 * T::S); + uint32_t words[2][GS / 32]; +#pragma unroll + for (int j = 0; j < 2; ++j) +#pragma unroll + for (int v = 0; v < GS / 32; ++v) words[j][v] = pw[((wn * 2 + j) * 32 + lane) * (GS / 32) + v]; + float d[TT][4]; +#pragma unroll + for (int kt = 0; kt < GS / 16; ++kt) { + uint32_t a[4]; + const int r = (lane & 7) + ((lane >> 3) & 1) * 8; + ldmatrix4(a, p + r * T::ROW + swz(r, kt * 2 + (lane >> 4)) * 16); + const int sh = (kt & 1) * 8; + const uint32_t wa[4] = {pair(words[0][kt / 2], sh), pair(words[1][kt / 2], sh), + pair(words[0][kt / 2], sh + 4), pair(words[1][kt / 2], sh + 4)}; +#pragma unroll + for (int u = 0; u < TT; ++u) { + if (kt == 0) mma0(d[u], wa, a[u], a[u + 2]); + else mma(d[u], wa, a[u], a[u + 2]); + } + } + const float s0 = __bfloat162float(ps[c0]), s1 = __bfloat162float(ps[c1]); + const float b0 = __bfloat162float(ps[BN + c0]), b1 = __bfloat162float(ps[BN + c1]); +#pragma unroll + for (int u = 0; u < TT; ++u) { + const float x0 = px[u * 8 + 2 * t], x1 = px[u * 8 + 2 * t + 1]; + acc[u][0] = __fmaf_rn(x0, b0, __fmaf_rn(d[u][0], s0, acc[u][0])); + acc[u][1] = __fmaf_rn(x1, b0, __fmaf_rn(d[u][1], s0, acc[u][1])); + acc[u][2] = __fmaf_rn(x0, b1, __fmaf_rn(d[u][2], s1, acc[u][2])); + acc[u][3] = __fmaf_rn(x1, b1, __fmaf_rn(d[u][3], s1, acc[u][3])); + } + } + wait<0>(); + sync(); +} + +// Several weights of one input in one launch, each at its own K slices; a tile's last block adds them in slice order. +constexpr int SEGS = 4; +struct Seg { + const uint32_t* w; + const __nv_bfloat16* scales; + const __nv_bfloat16* biases; + void* out; + float* part; + int N, npad, SK, blocks, sync; // blocks: tiles x SK; sync: its tiles' first counter +}; +struct Segs { + Seg s[SEGS]; + int count; +}; + +template +__device__ __forceinline__ void put(void* out, size_t i, float v) { + if (F32) reinterpret_cast(out)[i] = v; + else reinterpret_cast<__nv_bfloat16*>(out)[i] = __float2bfloat16_rn(v); +} + +// TT 0: tile_slice, 1-2: tile_slice_swap's 8-token tiles; G > 1: a block's G slices, summed in shared memory in order. +template +__device__ __forceinline__ void qmm_seg_kernel( + const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, const Segs segs, int* __restrict__ sync, + int M, int K, int ldx) { + using T = Tile; + extern __shared__ __align__(128) unsigned char buf[]; + __shared__ int last; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int wm = warp / WN, wn = warp % WN; + int b = blockIdx.x; + Seg sg = segs.s[0]; +#pragma unroll + for (int i = 1; i < SEGS; ++i) + if (i < segs.count && b >= sg.blocks / G) { // G > 1: a block a tile + b -= sg.blocks / G; + sg = segs.s[i]; + } + const int KG = K / GS, per = KG / sg.SK, tiles = sg.blocks / sg.SK; + const int slice = G > 1 ? static_cast(threadIdx.x) / T::THREADS : b / tiles; + const int tile = G > 1 ? b : b - slice * tiles, n0 = tile * BN; + constexpr int E = TT ? TT * 4 : T::MT * T::NT * 4; // the lane's outputs + float v[E]; + int rows[E], cols[E]; + if constexpr (TT == 0) { + float acc[T::MT][T::NT][4]; + tile_slice(acc, buf, x, xs, sg.w, sg.scales, sg.biases, M, KG, per, + slice * per, 0, n0, sg.npad, ldx); +#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) { + const int q = (i * T::NT + j) * 4 + e; + v[q] = acc[i][j][e]; + rows[q] = wm * (BM / WM) + i * 16 + (lane >> 2) + (e >> 1) * 8; + cols[q] = n0 + wn * (BN / WN) + j * 8 + (lane & 3) * 2 + (e & 1); + } + } else { + static_assert(WM == 1 && BM == 16, "the swapped form: one 16-row tile, a warp's 16 columns"); + float acc[TT][4]; + tile_slice_swap(acc, buf + (G > 1 ? slice * STAGES * T::STAGE : 0), x, xs, sg.w, + sg.scales, sg.biases, M, KG, per, slice * per, n0, sg.npad, ldx); +#pragma unroll + for (int u = 0; u < TT; ++u) +#pragma unroll + for (int e = 0; e < 4; ++e) { + v[u * 4 + e] = acc[u][e]; + rows[u * 4 + e] = u * 8 + 2 * (lane & 3) + (e & 1); + cols[u * 4 + e] = n0 + wn * 16 + (lane >> 2) + (e >> 1) * 8; + } + } + const int N = sg.N; + const size_t MN = static_cast(M) * N; + if constexpr (G > 1) { + float* red = reinterpret_cast(buf); // slice s's lane values at [s - 1][q][thread] + const int lt = threadIdx.x % T::THREADS; + __syncthreads(); // every slice's stages done + if (slice > 0) +#pragma unroll + for (int q = 0; q < E; ++q) red[((slice - 1) * E + q) * T::THREADS + lt] = v[q]; + __syncthreads(); + if (slice > 0) return; +#pragma unroll + for (int q = 0; q < E; ++q) { + float a = v[q]; +#pragma unroll + for (int s2 = 1; s2 < G; ++s2) a = a + red[((s2 - 1) * E + q) * T::THREADS + lt]; + if (rows[q] < M && cols[q] < N) put(sg.out, static_cast(rows[q]) * N + cols[q], a); + } + return; + } +#pragma unroll + for (int q = 0; q < E; ++q) { + if (rows[q] >= M || cols[q] >= N) continue; + const size_t at = static_cast(rows[q]) * N + cols[q]; + if (sg.SK == 1) put(sg.out, at, v[q]); + else sg.part[slice * MN + at] = v[q]; + } + if (sg.SK == 1) return; + __threadfence(); // this slice's partials, then its count + __syncthreads(); + if (tid == 0) last = atomicAdd(sync + sg.sync + tile, 1) == sg.SK - 1; + __syncthreads(); + if (!last) return; + __threadfence(); +#pragma unroll + for (int q = 0; q < E; ++q) { + if (rows[q] >= M || cols[q] >= N) continue; + const size_t at = static_cast(rows[q]) * N + cols[q]; + float p[8]; // every slice's load in flight, then added in order +#pragma unroll + for (int s = 0; s < 8; ++s) p[s] = s < sg.SK ? __ldcg(sg.part + s * MN + at) : 0.f; + float a = p[0]; +#pragma unroll + for (int s = 1; s < 8; ++s) + if (s < sg.SK) a = a + p[s]; + put(sg.out, at, a); + } + if (tid == 0) sync[sg.sync + tile] = 0; +} +} // namespace + +// Segmented decode launch (rows <= 16): TT 8-token tiles, G K slices a block, STAGES 3 or 4. +#define TF_SEG(ST, TT, G) \ + extern "C" __global__ void __launch_bounds__(128 * G) tf_qmm_seg_s##ST##_t##TT##_g##G( \ + const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, const Segs segs, int* __restrict__ sync, \ + int M, int K, int ldx) { \ + qmm_seg_kernel<64, 16, 64, 1, 4, ST, false, TT, G>(x, xs, segs, sync, M, K, ldx); \ + } +TF_SEG(3, 1, 1) TF_SEG(3, 2, 1) TF_SEG(3, 1, 2) TF_SEG(3, 2, 2) TF_SEG(3, 1, 4) TF_SEG(3, 2, 4) +TF_SEG(4, 1, 1) TF_SEG(4, 2, 1) TF_SEG(4, 1, 2) TF_SEG(4, 2, 2) TF_SEG(4, 1, 4) TF_SEG(4, 2, 4) +#undef TF_SEG + +// Wider rounds: BM-row tiles, K slices into part (SK, M, N) for tf_qmm_reduce, or straight to bf16 with one slice. +#define TF_LANE(BM) \ + extern "C" __global__ void __launch_bounds__(128) tf_qmm_lane_m##BM( \ + const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, const uint32_t* __restrict__ w, \ + const __nv_bfloat16* __restrict__ scales, const __nv_bfloat16* __restrict__ biases, \ + void* __restrict__ out, float* __restrict__ part, int M, int N, int K, int SK, int npad, int ldx, \ + int group) { \ + qmm_kernel<64, BM, 64, 1, 4, 4, false, false>(x, xs, w, scales, biases, out, part, M, N, K, SK, npad, ldx, group); \ + } +TF_LANE(16) TF_LANE(32) TF_LANE(64) +#undef TF_LANE + +extern "C" __global__ void tf_qmm_reduce(const float* __restrict__ part, void* __restrict__ out, long long total, int SK) { + reduce_kernel(part, out, total, SK); +} + diff --git a/zig/src/cli/qwen27_cuda.zig b/zig/src/cli/qwen27_cuda.zig index a92416ce7..aab471404 100644 --- a/zig/src/cli/qwen27_cuda.zig +++ b/zig/src/cli/qwen27_cuda.zig @@ -66,7 +66,8 @@ pub fn main(init: std.process.Init) !u8 { defer ctx.deinit(); const clock = std.Io.Clock.awake; var t0 = clock.now(init.io).toNanoseconds(); - const e = try qwen.Engine.init(gpa, init.io, &ctx, o.model, o.kernels, .{ .context = o.context }); + const rows_env = if (std.c.getenv("TF_QWEN27_PROMPT_ROWS")) |v| std.fmt.parseInt(usize, std.mem.span(v), 10) catch return error.BadPromptRows else qwen.engine.prompt_rows; + const e = try qwen.Engine.init(gpa, init.io, &ctx, o.model, o.kernels, .{ .context = o.context, .prompt_rows = rows_env }); defer e.deinit(); std.debug.print("loaded in {d:.1}s, weights {d:.2} GiB\n", .{ @as(f64, @floatFromInt(clock.now(init.io).toNanoseconds() - t0)) / 1e9, @as(f64, @floatFromInt(e.w.bytes())) / (1 << 30) }); if (windows) return checkWindows(gpa, arena, init.io, e, prompts, o.dump, o.cases); @@ -435,7 +436,7 @@ fn runMulti(gpa: std.mem.Allocator, a: std.mem.Allocator, io: std.Io, e: *qwen.E m.accepted += len - 1; } try f.commit(parts[0..n], paths[0..n]); - try qwen.drafts.absorb(d, f.ops, e.scratch.taps, e.scratch.copy_items, kept[0..n]); + try qwen.drafts.absorb(d, f.ops, e.scratch.taps, e.scratch.window_rows, e.scratch.copy_items, kept[0..n]); for (live[0..n]) |k| { const m = &ms[k]; if (m.out.items.len >= count or e.w.isEos(m.out.items[m.out.items.len - 1])) m.done = true; diff --git a/zig/src/cuda/aot.zig b/zig/src/cuda/aot.zig index b8e58d2e8..1a736ed7a 100644 --- a/zig/src/cuda/aot.zig +++ b/zig/src/cuda/aot.zig @@ -122,7 +122,12 @@ pub const Set = struct { pub fn find(self: *const Set, function: []const u8, args: []const Arg, consts: []const Const) !*const Variant { for (self.variants) |*v| { if (!std.mem.eql(u8, v.spec.@"fn", function)) continue; - if (matches(v.spec, args, consts)) return v; + if (matches(v.spec, args, consts, false)) return v; + } + // else one compiled without the 16-byte alignment hint: the same code, valid for any pointer + for (self.variants) |*v| { + if (!std.mem.eql(u8, v.spec.@"fn", function)) continue; + if (matches(v.spec, args, consts, true)) return v; } std.log.err("no captured Triton variant of {s} for this launch:", .{function}); for (args) |a| switch (a.value) { @@ -167,7 +172,8 @@ fn lookup(args: []const Arg, name: []const u8) ?Arg { return null; } -fn matches(k: KernelJson, args: []const Arg, consts: []const Const) bool { +/// `relaxed`: an argument divisible by 16 may also take a variant that doesn't assume it (never the reverse). +fn matches(k: KernelJson, args: []const Arg, consts: []const Const, relaxed: bool) bool { for (consts) |c| { const got = k.consts.map.get(c.name) orelse return false; if (c.int) |x| if (got.int == null or got.int.? != x) return false; @@ -189,11 +195,13 @@ fn matches(k: KernelJson, args: []const Arg, consts: []const Const) bool { const p = param.?; if (!std.mem.eql(u8, p.type, "i32")) return false; if (!p.nospec and x == 1) return false; - if (p.div16 != (!p.nospec and @mod(x, 16) == 0)) return false; + const div16 = !p.nospec and @mod(x, 16) == 0; + if (p.div16 != div16 and !(relaxed and div16 and !p.div16)) return false; }, .ptr => |x| { const p = param orelse return false; - if (!std.mem.eql(u8, p.type, x.ty) or p.div16 != (x.addr % 16 == 0)) return false; + const div16 = x.addr % 16 == 0; + if (!std.mem.eql(u8, p.type, x.ty) or (p.div16 != div16 and !(relaxed and div16 and !p.div16))) return false; }, .f32 => { const p = param orelse return false; @@ -201,7 +209,8 @@ fn matches(k: KernelJson, args: []const Arg, consts: []const Const) bool { }, .u64 => |x| { const p = param orelse return false; - if (!std.mem.eql(u8, p.type, "u64") or p.div16 != (x % 16 == 0)) return false; + const div16 = x % 16 == 0; + if (!std.mem.eql(u8, p.type, "u64") or (p.div16 != div16 and !(relaxed and div16 and !p.div16))) return false; }, } runtime += 1; diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index ac6c02628..5a56f6501 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -24,6 +24,7 @@ pub const prefill_attention: []const u8 = if (available) &Blob("fatbin_prefill_a pub const scan_rows: []const u8 = if (available) &Blob("fatbin_scan_rows").bytes else &.{}; pub const nemotron_ops: []const u8 = if (available) &Blob("fatbin_nemotron_ops").bytes else &.{}; pub const lane_gemv: []const u8 = if (available) &Blob("fatbin_lane_gemv").bytes else &.{}; +pub const qmm_lane: []const u8 = if (available) &Blob("fatbin_qmm_lane").bytes else &.{}; pub const affine4_pack: []const u8 = if (available) &Blob("fatbin_affine4_pack").bytes else &.{}; pub const qwen_ops: []const u8 = if (available) &Blob("fatbin_qwen_ops").bytes else &.{}; pub const sample: []const u8 = if (available) &Blob("fatbin_sample").bytes else &.{}; diff --git a/zig/src/families/qwen27/cuda/drafts.zig b/zig/src/families/qwen27/cuda/drafts.zig index 02e365952..08ad3fc9a 100644 --- a/zig/src/families/qwen27/cuda/drafts.zig +++ b/zig/src/families/qwen27/cuda/drafts.zig @@ -86,12 +86,12 @@ pub fn propose(d: *dr.DFlash2, asks: []const Ask) !void { pub const Kept = struct { x: *Stream, rows: []const u32 }; /// add_taps_streams: every stream's kept rows' taps gathered after the round's rows, then absorbed in one batch. -pub fn absorb(d: *dr.DFlash2, ops: kern.Ops, taps: u64, copy_items: u64, kept: []const Kept) !void { +pub fn absorb(d: *dr.DFlash2, ops: kern.Ops, taps: u64, round_rows: usize, copy_items: u64, kept: []const Kept) !void { var items: std.ArrayList(kern.Copy) = .empty; defer items.deinit(d.gpa); var absorbs: [st.max_streams]dr.Absorb = undefined; const row = taps_width * bf; - const gathered = taps + st.round_rows * row; + const gathered = taps + round_rows * row; // after the round's rows for (kept, 0..) |k, j| { for (k.rows) |r| try items.append(d.gpa, .{ .dst = gathered + items.items.len * row, .src = taps + r * row, .bytes = row }); absorbs[j] = .{ .ctx = &k.x.ctx, .rows = k.rows.len }; diff --git a/zig/src/families/qwen27/cuda/engine.zig b/zig/src/families/qwen27/cuda/engine.zig index 52576d7cd..2f24825a7 100644 --- a/zig/src/families/qwen27/cuda/engine.zig +++ b/zig/src/families/qwen27/cuda/engine.zig @@ -19,7 +19,8 @@ pub fn chunkBounds(start: usize, end: usize, size: usize, j: usize) [2]usize { return .{ start + (end - start) * j / n, start + (end - start) * (j + 1) / n }; } -pub const Options = struct { context: usize }; +/// `streams`: a shared round's windows at most (the scratch follows); `prompt_rows`: a prompt chunk's rows. +pub const Options = struct { context: usize, streams: usize = st.max_streams, prompt_rows: usize = prompt_rows }; pub const Engine = struct { gpa: std.mem.Allocator, @@ -29,6 +30,7 @@ pub const Engine = struct { w: wts.Weights, scratch: st.Scratch, context: usize, + prompt_rows: usize, pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: []const u8, o: Options) !*Engine { const cfg = try c.Config.read(gpa, io, dir); @@ -38,13 +40,15 @@ pub const Engine = struct { e.gpa = gpa; e.ctx = ctx; e.context = o.context; + if (o.prompt_rows < 1) return error.BadPromptRows; + e.prompt_rows = o.prompt_rows; e.s = try cuda.Stream.init(ctx.d, true); errdefer e.s.deinit(); e.k = try kern.Kernels.load(gpa, io, ctx, kernels, g.linear_v_heads); errdefer e.k.deinit(); e.w = try wts.load(gpa, io, e.ops(), dir, cfg); errdefer e.w.deinit(); - e.scratch = try st.Scratch.init(gpa, ctx.d, g, prompt_rows, o.context); + e.scratch = try st.Scratch.init(gpa, ctx.d, g, o.prompt_rows, o.context, o.streams); errdefer e.scratch.deinit(); try e.s.synchronize(); return e; @@ -77,9 +81,9 @@ pub const Engine = struct { if (prompt.len >= seq.capacity) return error.PromptTooLong; try seq.reset(e.ops(), e.w.g); var f = e.forward(); - const n = (prompt.len + prompt_rows - 1) / prompt_rows; + const n = (prompt.len + e.prompt_rows - 1) / e.prompt_rows; for (0..n) |j| { - const b = chunkBounds(0, prompt.len, prompt_rows, j); + const b = chunkBounds(0, prompt.len, e.prompt_rows, j); try f.chunk(prompt[b[0]..b[1]], seq, j == n - 1); } } diff --git a/zig/src/families/qwen27/cuda/forward.zig b/zig/src/families/qwen27/cuda/forward.zig index 4e4964be4..32914175a 100644 --- a/zig/src/families/qwen27/cuda/forward.zig +++ b/zig/src/families/qwen27/cuda/forward.zig @@ -97,10 +97,7 @@ pub const Forward = struct { switch (layer.mixer) { .delta => |d| { const rec = s.records[i]; - try f.ops.dense(s.y, s.xs, d.qkv, rec.qkv, W); - try f.ops.dense(s.y, s.xs, d.z, s.z, W); - try f.ops.dense(s.y, s.xs, d.b, s.b, W); - try f.ops.dense(s.y, s.xs, d.a, s.a, W); + try f.ops.denseMany(s.y, s.xs, &.{ d.qkv, d.z, d.b, d.a }, &.{ rec.qkv, s.z, s.b, s.a }, W); const sid: ?u64 = if (r.multi) r.sids else null; try f.t.gdnPre(rec.qkv, r.conv(parts, s, g_, i, lin), d.conv, r.windows, sid, s.a, s.b, d.a_log, d.dt_bias, rec.q, rec.k, rec.v, rec.g, rec.beta, W, f.gdnShape()); const state = if (r.multi) 0 else parts[0].seq.rec[i]; @@ -112,9 +109,7 @@ pub const Forward = struct { }, .attention => |a| { const rec = s.attn_records[i]; - try f.ops.dense(s.y, s.xs, a.q, s.qg, W); - try f.ops.dense(s.y, s.xs, a.k, s.key, W); - try f.ops.dense(s.y, s.xs, a.v, rec.v, W); + try f.ops.denseMany(s.y, s.xs, &.{ a.q, a.k, a.v }, &.{ s.qg, s.key, rec.v }, W); try f.t.attnPrep(s.qg, s.key, a.q_norm, a.k_norm, r.pos, w.inv_freq, s.q_rot, rec.k, W, f.attnShape(), c.eps); try f.t.attention(s.q_rot, rec.k, rec.v, s.origin, r.offsets(i, att), r.attn, .{ .o = s.part_o, .m = s.part_m, .l = s.part_l }, s.attn, f.attnShape()); try f.t.gateMul(s.attn, s.qg, s.gated, s.gated_xs, W, g_.query_heads, c.head_dim); @@ -123,8 +118,7 @@ pub const Forward = struct { }, } x = try f.norm(x, s.r, layer.post_norm, W); - try f.ops.dense(s.y, s.xs, layer.gate, s.gate, W); - try f.ops.dense(s.y, s.xs, layer.up, s.up, W); + try f.ops.denseMany(s.y, s.xs, &.{ layer.gate, layer.up }, &.{ s.gate, s.up }, W); try f.t.swiglu(s.gate, s.up, s.act, s.act_xs, W, g_.intermediate); try f.ops.dense(s.act, s.act_xs, layer.down, s.pending, W); pending = s.pending; @@ -295,7 +289,7 @@ pub const Forward = struct { const g_ = f.g(); const w = f.w; const S = pieces.len; - if (S == 0 or S > st.max_streams) return error.InvalidQwenWindows; + if (S == 0 or S > s.streams()) return error.InvalidQwenWindows; var starts: [st.max_streams + 1]usize = undefined; starts[0] = 0; for (pieces, 0..) |p, k| { @@ -329,7 +323,7 @@ pub const Forward = struct { defer items.deinit(f.gpa); var lin_n: usize = 0; for (0..g_.layers) |i| if (c.linear(i)) { - const cat = s.conv_cat + lin_n * st.max_streams * keep * row; + const cat = s.conv_cat + lin_n * s.streams() * keep * row; for (pieces, 0..) |p, k| try items.append(f.gpa, .{ .dst = cat + k * keep * row, .src = p.seq.conv[i], .bytes = keep * row }); lin_n += 1; }; @@ -340,7 +334,7 @@ pub const Forward = struct { at_layer[i] = items.items.len; if (c.linear(i)) { // the last `keep` rows of [kept rows; the stream's prompt rows] - const cat = s.conv_cat + lin_n * st.max_streams * keep * row; + const cat = s.conv_cat + lin_n * s.streams() * keep * row; for (pieces, 0..) |p, k| { const n = p.ids.len; const o = starts[k]; @@ -378,7 +372,7 @@ pub const Forward = struct { try f.ops.prefillDense(s.y, d.z, s.z, W); try f.ops.prefillDense(s.y, d.b, s.b, W); try f.ops.prefillDense(s.y, d.a, s.a, W); - const cat = s.conv_cat + lin_n * st.max_streams * keep * row; + const cat = s.conv_cat + lin_n * s.streams() * keep * row; try f.t.gdnPre(s.qkv, cat, d.conv, s.windows, s.sids, s.a, s.b, d.a_log, d.dt_bias, s.gq, s.gk, s.gv, s.gg, s.gbeta, W, f.gdnShape()); for (pieces, 0..) |p, k| { const o = starts[k]; diff --git a/zig/src/families/qwen27/cuda/kernels.zig b/zig/src/families/qwen27/cuda/kernels.zig index 4b486922f..04fe87dec 100644 --- a/zig/src/families/qwen27/cuda/kernels.zig +++ b/zig/src/families/qwen27/cuda/kernels.zig @@ -30,7 +30,12 @@ pub const GemvPart = extern struct { w: u64, scales: u64, biases: u64, out: u64, /// gdn.cu's Pending (the tree kernel's rows still to fold): all null here, commits replay instead. pub const Pending = extern struct { k: u64 = 0, v: u64 = 0, g: u64 = 0, beta: u64 = 0, rows: u64 = 0, row_stride: c_int = 0, counts: u64 = 0, count_stride: c_int = 0 }; +/// qmm_lane.cu's Seg and Segs (the segmented launch), passed by value: four weights at most. +pub const Seg = extern struct { w: u64, scales: u64, biases: u64, out: u64, part: u64, n: c_int, npad: c_int, sk: c_int, blocks: c_int, sync: c_int }; +pub const Segs = extern struct { s: [4]Seg, count: c_int }; + comptime { + std.debug.assert(@sizeOf(Seg) == 64 and @sizeOf(Segs) == 264); std.debug.assert(@sizeOf(Part) == 56 and @sizeOf(Parts) == 232 and @sizeOf(GemvPart) == 48 and @sizeOf(Pending) == 64); } @@ -38,6 +43,14 @@ pub const group_smem: u32 = 35328; // LaneTile<64, 16, 64, 1, 4, 8>::SMEM /// qmm_group's GB10 tiles 3-5 for wider rounds: rows a tile, columns a tile, threads, LaneTile::SMEM. pub const Wide = struct { bm: usize, bn: usize, threads: u32, smem: u32 }; pub const wide = [_]Wide{ .{ .bm = 32, .bn = 64, .threads = 128, .smem = 26112 }, .{ .bm = 64, .bn = 64, .threads = 128, .smem = 43008 }, .{ .bm = 64, .bn = 128, .threads = 256, .smem = 39168 } }; +/// qmm_lane.cu's Tile<64, BM, 64, 1, 4, STAGES>: a stage's bytes at 16 rows; SMEM at 4 stages for 16, 32, 64 rows. +pub const seg_stage: u32 = 16 * 128 + 64 * 32 + 2 * 64 * 2 + 16 * 4; +pub const lane_smem = [_]u32{ 17664, 26112, 43008 }; +/// Rows a wide lane launch takes at most (wider rounds run in turns), so K-slice partials fit `part`. +pub const lane_rows: usize = 128; +/// split_k stops once tiles x slices reach 192, so a weight's slices hold fewer than 384 column tiles of partials. +pub const part_bytes: usize = lane_rows * 384 * 64 * 4; +pub const sync_counters: usize = 1 << 14; pub const prefill_mm_smem: u32 = 41984; // Tile<64, 128, 128, 2, 2, 2>::SMEM: prompt_tile 9, two stages pub const pattn_smem: u32 = 98304; // two 32-key slots of 256 dims beside eight warps' staged queries @@ -49,11 +62,17 @@ pub const QLinear = struct { w: u64, s: u64, b: u64, n: usize, k: usize, npad: u pub const Kernels = struct { d: *const cuda.Driver, - mods: [15]cuda.Module, + mods: [16]cuda.Module, triton: cuda.aot.Set, group: cuda.Function, wide: [wide.len]cuda.Function, // tiles 3-5: rows <= 32, <= 64, more gemv: cuda.Function, + seg: [2][2][3]cuda.Function, // qmm_lane's segmented launch by stages (3, 4), 8-token tiles (1, 2), slices a block (1, 2, 4) + lane: [3]cuda.Function, // qmm_lane's row tiles of 16, 32 and 64 for wider rounds + reduce: cuda.Function, + part: cuda.DeviceBuffer, // K-slice partials of the pre-sm_90 launches + sync: cuda.DeviceBuffer, // the segmented launch's finished-slice counters, zero between launches + sms: usize, prefill_mm: cuda.Function, pattn: cuda.Function, trees: [cuda.kernels.gdn_variants.len]cuda.Function, // tree_kernel by gdn_variants: [0] the chain one @@ -74,6 +93,7 @@ pub const Kernels = struct { gemv_blocks: usize, // resident lane_gemv CTAs: per SM times SMs gemv_rows: usize, // windows of at most this many rows take lane_gemv (none; TF_QWEN27_GEMV_ROWS) gb10: bool, + clusters: bool, // sm_90 and later: K slices meet in a cluster; older chips take lane_gemv 16 rows at a time discrete: bool, // the card has its own memory: checkpoint bytes reach it through page-locked slots /// Loads every module; `triton_dir` holds the captured aot.json and cubins for this GPU. @@ -83,7 +103,7 @@ pub const Kernels = struct { var k: Kernels = undefined; k.d = d; const kk = cuda.kernels; - const images = [_][]const u8{ kk.qmm_group, kk.qmm_prefill, kk.lane_gemv, kk.prefill_attention, kk.gdn, kk.gdn_prefill, kk.affine4_pack, kk.torch_argmax, kk.torch_topk, kk.torch_pointwise, kk.torch_indexing, kk.torch_movement, kk.torch_nemotron_constants, kk.qwen_ops, kk.torch_norm_rotary }; + const images = [_][]const u8{ kk.qmm_group, kk.qmm_prefill, kk.lane_gemv, kk.prefill_attention, kk.gdn, kk.gdn_prefill, kk.affine4_pack, kk.torch_argmax, kk.torch_topk, kk.torch_pointwise, kk.torch_indexing, kk.torch_movement, kk.torch_nemotron_constants, kk.qwen_ops, kk.torch_norm_rotary, kk.qmm_lane }; var loaded: usize = 0; errdefer for (k.mods[0..loaded]) |*m| m.unload(); for (images, 0..) |img, i| { @@ -107,12 +127,31 @@ pub const Kernels = struct { k.inv_freq = try k.mods[13].function("tf_inv_freq"); k.rotary = try k.mods[13].function("tf_rotary"); k.torch_rms = try k.mods[14].function("tf_dflash_rms5120_bf16_fma_kernel"); + inline for (0..2) |si| inline for (0..2) |ti| inline for (0..3) |gi| { + const st: u32 = 3 + si; + const g: u32 = 1 << gi; + const f = &k.seg[si][ti][gi]; + f.* = try k.mods[15].function(std.fmt.comptimePrint("tf_qmm_seg_s{d}_t{d}_g{d}", .{ st, ti + 1, g })); + try f.allowDynamicShared(if (g > 1) g * st * seg_stage else @max(st * seg_stage, 4096)); + }; + for (&k.lane, [_][:0]const u8{ "tf_qmm_lane_m16", "tf_qmm_lane_m32", "tf_qmm_lane_m64" }, lane_smem) |*f, name, smem| { + f.* = try k.mods[15].function(name); + try f.allowDynamicShared(smem); + } + k.reduce = try k.mods[15].function("tf_qmm_reduce"); const sms: usize = @intCast(try ctx.attribute(.multiprocessor_count)); // gdn_prefill_cuda: a block takes 128 value rows, or 64 when there are fewer heads than SMs k.prompt_rows = if (heads >= sms) 128 else 64; k.prompt_chain = try k.mods[5].function(if (k.prompt_rows == 128) sym.chain128 else sym.chain64); k.triton = try cuda.aot.Set.load(gpa, io, d, ctx.device, triton_dir); errdefer k.triton.deinit(); + k.sms = sms; + k.clusters = try ctx.attribute(.compute_capability_major) >= 9; + k.part = try cuda.DeviceBuffer.alloc(d, if (k.clusters) 0 else part_bytes); + errdefer k.part.free(); + k.sync = try cuda.DeviceBuffer.alloc(d, if (k.clusters) 0 else sync_counters * 4); + errdefer k.sync.free(); + if (!k.clusters) try k.sync.fill32(0, null); try k.group.allowDynamicShared(group_smem); for (k.wide, wide) |f, t| try f.allowDynamicShared(t.smem); try k.gemv.allowDynamicShared(group_smem); @@ -131,6 +170,8 @@ pub const Kernels = struct { pub fn deinit(k: *Kernels) void { k.triton.deinit(); + k.part.free(); + k.sync.free(); for (&k.mods) |*m| m.unload(); } }; @@ -167,13 +208,44 @@ pub const Ops = struct { /// qmm.matmul on sm_12x: x (rows, k) bf16 with group sums xs -> out (rows, n) bf16, qmm_group's bits. pub fn dense(o: Ops, x: u64, xs: u64, q: QLinear, out: u64, rows: usize) !void { const sk = splitK(q.n, q.k); + if (!o.k.clusters and !(sk > 1 and rows <= o.k.gemv_rows)) return o.lane3090(x, xs, q, out, rows, sk); if (rows > 16) return o.wideTile(x, xs, q, out, rows, sk); // Python's clusters, a CTA a K slice; lane_gemv (byte-equal) when TF_QWEN27_GEMV_ROWS asks for it return if (sk > 1 and rows <= o.k.gemv_rows) o.gemv(x, xs, q, out, rows, sk) else o.cluster(x, xs, q, out, rows, sk); } + /// Weights over one input: before sm_90 and up to 16 rows, a segmented launch a K split, each weight at its own. + pub fn denseMany(o: Ops, x: u64, xs: u64, qs: []const QLinear, outs: []const u64, rows: usize) !void { + if (o.k.clusters or rows > 16 or (o.k.gemv_rows > 0 and rows <= o.k.gemv_rows)) { + for (qs, outs) |q, out| try o.dense(x, xs, q, out, rows); + return; + } + var done: u8 = 0; + for (qs, 0..) |q0, i| { + if (done & (@as(u8, 1) << @intCast(i)) != 0) continue; + const sk = splitK(q0.n, q0.k); + var segs: Segs = std.mem.zeroes(Segs); + var blocks: usize = 0; + var sync: usize = 0; + var part: u64 = o.k.part.ptr; + for (qs[i..], outs[i..], i..) |q, out, j| { + if (splitK(q.n, q.k) != sk or q.k != q0.k or segs.count == 4) continue; + done |= @as(u8, 1) << @intCast(j); + const tiles = (q.n + 63) / 64; + segs.s[@intCast(segs.count)] = .{ .w = q.w, .scales = q.s, .biases = q.b, .out = out, .part = if (sk > 1) part else 0, .n = int(q.n), .npad = int(q.npad), .sk = int(sk), .blocks = int(tiles * sk), .sync = int(sync) }; + segs.count += 1; + blocks += tiles * sk; + sync += tiles; + part += sk * rows * q.n * 4; + } + if (sk > 1 and sync > sync_counters) return error.TooManyTiles; + try o.segLaunch(x, xs, segs, blocks, sk, rows, q0.k); + } + } + /// qmm.matmul at a K split the caller names (matmul_rows takes the stacked weight's for each part). pub fn denseSplit(o: Ops, x: u64, xs: u64, q: QLinear, out: u64, rows: usize, sk: usize) !void { + if (!o.k.clusters and !(sk > 1 and rows <= o.k.gemv_rows)) return o.lane3090(x, xs, q, out, rows, sk); if (rows > 16) return o.wideTile(x, xs, q, out, rows, sk); return if (sk > 1 and rows <= o.k.gemv_rows) o.gemv(x, xs, q, out, rows, sk) else o.cluster(x, xs, q, out, rows, sk); } @@ -198,6 +270,60 @@ pub const Ops = struct { try cuda.launch.launch(o.k.gemv, cfg, o.s, &a); } + /// Before sm_90 (no clusters): qmm_lane's segmented launch to 16 rows, else row tiles 128 rows a turn. + fn lane3090(o: Ops, x: u64, xs: u64, q: QLinear, out: u64, rows: usize, sk: usize) !void { + if (rows <= 16) return o.seg(x, xs, q, out, rows, sk); + var r: usize = 0; + while (r < rows) : (r += lane_rows) { + try o.lane(x + r * q.k * 2, xs + r * (q.k / 64) * 4, q, out + r * q.n * 2, @min(lane_rows, rows - r), sk); + } + } + + fn seg(o: Ops, x: u64, xs: u64, q: QLinear, out: u64, rows: usize, sk: usize) !void { + const tiles = (q.n + 63) / 64; + const blocks = tiles * sk; + if (sk > 1 and tiles > sync_counters) return error.TooManyTiles; + var segs: Segs = std.mem.zeroes(Segs); + segs.count = 1; + segs.s[0] = .{ .w = q.w, .scales = q.s, .biases = q.b, .out = out, .part = if (sk > 1) o.k.part.ptr else 0, .n = int(q.n), .npad = int(q.npad), .sk = int(sk), .blocks = int(blocks), .sync = 0 }; + return o.segLaunch(x, xs, segs, blocks, sk, rows, q.k); + } + + fn segLaunch(o: Ops, x: u64, xs: u64, segs: Segs, blocks: usize, sk: usize, rows: usize, k: usize) !void { + // qmm_seg_cuda: past 4 blocks an SM, 3 stages; 2 or 4 slices side by side in a block + const st: u32 = if (blocks > 4 * o.k.sms) 3 else 4; + const g: u32 = if (sk == 2 or sk == 4) @intCast(sk) else 1; + const f = o.k.seg[st - 3][if (rows <= 8) 0 else 1][if (g == 1) 0 else if (g == 2) 1 else 2]; + var a: cuda.Args = .{}; + a.add(x); + a.add(xs); + a.add(segs); + a.add(o.k.sync.ptr); + for ([_]usize{ rows, k, k }) |v| a.add(int(v)); + const smem = if (g > 1) g * st * seg_stage else @max(st * seg_stage, 4096); + try cuda.launch.launch(f, .{ .grid = .{ .x = u(blocks / g) }, .block = .{ .x = 128 * g }, .shared = smem }, o.s, &a); + } + + fn lane(o: Ops, x: u64, xs: u64, q: QLinear, out: u64, rows: usize, sk: usize) !void { + const i: usize = if (rows <= 16) 0 else if (rows <= 32) 1 else 2; + const bm: usize = @as(usize, 16) << @intCast(i); + const rows_t = (rows + bm - 1) / bm; + // a group's inputs stay near 12 MB of L2 while its blocks sweep the column tiles + const group = @max(1, @min(rows_t, (12 << 20) / (bm * q.k * 2))); + var a: cuda.Args = .{}; + for ([_]u64{ x, xs, q.w, q.s, q.b, out, if (sk > 1) o.k.part.ptr else 0 }) |v| a.add(v); + for ([_]usize{ rows, q.n, q.k, sk, q.npad, q.k, group }) |v| a.add(int(v)); + try cuda.launch.launch(o.k.lane[i], .{ .grid = .{ .x = u(rows_t * ((q.n + 63) / 64)), .z = u(sk) }, .block = .{ .x = 128 }, .shared = lane_smem[i] }, o.s, &a); + if (sk == 1) return; + const total = rows * q.n; + var r: cuda.Args = .{}; + r.add(o.k.part.ptr); + r.add(out); + r.add(@as(i64, @intCast(total))); + r.add(int(sk)); + try cuda.launch.launch(o.k.reduce, .{ .grid = .{ .x = u((total + 255) / 256) }, .block = .{ .x = 256 } }, o.s, &r); + } + /// qmm_group's tile 2 with one K slice: a CTA a column tile. fn cluster(o: Ops, x: u64, xs: u64, q: QLinear, out: u64, rows: usize, sk: usize) !void { const tiles = (q.n + 63) / 64; diff --git a/zig/src/families/qwen27/cuda/lanes.zig b/zig/src/families/qwen27/cuda/lanes.zig index 8a5f86fad..e73796e3d 100644 --- a/zig/src/families/qwen27/cuda/lanes.zig +++ b/zig/src/families/qwen27/cuda/lanes.zig @@ -7,7 +7,6 @@ const st = @import("state.zig"); const Engine = @import("engine.zig").Engine; const Part = @import("forward.zig").Part; const Forward = @import("forward.zig").Forward; -const prompt_rows = @import("engine.zig").prompt_rows; const dl = @import("draft_load.zig"); const dr = @import("draft.zig"); const drafts = @import("drafts.zig"); @@ -78,7 +77,7 @@ pub const Cuda = struct { /// The facts the round loop reads: 16-row windows in one forward; with DFlash2, host trees with node chances. pub fn facts(x: *const Cuda) lanes.Model { - var m: lanes.Model = .{ .exact_width = st.max_rows, .hidden_rows = true, .batch_rows = st.round_rows, .max_streams = @min(x.streams, st.max_streams), .streams_exact = true }; + var m: lanes.Model = .{ .exact_width = st.max_rows, .hidden_rows = true, .batch_rows = @intCast(x.e.scratch.window_rows), .max_streams = @intCast(@min(x.streams, x.e.scratch.window_rows / st.max_rows)), .streams_exact = true }; if (x.drafter != null) { m.mtp = true; m.speculate = true; @@ -248,14 +247,14 @@ pub const Cuda = struct { continue; }; const ids = s.prompt(); - if (ids.len > prompt_rows) { + if (ids.len > x.e.prompt_rows) { x.prefillOpened(s, seq) catch |e| { x.dropSeq(s); err.* = e; }; continue; } - if (n == st.max_streams or rows + ids.len > prompt_rows) { + if (n == x.e.scratch.streams() or rows + ids.len > x.e.prompt_rows) { x.pass(streams, errs, pieces[0..n], which[0..n]); n = 0; rows = 0; @@ -471,7 +470,7 @@ pub const Cuda = struct { na += 1; } const f = x.e.forward(); - try drafts.absorb(d, f.ops, x.e.scratch.taps, x.e.scratch.copy_items, kept[0..nk]); + try drafts.absorb(d, f.ops, x.e.scratch.taps, x.e.scratch.window_rows, x.e.scratch.copy_items, kept[0..nk]); try drafts.propose(d, asks[0..na]); } diff --git a/zig/src/families/qwen27/cuda/lone.zig b/zig/src/families/qwen27/cuda/lone.zig index 0397ea706..5d75100bf 100644 --- a/zig/src/families/qwen27/cuda/lone.zig +++ b/zig/src/families/qwen27/cuda/lone.zig @@ -79,7 +79,7 @@ pub fn prefill(e: *Engine, seq: *st.Seq, ctx: *dr.Context, d: *dr.DFlash2, promp ctx.skip(tap_from); var f = e.forward(); f.taps = e.scratch.taps; - const rows = @import("engine.zig").prompt_rows; + const rows = e.prompt_rows; const chunks = (prompt.len + rows - 1) / rows; for (0..chunks) |j| { const b = @import("engine.zig").chunkBounds(0, prompt.len, rows, j); diff --git a/zig/src/families/qwen27/cuda/native.zig b/zig/src/families/qwen27/cuda/native.zig index 6e3f362fd..aa576026f 100644 --- a/zig/src/families/qwen27/cuda/native.zig +++ b/zig/src/families/qwen27/cuda/native.zig @@ -13,9 +13,11 @@ pub const model_type = "qwen3_5"; pub const formats: []const []const u8 = &.{"mlx-q4g64"}; pub const default_context: i64 = 32768; pub const max_segments: u32 = 1; -pub const prompt_rows: u32 = engine.prompt_rows; +/// A prompt chunk's rows: TF_QWEN27_PROMPT_ROWS, else engine.prompt_rows (set when the server opens the model). +pub var prompt_rows: u32 = engine.prompt_rows; -pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, drafter: ?[]const u8 = null }; +/// `streams`: --parallel, the shared round's windows at most (the scratch is sized for them). +pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, drafter: ?[]const u8 = null, streams: u32 = most_streams }; /// A lone drafted stream's own driver (none here: one stream decodes in the lane core too, as under --parallel). pub const LoneRun = *const fn (ctx: *anyopaque, s: *lanes.Stream, hooks: *anyopaque, committed: *const fn (*anyopaque) void, yield: *const fn (*anyopaque) bool) anyerror!bool; @@ -67,11 +69,14 @@ const most_streams = 32; pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { // the 27B's glue runs the Python engine's Triton kernels, captured on this chip (no own glue kernels yet) - const e = try engine.Engine.init(gpa, io, ctx, dir, kernels orelse return error.QwenNeedsCapturedKernels, .{ .context = o.context }); + const streams: u32 = @max(1, @min(o.streams, st.max_streams)); + if (getenv("TF_QWEN27_PROMPT_ROWS")) |v| prompt_rows = std.fmt.parseInt(u32, v, 10) catch return error.BadPromptRows; + if (prompt_rows < 16) return error.BadPromptRows; + const e = try engine.Engine.init(gpa, io, ctx, dir, kernels orelse return error.QwenNeedsCapturedKernels, .{ .context = o.context, .streams = streams, .prompt_rows = prompt_rows }); errdefer e.deinit(); const own = try gpa.create(Owned); errdefer gpa.destroy(own); - own.* = .{ .gpa = gpa, .e = e, .lanes = try Lanes.init(gpa, e, most_streams) }; + own.* = .{ .gpa = gpa, .e = e, .lanes = try Lanes.init(gpa, e, streams) }; errdefer own.lanes.deinit(); if (o.drafts) if (try drafterDir(gpa, io, o.drafter)) |draft_dir| { defer gpa.free(draft_dir); @@ -80,7 +85,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ own.lanes.drafter = own.draft; std.log.info("DFlash2 drafts every stream: {s}", .{draft_dir}); // the forward's ms by rows: a shared round's tree widths follow the curve - try own.lanes.calibrate(io, most_streams); + try own.lanes.calibrate(io, streams); }; return .{ .backend = own.lanes.backend(), diff --git a/zig/src/families/qwen27/cuda/round.zig b/zig/src/families/qwen27/cuda/round.zig index 2b2d4e269..9cc798dab 100644 --- a/zig/src/families/qwen27/cuda/round.zig +++ b/zig/src/families/qwen27/cuda/round.zig @@ -33,7 +33,7 @@ pub const Round = struct { /// The conv rows gdn_pre reads for layer `i` (its DeltaNet index `j`): the stream's own, or the stacked copy. pub fn conv(r: *const Round, parts: []const Part, s: *const st.Scratch, g: c.Geometry, i: usize, j: usize) u64 { if (!r.multi) return parts[0].seq.conv[i]; - return s.conv_cat + j * st.max_streams * (c.conv_taps - 1) * g.convDim() * bf; + return s.conv_cat + j * s.streams() * (c.conv_taps - 1) * g.convDim() * bf; } /// Layer `i`'s attention cache offsets: tree_forward's (one stream, by layer) or the packed (layer, streams) rows. @@ -57,6 +57,7 @@ pub fn parentsOf(p: Part, buf: []i32) []const i32 { /// attention.plan_host: rows, (first, rows, keys, slots) per stream, items, then window parents (global rows). fn attnPlan(gpa: std.mem.Allocator, out: *std.ArrayList(i32), parts: []const Part, group: usize) !usize { for (parts, 0..) |p, s| for (p.tokens) |_| try out.append(gpa, @intCast(s)); + while (out.items.len % 4 != 0) try out.append(gpa, 0); // each part from 16 bytes, as planAt places them var start: usize = 0; for (parts) |p| { const w = p.tokens.len; @@ -77,6 +78,7 @@ fn attnPlan(gpa: std.mem.Allocator, out: *std.ArrayList(i32), parts: []const Par } } } + while (out.items.len % 4 != 0) try out.append(gpa, 0); start = 0; var buf: [tree.max_nodes]i32 = undefined; for (parts) |p| { @@ -86,6 +88,13 @@ fn attnPlan(gpa: std.mem.Allocator, out: *std.ArrayList(i32), parts: []const Par return n_items; } +/// attnPlan's parts in i32s from its 16-byte start: rows, then streams, items and parents, each on 16 bytes. +const PlanAt = struct { streams: usize, items: usize, parents: usize }; +fn planAt(W: usize, S: usize, n_items: usize) PlanAt { + const rows = std.mem.alignForward(usize, W, 4); + return .{ .streams = rows, .items = rows + 4 * S, .parents = rows + 4 * S + std.mem.alignForward(usize, 3 * n_items, 4) }; +} + /// Each stream's key and value caches as bf16 element offsets from the scratch's origin. fn offsetOf(s: *const st.Scratch, ptr: u64) i64 { return @divExact(@as(i64, @bitCast(ptr -% s.origin)), bf); @@ -117,7 +126,7 @@ fn plans(gpa: std.mem.Allocator, parts: []const Part, starts: []const usize, mul /// Uploads the round's inputs and runs attention's _paths; the caller's layer loop reads the returned Round. pub fn stage(gpa: std.mem.Allocator, ops: kern.Ops, t: tri.Tri, s: *st.Scratch, g: c.Geometry, parts: []const Part) !Round { - if (parts.len == 0 or parts.len > st.max_streams) return error.InvalidQwenWindows; + if (parts.len == 0 or parts.len > s.streams()) return error.InvalidQwenWindows; var W: usize = 0; var r: Round = undefined; r.starts_host[0] = 0; @@ -128,7 +137,7 @@ pub fn stage(gpa: std.mem.Allocator, ops: kern.Ops, t: tri.Tri, s: *st.Scratch, W += p.tokens.len; r.starts_host[k + 1] = W; } - if (W > st.round_rows) return error.WindowTooWide; + if (W > s.window_rows) return error.WindowTooWide; r.rows = W; r.streams = parts.len; r.multi = parts.len > 1; @@ -175,30 +184,42 @@ pub fn stage(gpa: std.mem.Allocator, ops: kern.Ops, t: tri.Tri, s: *st.Scratch, r.starts = s.plan + 3 * W * 4; r.tables = 0; r.offs = s.offs; - r.attn = .{ .rows = s.attn_plan, .streams = s.attn_plan + W * 4, .items = s.attn_plan + (W + 4) * 4, .n_items = n_items, .paths = s.paths, .depths = s.depths, .width = W }; - try t.paths(s.attn_plan + (W + 4 + 3 * n_items) * 4, s.paths, s.depths, W); + const at = planAt(W, 1, n_items); + r.attn = .{ .rows = s.attn_plan, .streams = s.attn_plan + at.streams * 4, .items = s.attn_plan + at.items * 4, .n_items = n_items, .paths = s.paths, .depths = s.depths, .width = W }; + try t.paths(s.attn_plan + at.parents * 4, s.paths, s.depths, W); return r; } - // multi_tree_forward: positions, stream ids, plan entries, starts, ids, windows, attention's plan in one copy + // multi_tree_forward's tables in one copy, each part on 16 bytes: captured variants assume their pointer alignment + const pad = struct { + fn to16(gpa_: std.mem.Allocator, list: *std.ArrayList(i32)) !usize { + while (list.items.len % 4 != 0) try list.append(gpa_, 0); + return list.items.len; + } + }.to16; try host.appendSlice(gpa, h.pos.items); + const at_sids = try pad(gpa, &host); for (parts, 0..) |p, k| for (p.tokens) |_| try host.append(gpa, @intCast(k)); + const at_plan = try pad(gpa, &host); try host.appendSlice(gpa, h.entries.items); for (r.starts_host[0 .. S + 1]) |x| try host.append(gpa, @intCast(x)); + const at_ids = try pad(gpa, &host); for (parts) |p| for (p.tokens) |tok| try host.append(gpa, @intCast(tok)); + const at_windows = try pad(gpa, &host); try host.appendSlice(gpa, h.windows.items); - const at_attn = host.items.len; + const at_attn = try pad(gpa, &host); const n_items = try attnPlan(gpa, &host, parts, group); try ops.upload(s.round, std.mem.sliceAsBytes(host.items)); try ops.s.synchronize(); const base = s.round; r.pos = base; - r.sids = base + W * 4; - r.plan = base + 2 * W * 4; - r.starts = base + 5 * W * 4; - r.ids = base + (5 * W + S + 1) * 4; - r.windows = base + (6 * W + S + 1) * 4; + r.sids = base + at_sids * 4; + r.plan = base + at_plan * 4; + r.starts = r.plan + 3 * W * 4; + r.ids = base + at_ids * 4; + r.windows = base + at_windows * 4; const ab = base + at_attn * 4; - r.attn = .{ .rows = ab, .streams = ab + W * 4, .items = ab + (W + 4 * S) * 4, .n_items = n_items, .paths = s.paths, .depths = s.depths, .width = W }; + const at = planAt(W, S, n_items); + r.attn = .{ .rows = ab, .streams = ab + at.streams * 4, .items = ab + at.items * 4, .n_items = n_items, .paths = s.paths, .depths = s.depths, .width = W }; // the state pointer table and cache offsets, one copy each; every stream's conv rows stacked by layer var ptrs: std.ArrayList(u64) = .empty; defer ptrs.deinit(gpa); @@ -211,7 +232,7 @@ pub fn stage(gpa: std.mem.Allocator, ops: kern.Ops, t: tri.Tri, s: *st.Scratch, for (parts) |p| try offs.appendSlice(gpa, &.{ offsetOf(s, p.seq.keys[i]), offsetOf(s, p.seq.values[i]) }); continue; } - const cat = s.conv_cat + j * st.max_streams * keep * g.convDim() * bf; + const cat = s.conv_cat + j * s.streams() * keep * g.convDim() * bf; for (parts, 0..) |p, k| { try ptrs.append(gpa, p.seq.rec[i]); try stack.append(gpa, .{ .dst = cat + k * keep * g.convDim() * bf, .src = p.seq.conv[i], .bytes = keep * g.convDim() * bf }); @@ -225,6 +246,6 @@ pub fn stage(gpa: std.mem.Allocator, ops: kern.Ops, t: tri.Tri, s: *st.Scratch, try ops.s.synchronize(); r.tables = s.tables; r.offs = s.offs; - try t.paths(ab + (W + 4 * S + 3 * n_items) * 4, s.paths, s.depths, W); + try t.paths(ab + at.parents * 4, s.paths, s.depths, W); return r; } diff --git a/zig/src/families/qwen27/cuda/state.zig b/zig/src/families/qwen27/cuda/state.zig index 4cbe06c07..28366907e 100644 --- a/zig/src/families/qwen27/cuda/state.zig +++ b/zig/src/families/qwen27/cuda/state.zig @@ -122,6 +122,7 @@ pub const AttnRecord = struct { k: u64, v: u64 }; pub const Scratch = struct { arena: Arena, rows: usize, + window_rows: usize, // a shared round's rows at most: max_rows a stream // the residual stream and the normed rows with their group sums x: u64, h: u64, @@ -192,9 +193,11 @@ pub const Scratch = struct { attn_records: []AttnRecord, // by layer: the round's rotated keys and values gpa: std.mem.Allocator, - /// Scratch for prompt chunks of `prompt_rows` and windows of `max_rows`, over caches of `capacity` positions. - pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, g: c.Geometry, prompt_rows: usize, capacity: usize) !Scratch { + /// Scratch for `prompt_rows` prompt chunks and rounds of `streams` windows, over caches of `capacity` positions. + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, g: c.Geometry, prompt_rows: usize, capacity: usize, stream_count: usize) !Scratch { + if (stream_count < 1 or stream_count > max_streams) return error.BadStreams; var s: Scratch = undefined; + s.window_rows = max_rows * stream_count; s.gpa = gpa; s.records = try gpa.alloc(GdnRecord, g.layers); errdefer gpa.free(s.records); @@ -208,9 +211,14 @@ pub const Scratch = struct { return s; } + /// The streams a shared round or batched prompt may hold. + pub fn streams(s: *const Scratch) usize { + return s.window_rows / max_rows; + } + fn layout(s: *Scratch, g: c.Geometry, prompt_rows: usize, capacity: usize) !void { - const R = @max(prompt_rows, round_rows); - const W = round_rows; + const W = s.window_rows; + const R = @max(prompt_rows, W); const bf = 2; const chunks = (capacity + W + tri.chunk - 1) / tri.chunk; // attention.plan_host: a query tile's items per committed chunk and folded group @@ -263,7 +271,7 @@ pub const Scratch = struct { s.prompt_values = try A.take(R * g.kvInner() * bf); s.windows = try A.take(R * c.conv_taps * 4); s.plan = try A.take(R * 3 * 4); - s.attn_plan = try A.take((W + 4 + 3 * items + W) * 4); + s.attn_plan = try A.take((W + 4 + 3 * items + W + 6) * 4); // + its parts' 16-byte padding s.paths = try A.take(W * tri.max_nodes * 4); s.depths = try A.take(W * 4); s.offs = try A.take(g.layers * max_streams * 2 * 8); @@ -271,14 +279,14 @@ pub const Scratch = struct { s.part_m = try A.take(chunks * W * g.query_heads * 4); s.part_l = try A.take(chunks * W * g.query_heads * 4); s.conv_tmp = try A.take((c.conv_taps - 1 + R) * g.convDim() * bf); - s.conv_old = try A.take(max_streams * g.layers * (c.conv_taps - 1) * g.convDim() * bf); + s.conv_old = try A.take(s.streams() * g.layers * (c.conv_taps - 1) * g.convDim() * bf); s.copy_items = try A.take(2 * copy_items * @sizeOf(kern.Copy)); - s.taps = try A.take(R * 5 * g.hidden * bf); + s.taps = try A.take(@max(R, 2 * W) * 5 * g.hidden * bf); // a round's rows, then its kept rows gathered s.replay_table = try A.take(g.layers * (4 + max_streams) * 8); s.replay_rows = try A.take(max_streams * (W + 1) * 4); - s.round = try A.take((6 * W + max_streams + 1 + W * c.conv_taps + W + 4 * max_streams + 3 * items * max_streams + W) * 4); + s.round = try A.take((6 * W + max_streams + 1 + W * c.conv_taps + W + 4 * max_streams + 3 * items * max_streams + W + 40) * 4); // + each part's 16-byte padding s.tables = try A.take(g.layers * max_streams * 8); - s.conv_cat = try A.take(g.layers * max_streams * (c.conv_taps - 1) * g.convDim() * bf); + s.conv_cat = try A.take(g.layers * s.streams() * (c.conv_taps - 1) * g.convDim() * bf); for (0..g.layers) |i| { if (c.linear(i)) { s.records[i] = .{ diff --git a/zig/src/families/qwen27/cuda/weights.zig b/zig/src/families/qwen27/cuda/weights.zig index 5c796bd43..e7c2b9b6b 100644 --- a/zig/src/families/qwen27/cuda/weights.zig +++ b/zig/src/families/qwen27/cuda/weights.zig @@ -19,6 +19,7 @@ pub const Attention = struct { q: QLinear, k: QLinear, v: QLinear, o: QLinear, q pub const Mixer = union(enum) { delta: Delta, attention: Attention }; pub const Layer = struct { input_norm: u64, post_norm: u64, mixer: Mixer, gate: QLinear, up: QLinear, down: QLinear }; + pub const Weights = struct { gpa: std.mem.Allocator, g: c.Geometry, @@ -30,10 +31,13 @@ pub const Weights = struct { eos: [config_eos]u32 = undefined, // config.json's end tokens, then generation_config.json's eos_count: usize = 0, buffers: std.ArrayList(cuda.DeviceBuffer) = .empty, + host_buffers: std.ArrayList(cuda.HostBuffer) = .empty, // TF_QWEN27_HOST_EMBED=1: the token table, mapped pub fn deinit(w: *Weights) void { for (w.buffers.items) |*b| b.free(); w.buffers.deinit(w.gpa); + for (w.host_buffers.items) |*b| b.free(); + w.host_buffers.deinit(w.gpa); w.gpa.free(w.layers); w.* = undefined; } @@ -84,6 +88,15 @@ const Loader = struct { return L.host.items; } + /// A tensor in page-locked host memory the kernels read over PCIe (a row lookup reads a few KB a token). + fn rawHost(L: *Loader, t: Tensor) !u64 { + var b = try cuda.HostBuffer.allocMapped(L.ops.k.d, t.bytes.len); + errdefer b.free(); + try L.src.read(b.bytes, t.bytes); + try L.w.host_buffers.append(L.gpa, b); + return b.device(); + } + fn raw(L: *Loader, t: Tensor) !u64 { const ptr = try L.alloc(t.bytes.len); try L.src.upload(ptr, t.bytes); @@ -226,10 +239,13 @@ pub fn load(gpa: std.mem.Allocator, io: std.Io, ops: kern.Ops, dir: []const u8, layer.down = try L.tile(&ck, try join(&nb, p, "mlp.down_proj"), g.hidden, g.intermediate); layer.post_norm = try L.vector(&ck, try join(&nb, p, "post_attention_layernorm.weight"), g.hidden); } + // TF_QWEN27_HOST_EMBED=1: the token table in mapped host memory (prompt rows read it over PCIe), 0.67 GiB freed + const host_embed = if (std.c.getenv("TF_QWEN27_HOST_EMBED")) |v| std.mem.eql(u8, std.mem.span(v), "1") else false; + const table = if (host_embed) &Loader.rawHost else &Loader.raw; w.embed = .{ - .w = try L.raw(try ck.expect(prefix ++ "model.embed_tokens.weight", .u32, &.{ c.vocab, g.hidden / 8 })), - .s = try L.raw(try ck.expect(prefix ++ "model.embed_tokens.scales", .bf16, &.{ c.vocab, g.hidden / 64 })), - .b = try L.raw(try ck.expect(prefix ++ "model.embed_tokens.biases", .bf16, &.{ c.vocab, g.hidden / 64 })), + .w = try table(&L, try ck.expect(prefix ++ "model.embed_tokens.weight", .u32, &.{ c.vocab, g.hidden / 8 })), + .s = try table(&L, try ck.expect(prefix ++ "model.embed_tokens.scales", .bf16, &.{ c.vocab, g.hidden / 64 })), + .b = try table(&L, try ck.expect(prefix ++ "model.embed_tokens.biases", .bf16, &.{ c.vocab, g.hidden / 64 })), }; w.norm = try L.vector(&ck, "model.norm.weight", g.hidden); w.head = try L.tile(&ck, "lm_head", c.vocab, g.hidden); diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 032f64c60..1a9a70bfb 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -414,6 +414,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O var family_options: F.Options = .{ .context = @intCast(window), .drafts = o.drafts, .segments = segments }; if (@hasField(F.Options, "slide")) family_options.slide = o.slide; if (@hasField(F.Options, "drafter")) family_options.drafter = o.drafter; + if (@hasField(F.Options, "streams")) family_options.streams = o.lanes; const loaded = F.open(gpa, io, &g.ctx, o.dir, kernels, family_options) catch |e| { problem.* = try std.fmt.allocPrint(a, "the native CUDA engine cannot load {s} with {s} glue kernels ({s})", .{ o.dir, kernels orelse "its own", @errorName(e) }); return null;