From f264ecefaf75f77639fae93483518dc9a55f50bc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 27 Oct 2026 19:39:20 +0100 Subject: [PATCH 01/12] cuda: the paged shared-window attention at head dim 64 The shared-window chunk and merge bodies move to paged_attn_segs.cuh, included twice by nemotron_attention.cu: Nemotron's head-128 kernels keep their names and SASS, and head 64 adds tf_paged_attn_chunk_segs_64 and tf_paged_attn_merge_segs_64. prefill_attention.cu instantiates pattn_kernel<64, 8, 4, 8> over paged caches. sassid on sm_80/86/89/120/121: 0 functions changed, 15 added. --- zig/kernels/cuda/nemotron_attention.cu | 151 ++--------------------- zig/kernels/cuda/paged_attn_segs.cuh | 163 +++++++++++++++++++++++++ zig/kernels/cuda/prefill_attention.cu | 4 + 3 files changed, 180 insertions(+), 138 deletions(-) create mode 100644 zig/kernels/cuda/paged_attn_segs.cuh diff --git a/zig/kernels/cuda/nemotron_attention.cu b/zig/kernels/cuda/nemotron_attention.cu index 313b096f28..6dbd077325 100644 --- a/zig/kernels/cuda/nemotron_attention.cu +++ b/zig/kernels/cuda/nemotron_attention.cu @@ -225,144 +225,19 @@ extern "C" __global__ void tf_nemo_kv_write_segs(const AttnSegs sg, const __nv_b } } -extern "C" __global__ void __launch_bounds__(256) tf_nemo_attn_chunk_segs(const AttnSegs sg, const __nv_bfloat16* __restrict__ QKV, - float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, - int NQKV, int H, int HK, int G, int CH, int NCH, float scale) { - __shared__ __align__(16) float qs[GMAX][HD]; - __shared__ __align__(16) uint32_t kw[TILE][KW]; - __shared__ __align__(16) uint32_t vw[TILE][HD / 2]; - __shared__ float ps[GMAX][TILE]; - const int rq = blockIdx.x, k = seg_of(sg, rq), r = rq - sg.ROW0[k], hk = blockIdx.y, t = threadIdx.x; - const __nv_bfloat16* __restrict__ KC = sg.KC[k]; - const __nv_bfloat16* __restrict__ VC = sg.VC[k]; - const int* __restrict__ META = sg.META[k]; - const int limit = META[0] + r + 1; - if (static_cast(blockIdx.z) * CH >= limit) return; - const int g = t / 16, jj = t % 16; // softmax and values: head g, lanes jj of its 16 - const int key = t % TILE, quad = t / TILE; // scores: one key, heads 4 quad .. 4 quad + 3 - for (int i = t; i < GMAX * HD; i += 256) { - const int gg = i / HD, d = i % HD; - qs[gg][d] = gg < G ? bf(QKV[static_cast(rq) * NQKV + (hk * G + gg) * HD + d]) : 0.0f; - } - for (int c = blockIdx.z; c * CH < limit; c += gridDim.z) { // this block's chunks: every gridDim.z-th - float m = -INFINITY, den = 0.0f, o[8]; -#pragma unroll - for (int e = 0; e < 8; ++e) o[e] = 0.0f; - const int64_t cb = chunk_at(sg.PT[k], sg.PS, c, HK * HD) - static_cast(c) * CH * HK * HD; - for (int key0 = c * CH; key0 < (c + 1) * CH && key0 < limit; key0 += TILE) { - __syncthreads(); - for (int i = t; i < TILE * (HD / 2); i += 256) { - const int j = i / (HD / 2), w = i % (HD / 2); - const int64_t at = cb + (static_cast(key0 + j) * HK + hk) * HD; - const bool in = key0 + j < limit; - kw[j][w] = in ? reinterpret_cast(KC + at)[w] : 0u; - vw[j][w] = in ? reinterpret_cast(VC + at)[w] : 0u; - } - __syncthreads(); - float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; -#pragma unroll 4 - for (int w = 0; w < HD / 2; w += 2) { - const uint2 kk = *reinterpret_cast(&kw[key][w]); - const float k0 = lo(kk.x), k1 = hi(kk.x), k2 = lo(kk.y), k3 = hi(kk.y); -#pragma unroll - for (int h = 0; h < 4; ++h) { - const float4 q = *reinterpret_cast(&qs[4 * quad + h][2 * w]); - acc[h] = __fmaf_rn(q.x, k0, acc[h]); - acc[h] = __fmaf_rn(q.y, k1, acc[h]); - acc[h] = __fmaf_rn(q.z, k2, acc[h]); - acc[h] = __fmaf_rn(q.w, k3, acc[h]); - } - } -#pragma unroll - for (int h = 0; h < 4; ++h) ps[4 * quad + h][key] = key0 + key < limit ? __fmul_rn(acc[h], scale) : -INFINITY; - __syncthreads(); - float sc[4]; -#pragma unroll - for (int u = 0; u < 4; ++u) sc[u] = ps[g][jj + 16 * u]; - const float tile_m = half_max(fmaxf(fmaxf(sc[0], sc[1]), fmaxf(sc[2], sc[3]))); - const bool active = tile_m != -INFINITY; - const float next = active ? fmaxf(m, tile_m) : m; - const float alpha = active ? (m == -INFINITY ? 0.0f : tf_exp(__fsub_rn(m, next))) : 1.0f; - float local = 0.0f; -#pragma unroll - for (int u = 0; u < 4; ++u) { - const float p = active && key0 + jj + 16 * u < limit ? tf_exp(__fsub_rn(sc[u], next)) : 0.0f; - ps[g][jj + 16 * u] = p; - local = __fadd_rn(local, p); - } - den = __fmaf_rn(den, alpha, half_sum(local)); - m = next; - __syncthreads(); - float pv[8]; -#pragma unroll - for (int e = 0; e < 8; ++e) pv[e] = 0.0f; - for (int j = 0; j < TILE; ++j) { - const float p = ps[g][j]; - const uint4 v = *reinterpret_cast(&vw[j][4 * jj]); - const uint32_t vv[4] = {v.x, v.y, v.z, v.w}; -#pragma unroll - for (int e = 0; e < 4; ++e) { - pv[2 * e] = __fmaf_rn(p, lo(vv[e]), pv[2 * e]); - pv[2 * e + 1] = __fmaf_rn(p, hi(vv[e]), pv[2 * e + 1]); - } - } -#pragma unroll - for (int e = 0; e < 8; ++e) o[e] = __fmaf_rn(o[e], alpha, pv[e]); - } - if (g >= G) continue; - const int64_t base = (static_cast(rq) * NCH + c) * H + hk * G + g; - float4* out = reinterpret_cast(PO + base * HD + 8 * jj); - out[0] = make_float4(o[0], o[1], o[2], o[3]); - out[1] = make_float4(o[4], o[5], o[6], o[7]); - if (jj == 0) { - PM[base] = m; - PL[base] = den; - } - } -} - - -extern "C" __global__ void __launch_bounds__(256) tf_nemo_attn_merge_segs(const AttnSegs sg, const float* __restrict__ PO, const float* __restrict__ PM, - const float* __restrict__ PL, __nv_bfloat16* __restrict__ OUT, float* __restrict__ XS, - int H, int HK, int G, int CH, int NCH) { - __shared__ float ys[GMAX][HD]; - const int rq = blockIdx.x, k = seg_of(sg, rq), r = rq - sg.ROW0[k], hk = blockIdx.y, t = threadIdx.x, g = t / 16, jj = t % 16; - const int* __restrict__ META = sg.META[k]; - const int limit = META[0] + r + 1, nch = (limit + CH - 1) / CH, head = hk * G + g; - float m = -INFINITY, den = 0.0f, o[8]; -#pragma unroll - for (int e = 0; e < 8; ++e) o[e] = 0.0f; - if (g < G) { - for (int c = 0; c < nch; ++c) { - const int64_t base = (static_cast(rq) * NCH + c) * H + head; - const float cm = PM[base], cl = PL[base]; - const bool active = cl > 0.0f; - const float next = active ? fmaxf(m, cm) : m; - const float a = active ? (m == -INFINITY ? 0.0f : tf_exp(__fsub_rn(m, next))) : 1.0f; - const float b = active ? tf_exp(__fsub_rn(cm, next)) : 0.0f; - const float4* po = reinterpret_cast(PO + base * HD + 8 * jj); - const float4 p0 = po[0], p1 = po[1]; - const float co[8] = {p0.x, p0.y, p0.z, p0.w, p1.x, p1.y, p1.z, p1.w}; -#pragma unroll - for (int e = 0; e < 8; ++e) o[e] = __fmaf_rn(o[e], a, __fmul_rn(co[e], b)); - den = __fmaf_rn(den, a, __fmul_rn(cl, b)); - m = next; - } -#pragma unroll - for (int e = 0; e < 8; ++e) { - const __nv_bfloat16 y = __float2bfloat16_rn(__fdiv_rn(o[e], den)); - OUT[(static_cast(rq) * H + head) * HD + 8 * jj + e] = y; - ys[g][8 * jj + e] = bf(y); - } - } - __syncthreads(); - if (t >= GMAX * (HD / 64)) return; - const int gg = t / (HD / 64), gi = t % (HD / 64); - if (gg >= G) return; - float s = 0.0f; - for (int j = 0; j < 64; ++j) s = __fadd_rn(s, ys[gg][gi * 64 + j]); - XS[static_cast(rq) * (H * HD / 64) + (hk * G + gg) * (HD / 64) + gi] = s; -} +#define SEG_HD HD +#define SEG_KW KW +#define SEG_WIDE 1 +#define SEG_CHUNK tf_nemo_attn_chunk_segs +#define SEG_MERGE tf_nemo_attn_merge_segs +#include "paged_attn_segs.cuh" +// The head dim 64 forms (LFM2.5-24B-A2B: 32 query heads over 8 KV heads), four dims a lane. +#define SEG_HD 64 +#define SEG_KW (64 / 2 + 2) +#define SEG_WIDE 0 +#define SEG_CHUNK tf_paged_attn_chunk_segs_64 +#define SEG_MERGE tf_paged_attn_merge_segs_64 +#include "paged_attn_segs.cuh" // tf_nemo_attn_chunk with the next key tile loading into registers while this one computes (same values, same order). extern "C" __global__ void __launch_bounds__(256, 2) tf_nemo_attn_chunk_pf(const __nv_bfloat16* __restrict__ QKV, const __nv_bfloat16* __restrict__ KC, diff --git a/zig/kernels/cuda/paged_attn_segs.cuh b/zig/kernels/cuda/paged_attn_segs.cuh new file mode 100644 index 0000000000..56c3117a22 --- /dev/null +++ b/zig/kernels/cuda/paged_attn_segs.cuh @@ -0,0 +1,163 @@ +// Paged shared-window chunk/merge bodies shared by families (Nemotron head 128, LFM2 head 64); SEG_* set the form. + +#define SEG_E (SEG_WIDE ? 8 : 4) + +extern "C" __global__ void __launch_bounds__(256) SEG_CHUNK(const AttnSegs sg, const __nv_bfloat16* __restrict__ QKV, + float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, + int NQKV, int H, int HK, int G, int CH, int NCH, float scale) { + __shared__ __align__(16) float qs[GMAX][SEG_HD]; + __shared__ __align__(16) uint32_t kw[TILE][SEG_KW]; + __shared__ __align__(16) uint32_t vw[TILE][SEG_HD / 2]; + __shared__ float ps[GMAX][TILE]; + const int rq = blockIdx.x, k = seg_of(sg, rq), r = rq - sg.ROW0[k], hk = blockIdx.y, t = threadIdx.x; + const __nv_bfloat16* __restrict__ KC = sg.KC[k]; + const __nv_bfloat16* __restrict__ VC = sg.VC[k]; + const int* __restrict__ META = sg.META[k]; + const int limit = META[0] + r + 1; + if (static_cast(blockIdx.z) * CH >= limit) return; + const int g = t / 16, jj = t % 16; // softmax and values: head g, lanes jj of its 16 + const int key = t % TILE, quad = t / TILE; // scores: one key, heads 4 quad .. 4 quad + 3 + for (int i = t; i < GMAX * SEG_HD; i += 256) { + const int gg = i / SEG_HD, d = i % SEG_HD; + qs[gg][d] = gg < G ? bf(QKV[static_cast(rq) * NQKV + (hk * G + gg) * SEG_HD + d]) : 0.0f; + } + for (int c = blockIdx.z; c * CH < limit; c += gridDim.z) { // this block's chunks: every gridDim.z-th + float m = -INFINITY, den = 0.0f, o[SEG_E]; +#pragma unroll + for (int e = 0; e < SEG_E; ++e) o[e] = 0.0f; + const int64_t cb = chunk_at(sg.PT[k], sg.PS, c, HK * SEG_HD) - static_cast(c) * CH * HK * SEG_HD; + for (int key0 = c * CH; key0 < (c + 1) * CH && key0 < limit; key0 += TILE) { + __syncthreads(); + for (int i = t; i < TILE * (SEG_HD / 2); i += 256) { + const int j = i / (SEG_HD / 2), w = i % (SEG_HD / 2); + const int64_t at = cb + (static_cast(key0 + j) * HK + hk) * SEG_HD; + const bool in = key0 + j < limit; + kw[j][w] = in ? reinterpret_cast(KC + at)[w] : 0u; + vw[j][w] = in ? reinterpret_cast(VC + at)[w] : 0u; + } + __syncthreads(); + float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; +#pragma unroll 4 + for (int w = 0; w < SEG_HD / 2; w += 2) { + const uint2 kk = *reinterpret_cast(&kw[key][w]); + const float k0 = lo(kk.x), k1 = hi(kk.x), k2 = lo(kk.y), k3 = hi(kk.y); +#pragma unroll + for (int h = 0; h < 4; ++h) { + const float4 q = *reinterpret_cast(&qs[4 * quad + h][2 * w]); + acc[h] = __fmaf_rn(q.x, k0, acc[h]); + acc[h] = __fmaf_rn(q.y, k1, acc[h]); + acc[h] = __fmaf_rn(q.z, k2, acc[h]); + acc[h] = __fmaf_rn(q.w, k3, acc[h]); + } + } +#pragma unroll + for (int h = 0; h < 4; ++h) ps[4 * quad + h][key] = key0 + key < limit ? __fmul_rn(acc[h], scale) : -INFINITY; + __syncthreads(); + float sc[4]; +#pragma unroll + for (int u = 0; u < 4; ++u) sc[u] = ps[g][jj + 16 * u]; + const float tile_m = half_max(fmaxf(fmaxf(sc[0], sc[1]), fmaxf(sc[2], sc[3]))); + const bool active = tile_m != -INFINITY; + const float next = active ? fmaxf(m, tile_m) : m; + const float alpha = active ? (m == -INFINITY ? 0.0f : tf_exp(__fsub_rn(m, next))) : 1.0f; + float local = 0.0f; +#pragma unroll + for (int u = 0; u < 4; ++u) { + const float p = active && key0 + jj + 16 * u < limit ? tf_exp(__fsub_rn(sc[u], next)) : 0.0f; + ps[g][jj + 16 * u] = p; + local = __fadd_rn(local, p); + } + den = __fmaf_rn(den, alpha, half_sum(local)); + m = next; + __syncthreads(); + float pv[SEG_E]; +#pragma unroll + for (int e = 0; e < SEG_E; ++e) pv[e] = 0.0f; + for (int j = 0; j < TILE; ++j) { + const float p = ps[g][j]; +#if SEG_WIDE + const uint4 v = *reinterpret_cast(&vw[j][4 * jj]); + const uint32_t vv[4] = {v.x, v.y, v.z, v.w}; +#pragma unroll + for (int e = 0; e < 4; ++e) { +#else + const uint2 v = *reinterpret_cast(&vw[j][2 * jj]); + const uint32_t vv[2] = {v.x, v.y}; +#pragma unroll + for (int e = 0; e < 2; ++e) { +#endif + pv[2 * e] = __fmaf_rn(p, lo(vv[e]), pv[2 * e]); + pv[2 * e + 1] = __fmaf_rn(p, hi(vv[e]), pv[2 * e + 1]); + } + } +#pragma unroll + for (int e = 0; e < SEG_E; ++e) o[e] = __fmaf_rn(o[e], alpha, pv[e]); + } + if (g >= G) continue; + const int64_t base = (static_cast(rq) * NCH + c) * H + hk * G + g; + float4* out = reinterpret_cast(PO + base * SEG_HD + SEG_E * jj); + out[0] = make_float4(o[0], o[1], o[2], o[3]); +#if SEG_WIDE + out[1] = make_float4(o[4], o[5], o[6], o[7]); +#endif + if (jj == 0) { + PM[base] = m; + PL[base] = den; + } + } +} + + +extern "C" __global__ void __launch_bounds__(256) SEG_MERGE(const AttnSegs sg, const float* __restrict__ PO, const float* __restrict__ PM, + const float* __restrict__ PL, __nv_bfloat16* __restrict__ OUT, float* __restrict__ XS, + int H, int HK, int G, int CH, int NCH) { + __shared__ float ys[GMAX][SEG_HD]; + const int rq = blockIdx.x, k = seg_of(sg, rq), r = rq - sg.ROW0[k], hk = blockIdx.y, t = threadIdx.x, g = t / 16, jj = t % 16; + const int* __restrict__ META = sg.META[k]; + const int limit = META[0] + r + 1, nch = (limit + CH - 1) / CH, head = hk * G + g; + float m = -INFINITY, den = 0.0f, o[SEG_E]; +#pragma unroll + for (int e = 0; e < SEG_E; ++e) o[e] = 0.0f; + if (g < G) { + for (int c = 0; c < nch; ++c) { + const int64_t base = (static_cast(rq) * NCH + c) * H + head; + const float cm = PM[base], cl = PL[base]; + const bool active = cl > 0.0f; + const float next = active ? fmaxf(m, cm) : m; + const float a = active ? (m == -INFINITY ? 0.0f : tf_exp(__fsub_rn(m, next))) : 1.0f; + const float b = active ? tf_exp(__fsub_rn(cm, next)) : 0.0f; + const float4* po = reinterpret_cast(PO + base * SEG_HD + SEG_E * jj); +#if SEG_WIDE + const float4 p0 = po[0], p1 = po[1]; + const float co[8] = {p0.x, p0.y, p0.z, p0.w, p1.x, p1.y, p1.z, p1.w}; +#else + const float4 p0 = po[0]; + const float co[4] = {p0.x, p0.y, p0.z, p0.w}; +#endif +#pragma unroll + for (int e = 0; e < SEG_E; ++e) o[e] = __fmaf_rn(o[e], a, __fmul_rn(co[e], b)); + den = __fmaf_rn(den, a, __fmul_rn(cl, b)); + m = next; + } +#pragma unroll + for (int e = 0; e < SEG_E; ++e) { + const __nv_bfloat16 y = __float2bfloat16_rn(__fdiv_rn(o[e], den)); + OUT[(static_cast(rq) * H + head) * SEG_HD + SEG_E * jj + e] = y; + ys[g][SEG_E * jj + e] = bf(y); + } + } + __syncthreads(); + if (t >= GMAX * (SEG_HD / 64)) return; + const int gg = t / (SEG_HD / 64), gi = t % (SEG_HD / 64); + if (gg >= G) return; + float s = 0.0f; + for (int j = 0; j < 64; ++j) s = __fadd_rn(s, ys[gg][gi * 64 + j]); + XS[static_cast(rq) * (H * SEG_HD / 64) + (hk * G + gg) * (SEG_HD / 64) + gi] = s; +} + +#undef SEG_E +#undef SEG_HD +#undef SEG_KW +#undef SEG_WIDE +#undef SEG_CHUNK +#undef SEG_MERGE diff --git a/zig/kernels/cuda/prefill_attention.cu b/zig/kernels/cuda/prefill_attention.cu index 7fffbb9307..8d04c7d6bb 100644 --- a/zig/kernels/cuda/prefill_attention.cu +++ b/zig/kernels/cuda/prefill_attention.cu @@ -324,6 +324,10 @@ vision_kernel(const __nv_bfloat16* __restrict__ qkv, __nv_bfloat16* __restrict__ template __global__ void tf_prefill_attention::pattn_kernel<128, 8, 8, 8>(const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, int, int, int, int, int, float, const int* __restrict__, long long); +// Head dim 64, eight warps, four query heads a block, eight staging slots, over paged caches (LFM2.5-24B-A2B's chunks). +template __global__ void tf_prefill_attention::pattn_kernel<64, 8, 4, 8>(const __nv_bfloat16* __restrict__, + const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, int, int, int, + int, int, float, const int* __restrict__, long long); // Head dim 64, eight warps, four query heads a block, eight staging slots (LFM2's 4 heads a KV head, causal or trees). template __global__ void tf_prefill_attention::tree_kernel<64, 8, 4, 8>(const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, const int* __restrict__, int, int, int, int, int, float); From ac8f6e2d4cbdfa999c05741606eb246caaabbaea Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 27 Oct 2026 19:39:20 +0100 Subject: [PATCH 02/12] lfm2_moe: CUDA glue kernels under the contract lfm2_moe.cu holds what the shared operators do not: the short conv over each stream's state with every window row's snapshot, the commit of an accepted row's conv state, and the router (fp32 weights, top-4 of sigmoid + bias, ties to the lower id, a NaN ranked as -inf). Each has a planted control. Launch records cover these kernels and the head-64 paged attention on nvcc 13.0, with one 13.2 override. --- zig/build/cuda_kernels.zig | 1 + zig/kernels/cuda/contract_controls.cu | 20 + zig/kernels/cuda/lfm2_moe.cu | 25 + zig/kernels/cuda/lfm2_moe.cuh | 138 +++++ .../core/contract/lfm2_moe_cuda_records.zig | 24 + zig/src/core/contract/registry.zig | 4 +- zig/src/cuda/kernels.zig | 1 + zig/tests/cuda/contract_elementwise.zig | 4 + zig/tests/cuda/contract_lfm2_moe.zig | 493 ++++++++++++++++++ 9 files changed, 708 insertions(+), 2 deletions(-) create mode 100644 zig/kernels/cuda/lfm2_moe.cu create mode 100644 zig/kernels/cuda/lfm2_moe.cuh create mode 100644 zig/src/core/contract/lfm2_moe_cuda_records.zig create mode 100644 zig/tests/cuda/contract_lfm2_moe.zig diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index 887b2d84ba..bd0f03d169 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -83,6 +83,7 @@ pub const kernels = [_]Kernel{ .{ .name = "norms", .flags = &.{"-O3"} }, // ours: RMS norm rows for every family (cuda/norms.zig) .{ .name = "embed_rows", .flags = &.{"-O3"} }, // ours: MLX affine and bf16 token rows (cuda/embed_rows.zig) .{ .name = "lfm2", .flags = glue }, // ours: LFM2's short convolution, QK norm and rotary, gates and label rows + .{ .name = "lfm2_moe", .flags = glue }, // ours: LFM2.5-24B-A2B's conv over stream state, its commit and the router .{ .name = "vit", .flags = glue }, // ours: vision towers' LayerNorm, Linear epilogues, activations, unshuffle .{ .name = "tree_round", .flags = &.{"-O3"} }, // ours: verified tree rounds for every family (cuda/tree_round.zig) .{ .name = "act_rows", .flags = &.{"-O3"} }, // ours: SwiGLU, output gates, group sums (cuda/act_rows.zig) diff --git a/zig/kernels/cuda/contract_controls.cu b/zig/kernels/cuda/contract_controls.cu index 1cc1fd9c84..a6df0204e7 100644 --- a/zig/kernels/cuda/contract_controls.cu +++ b/zig/kernels/cuda/contract_controls.cu @@ -342,3 +342,23 @@ TF_CTL_SWIGLU_L2(tf_ctl_swiglu_l2_drop_gs, tf_act_l2::US | tf_act_l2::PROD) TF_CTL_SWIGLU_L2(tf_ctl_swiglu_l2_drop_us, tf_act_l2::GS | tf_act_l2::PROD) TF_CTL_SWIGLU_L2(tf_ctl_swiglu_l2_drop_prod, tf_act_l2::GS | tf_act_l2::US) #undef TF_CTL_SWIGLU_L2 + +// LFM2.5-24B-A2B's glue, one planted fault each: conv bx unrounded, commits to stream 0's slot, ties to the higher id. +#include "lfm2_moe.cuh" + +extern "C" __global__ void __launch_bounds__(256) tf_ctl_lfm2m_conv(const __nv_bfloat16* P, const __nv_bfloat16* Wc, const __nv_bfloat16* STATE, + const int* SLOT, const int* OFF, __nv_bfloat16* Y, __nv_bfloat16* SNAP, + int SNAP0, int D, long long SS) { + tf_lfm2m::lfm2m_conv<1>(P, Wc, STATE, SLOT, OFF, Y, SNAP, SNAP0, D, SS); +} + +extern "C" __global__ void __launch_bounds__(256) tf_ctl_lfm2m_commit(__nv_bfloat16* STATE, const __nv_bfloat16* SNAP, const int* SLOT, + const int* ROW, int D, long long state_stride, long long snap_stride, + long long SS) { + tf_lfm2m::lfm2m_commit<1>(STATE, SNAP, SLOT, ROW, D, state_stride, snap_stride, SS); +} + +extern "C" __global__ void __launch_bounds__(256) tf_ctl_lfm2m_route(const __nv_bfloat16* X, const float* W, const float* BIAS, int* PICK, + float* WTS, float* LOGITS, int rows, int D, float scale) { + tf_lfm2m::lfm2m_route<1>(X, W, BIAS, PICK, WTS, LOGITS, rows, D, scale); +} diff --git a/zig/kernels/cuda/lfm2_moe.cu b/zig/kernels/cuda/lfm2_moe.cu new file mode 100644 index 0000000000..dd32fcfdd8 --- /dev/null +++ b/zig/kernels/cuda/lfm2_moe.cu @@ -0,0 +1,25 @@ +// LFM2.5-24B-A2B's CUDA glue beside the shared operators (lfm2_moe.cuh): short conv, conv-state commit, router. + +#include "lfm2_moe.cuh" + +using namespace tf_lfm2m; + +extern "C" __global__ void __launch_bounds__(NT) tf_lfm2m_conv(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ Wc, + const __nv_bfloat16* __restrict__ STATE, const int* __restrict__ SLOT, + const int* __restrict__ OFF, __nv_bfloat16* __restrict__ Y, + __nv_bfloat16* __restrict__ SNAP, int SNAP0, int D, long long SS) { + lfm2m_conv<0>(P, Wc, STATE, SLOT, OFF, Y, SNAP, SNAP0, D, SS); +} + +extern "C" __global__ void __launch_bounds__(NT) tf_lfm2m_commit(__nv_bfloat16* __restrict__ STATE, const __nv_bfloat16* __restrict__ SNAP, + const int* __restrict__ SLOT, const int* __restrict__ ROW, int D, + long long state_stride, long long snap_stride, long long SS) { + lfm2m_commit<0>(STATE, SNAP, SLOT, ROW, D, state_stride, snap_stride, SS); +} + +extern "C" __global__ void __launch_bounds__(NT) tf_lfm2m_route(const __nv_bfloat16* __restrict__ X, const float* __restrict__ W, + const float* __restrict__ BIAS, int* __restrict__ PICK, + float* __restrict__ WTS, float* __restrict__ LOGITS, int rows, + int D, float scale) { + lfm2m_route<0>(X, W, BIAS, PICK, WTS, LOGITS, rows, D, scale); +} diff --git a/zig/kernels/cuda/lfm2_moe.cuh b/zig/kernels/cuda/lfm2_moe.cuh new file mode 100644 index 0000000000..ad8d3c3789 --- /dev/null +++ b/zig/kernels/cuda/lfm2_moe.cuh @@ -0,0 +1,138 @@ +// LFM2.5-24B-A2B's CUDA glue (short conv, its commit, router); each row's bits its own; CTL 1 is the planted control. + +#pragma once + +#include +#include +#include + +#include "glue_math.cuh" + +namespace tf_lfm2m { + +constexpr int NT = 256; // threads a block +constexpr int E = 64; // experts +constexpr int TOPK = 4; // picks a row +constexpr int RB = 8; // router rows a block: the weights are read once for them +constexpr int DMAX = 2048; // hidden width the router stages + +// (value, id) wins over (other value, other id): the larger value, the lower id among equals. +template +__device__ __forceinline__ bool beats(float v, int i, float w, int j) { return v > w || (v == w && (CTL == 1 ? i > j : i < j)); } + + +// Block (row, channel block): y = bf16(C * bf16(sum of w * bf16(B * x), oldest tap first)); SNAP rows keep state. +template // CTL 1: bx = B * x not rounded to bf16 before its tap. +__device__ __forceinline__ void lfm2m_conv(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ Wc, + const __nv_bfloat16* __restrict__ STATE, const int* __restrict__ SLOT, + const int* __restrict__ OFF, __nv_bfloat16* __restrict__ Y, + __nv_bfloat16* __restrict__ SNAP, int SNAP0, int D, long long SS) { + const int64_t r = blockIdx.x, D3 = 3 * static_cast(D); + const int c = blockIdx.y * NT + threadIdx.x; + if (c >= D) return; + const int o = OFF[r]; + const __nv_bfloat16* st = STATE + SLOT[r] * SS; + __nv_bfloat16 bx[3]; + float wide[3] = {0.0f, 0.0f, 0.0f}; +#pragma unroll + for (int j = 0; j < 3; ++j) { + const int back = 2 - j; // rows before this one + if (o >= back) { + const int64_t src = r - back; + const float v = __fmul_rn(tf_glue::bf(P[src * D3 + c]), tf_glue::bf(P[src * D3 + 2 * D + c])); + bx[j] = __float2bfloat16_rn(v); + if (CTL == 1) wide[j] = v; + } else { + bx[j] = st[static_cast(2 - back + o) * D + c]; + if (CTL == 1) wide[j] = tf_glue::bf(bx[j]); + } + } + float acc = 0.0f; +#pragma unroll + for (int j = 0; j < 3; ++j) acc = __fadd_rn(acc, __fmul_rn(tf_glue::bf(Wc[c * 3 + j]), CTL == 1 ? wide[j] : tf_glue::bf(bx[j]))); + Y[r * D + c] = __float2bfloat16_rn(__fmul_rn(tf_glue::bf(P[r * D3 + D + c]), tf_glue::rbf(acc))); + if (SNAP == nullptr || r < SNAP0) return; + SNAP[((r - SNAP0) * 2) * D + c] = bx[1]; + SNAP[((r - SNAP0) * 2 + 1) * D + c] = bx[2]; +} + +// Block (layer, stream): stream k's [2][D] conv state set to its snapshot row ROW[k]; ROW < 0 keeps the state. +template // CTL 1: every stream's state written to the first stream's slot. +__device__ __forceinline__ void lfm2m_commit(__nv_bfloat16* __restrict__ STATE, const __nv_bfloat16* __restrict__ SNAP, + const int* __restrict__ SLOT, const int* __restrict__ ROW, int D, + long long state_stride, long long snap_stride, long long SS) { + const int layer = blockIdx.x, k = blockIdx.y, row = ROW[k]; + if (row < 0) return; + __nv_bfloat16* dst = STATE + layer * state_stride + SLOT[CTL == 1 ? 0 : k] * SS; + const __nv_bfloat16* src = SNAP + layer * snap_stride + static_cast(row) * 2 * D; + for (int i = threadIdx.x; i < 2 * D; i += NT) dst[i] = src[i]; +} + +// Block of RB rows: fp32 logits (quarters of D in order, (q0 + q1) + (q2 + q3)), top-4 of sigmoid + bias, weights. +template // CTL 1: ties go to the higher expert id. +__device__ __forceinline__ void lfm2m_route(const __nv_bfloat16* __restrict__ X, const float* __restrict__ W, + const float* __restrict__ BIAS, int* __restrict__ PICK, + float* __restrict__ WTS, float* __restrict__ LOGITS, int rows, + int D, float scale) { + __shared__ __align__(16) __nv_bfloat16 xs[RB][DMAX]; + __shared__ float lg[RB][E]; + const int row0 = blockIdx.x * RB, t = threadIdx.x, e = t / 4, q = t % 4, n = min(RB, rows - row0); + for (int i = t; i < n * D; i += NT) xs[i / D][i % D] = X[static_cast(row0 + i / D) * D + i % D]; + __syncthreads(); + float acc[RB]; +#pragma unroll + for (int b = 0; b < RB; ++b) acc[b] = 0.0f; + const int per = D / 4, k0 = q * per; + const float* w = W + static_cast(e) * D + k0; + for (int k = 0; k < per; ++k) { + const float wk = w[k]; +#pragma unroll + for (int b = 0; b < RB; ++b) acc[b] = __fmaf_rn(tf_glue::bf(xs[b][k0 + k]), wk, acc[b]); + } +#pragma unroll + for (int b = 0; b < RB; ++b) { + float v = __fadd_rn(acc[b], __shfl_xor_sync(0xffffffffu, acc[b], 1)); + v = __fadd_rn(v, __shfl_xor_sync(0xffffffffu, v, 2)); + if (q == 0) lg[b][e] = v; + } + __syncthreads(); + const int b = t / 32, lane = t % 32; + if (b >= n) return; + const int64_t r = row0 + b; + float gate[2], sel[2]; +#pragma unroll + for (int h = 0; h < 2; ++h) { + const float l = lg[b][lane + 32 * h]; + if (LOGITS != nullptr) LOGITS[r * E + lane + 32 * h] = l; + gate[h] = tf_glue::tf_sigmoid(l); + const float v = __fadd_rn(gate[h], BIAS[lane + 32 * h]); + sel[h] = v != v ? -INFINITY : v; // a NaN logit is never picked ahead of a number + } + float total = 0.0f, mine = 0.0f; + int pick = 0; + for (int k = 0; k < TOPK; ++k) { + float bv = -INFINITY; + int bi = CTL == 1 ? -1 : 0x7fffffff; +#pragma unroll + for (int h = 0; h < 2; ++h) + if (beats(sel[h], lane + 32 * h, bv, bi)) bv = sel[h], bi = lane + 32 * h; +#pragma unroll + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffffu, bv, off); + const int oi = __shfl_xor_sync(0xffffffffu, bi, off); + if (beats(ov, oi, bv, bi)) bv = ov, bi = oi; + } + float p = 0.0f; +#pragma unroll + for (int h = 0; h < 2; ++h) + if (bi == lane + 32 * h) p = gate[h], sel[h] = NAN; // taken: NaN beats nothing, not even a -inf left + p = __shfl_sync(0xffffffffu, p, bi % 32); + total = __fadd_rn(total, p); + if (lane == k) mine = p, pick = bi; + } + if (lane >= TOPK) return; + PICK[r * TOPK + lane] = pick; + WTS[r * TOPK + lane] = __fmul_rn(__fdiv_rn(mine, __fadd_rn(total, 1e-6f)), scale); +} + +} // namespace tf_lfm2m diff --git a/zig/src/core/contract/lfm2_moe_cuda_records.zig b/zig/src/core/contract/lfm2_moe_cuda_records.zig new file mode 100644 index 0000000000..4615a01471 --- /dev/null +++ b/zig/src/core/contract/lfm2_moe_cuda_records.zig @@ -0,0 +1,24 @@ +//! LFM2.5-24B-A2B's CUDA records: its conv, commit and route, and the head-64 paged attention, every SM's image. +const std = @import("std"); +const registry = @import("registry.zig"); +const R = registry.CudaRecord; + +const pattn64 = "_ZN20tf_prefill_attention12pattn_kernelILi64ELi8ELi4ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifPKix"; + +pub const cuda = [_]R{ + .{ .launch = .{ .name = "lfm2_moe tf_lfm2m_conv", .image = "lfm2_moe", .symbol = "tf_lfm2m_conv", .threads = 256 }, .usage = .{ .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 26, .static = 0, .bound = 256 }, .{ .regs = 26, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "lfm2_moe tf_lfm2m_commit", .image = "lfm2_moe", .symbol = "tf_lfm2m_commit", .threads = 256 }, .usage = .{ .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "lfm2_moe tf_lfm2m_route", .image = "lfm2_moe", .symbol = "tf_lfm2m_route", .threads = 256 }, .usage = .{ .{ .regs = 55, .static = 34816, .bound = 256 }, .{ .regs = 50, .static = 34816, .bound = 256 }, .{ .regs = 50, .static = 34816, .bound = 256 }, .{ .regs = 58, .static = 34816, .bound = 256 }, .{ .regs = 59, .static = 34816, .bound = 256 } } }, + .{ .launch = .{ .name = "paging tf_paged_attn_chunk_segs_64", .image = "nemotron_attention", .symbol = "tf_paged_attn_chunk_segs_64", .threads = 256 }, .usage = .{ .{ .regs = 48, .static = 25088, .bound = 256 }, .{ .regs = 78, .static = 25088, .bound = 256 }, .{ .regs = 78, .static = 25088, .bound = 256 }, .{ .regs = 80, .static = 25088, .bound = 256 }, .{ .regs = 80, .static = 25088, .bound = 256 } } }, + .{ .launch = .{ .name = "paging tf_paged_attn_merge_segs_64", .image = "nemotron_attention", .symbol = "tf_paged_attn_merge_segs_64", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 4096, .bound = 256 }, .{ .regs = 40, .static = 4096, .bound = 256 }, .{ .regs = 40, .static = 4096, .bound = 256 }, .{ .regs = 40, .static = 4096, .bound = 256 }, .{ .regs = 40, .static = 4096, .bound = 256 } } }, + .{ .launch = .{ .name = "paging pattn64", .image = "prefill_attention", .symbol = pattn64, .threads = 256, .dynamic = 32768, .opt_in = true }, .usage = .{ .{ .regs = 148, .static = 0, .bound = 256 }, .{ .regs = 148, .static = 0, .bound = 256 }, .{ .regs = 148, .static = 0, .bound = 256 }, .{ .regs = 164, .static = 0, .bound = 256 }, .{ .regs = 164, .static = 0, .bound = 256 } } }, +}; + +test "LFM2.5-24B-A2B records: every launch fits every SM by its measured record" { + for (cuda) |r| try std.testing.expectEqual(@as(usize, 0), try registry.checkCuda(std.testing.allocator, &.{r}, true)); +} + +/// The launch nvcc 13.2 compiles differently: the router on sm_120 (one register more). +pub const cuda_132 = [_]registry.Override{ + .{ .name = "lfm2_moe tf_lfm2m_route", .usage = .{ .{ .regs = 55, .static = 34816, .bound = 256 }, .{ .regs = 50, .static = 34816, .bound = 256 }, .{ .regs = 50, .static = 34816, .bound = 256 }, .{ .regs = 59, .static = 34816, .bound = 256 }, .{ .regs = 59, .static = 34816, .bound = 256 } } }, +}; diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 7d83d71249..efa15e00ef 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -207,7 +207,7 @@ pub const cuda = @import("flashnext_cuda_records.zig").cuda ++ [_]CudaRecord{ .{ .launch = .{ .name = "experts3_prefill down", .symbol = "_ZN19tf_experts3_prefill14prefill_kernelILi1ELi3ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_PvifS6_i", .image = "experts3_prefill", .threads = 256, .dynamic = 47104, .opt_in = true }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 128, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 }, .{ .regs = 98, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "experts3 tf_experts3_widen", .symbol = "tf_experts3_widen", .image = "experts3", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 32, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "experts3 tf_experts3_planes", .symbol = "tf_experts3_planes", .image = "experts3", .threads = 256 }, .usage = .{ .{ .regs = 26, .static = 0, .bound = 0 }, .{ .regs = 28, .static = 0, .bound = 0 }, .{ .regs = 28, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 } } }, -} ++ @import("registry_media.zig").cuda ++ @import("nemotron_prompt_records.zig").cuda ++ @import("glm_cuda_records.zig").cuda ++ @import("registry_shared_attention.zig").records(CudaRecord) ++ @import("gemma_cuda_records.zig").cuda ++ @import("spark_cuda_records.zig").cuda ++ @import("row_cuda_records.zig").cuda ++ @import("flashnext_tile_records.zig").cuda ++ @import("window_rows_records.zig").cuda ++ @import("modelopt_lane_records.zig").cuda ++ @import("nvfp4_fact_records.zig").cuda ++ @import("lane_half_records.zig").cuda; +} ++ @import("registry_media.zig").cuda ++ @import("nemotron_prompt_records.zig").cuda ++ @import("glm_cuda_records.zig").cuda ++ @import("registry_shared_attention.zig").records(CudaRecord) ++ @import("gemma_cuda_records.zig").cuda ++ @import("spark_cuda_records.zig").cuda ++ @import("row_cuda_records.zig").cuda ++ @import("flashnext_tile_records.zig").cuda ++ @import("window_rows_records.zig").cuda ++ @import("modelopt_lane_records.zig").cuda ++ @import("nvfp4_fact_records.zig").cuda ++ @import("lane_half_records.zig").cuda ++ @import("lfm2_moe_cuda_records.zig").cuda; /// The nvcc release `cuda`'s usage records come from: the compiler that builds the release fatbins. pub const cuda_release: limits.Nvcc = .{ .major = 13, .minor = 0 }; @@ -319,7 +319,7 @@ pub const cuda_sets = [_]CudaSet{ .{ .name = "latent tf_super_topk_few", .usage = .{ .{ .regs = 32, .static = 4352, .bound = 1024 }, .{ .regs = 38, .static = 4352, .bound = 1024 }, .{ .regs = 38, .static = 4352, .bound = 1024 }, .{ .regs = 32, .static = 4352, .bound = 1024 }, .{ .regs = 32, .static = 4352, .bound = 1024 } } }, .{ .name = "ck tf_ck4_up_relu2_few", .usage = .{ .{ .regs = 19, .static = 0, .bound = 128 }, .{ .regs = 19, .static = 0, .bound = 128 }, .{ .regs = 19, .static = 0, .bound = 128 }, .{ .regs = 124, .static = 0, .bound = 128 }, .{ .regs = 124, .static = 0, .bound = 128 } } }, .{ .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) }, + } ++ learning.cuda_132 ++ @import("paging_pending.zig").cuda_132 ++ @import("lfm2_moe_cuda_records.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) }, .{ .release = .{ .major = 13, .minor = 3 }, .usage = &([_]Override{ .{ .name = "router tf_nemo_topk", .usage = .{ .{ .regs = 94, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 96, .static = 0, .bound = 32 }, .{ .regs = 96, .static = 0, .bound = 32 } } }, .{ .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 } } }, diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 33a3bbfe55..5e059a1b7e 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -34,6 +34,7 @@ pub const experts_pack: []const u8 = if (available) &Blob("fatbin_experts_pack") pub const experts_rows: []const u8 = if (available) &Blob("fatbin_experts_rows").bytes else &.{}; pub const prefill_attention: []const u8 = if (available) &Blob("fatbin_prefill_attention").bytes else &.{}; pub const lfm2: []const u8 = if (available) &Blob("fatbin_lfm2").bytes else &.{}; +pub const lfm2_moe: []const u8 = if (available) &Blob("fatbin_lfm2_moe").bytes else &.{}; pub const vit: []const u8 = if (available) &Blob("fatbin_vit").bytes else &.{}; 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 &.{}; diff --git a/zig/tests/cuda/contract_elementwise.zig b/zig/tests/cuda/contract_elementwise.zig index 0b6d8c46da..3ef8ff9980 100644 --- a/zig/tests/cuda/contract_elementwise.zig +++ b/zig/tests/cuda/contract_elementwise.zig @@ -5,6 +5,7 @@ const core = @import("core"); const check = @import("check.zig"); const embed_k = @import("contract_embed.zig"); const penalize_k = @import("contract_penalize.zig"); +const lfm2m_k = @import("contract_lfm2_moe.zig"); const g = core.contract.gates; @@ -15,6 +16,7 @@ pub const Set = struct { sample: ?cuda.Module = null, draws: cuda.sample.Draws = undefined, penalize: ?penalize_k.Penalize = null, + lfm2m: lfm2m_k.Set = .{}, // LFM2.5-24B-A2B's router, conv and commit (rows with sums, beside these) /// Every registration onto `kernels`; the embed controls' symbols come from the contract_controls image `ctl_mod`. pub fn add(self: *Set, gpu: check.Gpu, s: cuda.Stream, ctl_mod: cuda.Module, kernels: *std.ArrayList(g.Kernel)) !void { @@ -29,9 +31,11 @@ pub const Set = struct { self.draws = try cuda.sample.Draws.resolve(self.sample.?); self.penalize = .{ .rig = .{ .gpu = gpu, .s = s }, .draws = &self.draws }; try kernels.append(a, self.penalize.?.kernel()); + try self.lfm2m.add(gpu, s, ctl_mod, kernels); } pub fn deinit(self: *Set) void { + self.lfm2m.deinit(); if (self.penalize) |*p| p.deinit(); if (self.sample) |*m| m.unload(); for (self.embeds[0..self.n_embeds]) |*x| x.deinit(); diff --git a/zig/tests/cuda/contract_lfm2_moe.zig b/zig/tests/cuda/contract_lfm2_moe.zig new file mode 100644 index 0000000000..b686ab9ba3 --- /dev/null +++ b/zig/tests/cuda/contract_lfm2_moe.zig @@ -0,0 +1,493 @@ +//! lfm2_moe.cu under the contract: router logits and top-4 rule, short conv and snapshots, commit; each controlled. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const rows_ = contract.rows; +const Dense = contract.dense.Dense; + +const experts = 64; +const topk = 4; + +fn noSums(_: c.Shape) usize { + return 0; +} + +/// A quarter of the row an fma chain, then the quarters' two adds. +fn quarterChain(s: c.Shape) usize { + return s.k / 4 + 2; +} + +fn f32At(bytes: []const u8, i: usize) f32 { + return @bitCast(std.mem.readInt(u32, bytes[4 * i ..][0..4], .little)); +} + +fn routeArgs(x: u64, w: u64, bias: u64, picks: u64, wts: u64, logits: u64, rows: usize, d: usize) cuda.Args { + var a: cuda.Args = .{}; + for ([_]u64{ x, w, bias, picks, wts, logits }) |p| a.add(p); + a.add(@as(c_int, @intCast(rows))); + a.add(@as(c_int, @intCast(d))); + a.add(@as(f32, 1.0)); + return a; +} + +fn launch(f: cuda.Function, s: cuda.Stream, grid: [3]usize, a: *cuda.Args) !void { + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(grid[0]), .y = @intCast(grid[1]), .z = @intCast(grid[2]) }, .block = .{ .x = 256 } }, s, a); + try s.synchronize(); +} + +pub fn logitsSpec() c.Spec { + return .{ .name = "cuda lfm2_moe router logits tf_lfm2m_route (fp32 weights, quarters of the row)", .accum = .f32, .depth = quarterChain, .out = .f32, .tiles = .{ .{ .size = 8 }, .{}, .{ .size = 2048, .step = 4, .min = 4 } }, .served = &.{ .{ .rows = 1, .n = experts, .k = 2048 }, .{ .rows = 16, .n = experts, .k = 2048 }, .{ .rows = 1024, .n = experts, .k = 2048 } } }; +} + +/// x [rows, D] bf16 times fp32 router rows [64, D]: the logits (the picks and weights go to scratch). +pub const Logits = struct { + rig: rig_.Rig, + f: cuda.Function, + d: ?Dense = null, + dev: [4]u64 = @splat(0), + + pub fn deinit(self: *Logits) void { + if (self.d) |d| d.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Logits) g.Kernel { + return .{ .spec = logitsSpec(), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Logits); + + pub fn prepare(self: *Logits, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n != experts or shape.k % 4 != 0 or shape.k == 0 or shape.k > 2048 or shape.rows == 0) return false; + if (self.d) |d| d.deinit(); + self.d = null; + self.rig.reset(); + const a = self.rig.gpu.gpa; + const d = try Dense.init(a, shape, 1); + self.d = d; + d.fillX(fill, seed, 16); + const w = try a.alloc(f32, d.w.len); + defer a.free(w); + var prng = std.Random.DefaultPrng.init(seed ^ 0x51); + const rng = prng.random(); + const k = shape.k; + // stress: positive integers, exact in fp32 in any order; mirrored signs would cancel inside a 32-lane fp16 sum + for (w, 0..) |*v, i| v.* = if (fill != .stress) rng.float(f32) * 2 - 1 else @floatFromInt(1 + (i % k * 7 + i / k) % 6); + for (d.w, w) |*o, v| o.* = v; + const bias: [experts]f32 = @splat(0); + self.dev[0] = try self.rig.input(std.mem.sliceAsBytes(d.x)); + self.dev[1] = try self.rig.input(std.mem.sliceAsBytes(w)); + self.dev[2] = try self.rig.input(std.mem.sliceAsBytes(&bias)); + self.dev[3] = try self.rig.output(shape.rows * experts * 4); + return true; + } + + pub fn run(self: *Logits) !void { + const d = self.d.?; + const picks = try self.rig.scratch(d.shape.rows * topk * 4); + const wts = try self.rig.scratch(d.shape.rows * topk * 4); + var a = routeArgs(self.dev[0], self.dev[1], self.dev[2], picks, wts, self.dev[3], d.shape.rows, d.shape.k); + try launch(self.f, self.rig.s, .{ (d.shape.rows + 7) / 8, 1, 1 }, &a); + } + + pub fn count(self: *Logits) usize { + return self.d.?.outputs(); + } + + pub fn reference(self: *Logits, want: []f64, mag: []f64) !void { + return self.d.?.reference(want, mag); + } + + pub fn simulate(self: *Logits, acc: c.Accum, order: g.Order, out: []f64) !void { + return self.d.?.simulate(acc, order, out); + } +}; + +const pick_rule = "top 4 of sigmoid(logit) + bias, ties to the lower id, a NaN ranked as -inf; weights sigmoid / (sum + 1e-6)"; + +pub fn selectSpec(control: bool) c.Spec { + return .{ .name = if (control) "control: lfm2_moe router with ties to the higher expert" else "cuda lfm2_moe router top-4 tf_lfm2m_route (its rule on fp32 logits)", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = pick_rule }, .tiles = .{ .{ .size = 8 }, .{}, .{} }, .served = &.{ .{ .rows = 16, .n = experts, .k = experts }, .{ .rows = 300, .n = experts, .k = experts } } }; +} + +/// The rule's picks on logits the identity router passes through (logit e = x[e]): sixteenths, so ties are common. +pub const Select = struct { + rig: rig_.Rig, + f: cuda.Function, + control: bool = false, + shape: c.Shape = undefined, + x: []f32 = &.{}, + bias: [experts]f32 = undefined, + dev: [5]u64 = @splat(0), + msg: [200]u8 = undefined, + + pub fn deinit(self: *Select) void { + self.rig.gpu.gpa.free(self.x); + self.rig.deinit(); + } + + pub fn kernel(self: *Select) g.Kernel { + return .{ .spec = selectSpec(self.control), .ctx = self, .vt = &vt, .control = if (self.control) .shapes else null }; + } + + const vt = g.hooks(Select); + + pub fn prepare(self: *Select, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n != experts or shape.k != experts or shape.rows == 0) return false; + self.rig.reset(); + const a = self.rig.gpu.gpa; + self.shape = shape; + a.free(self.x); + self.x = &.{}; + self.x = try a.alloc(f32, shape.rows * experts); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (self.x) |*v| v.* = @as(f32, @floatFromInt(rng.uintLessThan(u32, 64))) / 16 - 2; + for (&self.bias) |*b| b.* = if (rng.boolean()) 0 else 0.0625; + if (fill == .nonfinite) { + const at = [_][2]usize{ .{ 0, 3 }, .{ 1, 5 }, .{ 2, 7 }, .{ 3, 2 }, .{ 3, experts - 1 } }; + const put = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32), std.math.inf(f32), -std.math.inf(f32) }; + for (at, put) |p, v| if (p[0] < shape.rows) { + self.x[p[0] * experts + p[1]] = v; + }; + } + const xb = try a.alloc(u16, self.x.len); + defer a.free(xb); + for (xb, self.x) |*b, v| b.* = c.toBf16(v); + var w: [experts * experts]f32 = @splat(0); + for (0..experts) |e| w[e * experts + e] = 1; + self.dev[0] = try self.rig.input(std.mem.sliceAsBytes(xb)); + self.dev[1] = try self.rig.input(std.mem.sliceAsBytes(&w)); + self.dev[2] = try self.rig.input(std.mem.sliceAsBytes(&self.bias)); + self.dev[3] = try self.rig.output(shape.rows * topk * 4); + self.dev[4] = try self.rig.output(shape.rows * topk * 4); + return true; + } + + pub fn run(self: *Select) !void { + const logits = try self.rig.scratch(self.shape.rows * experts * 4); + var a = routeArgs(self.dev[0], self.dev[1], self.dev[2], self.dev[3], self.dev[4], logits, self.shape.rows, experts); + try launch(self.f, self.rig.s, .{ (self.shape.rows + 7) / 8, 1, 1 }, &a); + } + + pub fn count(self: *Select) usize { + return self.shape.rows * topk * 8; + } + + pub fn reference(_: *Select, want: []f64, mag: []f64) !void { + @memset(want, 0); + @memset(mag, 0); + } + + /// Picks by the rule exactly, weights within 2^-20 of float64 (NaN where a picked logit is NaN). + pub fn check(self: *Select, _: g.Gate) !?[]const u8 { + const out = try self.rig.outputBytes(); + const n = self.shape.rows * topk; + for (0..self.shape.rows) |r| { + var lg: [experts]f64 = undefined; // the identity's logits as a dot: one non-finite x makes the whole row NaN + for (&lg, 0..) |*l, e| { + l.* = 0; + for (self.x[r * experts ..][0..experts], 0..) |v, i| l.* += @as(f64, v) * @as(f64, if (i == e) 1 else 0); + } + var taken: [experts]bool = @splat(false); + var gate: [topk]f64 = undefined; + var sum: f64 = 0; + for (0..topk) |k| { + var best: ?usize = null; + var bv: f64 = 0; + for (lg, 0..) |l, e| { + const s = 1 / (1 + @exp(-l)); + const sv: f64 = @as(f32, @floatCast(s)) + @as(f64, self.bias[e]); + const v = if (std.math.isNan(sv)) -std.math.inf(f64) else sv; + if (taken[e]) continue; + if (best == null or v > bv) { + best = e; + bv = v; + } + } + const id = best orelse return try std.fmt.bufPrint(&self.msg, "row {d}: fewer than 4 numbers", .{r}); + taken[id] = true; + const got = std.mem.readInt(u32, out[(r * topk + k) * 4 ..][0..4], .little); + if (got != id) return try std.fmt.bufPrint(&self.msg, "row {d} pick {d} is {d} where the rule gives {d}", .{ r, k, got, id }); + gate[k] = 1 / (1 + @exp(-lg[id])); + sum += gate[k]; + } + for (0..topk) |k| { + const want = gate[k] / (sum + 1e-6); + const got: f64 = f32At(out, n + r * topk + k); + if (std.math.isNan(want) and std.math.isNan(got)) continue; + if (!(@abs(got - want) <= @abs(want) * 0x1p-20 + 1e-30)) return try std.fmt.bufPrint(&self.msg, "row {d} slot {d}: weight {d:.8} where float64 gives {d:.8}", .{ r, k, got, want }); + } + } + return null; + } +}; + +const conv_rule = "the host's replay of each row's fp32 and bf16 steps, bit for bit (NaN as NaN)"; + +pub fn convSpec(control: bool) c.Spec { + return .{ .name = if (control) "control: lfm2_moe conv with bx unrounded" else "cuda lfm2_moe short conv tf_lfm2m_conv (stream state, every row's snapshot)", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = conv_rule }, .tiles = .{ .{}, .{ .size = 256, .step = 64, .min = 64 }, .{} }, .served = &.{ .{ .rows = 1, .n = 2048, .k = 3 }, .{ .rows = 16, .n = 2048, .k = 3 }, .{ .rows = 128, .n = 2048, .k = 3 } } }; +} + +/// Rows of two streams (the first rows/2 slot 1, the rest slot 0) over their states: y and every row's state after it. +pub const Conv = struct { + rig: rig_.Rig, + f: cuda.Function, + control: bool = false, + shape: c.Shape = undefined, + p: []u16 = &.{}, + w: []u16 = &.{}, + st: []u16 = &.{}, + slot: []i32 = &.{}, + off: []i32 = &.{}, + dev: [7]u64 = @splat(0), + msg: [200]u8 = undefined, + + pub fn deinit(self: *Conv) void { + self.free(); + self.rig.deinit(); + } + + fn free(self: *Conv) void { + const a = self.rig.gpu.gpa; + inline for (.{ "p", "w", "st", "slot", "off" }) |f| { + a.free(@field(self, f)); + @field(self, f) = &.{}; + } + } + + pub fn kernel(self: *Conv) g.Kernel { + return .{ .spec = convSpec(self.control), .ctx = self, .vt = &vt, .control = if (self.control) .shapes else null }; + } + + const vt = g.hooks(Conv); + + pub fn prepare(self: *Conv, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.n % 64 != 0 or shape.k != 3) return false; + self.rig.reset(); + self.free(); + const a = self.rig.gpu.gpa; + self.shape = shape; + const r = shape.rows; + const d = shape.n; + self.p = try a.alloc(u16, r * 3 * d); + self.w = try a.alloc(u16, d * 3); + self.st = try a.alloc(u16, 2 * 2 * d); + self.slot = try a.alloc(i32, r); + self.off = try a.alloc(i32, r); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (self.p) |*v| v.* = c.toBf16(rng.float(f32) * 4 - 2); + for (self.w) |*v| v.* = c.toBf16(rng.float(f32) - 0.5); + for (self.st) |*v| v.* = c.toBf16(rng.float(f32) * 4 - 2); + const half = r / 2; + for (self.slot, self.off, 0..) |*s, *o, i| { + s.* = if (i < half) 1 else 0; + o.* = @intCast(if (i < half) i else i - half); + } + if (fill == .nonfinite) { + const at = [_][2]usize{ .{ 0, d / 3 }, .{ 1, d + 1 }, .{ 2, 2 * d + 5 }, .{ 3, 0 }, .{ 3, 3 * d - 1 } }; + const put = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32), std.math.inf(f32), -std.math.inf(f32) }; + for (at, put) |q, v| if (q[0] < r) { + self.p[q[0] * 3 * d + q[1]] = c.toBf16(v); + }; + } + self.dev[0] = try self.rig.input(std.mem.sliceAsBytes(self.p)); + self.dev[1] = try self.rig.input(std.mem.sliceAsBytes(self.w)); + self.dev[2] = try self.rig.input(std.mem.sliceAsBytes(self.st)); + self.dev[3] = try self.rig.input(std.mem.sliceAsBytes(self.slot)); + self.dev[4] = try self.rig.input(std.mem.sliceAsBytes(self.off)); + self.dev[5] = try self.rig.output(r * d * 2); + self.dev[6] = try self.rig.output(r * 2 * d * 2); + return true; + } + + pub fn run(self: *Conv) !void { + const d = self.shape.n; + var a: cuda.Args = .{}; + for (self.dev) |p| a.add(p); + a.add(@as(c_int, 0)); + a.add(@as(c_int, @intCast(d))); + a.add(@as(i64, @intCast(2 * d))); + try launch(self.f, self.rig.s, .{ self.shape.rows, (d + 255) / 256, 1 }, &a); + } + + pub fn count(self: *Conv) usize { + return self.shape.rows * self.shape.n * 6; + } + + pub fn reference(_: *Conv, want: []f64, mag: []f64) !void { + @memset(want, 0); + @memset(mag, 0); + } + + fn val(h: u16) f32 { + return c.fromBf16(h); + } + + fn same(got: u16, want: u16) bool { + return got == want or (std.math.isNan(val(got)) and std.math.isNan(val(want))); + } + + pub fn check(self: *Conv, _: g.Gate) !?[]const u8 { + const out = try self.rig.outputBytes(); + const d = self.shape.n; + const r = self.shape.rows; + for (0..r) |row| for (0..d) |ch| { + const o: usize = @intCast(self.off[row]); + var bx: [3]u16 = undefined; + for (0..3) |j| { + const back = 2 - j; + bx[j] = if (o >= back) blk: { + const src = (row - back) * 3 * d; + break :blk c.toBf16(val(self.p[src + ch]) * val(self.p[src + 2 * d + ch])); + } else self.st[@as(usize, @intCast(self.slot[row])) * 2 * d + (2 - back + o) * d + ch]; + } + var acc: f32 = 0; + for (0..3) |j| acc += val(self.w[ch * 3 + j]) * val(bx[j]); + const y = c.toBf16(val(self.p[row * 3 * d + d + ch]) * val(c.toBf16(acc))); + const wants = [_][2]usize{ .{ row * d + ch, y }, .{ r * d + row * 2 * d + ch, bx[1] }, .{ r * d + row * 2 * d + d + ch, bx[2] } }; + for (wants, 0..) |wv, which| { + const got = std.mem.readInt(u16, out[2 * wv[0] ..][0..2], .little); + if (!same(got, @intCast(wv[1]))) return try std.fmt.bufPrint(&self.msg, "row {d} channel {d} output {d}: 0x{x:0>4} where the replay gives 0x{x:0>4}", .{ row, ch, which, got, wv[1] }); + } + }; + return null; + } +}; + +pub fn commitSpec(control: bool) c.Spec { + return .{ .name = if (control) "control: lfm2_moe commit into the first stream's slot" else "cuda lfm2_moe conv-state commit tf_lfm2m_commit", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = "a copy: every bit as stored" }, .served = &.{ .{ .rows = 1, .n = 2048, .k = 30 }, .{ .rows = 16, .n = 2048, .k = 30 } } }; +} + +/// Streams k = 0.. into slots rows - k of rows + 1, each its snapshot row 2k (stream 1: none, its slot as it was). +pub const Commit = struct { + rig: rig_.Rig, + f: cuda.Function, + control: bool = false, + shape: c.Shape = undefined, + snaps: []u16 = &.{}, + dev: [4]u64 = @splat(0), + msg: [200]u8 = undefined, + + pub fn deinit(self: *Commit) void { + self.rig.gpu.gpa.free(self.snaps); + self.rig.deinit(); + } + + pub fn kernel(self: *Commit) g.Kernel { + return .{ .spec = commitSpec(self.control), .ctx = self, .vt = &vt, .control = if (self.control) .shapes else null }; + } + + const vt = g.hooks(Commit); + + fn rowOf(k: usize) i32 { + return if (k == 1) -1 else @intCast(2 * k); + } + + pub fn prepare(self: *Commit, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.k == 0) return false; + self.rig.reset(); + const a = self.rig.gpu.gpa; + self.shape = shape; + a.free(self.snaps); + self.snaps = &.{}; + const r = shape.rows; + self.snaps = try a.alloc(u16, shape.k * 2 * r * 2 * shape.n); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (self.snaps) |*v| v.* = if (fill == .nonfinite and rng.uintLessThan(u32, 97) == 0) 0x7fc1 else rng.int(u16); + const slots = try a.alloc(i32, r); + defer a.free(slots); + const rows = try a.alloc(i32, r); + defer a.free(rows); + for (slots, rows, 0..) |*s, *w, k| { + s.* = @intCast(r - k); + w.* = rowOf(k); + } + self.dev[0] = try self.rig.output(shape.k * (r + 1) * 2 * shape.n * 2); + self.dev[1] = try self.rig.input(std.mem.sliceAsBytes(self.snaps)); + self.dev[2] = try self.rig.input(std.mem.sliceAsBytes(slots)); + self.dev[3] = try self.rig.input(std.mem.sliceAsBytes(rows)); + return true; + } + + pub fn run(self: *Commit) !void { + const d = self.shape.n; + const r = self.shape.rows; + var a: cuda.Args = .{}; + for (self.dev) |p| a.add(p); + a.add(@as(c_int, @intCast(d))); + a.add(@as(i64, @intCast((r + 1) * 2 * d))); + a.add(@as(i64, @intCast(2 * r * 2 * d))); + a.add(@as(i64, @intCast(2 * d))); + try launch(self.f, self.rig.s, .{ self.shape.k, r, 1 }, &a); + } + + pub fn count(self: *Commit) usize { + return self.shape.k * (self.shape.rows + 1) * 2 * self.shape.n * 2; + } + + pub fn reference(_: *Commit, want: []f64, mag: []f64) !void { + @memset(want, 0); + @memset(mag, 0); + } + + pub fn check(self: *Commit, _: g.Gate) !?[]const u8 { + const out = try self.rig.outputBytes(); + const d = self.shape.n; + const r = self.shape.rows; + for (0..self.shape.k) |layer| for (0..r + 1) |slot| for (0..2 * d) |i| { + const k = r - slot; // the stream that owns this slot (slot 0: none) + const want: u16 = if (slot == 0 or rowOf(k) < 0) 0xffff else self.snaps[layer * 2 * r * 2 * d + @as(usize, @intCast(rowOf(k))) * 2 * d + i]; + const got = std.mem.readInt(u16, out[2 * ((layer * (r + 1) + slot) * 2 * d + i) ..][0..2], .little); + if (got != want) return try std.fmt.bufPrint(&self.msg, "layer {d} slot {d} element {d}: 0x{x:0>4} where 0x{x:0>4}", .{ layer, slot, i, got, want }); + }; + return null; + } +}; + +/// The kernels and their controls, registered onto `kernels`; the controls from the contract_controls image `ctl`. +pub const Set = struct { + mod: ?cuda.Module = null, + logits: Logits = undefined, + select: [2]Select = undefined, + conv: [2]Conv = undefined, + commit: [2]Commit = undefined, + + pub fn add(self: *Set, gpu: check.Gpu, s: cuda.Stream, ctl: cuda.Module, kernels: *std.ArrayList(g.Kernel)) !void { + const a = gpu.gpa; + self.mod = try cuda.Module.load(gpu.d, cuda.kernels.lfm2_moe); + const m = self.mod.?; + const rig: rig_.Rig = .{ .gpu = gpu, .s = s }; + self.logits = .{ .rig = rig, .f = try m.function("tf_lfm2m_route") }; + for (&self.select, [_]bool{ false, true }) |*x, ctrl| x.* = .{ .rig = rig, .f = try (if (ctrl) ctl else m).function(if (ctrl) "tf_ctl_lfm2m_route" else "tf_lfm2m_route"), .control = ctrl }; + for (&self.conv, [_]bool{ false, true }) |*x, ctrl| x.* = .{ .rig = rig, .f = try (if (ctrl) ctl else m).function(if (ctrl) "tf_ctl_lfm2m_conv" else "tf_lfm2m_conv"), .control = ctrl }; + for (&self.commit, [_]bool{ false, true }) |*x, ctrl| x.* = .{ .rig = rig, .f = try (if (ctrl) ctl else m).function(if (ctrl) "tf_ctl_lfm2m_commit" else "tf_lfm2m_commit"), .control = ctrl }; + try kernels.append(a, self.logits.kernel()); + for (&self.select) |*x| try kernels.append(a, x.kernel()); + for (&self.conv) |*x| try kernels.append(a, x.kernel()); + for (&self.commit) |*x| try kernels.append(a, x.kernel()); + } + + pub fn deinit(self: *Set) void { + const m = self.mod orelse return; + self.logits.deinit(); + for (&self.select) |*x| x.deinit(); + for (&self.conv) |*x| x.deinit(); + for (&self.commit) |*x| x.deinit(); + var mm = m; + mm.unload(); + } +}; + +test "the commit's streams: stream 1 keeps its slot, the rest take row 2k" { + try std.testing.expectEqual(@as(i32, -1), Commit.rowOf(1)); + try std.testing.expectEqual(@as(i32, 6), Commit.rowOf(3)); + try std.testing.expectEqual(@as(usize, 514), quarterChain(.{ .rows = 1, .n = 64, .k = 2048 })); +} From 692cf2011c61d50366e179faf65f8dbc40931f0b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 27 Oct 2026 19:39:20 +0100 Subject: [PATCH 03/12] lfm2_moe: serve LFM2.5-24B-A2B on CUDA The MLX 4-bit checkpoint (embedding at 4, 6 or 8 bits) on the shared affine-4 projections and experts, with prompt chunks of 1024 rows, one-row decode and copy-drafted windows over --parallel streams. The ten attention layers keep their keys and values in 512-position pages, and each stream has its own conv state. There is no drafter and no prompt cache yet. --- zig/build/cuda.zig | 2 +- zig/build/cuda_engines.zig | 1 + zig/build/lfm2_moe_cuda.zig | 12 + zig/src/families/lfm2_moe/cuda_backend.zig | 318 +++++++++++++++++++++ zig/src/families/lfm2_moe/cuda_family.zig | 18 ++ zig/src/families/lfm2_moe/cuda_forward.zig | 228 +++++++++++++++ zig/src/families/lfm2_moe/cuda_kernels.zig | 193 +++++++++++++ zig/src/families/lfm2_moe/cuda_model.zig | 215 ++++++++++++++ zig/src/families/lfm2_moe/cuda_native.zig | 149 ++++++++++ zig/src/families/lfm2_moe/cuda_weights.zig | 239 ++++++++++++++++ zig/src/native/cuda.zig | 3 +- 11 files changed, 1376 insertions(+), 2 deletions(-) create mode 100644 zig/build/lfm2_moe_cuda.zig create mode 100644 zig/src/families/lfm2_moe/cuda_backend.zig create mode 100644 zig/src/families/lfm2_moe/cuda_family.zig create mode 100644 zig/src/families/lfm2_moe/cuda_forward.zig create mode 100644 zig/src/families/lfm2_moe/cuda_kernels.zig create mode 100644 zig/src/families/lfm2_moe/cuda_model.zig create mode 100644 zig/src/families/lfm2_moe/cuda_native.zig create mode 100644 zig/src/families/lfm2_moe/cuda_weights.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 054d85922b..61e81bab26 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -446,7 +446,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st @import("glm_family_host.zig").steps(b, cuda, glm, mods, step); @import("kolibri.zig").host(b, kolibri1Module(b, host, .debug, cuda, mods.core, mods.lanes), native_modules.api, step); @import("gemma4_cuda.zig").host(b, @import("gemma4_cuda.zig").module(b, host, .debug, cuda, mods.core, mods.lanes), step); - @import("spark25.zig").host(b, @import("spark25.zig").module(b, host, .debug, cuda, mods.core, mods.lanes), step); + inline for (.{ @import("spark25.zig"), @import("lfm2_moe_cuda.zig") }) |m| m.host(b, m.module(b, host, .debug, cuda, mods.core, mods.lanes), step); @import("glm_contracts.zig").tests(b, glm, mods.lanes, cuda, mods.core); // the shares a GLM-5.3-Flash rank takes (cluster/canon.zig against the Python family's) const shares = b.createModule(.{ .root_source_file = b.path("zig/src/cluster/canon_glm_test.zig"), .target = host, .optimize = .debug }); diff --git a/zig/build/cuda_engines.zig b/zig/build/cuda_engines.zig index bb929f7a23..e2f2c40b1a 100644 --- a/zig/build/cuda_engines.zig +++ b/zig/build/cuda_engines.zig @@ -52,6 +52,7 @@ pub fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu .{ .name = "kolibri1", .module = @import("kolibri.zig").module(b, target, optimize, cuda, core, lanes) }, .{ .name = "gemma4_cuda", .module = @import("gemma4_cuda.zig").module(b, target, optimize, cuda, core, lanes) }, .{ .name = "spark25", .module = @import("spark25.zig").module(b, target, optimize, cuda, core, lanes) }, + .{ .name = "lfm2_moe_cuda", .module = @import("lfm2_moe_cuda.zig").module(b, target, optimize, cuda, core, lanes) }, .{ .name = "dory", .module = @import("dory.zig").module(b, target, optimize, cuda, core, lanes) }, }, }); diff --git a/zig/build/lfm2_moe_cuda.zig b/zig/build/lfm2_moe_cuda.zig new file mode 100644 index 0000000000..e02fb69494 --- /dev/null +++ b/zig/build/lfm2_moe_cuda.zig @@ -0,0 +1,12 @@ +//! LFM2.5-24B-A2B's CUDA family over the maintained runtime, the core and the lane core; its host tests. +const std = @import("std"); + +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module) *std.Build.Module { + return b.createModule(.{ .root_source_file = b.path("zig/src/families/lfm2_moe/cuda_family.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "lanes", .module = lanes } } }); +} + +pub fn host(b: *std.Build, m: *std.Build.Module, all: *std.Build.Step) void { + const run = b.addRunArtifact(b.addTest(.{ .root_module = m })); + b.step("test-lfm2-moe-cuda-host", "LFM2.5-24B-A2B on CUDA: config admission, page and state sizes, RoPE table, refusals without a GPU").dependOn(&run.step); + all.dependOn(&run.step); +} diff --git a/zig/src/families/lfm2_moe/cuda_backend.zig b/zig/src/families/lfm2_moe/cuda_backend.zig new file mode 100644 index 0000000000..2ebc676d1d --- /dev/null +++ b/zig/src/families/lfm2_moe/cuda_backend.zig @@ -0,0 +1,318 @@ +//! LFM2.5-24B-A2B on CUDA behind lanes.Backend: a slot a stream, a chunk a prompt step, one forward a round of windows. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const cfg = @import("config.zig"); +const fwd = @import("cuda_forward.zig"); +const mdl = @import("cuda_model.zig"); +const kern = @import("cuda_kernels.zig"); + +const be = lanes.backend; +const Stream = lanes.Stream; +const Pick = cuda.tree_round.Pick; + +/// Prompt rows a step takes: the forward's chunk grid, from position 0. +pub const prompt_rows = mdl.prompt_rows; +/// Window rows a round takes over every stream. +pub const round_rows = mdl.round_rows; +/// The widest window a stream verifies: one-row steps and every wider window give a row the same bits. +pub const max_window = 16; +/// Drawn tokens a handle names (a ring; read before the round loop could wrap it). +const ring = 1024; + +const Lane = struct { + slot: usize, + len: u64 = 0, // the cache length: positions [0, len) hold this stream's keys + base: u64 = 0, // the cache length the last verify started from + row0: usize = 0, // the stream's first row in the last verify's forward (its conv snapshots) + rows: std.ArrayList(u32) = .empty, // the last verify's row tokens + last: u64, // the prompt's last row's bf16 logits [V] (the first draw reads them) + rule: u64, // sample.cu's Rule, when the stream samples + sampled: bool = false, + fill: ?usize = null, // a prompt pass going a chunk at a time: where its next chunk starts +}; + +pub const Lfm = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + m: *mdl.Model, + k: *const kern.Kernels, + s: cuda.Stream, + lanes: std.AutoHashMapUnmanaged(*const Stream, Lane) = .empty, + free: []bool, // slot i is free + tok: cuda.DeviceBuffer, // drawn tokens (u32): the handle ring, then a round's rows + picks: cuda.DeviceBuffer, // argmax picks (Pick), the same entries + greedy: []bool, // entry e of the ring (and the round's rows) is a pick, not a draw + meta: cuda.DeviceBuffer, // a draw's first position minus one, a window each + slot_bufs: []cuda.DeviceBuffer, // a slot's last-row logits, then its rule: two a slot + next: u64 = 0, + + /// The device bytes `init` allocates: the token and pick rings, metas, each slot's last row and rule. + pub fn openBytes(slots: usize) usize { + return (ring + round_rows) * (4 + @sizeOf(Pick)) + round_rows * 16 + slots * (cfg.vocab * 2 + @sizeOf(cuda.sample.Rule)); + } + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, m: *mdl.Model, k: *const kern.Kernels, s: cuda.Stream) !Lfm { + const free = try gpa.alloc(bool, m.slots); + errdefer gpa.free(free); + @memset(free, true); + const greedy = try gpa.alloc(bool, ring + round_rows); + errdefer gpa.free(greedy); + @memset(greedy, true); + var tok = try cuda.DeviceBuffer.alloc(d, (ring + round_rows) * 4); + errdefer tok.free(); + var picks = try cuda.DeviceBuffer.alloc(d, (ring + round_rows) * @sizeOf(Pick)); + errdefer picks.free(); + var meta = try cuda.DeviceBuffer.alloc(d, round_rows * 16); + errdefer meta.free(); + const bufs = try gpa.alloc(cuda.DeviceBuffer, 2 * m.slots); + errdefer gpa.free(bufs); + var made: usize = 0; + errdefer for (bufs[0..made]) |*b| b.free(); + for (bufs, 0..) |*b, i| { + b.* = try cuda.DeviceBuffer.alloc(d, if (i % 2 == 0) cfg.vocab * 2 else @sizeOf(cuda.sample.Rule)); + made += 1; + } + return .{ .gpa = gpa, .d = d, .m = m, .k = k, .s = s, .free = free, .tok = tok, .picks = picks, .greedy = greedy, .meta = meta, .slot_bufs = bufs }; + } + + pub fn deinit(x: *Lfm) void { + var it = x.lanes.valueIterator(); + while (it.next()) |l| l.rows.deinit(x.gpa); + x.lanes.deinit(x.gpa); + for (x.slot_bufs) |*b| b.free(); + x.gpa.free(x.slot_bufs); + x.gpa.free(x.free); + x.gpa.free(x.greedy); + x.tok.free(); + x.picks.free(); + x.meta.free(); + } + + pub fn backend(x: *Lfm) be.Backend { + return .{ .ptr = x, .vtable = &vtable }; + } + + const vtable: be.Backend.VTable = .{ .prefill = prefill, .prefill_step = prefillStep, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .release = release, .reports_fill = true }; + + /// The round loop's facts: no draft head; host copies verified in windows up to 16 rows, streams sharing a forward. + pub fn facts(x: *const Lfm) lanes.Model { + return .{ .exact_width = max_window, .gpu_tokens = false, .mtp = false, .speculate = false, .hidden_rows = true, .batch_rows = round_rows, .max_streams = @intCast(@min(x.m.slots, round_rows)) }; + } + + fn self(ptr: *anyopaque) *Lfm { + return @ptrCast(@alignCast(ptr)); + } + + fn upload(x: *Lfm, dst: u64, bytes: []const u8) !void { + try x.d.check(x.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, x.s.handle), "cuMemcpyHtoDAsync"); + } + + /// Rows [row0, row0 + rows) of bf16 `logits` into entries [at, at + rows), row r at `position0 + r`. + fn drawRows(x: *Lfm, l: *const Lane, logits: u64, row0: usize, rows: usize, position0: u64, at: usize, meta_at: usize) !void { + const v = cfg.vocab; + const src = logits + row0 * v * 2; + @memset(x.greedy[at..][0..rows], !l.sampled); + if (!l.sampled) return x.k.round.argmax(x.s, src, x.picks.ptr + at * @sizeOf(Pick), rows, v, v); + const meta = [4]i32{ @intCast(position0 - 1), 0, 0, 0 }; + const mp = x.meta.ptr + meta_at * 16; + try x.upload(mp, std.mem.asBytes(&meta)); + try x.k.draws.draw(x.s, src, v, l.rule, mp, 0, x.tok.ptr + at * 4, rows, null, null); + } + + /// Entries [at, at + out.len) as tokens, after the stream's work; a pick over non-finite logits is an error. + fn tokens(x: *Lfm, at: usize, out: []u32) !void { + try x.s.synchronize(); + var picks: [round_rows]Pick = undefined; + try x.tok.download(at * 4, std.mem.sliceAsBytes(out)); + try x.picks.download(at * @sizeOf(Pick), std.mem.sliceAsBytes(picks[0..out.len])); + for (out, picks[0..out.len], x.greedy[at..][0..out.len]) |*t, p, g| if (g) { + if (p.nonfinite != 0) return error.NonFiniteLogits; + t.* = p.token; + }; + } + + fn freeSlot(x: *const Lfm) ?usize { + return std.mem.indexOfScalar(bool, x.free, true); + } + + /// A prompt pass begins: a slot, the stream's rule, a kept state restored; where its first chunk starts. + fn opening(x: *Lfm, s: *Stream) !*Lane { + const p = s.prompt(); + if (p.len == 0) return error.EmptyPrompt; + if (s.images.len > 0) return error.LfmTextOnly; + if (p.len + @as(usize, s.max_new) + max_window > x.m.context) return error.PastContext; + const got = try x.lanes.getOrPut(x.gpa, s); + if (!got.found_existing) { + const slot = x.freeSlot() orelse { + x.lanes.removeByPtr(got.key_ptr); + return error.NoFreeSlot; + }; + x.free[slot] = false; + got.value_ptr.* = .{ .slot = slot, .last = x.slot_bufs[2 * slot].ptr, .rule = x.slot_bufs[2 * slot + 1].ptr }; + } + const l = got.value_ptr; + l.sampled = false; + if (s.sampling) |sp| if (sp.temperature > 0) { + const r = cuda.sample.Rule.from(lanes.gpu_rule.rule(sp)); + try x.upload(l.rule, std.mem.asBytes(&r)); + l.sampled = true; + }; + s.cached = 0; + if (s.reuse.saved != null) s.reuse_failed = true; // no prompt cache: every pass starts at 0 + try x.m.clearSlot(x.s, l.slot); + l.fill = 0; + s.fill_at.store(@intCast(l.fill.?), .monotonic); + return l; + } + + /// One chunk from `l.fill` to the next grid line or mark; true once the prompt's last row's logits are written. + fn chunk(x: *Lfm, s: *Stream, l: *Lane) !bool { + const p = s.prompt(); + const at = l.fill.?; + if (s.isCancelled()) return error.Cancelled; + var end = @min(p.len, (at / prompt_rows + 1) * prompt_rows); + for (s.reuse.marks) |mk| if (mk > at and mk < end) { + end = mk; + }; + const done = end == p.len; + try fwd.forward(x.m, x.k, x.s, &.{.{ .slot = l.slot, .p0 = at, .tokens = p[at..end] }}, true, end - at - 1, @intFromBool(done)); + if (done) try x.d.check(x.d.api.cuMemcpyDtoDAsync_v2(l.last, x.m.s.logits, cfg.vocab * 2, x.s.handle), "cuMemcpyDtoDAsync"); + s.fill_at.store(@intCast(end), .monotonic); + l.fill = if (done) null else end; + if (done) l.len = p.len; + if (s.reuse.hook) |hk| if (std.mem.indexOfScalar(u32, s.reuse.marks, @intCast(end)) != null) hk.at(hk.ptr, s, @intCast(end)); + return done; + } + + fn prefill(ptr: *anyopaque, s: *Stream) anyerror!void { + const x = self(ptr); + const l = try x.opening(s); + errdefer x.dropLane(s); + while (!try x.chunk(s, l)) {} + } + + fn prefillStep(ptr: *anyopaque, s: *Stream) anyerror!bool { + const x = self(ptr); + const going = if (x.lanes.getPtr(s)) |l| l.fill != null else false; + const l = if (going) x.lanes.getPtr(s).? else try x.opening(s); + errdefer x.dropLane(s); + return x.chunk(s, l); + } + + fn handle(x: *Lfm) usize { + const h = x.next; + x.next += 1; + return @intCast(h % ring); + } + + fn first(ptr: *anyopaque, s: *Stream, position: u64) anyerror!u64 { + const x = self(ptr); + const l = x.lanes.getPtr(s) orelse return error.UnknownStream; + if (position != l.len or l.fill != null) return error.PositionMismatch; + const h = x.next; + try x.drawRows(l, l.last, 0, 1, position, x.handle(), 0); + return h; + } + + fn value(x: *Lfm, feed: be.Feed) !u32 { + return switch (feed) { + .value => |v| v, + .handle => |h| read(x, h), + }; + } + + fn queue(ptr: *anyopaque, s: *Stream, feed: be.Feed, position: u64) anyerror!u64 { + const x = self(ptr); + const l = x.lanes.getPtr(s) orelse return error.UnknownStream; + const t = try x.value(feed); + if (position != l.len + 1) return error.PositionMismatch; + try fwd.forward(x.m, x.k, x.s, &.{.{ .slot = l.slot, .p0 = l.len, .tokens = &.{t} }}, false, 0, 1); + l.len += 1; + const h = x.next; + try x.drawRows(l, x.m.s.logits, 0, 1, position, x.handle(), 0); + return h; + } + + fn read(ptr: *anyopaque, h: u64) anyerror!u32 { + const x = self(ptr); + if (h >= x.next or x.next - h > ring) return error.NoSuchToken; + var out: [1]u32 = undefined; + try x.tokens(@intCast(h % ring), &out); + return out[0]; + } + + fn verify(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { + const x = self(ptr); + const chains = try x.gpa.alloc(fwd.Chain, windows.len); + defer x.gpa.free(chains); + var total: usize = 0; + for (windows, chains) |w, *c| { + if (w.held != 0) return error.NoHeldDrafts; // no draft head: every draft comes from the host + if (w.parents != null) return error.ChainsOnly; + if (w.masks.len != 0) return error.StructuresUnsupported; + const l = x.lanes.getPtr(w.stream) orelse return error.UnknownStream; + if (w.rows() > max_window) return error.WindowTooWide; + l.base = l.len; + l.rows.clearRetainingCapacity(); + try l.rows.append(x.gpa, w.pending); + try l.rows.appendSlice(x.gpa, w.tokens); + for (w.positions, 0..) |pos, r| if (pos != l.base + r + 1) return error.PositionMismatch; + c.* = .{ .slot = l.slot, .p0 = l.base, .tokens = l.rows.items }; + l.row0 = total; + total += l.rows.items.len; + } + if (total > round_rows) return error.RoundTooWide; + try fwd.forward(x.m, x.k, x.s, chains, false, 0, total); + var at: usize = 0; + for (windows, 0..) |w, i| { + const l = x.lanes.getPtr(w.stream).?; + try x.drawRows(l, x.m.s.logits, at, l.rows.items.len, w.positions[0], ring + at, i); + at += l.rows.items.len; + } + var got: [round_rows]u32 = undefined; + try x.tokens(ring, got[0..total]); + at = 0; + for (windows, out) |w, o| { + const l = x.lanes.getPtr(w.stream).?; + const n = l.rows.items.len; + @memcpy(o.sampled[0..n], got[at..][0..n]); + @memcpy(o.drafts[0 .. n - 1], l.rows.items[1..]); + l.len = l.base + n; // a chain keeps every row unless rolled back + at += n; + } + } + + fn keep(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) anyerror!void { + const x = self(ptr); + if (windows.len > round_rows) return error.RoundTooWide; + var slots: [round_rows]i32 = undefined; + var rows: [round_rows]i32 = undefined; + for (windows, paths, 0..) |w, path, i| { + const l = x.lanes.getPtr(w.stream) orelse return error.UnknownStream; + if (path.len == 0 or path.len > l.rows.items.len) return error.EmptyPath; + for (path, 0..) |r, j| if (r != j) return error.ChainsOnly; // a chain keeps a prefix of its rows + l.len = l.base + path.len; // caches are positional: later rows overwrite what the rejected ones wrote + slots[i] = @intCast(l.slot); + rows[i] = @intCast(l.row0 + path.len - 1); // the conv states after the last kept row + } + try fwd.commit(x.m, x.k, x.s, slots[0..windows.len], rows[0..windows.len]); + } + + /// No draft head: absorbing is a no-op, a request for drafts is refused by name (copy drafts come from the host). + fn draft(_: *anyopaque, requests: []const be.DraftRequest) anyerror!void { + for (requests) |r| if (r.depth != 0 or r.lanes != null) return error.NoDraftHead; + } + + fn dropLane(x: *Lfm, s: *const Stream) void { + if (x.lanes.fetchRemove(s)) |kv| { + var l = kv.value; + x.free[l.slot] = true; + l.rows.deinit(x.gpa); + } + } + + fn release(ptr: *anyopaque, s: *Stream) void { + self(ptr).dropLane(s); + } +}; diff --git a/zig/src/families/lfm2_moe/cuda_family.zig b/zig/src/families/lfm2_moe/cuda_family.zig new file mode 100644 index 0000000000..a938da0d0e --- /dev/null +++ b/zig/src/families/lfm2_moe/cuda_family.zig @@ -0,0 +1,18 @@ +//! LFM2.5-24B-A2B (lfm2_moe) on CUDA: the Metal family's config admission, the shared CUDA operators, lfm2_moe.cu. +pub const config = @import("config.zig"); +pub const weights = @import("cuda_weights.zig"); +pub const kernels = @import("cuda_kernels.zig"); +pub const model = @import("cuda_model.zig"); +pub const forward = @import("cuda_forward.zig"); +pub const backend = @import("cuda_backend.zig"); +pub const native = @import("cuda_native.zig"); + +test { + _ = config; + _ = weights; + _ = kernels; + _ = model; + _ = forward; + _ = backend; + _ = native; +} diff --git a/zig/src/families/lfm2_moe/cuda_forward.zig b/zig/src/families/lfm2_moe/cuda_forward.zig new file mode 100644 index 0000000000..6c2b7ef357 --- /dev/null +++ b/zig/src/families/lfm2_moe/cuda_forward.zig @@ -0,0 +1,228 @@ +//! LFM2.5-24B-A2B's forward on CUDA: a prompt chunk or a round of windows, each row's bits its own (solo == batched). +const std = @import("std"); +const cuda = @import("cuda"); +const c = @import("config.zig"); +const wts = @import("cuda_weights.zig"); +const kern = @import("cuda_kernels.zig"); +const mdl = @import("cuda_model.zig"); + +const Model = mdl.Model; +const Kernels = kern.Kernels; +const Stream = cuda.Stream; +const tail: cuda.norms.Tail = .{ .round = true }; // the norms' torch points: bf16(w * bf16(h * inv)) + +/// A stream's rows in one forward: its slot, first position, tokens. +pub const Chain = struct { slot: usize, p0: usize, tokens: []const u32 }; + +const Pass = struct { + m: *const Model, + k: *const Kernels, + s: Stream, + chains: []const Chain, + rows: usize, + prompt: bool, + + /// x [rows, k] through `l` into out [rows, n]: a prompt chunk's GEMM, or decode rows on their group sums. + fn project(p: Pass, l: wts.Proj, x: u64, xs: u64, out: u64, rows: usize, prompt: bool) !void { + const k = p.k; + if (prompt) return switch (l) { + .a4 => |q| k.affine.prompt(p.s, x, q, out, rows), + .an => |q| k.wide.prompt(p.s, x, q, out, rows), + }; + const tile: usize = switch (l) { + .a4 => k.affine.verifyRows(), + .an => 16, + }; + var at: usize = 0; + while (at < rows) : (at += tile) { + const n = @min(tile, rows - at); + const xa = x + at * l.k() * 2; + const sa = xs + at * (l.k() / 64) * 4; + const oa = out + at * l.n() * 2; + switch (l) { + .a4 => |q| try k.affine.verify(p.s, xa, sa, q, oa, n), + .an => |q| try k.wide.decode(p.s, xa, sa, q, oa, n), + } + } + } + + /// The 64-value sums of x [rows, width] a decode projection reads (prompt GEMMs take none). + fn sums(p: Pass, x: u64, width: usize) !void { + if (!p.prompt) try p.k.act.groupSums(p.s, x, width, p.m.s.xs, p.rows, width); + } + + /// h = bf16(h + r) (r null: as is), x = bf16(w * bf16(h * inv_rms)) and its sums. + fn norm(p: Pass, r: ?u64, w: u64) !void { + const x = p.m.s; + try p.k.norms.addRmsTail(p.s, x.h, r, w, x.h, x.x, x.xs, p.rows, c.hidden, c.eps, tail); + } + + fn conv(p: Pass, cv: wts.Conv, index: usize) !void { + const x = p.m.s; + try p.project(cv.in, x.x, x.xs, x.wide, p.rows, p.prompt); + const state = p.m.states + index * 2 * c.hidden * 2; + const snaps = x.snaps + index * mdl.round_rows * 2 * c.hidden * 2; + // a prompt chunk keeps its last row's state, a round each row's (a window keeps any prefix of its rows) + const snap0: usize = if (p.prompt) p.rows - 1 else 0; + try p.k.convRows(p.s, x.wide, cv.weight, state, mdl.slot_elems, x.slot, x.off, x.y, snaps, snap0, p.rows); + try p.sums(x.y, c.hidden); + try p.project(cv.out, x.y, x.xs, x.r, p.rows, p.prompt); + } + + fn attention(p: Pass, a: wts.Attention, index: usize) !void { + const x = p.m.s; + const m = p.m; + const qkv_w = c.hidden + 2 * mdl.kv_width; + try p.project(a.qkv, x.x, x.xs, x.wide, p.rows, p.prompt); + const ld: i64 = @intCast(qkv_w); + const dh: i64 = c.head_dim; + const src: cuda.qk_rows.Src = .{ .q = x.wide, .q_ld = ld, .q_hs = dh, .k = x.wide + c.hidden * 2, .k_ld = ld, .k_hs = dh, .v = x.wide + (c.hidden + mdl.kv_width) * 2, .v_ld = ld, .v_hs = dh }; + const shape: cuda.qk_rows.Shape = .{ .norms = a.norms, .norms_bf16 = true, .inv = m.inv, .rot = c.head_dim, .eps = c.eps, .round = true }; + const out: cuda.qk_rows.Out = .{ .q = x.q, .k = x.kv, .v = x.kv + mdl.kv_width * 2, .kv_ld = 2 * mdl.kv_width }; + try p.k.qk.rowsFrom(p.s, src, .{ .h = c.query_heads, .hk = c.kv_heads, .dim = c.head_dim }, shape, x.pos, null, out, p.rows); + const kc = m.cache(index, false); + const vc = m.cache(index, true); + var row0: usize = 0; + var at: usize = 0; + while (at < p.chains.len) : (at += kern.max_segs) { + const group = p.chains[at..@min(p.chains.len, at + kern.max_segs)]; + var sg = std.mem.zeroes(kern.AttnSegs); + var rows: usize = 0; + for (group, 0..) |ch, j| { + sg.kc[j] = kc; + sg.vc[j] = vc; + sg.meta[j] = x.meta + (at + j) * 4; + sg.pt[j] = m.tableOf(ch.slot); + sg.row0[j] = @intCast(rows); + rows += ch.tokens.len; + } + sg.ps = mdl.page_elems; + sg.n = @intCast(group.len); + try p.k.kvWrite(p.s, sg, x.kv + row0 * 2 * mdl.kv_width * 2, rows); + const qa = x.q + row0 * c.hidden * 2; + const ya = x.y + row0 * c.hidden * 2; + if (p.prompt) { + try p.k.prompt(p.s, qa, kc, vc, ya, group[0].p0, rows, m.tableOf(group[0].slot), mdl.page_elems); + } else { + const xa = x.xs + row0 * (c.hidden / 64) * 4; + try p.k.attend(p.s, sg, qa, x.po, x.pm, x.pl, ya, xa, rows, m.chunks(), m.chunkBlocks()); + } + row0 += rows; + } + try p.project(a.out, x.y, x.xs, x.r, p.rows, p.prompt); + } + + fn dense(p: Pass, dn: wts.Dense) !void { + const x = p.m.s; + try p.project(dn.gate_up, x.x, x.xs, x.wide, p.rows, p.prompt); + try p.k.siluMul(p.s, x.wide, x.y, p.rows, c.intermediate); + try p.sums(x.y, c.intermediate); + try p.project(dn.down, x.y, x.xs, x.r, p.rows, p.prompt); + } + + fn moe(p: Pass, e: wts.Moe) !void { + const x = p.m.s; + const k = p.k; + const pairs = p.rows * c.topk; + try k.routeRows(p.s, x.x, e.router, e.bias, x.picks, x.wts, null, p.rows, p.m.w.config.scale); + const tile: usize = if (p.prompt) cuda.experts.promptTile(.affine4) else 16; + const items = cuda.grouped.maxItems(pairs, c.experts, tile); + const plan: cuda.grouped.Plan = .{ .members = x.members, .items = x.items, .counts = x.counts, .rank = x.rank, .hist = x.hist }; + try k.plan.route(p.s, x.picks, pairs, c.experts, tile, plan); + const grouped: cuda.experts.Grouped = .{ .affine4 = &k.experts }; + const rows_in: cuda.experts.Rows = .{ .x = x.x, .stride = c.hidden, .slots = c.topk }; + const act_in: cuda.experts.Rows = .{ .x = x.eact, .stride = c.moe_intermediate }; + if (p.prompt) { + try grouped.prompt(p.s, .up, rows_in, e.experts, plan, x.eact, items, -1); + try grouped.prompt(p.s, .down_bf16, act_in, e.experts, plan, x.ey, items, -1); + } else { + try grouped.decode(p.s, .up, rows_in, e.experts, plan, x.eact, items, -1); + try grouped.decode(p.s, .down_f32, act_in, e.experts, plan, x.ey, items, -1); + } + try k.moe.combine(p.s, x.ey, !p.prompt, x.wts, x.r, p.rows, c.topk, c.hidden, 0); + } +}; + +/// Every chain's rows (a prompt chunk, or windows); logits rows [head0, head0 + heads); conv states end at last rows. +pub fn forward(m: *const Model, k: *const Kernels, s: Stream, chains: []const Chain, prompt: bool, head0: usize, heads: usize) !void { + var rows: usize = 0; + for (chains) |ch| { + if (ch.tokens.len == 0 or ch.slot >= m.slots or ch.p0 + ch.tokens.len > m.context) return error.BadChain; + rows += ch.tokens.len; + } + if (rows > @as(usize, if (prompt) mdl.prompt_rows else mdl.round_rows) or (prompt and chains.len != 1)) return error.BadWindow; + if (chains.len > mdl.round_rows or heads > mdl.round_rows or head0 + heads > rows) return error.BadWindow; + try upload(m, s, chains, rows); + const p: Pass = .{ .m = m, .k = k, .s = s, .chains = chains, .rows = rows, .prompt = prompt }; + const x = m.s; + try k.embed.affine(s, x.ids, m.w.embed_w, m.w.embed_s, m.w.embed_b, x.h, rows, c.hidden, .{ .bits = @intCast(m.w.config.embed_bits), .group = c.group }); + var conv_at: usize = 0; + var attn_at: usize = 0; + for (m.w.layers, 0..) |l, i| { + try p.norm(if (i == 0) null else x.r, l.operator_norm); + switch (l.mixer) { + .conv => |cv| { + try p.conv(cv, conv_at); + conv_at += 1; + }, + .attention => |a| { + try p.attention(a, attn_at); + attn_at += 1; + }, + } + try p.norm(x.r, l.ffn_norm); + switch (l.ffn) { + .dense => |dn| try p.dense(dn), + .moe => |e| try p.moe(e), + } + } + try p.norm(x.r, m.w.final_norm); + try commitAll(m, k, s, chains, prompt); + if (heads == 0) return; + try p.project(m.w.head, x.x + head0 * c.hidden * 2, x.xs + head0 * (c.hidden / 64) * 4, x.logits, heads, false); +} + +/// Each stream's conv states (slots[i]) set to its state after snapshot row rows[i]; a row < 0 keeps them. +pub fn commit(m: *const Model, k: *const Kernels, s: Stream, slots: []const i32, rows: []const i32) !void { + if (slots.len != rows.len or slots.len > mdl.round_rows) return error.BadWindow; + try put(m, s, m.s.commit_slot, std.mem.sliceAsBytes(slots)); + try put(m, s, m.s.commit_row, std.mem.sliceAsBytes(rows)); + try k.commitRows(s, m.states, m.s.snaps, m.s.commit_slot, m.s.commit_row, slots.len, mdl.conv_layers, 2 * c.hidden, mdl.round_rows * 2 * c.hidden, mdl.slot_elems); +} + +/// Each chain's conv states at its last row. +fn commitAll(m: *const Model, k: *const Kernels, s: Stream, chains: []const Chain, prompt: bool) !void { + var slots: [mdl.round_rows]i32 = undefined; + var rows: [mdl.round_rows]i32 = undefined; + var base: usize = 0; + for (chains, 0..) |ch, i| { + base += ch.tokens.len; + slots[i] = @intCast(ch.slot); + rows[i] = if (prompt) 0 else @intCast(base - 1); + } + try commit(m, k, s, slots[0..chains.len], rows[0..chains.len]); +} + +fn put(m: *const Model, s: Stream, dst: u64, bytes: []const u8) !void { + try m.d.check(m.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, s.handle), "cuMemcpyHtoDAsync"); +} + +/// Tokens, positions, slots and offsets a row, each chain's first position, staged for the launches. +fn upload(m: *const Model, s: Stream, chains: []const Chain, rows: usize) !void { + const ids = try m.gpa.alloc(i32, 4 * rows + chains.len); + defer m.gpa.free(ids); + var at: usize = 0; + for (chains, 0..) |ch, j| { + for (ch.tokens, 0..) |t, i| { + ids[at + i] = @intCast(t); + ids[rows + at + i] = @intCast(ch.p0 + i); + ids[2 * rows + at + i] = @intCast(ch.slot); + ids[3 * rows + at + i] = @intCast(i); + } + ids[4 * rows + j] = @intCast(ch.p0); + at += ch.tokens.len; + } + const x = m.s; + for ([_]u64{ x.ids, x.pos, x.slot, x.off }, 0..) |dst, i| try put(m, s, dst, std.mem.sliceAsBytes(ids[i * rows ..][0..rows])); + try put(m, s, x.meta, std.mem.sliceAsBytes(ids[4 * rows ..])); +} diff --git a/zig/src/families/lfm2_moe/cuda_kernels.zig b/zig/src/families/lfm2_moe/cuda_kernels.zig new file mode 100644 index 0000000000..2429f7d02a --- /dev/null +++ b/zig/src/families/lfm2_moe/cuda_kernels.zig @@ -0,0 +1,193 @@ +//! LFM2.5-24B-A2B's launches on CUDA: the shared operators and lfm2_moe.cu's conv, conv-state commit and router. +const std = @import("std"); +const cuda = @import("cuda"); +const c = @import("config.zig"); + +const Args = cuda.Args; +const Function = cuda.Function; +const Stream = cuda.Stream; + +/// prefill_attention.cu's pattn_kernel<64, 8, 4, 8> over paged caches: eight 32-key slots of 64 dims, 32 rows a block. +const pattn_sym = "_ZN20tf_prefill_attention12pattn_kernelILi64ELi8ELi4ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifPKix"; +const pattn_smem: u32 = 8 * 32 * 8 * 16; +const pattn_rows = 32; +/// Positions a KV page holds (nemotron_attention.cu's PAGE and prefill_attention_rows.cuh's 512). +pub const page = 512; +/// Streams one shared-window attention launch takes (nemotron_attention.cu's MAXSEG). +pub const max_segs = 8; + +/// nemotron_attention.cu's AttnSegs: each stream's caches, position, page table and first row. +pub const AttnSegs = extern struct { kc: [max_segs]u64, vc: [max_segs]u64, meta: [max_segs]u64, pt: [max_segs]u64, ps: i64, row0: [max_segs]c_int, n: c_int }; + +const images = .{ "qmm_group", "qmm_prefill", "lane_gemv", "affine4_pack", "sample", "tree_round", "affine8", "affine_prefill", "experts", "experts_prefill", "experts_pack", "moe", "prefill_attention", "nemotron_attention", "lfm2", "lfm2_moe" }; + +fn int(x: usize) c_int { + return @intCast(x); +} + +pub const Kernels = struct { + mods: [images.len]cuda.Module, + affine: cuda.qlinear.Affine4Kernels, + wide: cuda.qlinear.AffineNKernels, + draws: cuda.sample.Draws, + round: cuda.tree_round.Kernels, + embed: cuda.embed_rows.EmbedRows, + norms: cuda.norms.Norms, + qk: cuda.qk_rows.QkRows, + act: cuda.act_rows.ActRows, + experts: cuda.experts.Affine4Experts, + plan: cuda.grouped.Router, + moe: cuda.moe.Fns, + pattn: Function, + kv_write: Function, + chunk: Function, + merge: Function, + silu_mul: Function, + conv: Function, + commit: Function, + route: Function, + + pub fn load(ctx: *const cuda.Context) !Kernels { + const d = ctx.d; + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var k: Kernels = undefined; + var loaded: usize = 0; + errdefer for (k.mods[0..loaded]) |*m| m.unload(); + inline for (images, 0..) |name, i| { + k.mods[i] = try cuda.Module.load(d, @field(cuda.kernels, name)); + loaded += 1; + } + const m = &k.mods; + const sms: usize = @intCast(try ctx.attribute(.multiprocessor_count)); + const gb10 = try ctx.capability() == 121; + k.affine = try cuda.qlinear.Affine4Kernels.resolve(d, m[0], m[1], m[2], m[3], sms, gb10); + errdefer k.affine.deinit(); + k.wide = try cuda.qlinear.AffineNKernels.resolve(d, m[6], m[7], m[3], sms, gb10); + errdefer k.wide.deinit(); + k.draws = try cuda.sample.Draws.resolve(m[4]); + k.round = try cuda.tree_round.Kernels.resolve(m[5]); + k.experts = try cuda.experts.Affine4Experts.resolve(m[8], m[9], m[10], sms); + k.plan = try cuda.grouped.Router.resolve(m[8]); + k.moe = try cuda.moe.Fns.resolve(m[11]); + k.pattn = try m[12].function(pattn_sym); + try k.pattn.allowDynamicShared(pattn_smem); + k.kv_write = try m[13].function("tf_nemo_kv_write_segs"); + k.chunk = try m[13].function("tf_paged_attn_chunk_segs_64"); + k.merge = try m[13].function("tf_paged_attn_merge_segs_64"); + k.silu_mul = try m[14].function("tf_lfm2_silu_mul"); + k.conv = try m[15].function("tf_lfm2m_conv"); + k.commit = try m[15].function("tf_lfm2m_commit"); + k.route = try m[15].function("tf_lfm2m_route"); + k.embed = try cuda.embed_rows.EmbedRows.load(d); + errdefer k.embed.unload(); + k.norms = try cuda.norms.Norms.load(d); + errdefer k.norms.unload(); + k.qk = try cuda.qk_rows.QkRows.load(d); + errdefer k.qk.unload(); + k.act = try cuda.act_rows.ActRows.load(d); + return k; + } + + pub fn deinit(k: *Kernels) void { + k.act.unload(); + k.qk.unload(); + k.norms.unload(); + k.embed.unload(); + k.wide.deinit(); + k.affine.deinit(); + for (&k.mods) |*m| m.unload(); + } + + fn go(f: Function, s: Stream, grid: [3]usize, block: usize, smem: u32, a: *Args) !void { + const g: cuda.Dim3 = .{ .x = @intCast(grid[0]), .y = @intCast(grid[1]), .z = @intCast(grid[2]) }; + try cuda.launch.launch(f, .{ .grid = g, .block = .{ .x = @intCast(block) }, .shared = smem }, s, a); + } + + /// Short conv over bcx [rows, 3 hidden]: each row's slot and offset; a slot's [2, hidden] state every `ss`. + pub fn convRows(k: *const Kernels, s: Stream, bcx: u64, w: u64, state: u64, ss: usize, slot: u64, off: u64, y: u64, snap: ?u64, snap0: usize, rows: usize) !void { + var a: Args = .{}; + for ([_]u64{ bcx, w, state, slot, off, y, snap orelse 0 }) |v| a.add(v); + a.add(int(snap0)); + a.add(int(c.hidden)); + a.add(@as(i64, @intCast(ss))); + try go(k.conv, s, .{ rows, c.hidden / 256, 1 }, 256, 0, &a); + } + + /// Each stream k's conv states (every conv layer) set to its snapshot row rows[k] (int32; < 0: kept). + pub fn commitRows(k: *const Kernels, s: Stream, state: u64, snap: u64, slots: u64, rows: u64, streams: usize, layers: usize, state_stride: usize, snap_stride: usize, ss: usize) !void { + if (streams == 0) return; + var a: Args = .{}; + for ([_]u64{ state, snap, slots, rows }) |v| a.add(v); + a.add(int(c.hidden)); + a.add(@as(i64, @intCast(state_stride))); + a.add(@as(i64, @intCast(snap_stride))); + a.add(@as(i64, @intCast(ss))); + try go(k.commit, s, .{ layers, streams, 1 }, 256, 0, &a); + } + + /// Router picks and weights ([rows, 4] int32, fp32) of x [rows, hidden] against fp32 rows [experts, hidden]. + pub fn routeRows(k: *const Kernels, s: Stream, x: u64, w: u64, bias: u64, picks: u64, wts: u64, logits: ?u64, rows: usize, scale: f32) !void { + var a: Args = .{}; + for ([_]u64{ x, w, bias, picks, wts, logits orelse 0 }) |v| a.add(v); + a.add(int(rows)); + a.add(int(c.hidden)); + a.add(scale); + try go(k.route, s, .{ (rows + 7) / 8, 1, 1 }, 256, 0, &a); + } + + /// out [rows, f] = bf16(bf16(silu(gate)) * up) from rows gate | up [rows, 2 f] (lfm2.cu). + pub fn siluMul(k: *const Kernels, s: Stream, g: u64, out: u64, rows: usize, f: usize) !void { + var a: Args = .{}; + a.add(g); + a.add(out); + a.add(int(f)); + try go(k.silu_mul, s, .{ rows, (f + 255) / 256, 1 }, 256, 0, &a); + } + + /// Rows k, v (kv [rows, 2 kv width]) into each stream's pages at its position on. + pub fn kvWrite(k: *const Kernels, s: Stream, sg: AttnSegs, kv: u64, rows: usize) !void { + var a: Args = .{}; + a.add(sg); + a.add(kv); + a.add(int(2 * c.kv_heads * c.head_dim)); + a.add(@as(c_int, 0)); + a.add(int(c.kv_heads * c.head_dim)); + try go(k.kv_write, s, .{ rows, 1, 1 }, 256, 0, &a); + } + + /// Decode rows: q [rows, heads, 64] over each stream's pages in 512-key chunks, merged in order into out and xs. + pub fn attend(k: *const Kernels, s: Stream, sg: AttnSegs, q: u64, po: u64, pm: u64, pl: u64, out: u64, xs: u64, rows: usize, nch: usize, z: usize) !void { + const h = c.query_heads; + const hk = c.kv_heads; + const g = h / hk; + var a: Args = .{}; + a.add(sg); + a.add(q); + for ([_]u64{ po, pm, pl }) |v| a.add(v); + for ([_]usize{ h * c.head_dim, h, hk, g, page, nch }) |v| a.add(int(v)); + a.add(1.0 / @sqrt(@as(f32, c.head_dim))); + try go(k.chunk, s, .{ rows, hk, z }, 256, 0, &a); + var b: Args = .{}; + b.add(sg); + for ([_]u64{ po, pm, pl, out, xs }) |v| b.add(v); + for ([_]usize{ h, hk, g, page, nch }) |v| b.add(int(v)); + try go(k.merge, s, .{ rows, hk, 1 }, 256, 0, &b); + } + + /// A prompt chunk's rows at positions p0.. against its stream's pages (filled through them): out [rows, heads, 64]. + pub fn prompt(k: *const Kernels, s: Stream, q: u64, kc: u64, vc: u64, out: u64, p0: usize, rows: usize, pt: u64, ps: usize) !void { + const g = c.query_heads / c.kv_heads; + var a: Args = .{}; + for ([_]u64{ q, kc, vc, out }) |v| a.add(v); + for ([_]usize{ p0, rows, c.query_heads, c.kv_heads, g }) |v| a.add(int(v)); + a.add(1.0 / @sqrt(@as(f32, c.head_dim))); + a.add(pt); + a.add(@as(i64, @intCast(ps))); + try go(k.pattn, s, .{ (rows + pattn_rows - 1) / pattn_rows, c.kv_heads * (g / 4), 1 }, 256, pattn_smem, &a); + } +}; + +test "the attention launches' geometry: 4 query heads a KV head, a pattn block's 4 heads" { + try std.testing.expectEqual(@as(usize, 4), c.query_heads / c.kv_heads); + try std.testing.expectEqual(@as(usize, 304), @sizeOf(AttnSegs)); // nemotron_attention.cu's struct, padded to 8 +} diff --git a/zig/src/families/lfm2_moe/cuda_model.zig b/zig/src/families/lfm2_moe/cuda_model.zig new file mode 100644 index 0000000000..50f2c51277 --- /dev/null +++ b/zig/src/families/lfm2_moe/cuda_model.zig @@ -0,0 +1,215 @@ +//! LFM2.5-24B-A2B's device state on CUDA: KV in 512-position pages, each slot's conv states, the scratch. +const std = @import("std"); +const cuda = @import("cuda"); +const c = @import("config.zig"); +const wts = @import("cuda_weights.zig"); +const kern = @import("cuda_kernels.zig"); + +const DeviceBuffer = cuda.DeviceBuffer; + +/// Rows a prompt chunk takes. +pub const prompt_rows = 1024; +/// Rows a round of windows takes over every stream. +pub const round_rows = 128; +/// Conv layers (every layer that is not attention) and attention layers of the admitted geometry. +pub const conv_layers = 30; +pub const attention_layers = c.layers - conv_layers; +/// Elements a KV row holds (8 heads of 64) and a page holds over every attention layer's keys and values. +pub const kv_width = c.kv_heads * c.head_dim; +pub const page_elems = attention_layers * 2 * kern.page * kv_width; +/// Elements a slot's conv states hold: two bx rows a conv layer. +pub const slot_elems = conv_layers * 2 * c.hidden; +/// Key chunks one attention launch walks in parallel at most (nemotron's chunk_blocks). +const chunk_blocks = 64; + +pub const Scratch = struct { + ids: u64, // int32 [R]: tokens, positions, slots, offsets in their stream's rows, R each + pos: u64, + slot: u64, + off: u64, + meta: u64, // int32 [round_rows]: each chain's first position (AttnSegs' META), its commit slot and row + commit_slot: u64, + commit_row: u64, + h: u64, // bf16 [R, hidden]: the residual + x: u64, // bf16 [R, hidden]: a norm's output + xs: u64, // fp32 [R, intermediate / 64]: a projection input's 64-value sums + r: u64, // bf16 [R, hidden]: a block's output + wide: u64, // bf16 [R, 2 intermediate]: conv in (3 hidden), q|k|v, dense gate|up + y: u64, // bf16 [R, intermediate]: conv out, attention out, dense act + q: u64, // bf16 [R, hidden]: normed and rotated queries + kv: u64, // bf16 [R, 2 kv_width]: keys then values a row + picks: u64, // int32 [R, 4] + wts: u64, // fp32 [R, 4] + members: u64, + items: u64, + counts: u64, + rank: u64, + hist: u64, + eact: u64, // bf16 [R * 4, moe_intermediate] + ey: u64, // fp32 [R * 4, hidden] (bf16 at prompt rows) + po: u64, // decode attention partials: o, m, l + pm: u64, + pl: u64, + snaps: u64, // bf16 [conv_layers][round_rows][2][hidden]: each window row's conv state after it + logits: u64, // bf16 [round_rows, vocab] +}; + +pub const Model = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + w: *const wts.Weights, + context: usize, // positions a slot holds: whole pages + slots: usize, + arena: u64 = 0, // [pages][attention layer][K, V][512][kv_width] bf16 + table: u64 = 0, // int32 [slots][pages a slot]: slot s's pages s * per .. (s + 1) * per + states: u64 = 0, // bf16 [slots][conv layer][2][hidden] + inv: u64 = 0, // fp32 [head_dim / 2]: RoPE's inverse frequencies + s: Scratch = undefined, + buffers: std.ArrayList(DeviceBuffer) = .empty, + planned: ?usize = null, // a plan: the bytes counted, nothing allocated + + pub fn perSlot(m: *const Model) usize { + return m.context / kern.page; + } + + /// Positions a slot holds for a window of `context`: whole pages. + pub fn positions(context: usize) usize { + return (context + kern.page - 1) / kern.page * kern.page; + } + + /// Bytes one stream's caches take at `context` positions: its pages and its conv states. + pub fn streamBytes(context: usize) usize { + return positions(context) / kern.page * page_elems * 2 + slot_elems * 2; + } + + /// Decode attention's key chunks a row walks at most, and the blocks that walk them. + pub fn chunks(m: *const Model) usize { + return m.perSlot(); + } + + pub fn chunkBlocks(m: *const Model) usize { + return @min(m.chunks(), chunk_blocks); + } + + fn alloc(m: *Model, bytes: usize) !u64 { + if (m.planned) |*p| { + p.* += (@max(bytes, 256) + 255) / 256 * 256; + return 0; + } + const b = try DeviceBuffer.alloc(m.d, @max(bytes, 256)); + errdefer { + var x = b; + x.free(); + } + try m.buffers.append(m.gpa, b); + return b.ptr; + } + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, w: *const wts.Weights, context: usize, slots: usize) !Model { + var m: Model = .{ .gpa = gpa, .d = d, .w = w, .context = positions(context), .slots = slots }; + errdefer m.deinit(); + try m.build(s); + return m; + } + + /// The device bytes `init` allocates for these sizes, counted by the same path without allocating. + pub fn openBytes(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, w: *const wts.Weights, context: usize, slots: usize) !usize { + var m: Model = .{ .gpa = gpa, .d = d, .w = w, .context = positions(context), .slots = slots, .planned = 0 }; + defer m.deinit(); + try m.build(s); + return m.planned.?; + } + + fn build(m: *Model, s: cuda.Stream) !void { + const per = m.perSlot(); + m.arena = try m.alloc(m.slots * per * page_elems * 2); + m.table = try m.alloc(m.slots * per * 4); + m.states = try m.alloc(m.slots * slot_elems * 2); + m.inv = try m.alloc(c.head_dim / 2 * 4); + const r = prompt_rows; + const h = c.hidden; + const pairs = r * c.topk; + const items = cuda.grouped.maxItems(pairs, c.experts, 16); + const parts = round_rows * m.chunks() * c.query_heads; + m.s = .{ + .ids = try m.alloc(r * 4), + .pos = try m.alloc(r * 4), + .slot = try m.alloc(r * 4), + .off = try m.alloc(r * 4), + .meta = try m.alloc(round_rows * 4), + .commit_slot = try m.alloc(round_rows * 4), + .commit_row = try m.alloc(round_rows * 4), + .h = try m.alloc(r * h * 2), + .x = try m.alloc(r * h * 2), + .xs = try m.alloc(r * (c.intermediate / 64) * 4), + .r = try m.alloc(r * h * 2), + .wide = try m.alloc(r * 2 * c.intermediate * 2), + .y = try m.alloc(r * c.intermediate * 2), + .q = try m.alloc(r * h * 2), + .kv = try m.alloc(r * 2 * kv_width * 2), + .picks = try m.alloc(pairs * 4), + .wts = try m.alloc(pairs * 4), + .members = try m.alloc(pairs * 4), + .items = try m.alloc(3 * items * 4), + .counts = try m.alloc(2 * 4), + .rank = try m.alloc(pairs * 4), + .hist = try m.alloc(((pairs + 1023) / 1024) * c.experts * 4), + .eact = try m.alloc(pairs * c.moe_intermediate * 2), + .ey = try m.alloc(pairs * h * 4), + .po = try m.alloc(parts * c.head_dim * 4), + .pm = try m.alloc(parts * 4), + .pl = try m.alloc(parts * 4), + .snaps = try m.alloc(conv_layers * round_rows * 2 * h * 2), + .logits = try m.alloc(round_rows * c.vocab * 2), + }; + if (m.planned != null) return; + const ids = try m.gpa.alloc(i32, m.slots * per); + defer m.gpa.free(ids); + for (ids, 0..) |*v, i| v.* = @intCast(i); + try (DeviceBuffer{ .d = m.d, .ptr = m.table, .len = ids.len * 4, .borrowed = true }).upload(0, std.mem.sliceAsBytes(ids)); + const inv = inverseFrequencies(m.w.config.theta); + try (DeviceBuffer{ .d = m.d, .ptr = m.inv, .len = inv.len * 4, .borrowed = true }).upload(0, std.mem.sliceAsBytes(&inv)); + try (DeviceBuffer{ .d = m.d, .ptr = m.states, .len = m.slots * slot_elems * 2, .borrowed = true }).fill8(0, s.handle); + try s.synchronize(); + } + + pub fn deinit(m: *Model) void { + for (m.buffers.items) |*b| b.free(); + m.buffers.deinit(m.gpa); + } + + /// Attention layer a's key (or value) cache as the kernels take it: its part of page 0. + pub fn cache(m: *const Model, a: usize, value: bool) u64 { + return m.arena + ((2 * a + @intFromBool(value)) * kern.page * kv_width) * 2; + } + + /// Slot s's page table row. + pub fn tableOf(m: *const Model, slot: usize) u64 { + return m.table + slot * m.perSlot() * 4; + } + + /// Zero slot s's conv states, queued on `s` (a new stream starts from no history). + pub fn clearSlot(m: *const Model, s: cuda.Stream, slot: usize) !void { + try (DeviceBuffer{ .d = m.d, .ptr = m.states + slot * slot_elems * 2, .len = slot_elems * 2, .borrowed = true }).fill8(0, s.handle); + } +}; + +/// RoPE's inverse frequencies as torch's float32 rows: 1 / theta^(2i / 64), the power in float64 rounded once. +pub fn inverseFrequencies(theta: f32) [c.head_dim / 2]f32 { + var out: [c.head_dim / 2]f32 = undefined; + for (&out, 0..) |*v, i| { + const e: f32 = @as(f32, @floatFromInt(2 * i)) / @as(f32, c.head_dim); + const pw: f32 = @floatCast(std.math.pow(f64, theta, e)); + v.* = 1.0 / pw; + } + return out; +} + +test "a page holds 512 positions of the ten attention layers' keys and values: 10 MiB" { + try std.testing.expectEqual(@as(usize, 10 << 20), page_elems * 2); + try std.testing.expectEqual(@as(usize, 4 * (10 << 20) + slot_elems * 2), Model.streamBytes(1537)); + try std.testing.expectEqual(@as(usize, 2048), Model.positions(1537)); + const inv = inverseFrequencies(1e7); + try std.testing.expectEqual(@as(f32, 1), inv[0]); + try std.testing.expect(@abs(inv[31] - 1.0 / std.math.pow(f64, 1e7, 62.0 / 64.0)) < 1e-12); +} diff --git a/zig/src/families/lfm2_moe/cuda_native.zig b/zig/src/families/lfm2_moe/cuda_native.zig new file mode 100644 index 0000000000..b5699b5795 --- /dev/null +++ b/zig/src/families/lfm2_moe/cuda_native.zig @@ -0,0 +1,149 @@ +//! LFM2.5-24B-A2B (MLX affine 4-bit) behind the native CUDA registry: plain and copy-drafted decoding, no prompt cache. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const cfg = @import("config.zig"); +const wts = @import("cuda_weights.zig"); +const kern = @import("cuda_kernels.zig"); +const mdl = @import("cuda_model.zig"); +const be = @import("cuda_backend.zig"); + +pub const model_type = "lfm2_moe"; +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 = be.prompt_rows; +pub const snap_lookahead: u32 = 0; +pub const explains_open = true; + +/// The server's reading of the room left after the weights; null if unreadable. +pub const Room = struct { ctx: *anyopaque, read: *const fn (*anyopaque) ?u64 }; +/// No `precision` field: the server refuses --precision for this family by name. +pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false, lanes: u32 = 8, lanes_fixed: bool = false, room: ?Room = null }; + +pub const Loaded = struct { + backend: lanes.backend.Backend, + facts: lanes.Model, + rows: u32, + stream_bytes: usize, + free_streams: u32, + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, + target: *be.Lfm, +}; + +const Owned = struct { gpa: std.mem.Allocator, w: wts.Weights, k: kern.Kernels, s: cuda.Stream, m: mdl.Model, lb: be.Lfm }; + +/// Streams to allocate in `room`: as asked, else (auto) the most that fit; a fixed count past them or a round refuses. +fn fit(room: u64, asked: u32, fixed: bool, plan: anytype) !usize { + if (fixed and asked > be.round_rows) return error.ParallelPastRound; + var n = std.math.clamp(@as(usize, asked), 1, be.round_rows); + while (try plan.bytes(n) > room) : (n -= 1) { + if (fixed) return error.ParallelPastMemory; + if (n == 1) return error.NoStreamFits; + } + return n; +} + +const Plan = struct { + own: *const Owned, + d: *const cuda.Driver, + context: usize, + + fn bytes(p: Plan, n: usize) !u64 { + const o = p.own; + return try mdl.Model.openBytes(o.gpa, p.d, o.s, &o.w, p.context, n) + be.Lfm.openBytes(n); + } +}; + +/// The attention layers the kernels' page arena holds: the admitted 10 of 40 (conv the rest). +fn checkLayers(config: *const cfg.Config) !void { + var n: usize = 0; + for (config.attention) |a| n += @intFromBool(a); + if (n != mdl.attention_layers) return error.UnsupportedLfmGeometry; +} + +pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { + if (kernels != null) return error.CapturedKernels; + if (o.lanes_fixed and o.lanes > be.round_rows) return error.ParallelPastRound; + const own = try gpa.create(Owned); + errdefer gpa.destroy(own); + own.* = .{ .gpa = gpa, .w = undefined, .k = undefined, .s = undefined, .m = undefined, .lb = undefined }; + own.k = try kern.Kernels.load(ctx); + errdefer own.k.deinit(); + own.s = try cuda.Stream.init(ctx.d, false); + errdefer own.s.deinit(); + own.w = try wts.load(gpa, io, ctx.d, own.s, &own.k, dir, null); + errdefer own.w.deinit(); + try checkLayers(&own.w.config); + const context = @min(mdl.Model.positions(o.context), own.w.config.context); + const plan: Plan = .{ .own = own, .d = ctx.d, .context = context }; + const room: u64 = if (o.room) |r| r.read(r.ctx) orelse 0 else std.math.maxInt(u64); + const n = fit(room, o.lanes, o.lanes_fixed, plan) catch |e| { + if (e == error.ParallelPastMemory) { + const most = fit(room, o.lanes, false, plan) catch 0; + std.log.err("--parallel {d} needs {d:.2} GiB at a {d}-token window and the CUDA memory budget leaves {d:.2} GiB after the weights: {d} stream{s} fit", .{ o.lanes, gib(try plan.bytes(o.lanes)), context, gib(room), most, if (most == 1) "" else "s" }); + } + return e; + }; + own.m = try mdl.Model.init(gpa, ctx.d, own.s, &own.w, context, n); + errdefer own.m.deinit(); + own.lb = try be.Lfm.init(gpa, ctx.d, &own.m, &own.k, own.s); + errdefer own.lb.deinit(); + return .{ + .backend = own.lb.backend(), + .facts = own.lb.facts(), + .rows = be.max_window, + .stream_bytes = mdl.Model.streamBytes(context), + .free_streams = @intCast(n), + .ctx = own, + .deinit = release, + .target = &own.lb, + }; +} + +fn gib(bytes: u64) f64 { + return @as(f64, @floatFromInt(bytes)) / (1 << 30); +} + +pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { + return switch (err) { + error.PastContext => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", + error.NoFreeSlot => "every stream's caches are in use: retry once another request ends", + error.NoDraftHead => "LFM2.5-24B-A2B on CUDA has no draft head: it drafts from the prompt's own text", + error.LfmTextOnly => "LFM2.5-24B-A2B is a text model: this request has images", + error.StructuresUnsupported => "structured output is not built for LFM2.5-24B-A2B on CUDA yet", + error.NonFiniteLogits => "the model's logits went non-finite at this row: the request is stopped rather than answered from them", + error.CapturedKernels => "LFM2.5-24B-A2B runs on the engine's own kernels: unset TENSORFOLD_CUDA_KERNELS", + error.UnsupportedLfmQuantization => "LFM2.5-24B-A2B on CUDA reads MLX affine 4-bit checkpoints in groups of 64 (the embedding at 4, 6 or 8 bits): the 6- and 8-bit exports are served on Metal", + error.UnsupportedLfmGeometry, error.UnsupportedLfmConfig, error.UnsupportedLfmRotary => "this lfm2_moe checkpoint's geometry is not LFM2.5-24B-A2B's, which this engine serves", + error.ParallelPastRound => std.fmt.comptimePrint("--parallel: LFM2.5-24B-A2B on CUDA holds at most {d} streams (a round's rows): serve --parallel {d} or fewer, or auto", .{ be.round_rows, be.round_rows }), + error.ParallelPastMemory => "--parallel: the CUDA memory budget fits fewer LFM2.5-24B-A2B streams than asked (each stream's caches are allocated at load): lower --parallel or --context, serve --parallel auto, or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", + error.NoStreamFits => "the CUDA memory budget fits no LFM2.5-24B-A2B stream after the weights: lower --context, free device memory, or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", + else => null, + }; +} + +fn release(p: *anyopaque) void { + const own: *Owned = @ptrCast(@alignCast(p)); + own.lb.deinit(); + own.m.deinit(); + own.w.deinit(); + own.s.deinit(); + own.k.deinit(); + own.gpa.destroy(own); +} + +test "every refusal of its own is worded, the rest left to the server" { + for ([_]anyerror{ error.PastContext, error.NoFreeSlot, error.NoDraftHead, error.LfmTextOnly, error.StructuresUnsupported, error.NonFiniteLogits, error.CapturedKernels, error.UnsupportedLfmQuantization, error.UnsupportedLfmGeometry, error.ParallelPastRound, error.ParallelPastMemory, error.NoStreamFits }) |e| { + try std.testing.expect(explain(null, e) != null); + } + try std.testing.expect(explain(null, error.OutOfMemory) == null); +} + +test "the admitted layer pattern has the arena's ten attention layers" { + var config = try cfg.parse(std.testing.allocator, cfg.fixture); + try checkLayers(&config); + config.attention[0] = true; + try std.testing.expectError(error.UnsupportedLfmGeometry, checkLayers(&config)); +} diff --git a/zig/src/families/lfm2_moe/cuda_weights.zig b/zig/src/families/lfm2_moe/cuda_weights.zig new file mode 100644 index 0000000000..05f00a9aff --- /dev/null +++ b/zig/src/families/lfm2_moe/cuda_weights.zig @@ -0,0 +1,239 @@ +//! LFM2.5-24B-A2B's weights on CUDA: affine-4 projections, expert pair/down tables, the tied head, bf16 norms. +const std = @import("std"); +const core = @import("core"); +const cuda = @import("cuda"); +const c = @import("config.zig"); +const kern = @import("cuda_kernels.zig"); + +const DeviceBuffer = cuda.DeviceBuffer; +const Affine4 = cuda.qlinear.Affine4; +const AffineN = cuda.qlinear.AffineN; +const A4 = cuda.experts.Affine4Experts; + +/// A projection at its MLX width: affine-4 tiles, or the tiles of every other width (the head). +pub const Proj = union(enum) { + a4: Affine4, + an: AffineN, + + pub fn n(p: Proj) usize { + return switch (p) { + inline else => |q| q.n, + }; + } + + pub fn k(p: Proj) usize { + return switch (p) { + inline else => |q| q.k, + }; + } +}; + +pub const Conv = struct { in: Proj, out: Proj, weight: u64 }; +/// q, k and v stacked into one projection; the q and k norms bf16 [64] each, q's then k's. +pub const Attention = struct { qkv: Proj, out: Proj, norms: u64 }; +pub const Dense = struct { gate_up: Proj, down: Proj }; +pub const Moe = struct { router: u64, bias: u64, experts: cuda.experts.Layer }; + +pub const Layer = struct { + operator_norm: u64, + ffn_norm: u64, + mixer: union(enum) { conv: Conv, attention: Attention }, + ffn: union(enum) { dense: Dense, moe: Moe }, +}; + +pub const Weights = struct { + gpa: std.mem.Allocator, + config: c.Config, + /// The embedding as MLX stores it: words u32 [V, H bits / 32], scales and biases [V, H / 64]; also the head. + embed_w: u64 = 0, + embed_s: u64 = 0, + embed_b: u64 = 0, + head: Proj = undefined, + final_norm: u64 = 0, + layers: []Layer = &.{}, + buffers: std.ArrayList(DeviceBuffer) = .empty, + + pub fn deinit(w: *Weights) void { + for (w.buffers.items) |*b| b.free(); + w.buffers.deinit(w.gpa); + w.gpa.free(w.layers); + w.* = undefined; + } + + fn alloc(w: *Weights, d: *const cuda.Driver, bytes: usize) !u64 { + const b = try DeviceBuffer.alloc(d, @max(bytes, 256)); + errdefer { + var x = b; + x.free(); + } + try w.buffers.append(w.gpa, b); + return b.ptr; + } + + fn upload(w: *Weights, d: *const cuda.Driver, bytes: []const u8) !u64 { + const b = try DeviceBuffer.fromHost(d, bytes); + errdefer { + var x = b; + x.free(); + } + try w.buffers.append(w.gpa, b); + return b.ptr; + } +}; + +const Ctx = struct { + gpa: std.mem.Allocator, + w: *Weights, + d: *const cuda.Driver, + s: cuda.Stream, + k: *const kern.Kernels, + ck: *core.checkpoint.Checkpoint, + stage: DeviceBuffer, + + fn tensor(x: *Ctx, name: []const u8, dtype: core.checkpoint.DType, shape: []const usize) !u64 { + return x.w.upload(x.d, (try x.ck.expect(name, dtype, shape)).bytes); + } + + /// Parts' bytes one after another into the stage; their device addresses. + fn staged(x: *Ctx, parts: []const []const u8, out: []u64) !void { + var at: usize = 0; + for (parts, out) |p, *o| { + if (at + p.len > x.stage.len) return error.StageTooSmall; + try x.stage.upload(at, p); + o.* = x.stage.ptr + at; + at += (p.len + 255) / 256 * 256; + } + } + + /// MLX affine projections stacked along n (each [n_i, k] at `bits`) packed as one into qlinear's tiles. + fn linear(x: *Ctx, parts: []const []const u8, ns: []const usize, k: usize, bits: u8) !Proj { + var n: usize = 0; + for (ns) |v| n += v; + var buf: [192]u8 = undefined; + const specs = [_]struct { []const u8, core.checkpoint.DType, usize }{ .{ "weight", .u32, k * bits / 32 }, .{ "scales", .bf16, k / 64 }, .{ "biases", .bf16, k / 64 } }; + var raw: [3]u64 = undefined; + var at: usize = 0; + for (specs, &raw) |sp, *r| { + r.* = x.stage.ptr + at; + for (parts, ns) |prefix, rows| { + const t = try x.ck.expect(try std.fmt.bufPrint(&buf, "{s}.{s}", .{ prefix, sp[0] }), sp[1], &.{ rows, sp[2] }); + if (at + t.bytes.len > x.stage.len) return error.StageTooSmall; + try x.stage.upload(at, t.bytes); + at += t.bytes.len; + } + at = (at + 255) / 256 * 256; + } + const served = cuda.qlinear.servedBits(bits); + const out: Proj = if (served == 4) blk: { + const l = Affine4.layout(n, k); + const q: Affine4 = .{ .w = try x.w.alloc(x.d, l.words), .s = try x.w.alloc(x.d, l.scales), .b = try x.w.alloc(x.d, l.scales), .n = n, .k = k, .npad = l.npad }; + try x.k.affine.pack(x.s, raw[0], raw[1], raw[2], q); + break :blk .{ .a4 = q }; + } else blk: { + const l = AffineN.layout(n, k, served); + const q: AffineN = .{ .w = try x.w.alloc(x.d, l.words), .s = try x.w.alloc(x.d, l.scales), .b = try x.w.alloc(x.d, l.scales), .n = n, .k = k, .npad = l.npad, .bits = served }; + try x.k.wide.packWords(x.s, raw[0], bits, n, k, q.w); + try x.k.wide.packTable(x.s, raw[1], n, k, q.npad, q.s); + try x.k.wide.packTable(x.s, raw[2], n, k, q.npad, q.b); + break :blk .{ .an = q }; + }; + try x.s.synchronize(); // the next projection's staging overwrites this one's + return out; + } + + /// The 64 experts' `name` ([64, n, k] at 4 bits) packed as `out`'s mi-th of m interleaved tables. + fn experts(x: *Ctx, prefix: []const u8, name: []const u8, n: usize, k: usize, out: u64, m: usize, mi: usize) !void { + var buf: [192]u8 = undefined; + var parts: [3][]const u8 = undefined; + const specs = [_]struct { []const u8, core.checkpoint.DType, usize }{ .{ "weight", .u32, k / 8 }, .{ "scales", .bf16, k / 64 }, .{ "biases", .bf16, k / 64 } }; + for (specs, &parts) |sp, *p| p.* = (try x.ck.expect(try std.fmt.bufPrint(&buf, "{s}{s}.{s}", .{ prefix, name, sp[0] }), sp[1], &.{ c.experts, n, sp[2] })).bytes; + var raw: [3]u64 = undefined; + try x.staged(&parts, &raw); + try x.k.experts.packInto(x.s, raw[0], raw[1], raw[2], out, c.experts, n, k, m, mi); + try x.s.synchronize(); + } +}; + +/// Bytes the stage needs: the largest of the head's table as stored and one expert table (words, scales, biases). +pub fn stageBytes(embed_bits: u8) usize { + const head = c.vocab * c.hidden * @as(usize, embed_bits) / 8 + 2 * c.vocab * (c.hidden / 64) * 2; + const table = c.experts * c.moe_intermediate * (c.hidden / 2 + 2 * (c.hidden / 64) * 2); + return @max(head, table) + 4 * 256; +} + +/// Every tensor checked against the admitted geometry; projections and experts at 4 bits. `layers` caps them. +pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, s: cuda.Stream, k: *const kern.Kernels, dir: []const u8, layers: ?usize) !Weights { + const config = try c.Config.read(gpa, io, dir); + if (config.bits != 4) return error.UnsupportedLfmQuantization; // the experts' kernels read 4-bit tables + var w: Weights = .{ .gpa = gpa, .config = config }; + errdefer w.deinit(); + var ck = try core.checkpoint.Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + const h = c.hidden; + const eb = config.embed_bits; + var stage = try DeviceBuffer.alloc(d, stageBytes(eb)); + defer stage.free(); + var x: Ctx = .{ .gpa = gpa, .w = &w, .d = d, .s = s, .k = k, .ck = &ck, .stage = stage }; + w.embed_w = try x.tensor("model.embed_tokens.weight", .u32, &.{ c.vocab, h * @as(usize, eb) / 32 }); + w.embed_s = try x.tensor("model.embed_tokens.scales", .bf16, &.{ c.vocab, h / 64 }); + w.embed_b = try x.tensor("model.embed_tokens.biases", .bf16, &.{ c.vocab, h / 64 }); + w.head = try x.linear(&.{"model.embed_tokens"}, &.{c.vocab}, h, eb); + w.final_norm = try x.tensor("model.embedding_norm.weight", .bf16, &.{h}); + const n = @min(layers orelse c.layers, c.layers); + w.layers = try gpa.alloc(Layer, n); + var b1: [128]u8 = undefined; + var b2: [128]u8 = undefined; + var b3: [128]u8 = undefined; + for (w.layers, 0..) |*l, i| { + var pb: [64]u8 = undefined; + const p = try std.fmt.bufPrint(&pb, "model.layers.{d}.", .{i}); + const name = struct { + fn of(b: []u8, prefix: []const u8, suffix: []const u8) ![]const u8 { + return std.fmt.bufPrint(b, "{s}{s}", .{ prefix, suffix }); + } + }.of; + l.operator_norm = try x.tensor(try name(&b1, p, "operator_norm.weight"), .bf16, &.{h}); + l.ffn_norm = try x.tensor(try name(&b1, p, "ffn_norm.weight"), .bf16, &.{h}); + if (config.attention[i]) { + const qw = c.query_heads * c.head_dim; + const kw = c.kv_heads * c.head_dim; + const norms = try w.alloc(d, 2 * c.head_dim * 2); + const qn = try ck.expect(try name(&b1, p, "self_attn.q_layernorm.weight"), .bf16, &.{c.head_dim}); + const kn = try ck.expect(try name(&b1, p, "self_attn.k_layernorm.weight"), .bf16, &.{c.head_dim}); + try (DeviceBuffer{ .d = d, .ptr = norms, .len = 4 * c.head_dim, .borrowed = true }).upload(0, qn.bytes); + try (DeviceBuffer{ .d = d, .ptr = norms, .len = 4 * c.head_dim, .borrowed = true }).upload(2 * c.head_dim, kn.bytes); + l.mixer = .{ .attention = .{ + .qkv = try x.linear(&.{ try name(&b1, p, "self_attn.q_proj"), try name(&b2, p, "self_attn.k_proj"), try name(&b3, p, "self_attn.v_proj") }, &.{ qw, kw, kw }, h, 4), + .out = try x.linear(&.{try name(&b1, p, "self_attn.out_proj")}, &.{h}, qw, 4), + .norms = norms, + } }; + } else l.mixer = .{ .conv = .{ + .in = try x.linear(&.{try name(&b1, p, "conv.in_proj")}, &.{3 * h}, h, 4), + .out = try x.linear(&.{try name(&b1, p, "conv.out_proj")}, &.{h}, h, 4), + .weight = try x.tensor(try name(&b1, p, "conv.conv.weight"), .bf16, &.{ h, c.conv_taps, 1 }), + } }; + if (config.moe(i)) { + const ni = c.moe_intermediate; + const pair = try w.alloc(d, A4.pairBytes(c.experts, ni, h)); + const down = try w.alloc(d, A4.bytes(c.experts, h, ni)); + const ff = try name(&b2, p, "feed_forward.switch_mlp."); + try x.experts(ff, "gate_proj", ni, h, pair, 2, 0); + try x.experts(ff, "up_proj", ni, h, pair, 2, 1); + try x.experts(ff, "down_proj", h, ni, down, 1, 0); + l.ffn = .{ .moe = .{ + .router = try x.tensor(try name(&b1, p, "feed_forward.gate.weight"), .f32, &.{ c.experts, h }), + .bias = try x.tensor(try name(&b1, p, "feed_forward.expert_bias"), .f32, &.{c.experts}), + .experts = .{ .format = .affine4, .up = pair, .down = down, .width = ni, .dims = h, .experts = c.experts, .swiglu = true }, + } }; + } else l.ffn = .{ .dense = .{ + .gate_up = try x.linear(&.{ try name(&b1, p, "feed_forward.gate_proj"), try name(&b2, p, "feed_forward.up_proj") }, &.{ c.intermediate, c.intermediate }, h, 4), + .down = try x.linear(&.{try name(&b1, p, "feed_forward.down_proj")}, &.{h}, c.intermediate, 4), + } }; + } + return w; +} + +test "the stage holds the 6-bit head's table and one expert table" { + try std.testing.expect(stageBytes(6) >= c.vocab * c.hidden * 6 / 8); + try std.testing.expect(stageBytes(4) >= c.experts * c.moe_intermediate * c.hidden / 2); +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index ab6a6e9443..c808cb8d84 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -18,6 +18,7 @@ const kolibri1 = @import("kolibri1"); const gemma4 = @import("gemma4_cuda"); const spark25 = @import("spark25"); const dory = @import("dory"); +const lfm2_moe = @import("lfm2_moe_cuda"); const Allocator = std.mem.Allocator; pub const budget = @import("cuda_memory.zig"); const lane_backend = @import("cuda_bound.zig"); @@ -28,7 +29,7 @@ const vision_offer = @import("cuda_vision_offer.zig"); const drafted = @import("cuda_drafted.zig"); const kv_pages = @import("cuda_kv_pool.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. -const registry = .{ nemotron.native, nemotron.omni, glm.native, qwen35.native, qwen35.native_moe, flashnext.native, kolibri1.native, gemma4.native, gemma4.native31, spark25.native, dory.native }; +const registry = .{ nemotron.native, nemotron.omni, glm.native, qwen35.native, qwen35.native_moe, flashnext.native, kolibri1.native, gemma4.native, gemma4.native31, spark25.native, dory.native, lfm2_moe.native }; /// --record opens here (engines.zig refuses it for a backend without this). pub const records = true; pub const backends: []const []const u8 = &.{"cuda"}; From 173e06901ae80a8226753f79f8480ba0591c65c9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 27 Oct 2026 21:12:08 +0100 Subject: [PATCH 04/12] cuda: Liquid's short conv and its commit in liquid_conv.cuh, by tap count The gated short conv over each stream's state and the conv-state commit move out of lfm2_moe.cuh into liquid_conv.cuh, templated on the tap count, so LFM3 (7 taps) can take the same bodies as LFM2.5 (3). lfm2_moe.cu keeps the router and instantiates the 3-tap pair as tf_liquid_conv_3 and tf_liquid_commit_3; the controls become tf_ctl_liquid_conv_3 and tf_ctl_liquid_commit_3. cuobjdump -sass on sm_80/86/89/120/121, nvcc 13.0 and 13.2: the moved bodies are identical under their new names, every other function unchanged. --- zig/build/cuda_kernels.zig | 2 +- zig/kernels/cuda/contract_controls.cu | 23 ++++--- zig/kernels/cuda/lfm2_moe.cu | 21 +++--- zig/kernels/cuda/lfm2_moe.cuh | 49 +------------- zig/kernels/cuda/liquid_conv.cuh | 65 +++++++++++++++++++ .../core/contract/lfm2_moe_cuda_records.zig | 6 +- zig/src/families/lfm2_moe/cuda_kernels.zig | 6 +- zig/tests/cuda/contract_elementwise.zig | 2 +- zig/tests/cuda/contract_lfm2_moe.zig | 10 +-- 9 files changed, 103 insertions(+), 81 deletions(-) create mode 100644 zig/kernels/cuda/liquid_conv.cuh diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index bd0f03d169..38a3d44ae3 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -83,7 +83,7 @@ pub const kernels = [_]Kernel{ .{ .name = "norms", .flags = &.{"-O3"} }, // ours: RMS norm rows for every family (cuda/norms.zig) .{ .name = "embed_rows", .flags = &.{"-O3"} }, // ours: MLX affine and bf16 token rows (cuda/embed_rows.zig) .{ .name = "lfm2", .flags = glue }, // ours: LFM2's short convolution, QK norm and rotary, gates and label rows - .{ .name = "lfm2_moe", .flags = glue }, // ours: LFM2.5-24B-A2B's conv over stream state, its commit and the router + .{ .name = "lfm2_moe", .flags = glue }, // ours: LFM2.5-24B-A2B's router, Liquid's 3-tap conv and commit (liquid_conv.cuh) .{ .name = "vit", .flags = glue }, // ours: vision towers' LayerNorm, Linear epilogues, activations, unshuffle .{ .name = "tree_round", .flags = &.{"-O3"} }, // ours: verified tree rounds for every family (cuda/tree_round.zig) .{ .name = "act_rows", .flags = &.{"-O3"} }, // ours: SwiGLU, output gates, group sums (cuda/act_rows.zig) diff --git a/zig/kernels/cuda/contract_controls.cu b/zig/kernels/cuda/contract_controls.cu index a6df0204e7..da3f8243d4 100644 --- a/zig/kernels/cuda/contract_controls.cu +++ b/zig/kernels/cuda/contract_controls.cu @@ -343,21 +343,24 @@ TF_CTL_SWIGLU_L2(tf_ctl_swiglu_l2_drop_us, tf_act_l2::GS | tf_act_l2::PROD) TF_CTL_SWIGLU_L2(tf_ctl_swiglu_l2_drop_prod, tf_act_l2::GS | tf_act_l2::US) #undef TF_CTL_SWIGLU_L2 -// LFM2.5-24B-A2B's glue, one planted fault each: conv bx unrounded, commits to stream 0's slot, ties to the higher id. -#include "lfm2_moe.cuh" +// Liquid's 3-tap short conv and commit, one planted fault each: conv bx unrounded, commits to stream 0's slot. +#include "liquid_conv.cuh" -extern "C" __global__ void __launch_bounds__(256) tf_ctl_lfm2m_conv(const __nv_bfloat16* P, const __nv_bfloat16* Wc, const __nv_bfloat16* STATE, - const int* SLOT, const int* OFF, __nv_bfloat16* Y, __nv_bfloat16* SNAP, - int SNAP0, int D, long long SS) { - tf_lfm2m::lfm2m_conv<1>(P, Wc, STATE, SLOT, OFF, Y, SNAP, SNAP0, D, SS); +extern "C" __global__ void __launch_bounds__(256) tf_ctl_liquid_conv_3(const __nv_bfloat16* P, const __nv_bfloat16* Wc, const __nv_bfloat16* STATE, + const int* SLOT, const int* OFF, __nv_bfloat16* Y, __nv_bfloat16* SNAP, + int SNAP0, int D, long long SS) { + tf_liquid::liquid_conv<3, 1>(P, Wc, STATE, SLOT, OFF, Y, SNAP, SNAP0, D, SS); } -extern "C" __global__ void __launch_bounds__(256) tf_ctl_lfm2m_commit(__nv_bfloat16* STATE, const __nv_bfloat16* SNAP, const int* SLOT, - const int* ROW, int D, long long state_stride, long long snap_stride, - long long SS) { - tf_lfm2m::lfm2m_commit<1>(STATE, SNAP, SLOT, ROW, D, state_stride, snap_stride, SS); +extern "C" __global__ void __launch_bounds__(256) tf_ctl_liquid_commit_3(__nv_bfloat16* STATE, const __nv_bfloat16* SNAP, const int* SLOT, + const int* ROW, int D, long long state_stride, long long snap_stride, + long long SS) { + tf_liquid::liquid_commit<3, 1>(STATE, SNAP, SLOT, ROW, D, state_stride, snap_stride, SS); } +// LFM2.5-24B-A2B's router, one planted fault: ties go to the higher expert id. +#include "lfm2_moe.cuh" + extern "C" __global__ void __launch_bounds__(256) tf_ctl_lfm2m_route(const __nv_bfloat16* X, const float* W, const float* BIAS, int* PICK, float* WTS, float* LOGITS, int rows, int D, float scale) { tf_lfm2m::lfm2m_route<1>(X, W, BIAS, PICK, WTS, LOGITS, rows, D, scale); diff --git a/zig/kernels/cuda/lfm2_moe.cu b/zig/kernels/cuda/lfm2_moe.cu index dd32fcfdd8..d7b3783388 100644 --- a/zig/kernels/cuda/lfm2_moe.cu +++ b/zig/kernels/cuda/lfm2_moe.cu @@ -1,20 +1,21 @@ -// LFM2.5-24B-A2B's CUDA glue beside the shared operators (lfm2_moe.cuh): short conv, conv-state commit, router. +// LFM2.5-24B-A2B's CUDA glue: Liquid's 3-tap short conv and commit (liquid_conv.cuh), the router (lfm2_moe.cuh). +#include "liquid_conv.cuh" #include "lfm2_moe.cuh" using namespace tf_lfm2m; -extern "C" __global__ void __launch_bounds__(NT) tf_lfm2m_conv(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ Wc, - const __nv_bfloat16* __restrict__ STATE, const int* __restrict__ SLOT, - const int* __restrict__ OFF, __nv_bfloat16* __restrict__ Y, - __nv_bfloat16* __restrict__ SNAP, int SNAP0, int D, long long SS) { - lfm2m_conv<0>(P, Wc, STATE, SLOT, OFF, Y, SNAP, SNAP0, D, SS); +extern "C" __global__ void __launch_bounds__(tf_liquid::NT) tf_liquid_conv_3(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ Wc, + const __nv_bfloat16* __restrict__ STATE, const int* __restrict__ SLOT, + const int* __restrict__ OFF, __nv_bfloat16* __restrict__ Y, + __nv_bfloat16* __restrict__ SNAP, int SNAP0, int D, long long SS) { + tf_liquid::liquid_conv<3, 0>(P, Wc, STATE, SLOT, OFF, Y, SNAP, SNAP0, D, SS); } -extern "C" __global__ void __launch_bounds__(NT) tf_lfm2m_commit(__nv_bfloat16* __restrict__ STATE, const __nv_bfloat16* __restrict__ SNAP, - const int* __restrict__ SLOT, const int* __restrict__ ROW, int D, - long long state_stride, long long snap_stride, long long SS) { - lfm2m_commit<0>(STATE, SNAP, SLOT, ROW, D, state_stride, snap_stride, SS); +extern "C" __global__ void __launch_bounds__(tf_liquid::NT) tf_liquid_commit_3(__nv_bfloat16* __restrict__ STATE, const __nv_bfloat16* __restrict__ SNAP, + const int* __restrict__ SLOT, const int* __restrict__ ROW, int D, + long long state_stride, long long snap_stride, long long SS) { + tf_liquid::liquid_commit<3, 0>(STATE, SNAP, SLOT, ROW, D, state_stride, snap_stride, SS); } extern "C" __global__ void __launch_bounds__(NT) tf_lfm2m_route(const __nv_bfloat16* __restrict__ X, const float* __restrict__ W, diff --git a/zig/kernels/cuda/lfm2_moe.cuh b/zig/kernels/cuda/lfm2_moe.cuh index ad8d3c3789..be336c21e7 100644 --- a/zig/kernels/cuda/lfm2_moe.cuh +++ b/zig/kernels/cuda/lfm2_moe.cuh @@ -1,4 +1,4 @@ -// LFM2.5-24B-A2B's CUDA glue (short conv, its commit, router); each row's bits its own; CTL 1 is the planted control. +// LFM2.5-24B-A2B's router; each row's bits its own; CTL 1 is the planted control (short conv: liquid_conv.cuh). #pragma once @@ -21,53 +21,6 @@ template __device__ __forceinline__ bool beats(float v, int i, float w, int j) { return v > w || (v == w && (CTL == 1 ? i > j : i < j)); } -// Block (row, channel block): y = bf16(C * bf16(sum of w * bf16(B * x), oldest tap first)); SNAP rows keep state. -template // CTL 1: bx = B * x not rounded to bf16 before its tap. -__device__ __forceinline__ void lfm2m_conv(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ Wc, - const __nv_bfloat16* __restrict__ STATE, const int* __restrict__ SLOT, - const int* __restrict__ OFF, __nv_bfloat16* __restrict__ Y, - __nv_bfloat16* __restrict__ SNAP, int SNAP0, int D, long long SS) { - const int64_t r = blockIdx.x, D3 = 3 * static_cast(D); - const int c = blockIdx.y * NT + threadIdx.x; - if (c >= D) return; - const int o = OFF[r]; - const __nv_bfloat16* st = STATE + SLOT[r] * SS; - __nv_bfloat16 bx[3]; - float wide[3] = {0.0f, 0.0f, 0.0f}; -#pragma unroll - for (int j = 0; j < 3; ++j) { - const int back = 2 - j; // rows before this one - if (o >= back) { - const int64_t src = r - back; - const float v = __fmul_rn(tf_glue::bf(P[src * D3 + c]), tf_glue::bf(P[src * D3 + 2 * D + c])); - bx[j] = __float2bfloat16_rn(v); - if (CTL == 1) wide[j] = v; - } else { - bx[j] = st[static_cast(2 - back + o) * D + c]; - if (CTL == 1) wide[j] = tf_glue::bf(bx[j]); - } - } - float acc = 0.0f; -#pragma unroll - for (int j = 0; j < 3; ++j) acc = __fadd_rn(acc, __fmul_rn(tf_glue::bf(Wc[c * 3 + j]), CTL == 1 ? wide[j] : tf_glue::bf(bx[j]))); - Y[r * D + c] = __float2bfloat16_rn(__fmul_rn(tf_glue::bf(P[r * D3 + D + c]), tf_glue::rbf(acc))); - if (SNAP == nullptr || r < SNAP0) return; - SNAP[((r - SNAP0) * 2) * D + c] = bx[1]; - SNAP[((r - SNAP0) * 2 + 1) * D + c] = bx[2]; -} - -// Block (layer, stream): stream k's [2][D] conv state set to its snapshot row ROW[k]; ROW < 0 keeps the state. -template // CTL 1: every stream's state written to the first stream's slot. -__device__ __forceinline__ void lfm2m_commit(__nv_bfloat16* __restrict__ STATE, const __nv_bfloat16* __restrict__ SNAP, - const int* __restrict__ SLOT, const int* __restrict__ ROW, int D, - long long state_stride, long long snap_stride, long long SS) { - const int layer = blockIdx.x, k = blockIdx.y, row = ROW[k]; - if (row < 0) return; - __nv_bfloat16* dst = STATE + layer * state_stride + SLOT[CTL == 1 ? 0 : k] * SS; - const __nv_bfloat16* src = SNAP + layer * snap_stride + static_cast(row) * 2 * D; - for (int i = threadIdx.x; i < 2 * D; i += NT) dst[i] = src[i]; -} - // Block of RB rows: fp32 logits (quarters of D in order, (q0 + q1) + (q2 + q3)), top-4 of sigmoid + bias, weights. template // CTL 1: ties go to the higher expert id. __device__ __forceinline__ void lfm2m_route(const __nv_bfloat16* __restrict__ X, const float* __restrict__ W, diff --git a/zig/kernels/cuda/liquid_conv.cuh b/zig/kernels/cuda/liquid_conv.cuh new file mode 100644 index 0000000000..3e6b511159 --- /dev/null +++ b/zig/kernels/cuda/liquid_conv.cuh @@ -0,0 +1,65 @@ +// Liquid's gated short conv over stream state and its commit, by tap count (LFM2.5 3, LFM3 7); CTL 1: control. + +#pragma once + +#include +#include + +#include "glue_math.cuh" + +namespace tf_liquid { + +constexpr int NT = 256; // threads a block + +// Block (row, channel block): y = bf16(C * bf16(sum of w * bf16(B * x), oldest tap first)); SNAP rows keep state. +template // CTL 1: bx = B * x not rounded to bf16 before its tap. +__device__ __forceinline__ void liquid_conv(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ Wc, + const __nv_bfloat16* __restrict__ STATE, const int* __restrict__ SLOT, + const int* __restrict__ OFF, __nv_bfloat16* __restrict__ Y, + __nv_bfloat16* __restrict__ SNAP, int SNAP0, int D, long long SS) { + constexpr int S = TAPS - 1; // state rows: the earlier taps come from the round (OFF before the row), else the slot + const int64_t r = blockIdx.x, D3 = 3 * static_cast(D); + const int c = blockIdx.y * NT + threadIdx.x; + if (c >= D) return; + const int o = OFF[r]; + const __nv_bfloat16* st = STATE + SLOT[r] * SS; + __nv_bfloat16 bx[TAPS]; + float wide[TAPS]; +#pragma unroll + for (int j = 0; j < TAPS; ++j) wide[j] = 0.0f; +#pragma unroll + for (int j = 0; j < TAPS; ++j) { + const int back = S - j; // rows before this one + if (o >= back) { + const int64_t src = r - back; + const float v = __fmul_rn(tf_glue::bf(P[src * D3 + c]), tf_glue::bf(P[src * D3 + 2 * D + c])); + bx[j] = __float2bfloat16_rn(v); + if (CTL == 1) wide[j] = v; + } else { + bx[j] = st[static_cast(S - back + o) * D + c]; + if (CTL == 1) wide[j] = tf_glue::bf(bx[j]); + } + } + float acc = 0.0f; +#pragma unroll + for (int j = 0; j < TAPS; ++j) acc = __fadd_rn(acc, __fmul_rn(tf_glue::bf(Wc[c * TAPS + j]), CTL == 1 ? wide[j] : tf_glue::bf(bx[j]))); + Y[r * D + c] = __float2bfloat16_rn(__fmul_rn(tf_glue::bf(P[r * D3 + D + c]), tf_glue::rbf(acc))); + if (SNAP == nullptr || r < SNAP0) return; +#pragma unroll + for (int j = 1; j < TAPS; ++j) SNAP[((r - SNAP0) * S + (j - 1)) * D + c] = bx[j]; +} + +// Block (layer, stream): stream k's [TAPS - 1][D] conv state set to its snapshot row ROW[k]; ROW < 0 keeps the state. +template // CTL 1: every stream's state written to the first stream's slot. +__device__ __forceinline__ void liquid_commit(__nv_bfloat16* __restrict__ STATE, const __nv_bfloat16* __restrict__ SNAP, + const int* __restrict__ SLOT, const int* __restrict__ ROW, int D, + long long state_stride, long long snap_stride, long long SS) { + constexpr int S = TAPS - 1; // state rows + const int layer = blockIdx.x, k = blockIdx.y, row = ROW[k]; + if (row < 0) return; + __nv_bfloat16* dst = STATE + layer * state_stride + SLOT[CTL == 1 ? 0 : k] * SS; + const __nv_bfloat16* src = SNAP + layer * snap_stride + static_cast(row) * S * D; + for (int i = threadIdx.x; i < S * D; i += NT) dst[i] = src[i]; +} + +} // namespace tf_liquid diff --git a/zig/src/core/contract/lfm2_moe_cuda_records.zig b/zig/src/core/contract/lfm2_moe_cuda_records.zig index 4615a01471..e517e0f33f 100644 --- a/zig/src/core/contract/lfm2_moe_cuda_records.zig +++ b/zig/src/core/contract/lfm2_moe_cuda_records.zig @@ -1,4 +1,4 @@ -//! LFM2.5-24B-A2B's CUDA records: its conv, commit and route, and the head-64 paged attention, every SM's image. +//! LFM2.5-24B-A2B's CUDA records: the 3-tap liquid conv and commit, its route, head-64 paged attention, every SM. const std = @import("std"); const registry = @import("registry.zig"); const R = registry.CudaRecord; @@ -6,8 +6,8 @@ const R = registry.CudaRecord; const pattn64 = "_ZN20tf_prefill_attention12pattn_kernelILi64ELi8ELi4ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifPKix"; pub const cuda = [_]R{ - .{ .launch = .{ .name = "lfm2_moe tf_lfm2m_conv", .image = "lfm2_moe", .symbol = "tf_lfm2m_conv", .threads = 256 }, .usage = .{ .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 26, .static = 0, .bound = 256 }, .{ .regs = 26, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 } } }, - .{ .launch = .{ .name = "lfm2_moe tf_lfm2m_commit", .image = "lfm2_moe", .symbol = "tf_lfm2m_commit", .threads = 256 }, .usage = .{ .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "liquid_conv tf_liquid_conv_3", .image = "lfm2_moe", .symbol = "tf_liquid_conv_3", .threads = 256 }, .usage = .{ .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 26, .static = 0, .bound = 256 }, .{ .regs = 26, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "liquid_conv tf_liquid_commit_3", .image = "lfm2_moe", .symbol = "tf_liquid_commit_3", .threads = 256 }, .usage = .{ .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 24, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "lfm2_moe tf_lfm2m_route", .image = "lfm2_moe", .symbol = "tf_lfm2m_route", .threads = 256 }, .usage = .{ .{ .regs = 55, .static = 34816, .bound = 256 }, .{ .regs = 50, .static = 34816, .bound = 256 }, .{ .regs = 50, .static = 34816, .bound = 256 }, .{ .regs = 58, .static = 34816, .bound = 256 }, .{ .regs = 59, .static = 34816, .bound = 256 } } }, .{ .launch = .{ .name = "paging tf_paged_attn_chunk_segs_64", .image = "nemotron_attention", .symbol = "tf_paged_attn_chunk_segs_64", .threads = 256 }, .usage = .{ .{ .regs = 48, .static = 25088, .bound = 256 }, .{ .regs = 78, .static = 25088, .bound = 256 }, .{ .regs = 78, .static = 25088, .bound = 256 }, .{ .regs = 80, .static = 25088, .bound = 256 }, .{ .regs = 80, .static = 25088, .bound = 256 } } }, .{ .launch = .{ .name = "paging tf_paged_attn_merge_segs_64", .image = "nemotron_attention", .symbol = "tf_paged_attn_merge_segs_64", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 4096, .bound = 256 }, .{ .regs = 40, .static = 4096, .bound = 256 }, .{ .regs = 40, .static = 4096, .bound = 256 }, .{ .regs = 40, .static = 4096, .bound = 256 }, .{ .regs = 40, .static = 4096, .bound = 256 } } }, diff --git a/zig/src/families/lfm2_moe/cuda_kernels.zig b/zig/src/families/lfm2_moe/cuda_kernels.zig index 2429f7d02a..2978245e6f 100644 --- a/zig/src/families/lfm2_moe/cuda_kernels.zig +++ b/zig/src/families/lfm2_moe/cuda_kernels.zig @@ -1,4 +1,4 @@ -//! LFM2.5-24B-A2B's launches on CUDA: the shared operators and lfm2_moe.cu's conv, conv-state commit and router. +//! LFM2.5-24B-A2B's launches on CUDA: the shared operators, liquid_conv.cuh's 3-tap conv and commit, the router. const std = @import("std"); const cuda = @import("cuda"); const c = @import("config.zig"); @@ -75,8 +75,8 @@ pub const Kernels = struct { k.chunk = try m[13].function("tf_paged_attn_chunk_segs_64"); k.merge = try m[13].function("tf_paged_attn_merge_segs_64"); k.silu_mul = try m[14].function("tf_lfm2_silu_mul"); - k.conv = try m[15].function("tf_lfm2m_conv"); - k.commit = try m[15].function("tf_lfm2m_commit"); + k.conv = try m[15].function("tf_liquid_conv_3"); + k.commit = try m[15].function("tf_liquid_commit_3"); k.route = try m[15].function("tf_lfm2m_route"); k.embed = try cuda.embed_rows.EmbedRows.load(d); errdefer k.embed.unload(); diff --git a/zig/tests/cuda/contract_elementwise.zig b/zig/tests/cuda/contract_elementwise.zig index 3ef8ff9980..f4f49d4c70 100644 --- a/zig/tests/cuda/contract_elementwise.zig +++ b/zig/tests/cuda/contract_elementwise.zig @@ -16,7 +16,7 @@ pub const Set = struct { sample: ?cuda.Module = null, draws: cuda.sample.Draws = undefined, penalize: ?penalize_k.Penalize = null, - lfm2m: lfm2m_k.Set = .{}, // LFM2.5-24B-A2B's router, conv and commit (rows with sums, beside these) + lfm2m: lfm2m_k.Set = .{}, // LFM2.5-24B-A2B's router, Liquid's 3-tap conv and commit (rows with sums, beside these) /// Every registration onto `kernels`; the embed controls' symbols come from the contract_controls image `ctl_mod`. pub fn add(self: *Set, gpu: check.Gpu, s: cuda.Stream, ctl_mod: cuda.Module, kernels: *std.ArrayList(g.Kernel)) !void { diff --git a/zig/tests/cuda/contract_lfm2_moe.zig b/zig/tests/cuda/contract_lfm2_moe.zig index b686ab9ba3..7b30139e39 100644 --- a/zig/tests/cuda/contract_lfm2_moe.zig +++ b/zig/tests/cuda/contract_lfm2_moe.zig @@ -1,4 +1,4 @@ -//! lfm2_moe.cu under the contract: router logits and top-4 rule, short conv and snapshots, commit; each controlled. +//! lfm2_moe.cu under the contract: router logits and top-4 rule, 3-tap liquid conv and snapshots, commit; controlled. const std = @import("std"); const cuda = @import("cuda"); const core = @import("core"); @@ -231,7 +231,7 @@ pub const Select = struct { const conv_rule = "the host's replay of each row's fp32 and bf16 steps, bit for bit (NaN as NaN)"; pub fn convSpec(control: bool) c.Spec { - return .{ .name = if (control) "control: lfm2_moe conv with bx unrounded" else "cuda lfm2_moe short conv tf_lfm2m_conv (stream state, every row's snapshot)", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = conv_rule }, .tiles = .{ .{}, .{ .size = 256, .step = 64, .min = 64 }, .{} }, .served = &.{ .{ .rows = 1, .n = 2048, .k = 3 }, .{ .rows = 16, .n = 2048, .k = 3 }, .{ .rows = 128, .n = 2048, .k = 3 } } }; + return .{ .name = if (control) "control: liquid 3-tap conv with bx unrounded" else "cuda liquid 3-tap short conv tf_liquid_conv_3 (stream state, every row's snapshot)", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = conv_rule }, .tiles = .{ .{}, .{ .size = 256, .step = 64, .min = 64 }, .{} }, .served = &.{ .{ .rows = 1, .n = 2048, .k = 3 }, .{ .rows = 16, .n = 2048, .k = 3 }, .{ .rows = 128, .n = 2048, .k = 3 } } }; } /// Rows of two streams (the first rows/2 slot 1, the rest slot 0) over their states: y and every row's state after it. @@ -362,7 +362,7 @@ pub const Conv = struct { }; pub fn commitSpec(control: bool) c.Spec { - return .{ .name = if (control) "control: lfm2_moe commit into the first stream's slot" else "cuda lfm2_moe conv-state commit tf_lfm2m_commit", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = "a copy: every bit as stored" }, .served = &.{ .{ .rows = 1, .n = 2048, .k = 30 }, .{ .rows = 16, .n = 2048, .k = 30 } } }; + return .{ .name = if (control) "control: liquid 3-tap commit into the first stream's slot" else "cuda liquid 3-tap conv-state commit tf_liquid_commit_3", .accum = .none, .depth = noSums, .out = .raw, .nan = .{ .stated = "a copy: every bit as stored" }, .served = &.{ .{ .rows = 1, .n = 2048, .k = 30 }, .{ .rows = 16, .n = 2048, .k = 30 } } }; } /// Streams k = 0.. into slots rows - k of rows + 1, each its snapshot row 2k (stream 1: none, its slot as it was). @@ -467,8 +467,8 @@ pub const Set = struct { const rig: rig_.Rig = .{ .gpu = gpu, .s = s }; self.logits = .{ .rig = rig, .f = try m.function("tf_lfm2m_route") }; for (&self.select, [_]bool{ false, true }) |*x, ctrl| x.* = .{ .rig = rig, .f = try (if (ctrl) ctl else m).function(if (ctrl) "tf_ctl_lfm2m_route" else "tf_lfm2m_route"), .control = ctrl }; - for (&self.conv, [_]bool{ false, true }) |*x, ctrl| x.* = .{ .rig = rig, .f = try (if (ctrl) ctl else m).function(if (ctrl) "tf_ctl_lfm2m_conv" else "tf_lfm2m_conv"), .control = ctrl }; - for (&self.commit, [_]bool{ false, true }) |*x, ctrl| x.* = .{ .rig = rig, .f = try (if (ctrl) ctl else m).function(if (ctrl) "tf_ctl_lfm2m_commit" else "tf_lfm2m_commit"), .control = ctrl }; + for (&self.conv, [_]bool{ false, true }) |*x, ctrl| x.* = .{ .rig = rig, .f = try (if (ctrl) ctl else m).function(if (ctrl) "tf_ctl_liquid_conv_3" else "tf_liquid_conv_3"), .control = ctrl }; + for (&self.commit, [_]bool{ false, true }) |*x, ctrl| x.* = .{ .rig = rig, .f = try (if (ctrl) ctl else m).function(if (ctrl) "tf_ctl_liquid_commit_3" else "tf_liquid_commit_3"), .control = ctrl }; try kernels.append(a, self.logits.kernel()); for (&self.select) |*x| try kernels.append(a, x.kernel()); for (&self.conv) |*x| try kernels.append(a, x.kernel()); From 064089ddd478cdbb3600163e5fec1ef0fcc55fa7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 27 Oct 2026 22:56:54 +0100 Subject: [PATCH 05/12] kernel contract: the limits check reads lfm2_moe's fatbin --- zig/tests/cuda/contract_limits.zig | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zig/tests/cuda/contract_limits.zig b/zig/tests/cuda/contract_limits.zig index 2a16afa5b6..1fd1a1cc3f 100644 --- a/zig/tests/cuda/contract_limits.zig +++ b/zig/tests/cuda/contract_limits.zig @@ -25,7 +25,7 @@ pub fn image(name: []const u8) error{UnknownContractImage}![]const u8 { if (std.mem.eql(u8, name, "h3")) return &h3_image.bytes; inline for (.{ "qwen4_projection", "qwen4_hc", "qwen4_gdn", "qwen4_ple" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); - inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); + inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows", "lfm2_moe" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); if (std.mem.eql(u8, name, "contract_controls")) return if (options.controls) &controls.bytes else &.{}; std.debug.print("FAIL kernel-contract: unknown image name: {s}\n", .{name}); return error.UnknownContractImage; From 1b4893416f8cb63561b87aa0de0128f76f3f8721 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 28 Oct 2026 09:32:42 +0100 Subject: [PATCH 06/12] kernel contract: the head-64 paged attention, decode and prompt rows, against float64 tf_paged_attn_chunk_segs_64 with tf_paged_attn_merge_segs_64 and pattn_kernel<64, 8, 4, 8> run over shuffled 512-key pages, decode rows over up to eight streams a launch, and are held head by head to contract_attn's float64 bound. Two controls in contract_controls.cu read the values from the neighbouring lane's dims (decode) and from the neighbouring 8 dims (prompt rows). To build them from the same bodies, nemotron_attention.cu's page and tile helpers move to paged_attn.cuh and prefill_attention.cu's to prefill_attention.cuh; every existing function's SASS is unchanged on sm_80/86/89/120/121 under nvcc 13.0 and 13.2. --- zig/kernels/cuda/contract_controls.cu | 33 ++ zig/kernels/cuda/lfm2_moe.cuh | 2 +- zig/kernels/cuda/liquid_conv.cuh | 2 +- zig/kernels/cuda/nemotron_attention.cu | 49 +-- zig/kernels/cuda/paged_attn.cuh | 62 ++++ zig/kernels/cuda/paged_attn_segs.cuh | 6 +- zig/kernels/cuda/prefill_attention.cu | 57 +--- zig/kernels/cuda/prefill_attention.cuh | 67 ++++ zig/kernels/cuda/prefill_attention_rows.cuh | 4 +- zig/tests/cuda/contract_elementwise.zig | 4 + zig/tests/cuda/contract_paged_attn.zig | 344 ++++++++++++++++++++ 11 files changed, 521 insertions(+), 109 deletions(-) create mode 100644 zig/kernels/cuda/paged_attn.cuh create mode 100644 zig/kernels/cuda/prefill_attention.cuh create mode 100644 zig/tests/cuda/contract_paged_attn.zig diff --git a/zig/kernels/cuda/contract_controls.cu b/zig/kernels/cuda/contract_controls.cu index da3f8243d4..c3e58c102c 100644 --- a/zig/kernels/cuda/contract_controls.cu +++ b/zig/kernels/cuda/contract_controls.cu @@ -365,3 +365,36 @@ extern "C" __global__ void __launch_bounds__(256) tf_ctl_lfm2m_route(const __nv_ float* WTS, float* LOGITS, int rows, int D, float scale) { tf_lfm2m::lfm2m_route<1>(X, W, BIAS, PICK, WTS, LOGITS, rows, D, scale); } + +// Paged attention at head 64, one planted fault each: values from the next lane's dims (decode) or next 8 (prompt). +namespace tf_ctl_paged { +using tf_glue::bf; +using tf_glue::tf_exp; +#include "paged_attn.cuh" +#define SEG_HD 64 +#define SEG_KW (64 / 2 + 2) +#define SEG_WIDE 0 +#define SEG_VLANE 1 +#define SEG_CHUNK tf_ctl_paged_attn_chunk_segs_64 +#define SEG_MERGE tf_ctl_paged_attn_merge_segs_64 +#include "paged_attn_segs.cuh" +} // namespace tf_ctl_paged + +#define PATTN_VLANE 1 +#include "prefill_attention.cuh" + +namespace tf_prefill_attention { +extern "C" __global__ void __launch_bounds__(256, 1) tf_ctl_pattn_64(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc, + const __nv_bfloat16* __restrict__ vc, __nv_bfloat16* __restrict__ out, int p0, + int W, int H, int HK, int G, float scale, const int* __restrict__ PT, long long PS) { + constexpr int D = 64, WARPS = 8, HPC = 4, NS = 8; + constexpr bool causal = true; + constexpr int skip = 0; + constexpr bool HPAGED = false, F8 = false, HM = false; + constexpr int CAP = 0, PTOK = 1; + const float *ks = nullptr, *vs = nullptr; + const int HO = H; + auto first_row = [] { return blockIdx.x * 16 * (WARPS / HPC); }; +#include "prefill_attention_rows.cuh" +} +} // namespace tf_prefill_attention diff --git a/zig/kernels/cuda/lfm2_moe.cuh b/zig/kernels/cuda/lfm2_moe.cuh index be336c21e7..e6f390f05c 100644 --- a/zig/kernels/cuda/lfm2_moe.cuh +++ b/zig/kernels/cuda/lfm2_moe.cuh @@ -6,7 +6,7 @@ #include #include -#include "glue_math.cuh" +#include "common/glue_math.cuh" namespace tf_lfm2m { diff --git a/zig/kernels/cuda/liquid_conv.cuh b/zig/kernels/cuda/liquid_conv.cuh index 3e6b511159..51d1ec0e7e 100644 --- a/zig/kernels/cuda/liquid_conv.cuh +++ b/zig/kernels/cuda/liquid_conv.cuh @@ -5,7 +5,7 @@ #include #include -#include "glue_math.cuh" +#include "common/glue_math.cuh" namespace tf_liquid { diff --git a/zig/kernels/cuda/nemotron_attention.cu b/zig/kernels/cuda/nemotron_attention.cu index 6dbd077325..6c66d56704 100644 --- a/zig/kernels/cuda/nemotron_attention.cu +++ b/zig/kernels/cuda/nemotron_attention.cu @@ -5,45 +5,15 @@ #include #include "common/glue_math.cuh" +#include "paged_attn.cuh" using namespace tf_glue; namespace { constexpr int HD = 128; // head dim the chunk kernels take -constexpr int GMAX = 16; // query heads a KV head -constexpr int TILE = 64; // keys a tile: the online softmax's unit constexpr int KW = HD / 2 + 2; // a key row's words in shared memory, padded so 64-bit reads of 16 rows miss no bank -__device__ __forceinline__ float lo(uint32_t w) { return __uint_as_float(w << 16); } -__device__ __forceinline__ float hi(uint32_t w) { return __uint_as_float(w & 0xFFFF0000u); } - -// Max over the 16 lanes of a half warp (xor 8, 4, 2, 1). -__device__ __forceinline__ float half_max(float v) { -#pragma unroll - for (int off = 8; off > 0; off >>= 1) v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, off)); - return v; -} - -// Sum over the 16 lanes of a half warp (xor 8, 4, 2, 1): each lane ends with the same sum. -__device__ __forceinline__ float half_sum(float v) { -#pragma unroll - for (int off = 8; off > 0; off >>= 1) v = __fadd_rn(v, __shfl_xor_sync(0xffffffffu, v, off)); - return v; -} - -constexpr int PAGE = 512; // rows a KV page holds: one attention chunk - -// Element offset of chunk c's first key row: contiguous rows, or the page the stream's table names (PT null: rows). -__device__ __forceinline__ int64_t chunk_at(const int* PT, int64_t PS, int c, int KVD) { - return PT ? static_cast(PT[c]) * PS : static_cast(c) * PAGE * KVD; -} - -// Element offset of key row k: its chunk's first row, then its row inside the chunk. -__device__ __forceinline__ int64_t key_at(const int* PT, int64_t PS, int64_t k, int KVD) { - return chunk_at(PT, PS, static_cast(k / PAGE), KVD) + (k % PAGE) * KVD; -} - } // namespace // Row r's keys and values written at its position: K and V follow the query columns of each qkv row. @@ -197,23 +167,6 @@ extern "C" __global__ void __launch_bounds__(256) tf_nemo_attn_merge(const float } // Shared-window copies of the kernels above (a shared body slowed a lone stream): stream k's rows from ROW0[k] on. -constexpr int MAXSEG = 8; -struct AttnSegs { - __nv_bfloat16* KC[MAXSEG]; - __nv_bfloat16* VC[MAXSEG]; - const int* META[MAXSEG]; - const int* PT[MAXSEG]; // each stream's page table (null: its rows are contiguous) - long long PS; // elements between pages - int ROW0[MAXSEG]; - int N; -}; - -__device__ __forceinline__ int seg_of(const AttnSegs& sg, int rq) { - int k = 0; - while (k + 1 < sg.N && sg.ROW0[k + 1] <= rq) ++k; - return k; -} - extern "C" __global__ void tf_nemo_kv_write_segs(const AttnSegs sg, const __nv_bfloat16* __restrict__ QKV, int NQKV, int QD, int KVD) { const int rq = blockIdx.x, k = seg_of(sg, rq); __nv_bfloat16* __restrict__ KC = sg.KC[k]; diff --git a/zig/kernels/cuda/paged_attn.cuh b/zig/kernels/cuda/paged_attn.cuh new file mode 100644 index 0000000000..f5bba2e249 --- /dev/null +++ b/zig/kernels/cuda/paged_attn.cuh @@ -0,0 +1,62 @@ +// Paged attention's shared pieces (Nemotron's chunk/merge kernels, paged_attn_segs.cuh): tiles, pages, stream segments. +#pragma once + +#include +#include +#include + +#include "common/glue_math.cuh" + +namespace { + +constexpr int GMAX = 16; // query heads a KV head +constexpr int TILE = 64; // keys a tile: the online softmax's unit + +__device__ __forceinline__ float lo(uint32_t w) { return __uint_as_float(w << 16); } +__device__ __forceinline__ float hi(uint32_t w) { return __uint_as_float(w & 0xFFFF0000u); } + +// Max over the 16 lanes of a half warp (xor 8, 4, 2, 1). +__device__ __forceinline__ float half_max(float v) { +#pragma unroll + for (int off = 8; off > 0; off >>= 1) v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, off)); + return v; +} + +// Sum over the 16 lanes of a half warp (xor 8, 4, 2, 1): each lane ends with the same sum. +__device__ __forceinline__ float half_sum(float v) { +#pragma unroll + for (int off = 8; off > 0; off >>= 1) v = __fadd_rn(v, __shfl_xor_sync(0xffffffffu, v, off)); + return v; +} + +constexpr int PAGE = 512; // rows a KV page holds: one attention chunk + +// Element offset of chunk c's first key row: contiguous rows, or the page the stream's table names (PT null: rows). +__device__ __forceinline__ int64_t chunk_at(const int* PT, int64_t PS, int c, int KVD) { + return PT ? static_cast(PT[c]) * PS : static_cast(c) * PAGE * KVD; +} + +// Element offset of key row k: its chunk's first row, then its row inside the chunk. +__device__ __forceinline__ int64_t key_at(const int* PT, int64_t PS, int64_t k, int KVD) { + return chunk_at(PT, PS, static_cast(k / PAGE), KVD) + (k % PAGE) * KVD; +} + +} // namespace + +// Stream k's rows from ROW0[k] on, for the shared-window kernels. +constexpr int MAXSEG = 8; +struct AttnSegs { + __nv_bfloat16* KC[MAXSEG]; + __nv_bfloat16* VC[MAXSEG]; + const int* META[MAXSEG]; + const int* PT[MAXSEG]; // each stream's page table (null: its rows are contiguous) + long long PS; // elements between pages + int ROW0[MAXSEG]; + int N; +}; + +__device__ __forceinline__ int seg_of(const AttnSegs& sg, int rq) { + int k = 0; + while (k + 1 < sg.N && sg.ROW0[k + 1] <= rq) ++k; + return k; +} diff --git a/zig/kernels/cuda/paged_attn_segs.cuh b/zig/kernels/cuda/paged_attn_segs.cuh index 56c3117a22..5a18303cb4 100644 --- a/zig/kernels/cuda/paged_attn_segs.cuh +++ b/zig/kernels/cuda/paged_attn_segs.cuh @@ -1,6 +1,9 @@ // Paged shared-window chunk/merge bodies shared by families (Nemotron head 128, LFM2 head 64); SEG_* set the form. #define SEG_E (SEG_WIDE ? 8 : 4) +#ifndef SEG_VLANE +#define SEG_VLANE 0 // 1 in contract_controls.cu only: the narrow form reads its values from the neighbouring lane's dims +#endif extern "C" __global__ void __launch_bounds__(256) SEG_CHUNK(const AttnSegs sg, const __nv_bfloat16* __restrict__ QKV, float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, @@ -81,7 +84,7 @@ extern "C" __global__ void __launch_bounds__(256) SEG_CHUNK(const AttnSegs sg, c #pragma unroll for (int e = 0; e < 4; ++e) { #else - const uint2 v = *reinterpret_cast(&vw[j][2 * jj]); + const uint2 v = *reinterpret_cast(&vw[j][2 * (jj ^ SEG_VLANE)]); const uint32_t vv[2] = {v.x, v.y}; #pragma unroll for (int e = 0; e < 2; ++e) { @@ -156,6 +159,7 @@ extern "C" __global__ void __launch_bounds__(256) SEG_MERGE(const AttnSegs sg, c } #undef SEG_E +#undef SEG_VLANE #undef SEG_HD #undef SEG_KW #undef SEG_WIDE diff --git a/zig/kernels/cuda/prefill_attention.cu b/zig/kernels/cuda/prefill_attention.cu index 8d04c7d6bb..db33c34253 100644 --- a/zig/kernels/cuda/prefill_attention.cu +++ b/zig/kernels/cuda/prefill_attention.cu @@ -1,63 +1,8 @@ // Device code of src/tensorfold/cuda/kernels/prefill_attention.cu (comments, ATen includes dropped), paged keys added. -#include -#include -#include -#include -#include - -#include "common/qmm_frag.cuh" +#include "prefill_attention.cuh" namespace tf_prefill_attention { -using qmm_frag::commit; -using qmm_frag::cp16z; -using qmm_frag::ldmatrix4; -using qmm_frag::mma; -using qmm_frag::mma0; - -constexpr int BN = 64, HALF = 32; // keys a tile (the fold's unit), keys a staging slot -constexpr float LOG2E = 1.4426950408889634f; -constexpr int SMEM_MAX = 96 * 1024; - -template -__device__ __forceinline__ void wait_group() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); } - -__device__ __forceinline__ void ldmatrix4t(uint32_t (&r)[4], const void* p) { - asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" - : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) : "r"(qmm_frag::smem(p))); -} - -__device__ __forceinline__ float texp(float x) { - float y; - asm("ex2.approx.f32 %0, %1;\n" : "=f"(y) : "f"(__fmul_rn(x, LOG2E))); - return y; -} - -__device__ __forceinline__ float tdiv(float a, float b) { - float y; - asm("div.full.f32 %0, %1, %2;\n" : "=f"(y) : "f"(a), "f"(b)); - return y; -} - -__device__ __forceinline__ uint32_t pack(float lo, float hi) { - __nv_bfloat162 v = __floats2bfloat162_rn(lo, hi); - return *reinterpret_cast(&v); -} - -template -__device__ __forceinline__ int sw(int r, int c) { return r * (D / 8) + (c ^ (r & 7)); } - -__device__ __forceinline__ float rowsum(const float (&c)[8]) { - float w[8]; -#pragma unroll - for (int j = 0; j < 8; ++j) { - const float x = __fadd_rn(c[j], __shfl_xor_sync(0xffffffffu, c[j], 2)); - w[j] = __fadd_rn(x, __shfl_xor_sync(0xffffffffu, x, 1)); - } - return __fadd_rn(__fadd_rn(__fadd_rn(w[0], w[4]), __fadd_rn(w[2], w[6])), - __fadd_rn(__fadd_rn(w[1], w[5]), __fadd_rn(w[3], w[7]))); -} - template __global__ void __launch_bounds__(32 * WARPS, 1) pattn_kernel(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc, diff --git a/zig/kernels/cuda/prefill_attention.cuh b/zig/kernels/cuda/prefill_attention.cuh new file mode 100644 index 0000000000..2bebef02aa --- /dev/null +++ b/zig/kernels/cuda/prefill_attention.cuh @@ -0,0 +1,67 @@ +// prefill_attention.cu's staging, MMA and fold helpers, shared with its controls (contract_controls.cu). +#pragma once + +#include +#include +#include +#include +#include + +#include "common/qmm_frag.cuh" + +namespace tf_prefill_attention { + +using qmm_frag::commit; +using qmm_frag::cp16z; +using qmm_frag::ldmatrix4; +using qmm_frag::mma; +using qmm_frag::mma0; + +constexpr int BN = 64, HALF = 32; // keys a tile (the fold's unit), keys a staging slot +constexpr float LOG2E = 1.4426950408889634f; +constexpr int SMEM_MAX = 96 * 1024; + +template +__device__ __forceinline__ void wait_group() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); } + +__device__ __forceinline__ void ldmatrix4t(uint32_t (&r)[4], const void* p) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) : "r"(qmm_frag::smem(p))); +} + +__device__ __forceinline__ float texp(float x) { + float y; + asm("ex2.approx.f32 %0, %1;\n" : "=f"(y) : "f"(__fmul_rn(x, LOG2E))); + return y; +} + +__device__ __forceinline__ float tdiv(float a, float b) { + float y; + asm("div.full.f32 %0, %1, %2;\n" : "=f"(y) : "f"(a), "f"(b)); + return y; +} + +__device__ __forceinline__ uint32_t pack(float lo, float hi) { + __nv_bfloat162 v = __floats2bfloat162_rn(lo, hi); + return *reinterpret_cast(&v); +} + +template +__device__ __forceinline__ int sw(int r, int c) { return r * (D / 8) + (c ^ (r & 7)); } + +__device__ __forceinline__ float rowsum(const float (&c)[8]) { + float w[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const float x = __fadd_rn(c[j], __shfl_xor_sync(0xffffffffu, c[j], 2)); + w[j] = __fadd_rn(x, __shfl_xor_sync(0xffffffffu, x, 1)); + } + return __fadd_rn(__fadd_rn(__fadd_rn(w[0], w[4]), __fadd_rn(w[2], w[6])), + __fadd_rn(__fadd_rn(w[1], w[5]), __fadd_rn(w[3], w[7]))); +} + +#ifndef PATTN_VLANE +#define PATTN_VLANE 0 // 1 in contract_controls.cu only: values read from the neighbouring 8 dims +#endif + +} // namespace tf_prefill_attention diff --git a/zig/kernels/cuda/prefill_attention_rows.cuh b/zig/kernels/cuda/prefill_attention_rows.cuh index 1986522ed2..0f4a5e9953 100644 --- a/zig/kernels/cuda/prefill_attention_rows.cuh +++ b/zig/kernels/cuda/prefill_attention_rows.cuh @@ -165,14 +165,14 @@ #pragma unroll for (int i = 0; i < NB; i += 2) { uint32_t b[4]; - ldmatrix4t(b, sl + sw<8 * CS>(8 * (mt % 2) + lane % 8, i + mt / 2)); + ldmatrix4t(b, sl + sw<8 * CS>(8 * (mt % 2) + lane % 8, (i + mt / 2) ^ PATTN_VLANE)); mma(o[i], pa0, b[0], b[1]); mma(o[i + 1], pa0, b[2], b[3]); } #pragma unroll for (int i = 0; i < NB; i += 2) { uint32_t b[4]; - ldmatrix4t(b, sl + sw<8 * CS>(16 + 8 * (mt % 2) + lane % 8, i + mt / 2)); + ldmatrix4t(b, sl + sw<8 * CS>(16 + 8 * (mt % 2) + lane % 8, (i + mt / 2) ^ PATTN_VLANE)); mma(o[i], pa1, b[0], b[1]); mma(o[i + 1], pa1, b[2], b[3]); } diff --git a/zig/tests/cuda/contract_elementwise.zig b/zig/tests/cuda/contract_elementwise.zig index f4f49d4c70..fc949f0032 100644 --- a/zig/tests/cuda/contract_elementwise.zig +++ b/zig/tests/cuda/contract_elementwise.zig @@ -6,6 +6,7 @@ const check = @import("check.zig"); const embed_k = @import("contract_embed.zig"); const penalize_k = @import("contract_penalize.zig"); const lfm2m_k = @import("contract_lfm2_moe.zig"); +const paged_k = @import("contract_paged_attn.zig"); const g = core.contract.gates; @@ -17,6 +18,7 @@ pub const Set = struct { draws: cuda.sample.Draws = undefined, penalize: ?penalize_k.Penalize = null, lfm2m: lfm2m_k.Set = .{}, // LFM2.5-24B-A2B's router, Liquid's 3-tap conv and commit (rows with sums, beside these) + paged: paged_k.Set = .{}, // the Liquid families' paged attention at head 64, decode and prompt rows /// Every registration onto `kernels`; the embed controls' symbols come from the contract_controls image `ctl_mod`. pub fn add(self: *Set, gpu: check.Gpu, s: cuda.Stream, ctl_mod: cuda.Module, kernels: *std.ArrayList(g.Kernel)) !void { @@ -32,9 +34,11 @@ pub const Set = struct { self.penalize = .{ .rig = .{ .gpu = gpu, .s = s }, .draws = &self.draws }; try kernels.append(a, self.penalize.?.kernel()); try self.lfm2m.add(gpu, s, ctl_mod, kernels); + try self.paged.add(gpu, s, ctl_mod, kernels); } pub fn deinit(self: *Set) void { + self.paged.deinit(); self.lfm2m.deinit(); if (self.penalize) |*p| p.deinit(); if (self.sample) |*m| m.unload(); diff --git a/zig/tests/cuda/contract_paged_attn.zig b/zig/tests/cuda/contract_paged_attn.zig new file mode 100644 index 0000000000..ff5fceaa19 --- /dev/null +++ b/zig/tests/cuda/contract_paged_attn.zig @@ -0,0 +1,344 @@ +//! Liquid's paged attention, decode (segs chunk and merge) and prompt rows, over shuffled pages against float64. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); +const attn = @import("contract_attn.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; + +const page = 512; // positions a page holds (paged_attn.cuh's PAGE) +const max_segs = 8; // streams a decode launch takes (paged_attn.cuh's MAXSEG) +const chunk_blocks = 64; // key chunks a decode launch walks in parallel at most (the families' chunk_blocks) +const pad = 64; // elements between a page's values and the next page: a stride that ignores PS reads them + +/// Decode: the chunk and merge kernels' symbols; prompt: pattn_kernel's symbol, rows a block, dynamic shared bytes. +pub const Form = union(enum) { + decode: struct { chunk: [:0]const u8, merge: [:0]const u8 }, + prompt: struct { sym: [:0]const u8, rows: usize, smem: u32 }, +}; + +/// A launch's fixed heads and head dim, its kernels, and whether they come from the controls image. +pub const Cfg = struct { name: []const u8, h: usize, hk: usize, d: usize, form: Form, control: bool = false, served: []const c.Shape }; + +const pattn64 = "_ZN20tf_prefill_attention12pattn_kernelILi64ELi8ELi4ELi8EEEvPK13__nv_bfloat16S3_S3_PS1_iiiiifPKix"; +const decode64: Form = .{ .decode = .{ .chunk = "tf_paged_attn_chunk_segs_64", .merge = "tf_paged_attn_merge_segs_64" } }; +const prompt64: Form = .{ .prompt = .{ .sym = pattn64, .rows = 32, .smem = 8 * 32 * 8 * 16 } }; +const decode_served = [_]c.Shape{ .{ .rows = 1, .n = 32, .k = 3000 }, .{ .rows = 16, .n = 32, .k = 1500 }, .{ .rows = 128, .n = 32, .k = 700 }, .{ .rows = 8, .n = 32, .k = 4000 } }; +const prompt_served = [_]c.Shape{ .{ .rows = 256, .n = 32, .k = 0 }, .{ .rows = 64, .n = 32, .k = 3000 }, .{ .rows = 100, .n = 32, .k = 500 } }; + +pub const cfgs = [_]Cfg{ + // LFM2.5-24B-A2B: 32 query heads over 8 KV heads of 64 + .{ .name = "cuda paged decode tf_paged_attn_chunk_segs_64 + tf_paged_attn_merge_segs_64, 32/8 heads of 64, 512-key pages", .h = 32, .hk = 8, .d = 64, .form = decode64, .served = &decode_served }, + .{ .name = "control: paged decode at head 64 reading its values from the neighbouring lane's dims", .h = 32, .hk = 8, .d = 64, .form = .{ .decode = .{ .chunk = "tf_ctl_paged_attn_chunk_segs_64", .merge = "tf_paged_attn_merge_segs_64" } }, .control = true, .served = &decode_served }, + .{ .name = "cuda paged prompt pattn_kernel<64, 8, 4, 8>, 32/8 heads of 64, 512-key pages", .h = 32, .hk = 8, .d = 64, .form = prompt64, .served = &prompt_served }, + .{ .name = "control: paged prompt rows at head 64 reading their values from the neighbouring 8 dims", .h = 32, .hk = 8, .d = 64, .form = .{ .prompt = .{ .sym = "tf_ctl_pattn_64", .rows = 32, .smem = 8 * 32 * 8 * 16 } }, .control = true, .served = &prompt_served }, +}; + +/// Gates 1 and 3 run attention_ref's bound: decode weights are fp32 through fast exp, prompt P enters the mma as bf16. +const decode_why = "a score's absolute rounding error is its weight's relative error: gates 1 and 3 hold outputs to attention_ref's bound"; + +/// The P x V chain of a row: every key it sees (positions below k, then the rows before it), a tile, then the fold. +fn keyChain(s: c.Shape) usize { + return s.k + s.rows + 64 + 64; +} + +pub fn spec(f: Cfg) c.Spec { + return switch (f.form) { + .decode => .{ .name = f.name, .accum = .f32, .depth = keyChain, .ops = 1, .out = .bf16, .tiles = .{ .{}, .{}, .{ .size = 64 } }, .served = f.served, .sums = .{ .why = decode_why } }, + .prompt => |p| .{ .name = f.name, .accum = .f32, .depth = keyChain, .ops = 1, .out = .bf16, .tiles = .{ .{ .size = p.rows }, .{}, .{ .size = 64, .min = 0 } }, .served = f.served, .sums = .{ .why = attn.why } }, + }; +} + +/// nemotron_attention.cu's AttnSegs (paged_attn.cuh): each stream's caches, position, page table and first row. +const AttnSegs = extern struct { kc: [max_segs]u64, vc: [max_segs]u64, meta: [max_segs]u64, pt: [max_segs]u64, ps: i64, row0: [max_segs]c_int, n: c_int }; + +/// One stream: its rows from row0 on at positions pos0.., and its pages (the first of them at `first` in `ids`). +const Seg = struct { row0: usize, rows: usize, pos0: usize, first: usize, pages: usize }; + +pub const Paged = struct { + rig: rig_.Rig, + f: Cfg, + fns: [2]cuda.Function, // decode: chunk, merge; prompt: the kernel twice + r: ?attn.Ref = null, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + segs: std.ArrayList(Seg) = .empty, + ids: std.ArrayList(i32) = .empty, // every stream's page table, one after another + arena: std.ArrayList(u16) = .empty, // pages of PS elements: page rows of keys, then of values, then pad + q: std.ArrayList(u16) = .empty, + dev: struct { q: u64 = 0, arena: u64 = 0, meta: u64 = 0, ids: u64 = 0, out: u64 = 0, xs: u64 = 0, po: u64 = 0, pm: u64 = 0, pl: u64 = 0 } = .{}, + nch: usize = 0, + + pub fn deinit(self: *Paged) void { + const a = self.rig.gpu.gpa; + if (self.r) |*r| r.deinit(); + inline for (.{ "segs", "ids", "arena", "q" }) |n| @field(self, n).deinit(a); + self.rig.deinit(); + } + + pub fn kernel(self: *Paged) g.Kernel { + return .{ .spec = spec(self.f), .ctx = self, .vt = &vt, .control = if (self.f.control) .shapes else null }; + } + + const vt = g.hooks(Paged); + + fn kvd(self: *const Paged) usize { + return self.f.hk * self.f.d; + } + + fn ps(self: *const Paged) usize { + return 2 * page * self.kvd() + pad; + } + + fn scale(self: *const Paged) f64 { + return @floatCast(@as(f32, 1) / @sqrt(@as(f32, @floatFromInt(self.f.d)))); + } + + /// Row r's stream and position. + fn at(self: *const Paged, r: usize) struct { seg: Seg, p: usize } { + for (self.segs.items) |s| if (r < s.row0 + s.rows) return .{ .seg = s, .p = s.pos0 + r - s.row0 }; + unreachable; + } + + /// Element of key `key`'s row (key head kh) in stream s's pages; values sit page * kvd further on. + fn keyAt(self: *const Paged, s: Seg, key: usize, kh: usize) usize { + const id: usize = @intCast(self.ids.items[s.first + key / page]); + return id * self.ps() + (key % page) * self.kvd() + kh * self.f.d; + } + + fn value(self: *const Paged, s: Seg, key: usize, kh: usize) usize { + return self.keyAt(s, key, kh) + page * self.kvd(); + } + + pub fn prepare(self: *Paged, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + self.rig.reset(); + const f = self.f; + const n = shape.rows; + if (shape.n != f.h or n == 0) return false; + const a = self.rig.gpu.gpa; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + self.shape = shape; + // streams: decode rows split over up to 8 (one under stress), the first ending at k; a prompt chunk is one at k + const nseg: usize = if (f.form == .prompt or fill == .stress) 1 else @min(n, max_segs); + self.segs.clearRetainingCapacity(); + var row0: usize = 0; + var pages: usize = 0; + for (0..nseg) |j| { + const rows = n / nseg + @intFromBool(j < n % nseg); + const pos0 = switch (f.form) { + .prompt => shape.k, + .decode => blk: { + if (shape.k < rows) return false; + const room = shape.k - rows + 1; + break :blk if (j == 0) room - 1 else (j * 7919 + seed % 97) % room; + }, + }; + const need = (pos0 + rows + page - 1) / page; + try self.segs.append(a, .{ .row0 = row0, .rows = rows, .pos0 = pos0, .first = pages, .pages = need }); + row0 += rows; + pages += need; + } + // page ids: every stream's pages shuffled over the arena, two spare pages left random + const total = pages + 2; + try self.ids.resize(a, total); + for (self.ids.items, 0..) |*x, i| x.* = @intCast(i); + rng.shuffle(i32, self.ids.items); + try self.arena.resize(a, total * self.ps()); + for (self.arena.items) |*e| e.* = attn.normal(rng); + try self.q.resize(a, n * f.h * f.d); + for (self.q.items) |*e| e.* = attn.normal(rng); + var most: usize = 1; + var top: usize = 0; + for (self.segs.items) |s| { + most = @max(most, s.pos0 + s.rows); + top = @max(top, s.pos0 + s.rows); + } + self.nch = (top + page - 1) / page + 1; + if (self.r) |*r| r.deinit(); + self.r = null; + self.r = try attn.Ref.init(a, f.d, most); + switch (fill) { + .random => {}, + .nonfinite => self.plant(), + .stress => try self.stress(rng), + } + const metas = try a.alloc(i32, nseg); + defer a.free(metas); + for (metas, self.segs.items) |*m, s| m.* = @intCast(s.pos0); + self.dev.q = try self.rig.input(std.mem.sliceAsBytes(self.q.items)); + self.dev.arena = try self.rig.input(std.mem.sliceAsBytes(self.arena.items)); + self.dev.meta = try self.rig.input(std.mem.sliceAsBytes(metas)); + self.dev.ids = try self.rig.input(std.mem.sliceAsBytes(self.ids.items)); + self.dev.out = try self.rig.output(n * f.h * f.d * 2); + if (f.form == .decode) { + self.dev.xs = try self.rig.output(n * f.h * f.d / 64 * 4); + const parts = n * self.nch * f.h; + self.dev.po = try self.rig.scratch(parts * f.d * 4); + self.dev.pm = try self.rig.scratch(parts * 4); + self.dev.pl = try self.rig.scratch(parts * 4); + } + return true; + } + + /// NaN query (row 0, head 0); +inf and -inf values at row 0's key, NaN where they meet +inf a key earlier; NaN key. + fn plant(self: *Paged) void { + const f = self.f; + const s = self.segs.items[0]; + const p = s.pos0; // row 0's own key: every row of its stream sees it + self.q.items[3] = c.toBf16(std.math.nan(f32)); + self.arena.items[self.value(s, p, 0) + 7] = c.toBf16(std.math.inf(f32)); + self.arena.items[self.value(s, p, 1 % f.hk) + 9] = c.toBf16(-std.math.inf(f32)); + if (p > 0) self.arena.items[self.value(s, p - 1, 1 % f.hk) + 9] = c.toBf16(std.math.inf(f32)); + self.arena.items[self.keyAt(s, p, 2 % f.hk)] = c.toBf16(std.math.nan(f32)); + } + + /// Gate 4: every key of the stream alike (each p is 1); values mirrored over the keys row 0 sees, 0 after them. + fn stress(self: *Paged, rng: std.Random) !void { + const a = self.rig.gpu.gpa; + const f = self.f; + const s = self.segs.items[0]; + const key = try a.alloc(u16, self.kvd()); + defer a.free(key); + for (key) |*e| e.* = attn.normal(rng); + const last = s.pos0 + s.rows; + for (0..last) |k| { + @memcpy(self.arena.items[self.keyAt(s, k, 0)..][0..self.kvd()], key); + @memset(self.arena.items[self.value(s, k, 0)..][0..self.kvd()], 0); + } + const col = try a.alloc(f32, s.pos0 + 1); + defer a.free(col); + for (0..f.hk) |kh| for (0..f.d) |e| { + attn.mirrored(rng, col); + for (col, 0..) |y, k| self.arena.items[self.value(s, k, kh) + e] = c.toBf16(y); + }; + } + + pub fn run(self: *Paged) !void { + const f = self.f; + const n = self.shape.rows; + const g_: usize = f.h / f.hk; + const s = self.rig.s; + switch (f.form) { + .decode => { + var sg = std.mem.zeroes(AttnSegs); + for (self.segs.items, 0..) |x, j| { + sg.kc[j] = self.dev.arena; + sg.vc[j] = self.dev.arena + page * self.kvd() * 2; + sg.meta[j] = self.dev.meta + j * 4; + sg.pt[j] = self.dev.ids + x.first * 4; + sg.row0[j] = @intCast(x.row0); + } + sg.ps = @intCast(self.ps()); + sg.n = @intCast(self.segs.items.len); + var a: cuda.Args = .{}; + a.add(sg); + inline for (.{ "q", "po", "pm", "pl" }) |name| a.add(@field(self.dev, name)); + for ([_]usize{ f.h * f.d, f.h, f.hk, g_, page, self.nch }) |v| a.add(@as(c_int, @intCast(v))); + a.add(@as(f32, @floatCast(self.scale()))); + try cuda.launch.launch(self.fns[0], .{ .grid = .{ .x = @intCast(n), .y = @intCast(f.hk), .z = @intCast(@min(self.nch, chunk_blocks)) }, .block = .{ .x = 256 } }, s, &a); + var b: cuda.Args = .{}; + b.add(sg); + inline for (.{ "po", "pm", "pl", "out", "xs" }) |name| b.add(@field(self.dev, name)); + for ([_]usize{ f.h, f.hk, g_, page, self.nch }) |v| b.add(@as(c_int, @intCast(v))); + try cuda.launch.launch(self.fns[1], .{ .grid = .{ .x = @intCast(n), .y = @intCast(f.hk), .z = 1 }, .block = .{ .x = 256 } }, s, &b); + }, + .prompt => |p| { + const x = self.segs.items[0]; + var a: cuda.Args = .{}; + a.add(self.dev.q); + a.add(self.dev.arena); + a.add(self.dev.arena + page * self.kvd() * 2); + a.add(self.dev.out); + for ([_]usize{ x.pos0, n, f.h, f.hk, g_ }) |v| a.add(@as(c_int, @intCast(v))); + a.add(@as(f32, @floatCast(self.scale()))); + a.add(self.dev.ids + x.first * 4); + a.add(@as(i64, @intCast(self.ps()))); + const hpc = 16 * 8 / p.rows; // a block's query heads: 16 rows a warp, eight warps + try cuda.launch.launch(self.fns[0], .{ .grid = .{ .x = @intCast((n + p.rows - 1) / p.rows), .y = @intCast(f.hk * (g_ / hpc)) }, .block = .{ .x = 256 }, .shared = p.smem }, s, &a); + }, + } + try s.synchronize(); + } + + pub fn count(self: *Paged) usize { + return self.shape.rows * self.f.h * self.f.d; + } + + /// Output head i (row i / h, head i % h): its stream's keys [0, p] through the page table. + pub fn load(self: *Paged, i: usize, hd: *attn.Head) !void { + const f = self.f; + const d = f.d; + const r = i / f.h; + const head = i % f.h; + const kh = head / (f.h / f.hk); + const w = self.at(r); + for (hd.q, 0..) |*x, e| x.* = c.fromBf16(self.q.items[(r * f.h + head) * d + e]); + for (0..w.p + 1) |key| { + const kat = self.keyAt(w.seg, key, kh); + const vat = self.value(w.seg, key, kh); + for (0..d) |e| { + hd.k[key * d + e] = c.fromBf16(self.arena.items[kat + e]); + hd.v[key * d + e] = c.fromBf16(self.arena.items[vat + e]); + } + } + hd.n = w.p + 1; + } + + fn heads(self: *Paged) usize { + return self.shape.rows * self.f.h; + } + + pub fn check(self: *Paged, gate: g.Gate) !?[]const u8 { + return self.r.?.check(self, self.heads(), self.scale(), try self.rig.outputBytes(), gate); + } + + pub fn reference(self: *Paged, want: []f64, mag: []f64) !void { + return self.r.?.reference(self, self.heads(), self.scale(), want, mag); + } + + pub fn simulate(self: *Paged, acc: c.Accum, order: g.Order, out: []f64) !void { + return self.r.?.simulate(self, self.heads(), self.scale(), acc, order, out); + } +}; + +/// The kernels and their controls, registered onto `kernels`; the controls from the contract_controls image `ctl`. +pub const Set = struct { + mods: [2]?cuda.Module = .{ null, null }, + items: [cfgs.len]Paged = undefined, + n: usize = 0, + + pub fn add(self: *Set, gpu: check.Gpu, s: cuda.Stream, ctl: cuda.Module, kernels: *std.ArrayList(g.Kernel)) !void { + self.mods[0] = try cuda.Module.load(gpu.d, cuda.kernels.nemotron_attention); + self.mods[1] = try cuda.Module.load(gpu.d, cuda.kernels.prefill_attention); + const segs = self.mods[0].?; + const prompt = self.mods[1].?; + for (cfgs, &self.items) |f, *x| { + const fns: [2]cuda.Function = switch (f.form) { + .decode => |k| .{ try (if (f.control) ctl else segs).function(k.chunk), try segs.function(k.merge) }, + .prompt => |k| blk: { + const fn_ = try (if (f.control) ctl else prompt).function(k.sym); + try fn_.allowDynamicShared(k.smem); + break :blk .{ fn_, fn_ }; + }, + }; + x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = f, .fns = fns }; + self.n += 1; + } + for (self.items[0..self.n]) |*x| try kernels.append(gpu.gpa, x.kernel()); + } + + pub fn deinit(self: *Set) void { + for (self.items[0..self.n]) |*x| x.deinit(); + for (&self.mods) |*m| if (m.*) |*x| x.unload(); + } +}; + +test "the paged forms' geometry: 4 query heads a KV head, AttnSegs as nemotron_attention.cu pads it" { + try std.testing.expectEqual(@as(usize, 304), @sizeOf(AttnSegs)); + for (cfgs) |f| try std.testing.expectEqual(@as(usize, 4), f.h / f.hk); + try std.testing.expectEqual(@as(usize, 32), prompt64.prompt.rows); +} From 68b48aaafb066915c51da3102f0827e6acbea30e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 28 Oct 2026 09:32:42 +0100 Subject: [PATCH 07/12] tools: LFM2.5-24B-A2B's float64 reference on the CPU, and its receipt of every top-1 miss ref64.py dequantizes the MLX 4-bit weights exactly and runs the forward in float64; --bf16 adds the engine's bf16 storage points. It checks the CUDA engine's greedy picks on teacher-forced prompt rows and on reply rows; REF64.md lists each miss with the prompt, the position, both tokens and both margins. --- tools/lfm2_moe/REF64.md | 65 ++++++++++ tools/lfm2_moe/ref64.py | 254 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 319 insertions(+) create mode 100644 tools/lfm2_moe/REF64.md create mode 100644 tools/lfm2_moe/ref64.py diff --git a/tools/lfm2_moe/REF64.md b/tools/lfm2_moe/REF64.md new file mode 100644 index 0000000000..3f5f4be85b --- /dev/null +++ b/tools/lfm2_moe/REF64.md @@ -0,0 +1,65 @@ +# Float64 reference receipt: LFM2.5-24B-A2B on CUDA + +- Checkpoint: [LiquidAI/LFM2.5-24B-A2B-MLX-4bit](https://huggingface.co/LiquidAI/LFM2.5-24B-A2B-MLX-4bit) at revision + `667408d494bde191c2cb9dd4f4549827816b9275`. +- Engine: this PR's serving commit, before the rebase onto main (`c6645063a`). The rebased head serves the same + tokens: the token hashes of 8 prompts and of a chat conversation are equal. + Served on a GB10 (sm_121), greedy, drafts off, `--parallel 8 --context 16384`. +- Reference: `ref64.py` on the CPU in float64 (Python 3.12, torch 2.14, safetensors 0.8, tokenizers 0.23), 2026-10-28. + +Reproduce (the reference pass takes about 10 minutes and 17 GB of host memory): + +``` +zig-out/native/bin/tensorfold-native serve MODEL_DIR --port 8080 --parallel 8 --context 16384 +python -I tools/lfm2_moe/ref64.py collect http://127.0.0.1:8080 MODEL_DIR engine.json +python -I tools/lfm2_moe/ref64.py MODEL_DIR engine.json --bf16 +``` + +Rows. A prompt row is the engine's one-token greedy pick after the prefix [0, t] of a prompt, one request per prefix. +A reply row is a token of the engine's 32-token greedy reply, teacher-forced on that reply. `t` is the logits row after +token t, so prompt row P - 1 and reply token 0 are the same position. The prompts (`prompts()` in the script) are +capitals (13 tokens), code (19), a chat turn (19) and items (384, a list of 30 items and a question). + +The float64 pass dequantizes the MLX weights exactly. The `--bf16` pass adds the engine's bf16 storage points (see the +script's docstring). It shows which side of a miss rounding alone lands on; it is not bit-exact. + +## Result + +- f64: prompt rows 421/435, reply rows 123/128 +- bf16: prompt rows 424/435, reply rows 123/128 + +Float64 top-1 misses (margin = f64 logit of its top minus f64 logit of the engine's token; bf16 margin = the same difference in the bf16-points pass, < 0 means it prefers the engine's token): + +| case | row | t | engine | float64 | f64 margin | bf16 margin | bf16 picks | +|---|---|---|---|---|---|---|---| +| capitals | prompt | 1 | 589 "'s" | 124901 "" | 0.3389 | -0.1250 (pre-round -0.1026) | engine's | +| capitals | prompt | 3 | 1394 " each" | 278 " the" | 0.0704 | -1.0625 (pre-round -1.0871) | engine's | +| capitals | reply | 39 | 278 " the" | 17755 " Greece" | 0.0095 | +0.0000 (pre-round -0.0276) | engine's | +| code | prompt | 7 | 1288 " return" | 21744 " \"\"\"" | 0.3008 | +0.2500 (pre-round +0.2923) | f64's | +| code | reply | 26 | 4876 "_m" | 71 "_" | 0.3628 | +0.3125 (pre-round +0.2956) | f64's | +| code | reply | 35 | 73 "a" | 2762 "add" | 0.0687 | -0.3750 (pre-round -0.4420) | engine's | +| chat | prompt | 13 | 90978 " Explanation" | 3103 " \n" | 0.7478 | -0.8125 (pre-round -0.7755) | engine's | +| chat | prompt | 15 | 25 "1" | 3103 " \n" | 0.1301 | +0.0625 (pre-round +0.0216) | f64's | +| items | prompt | 3 | 22 "." | 34 ":" | 0.9143 | +0.8750 (pre-round +0.8393) | f64's | +| items | prompt | 5 | 2719 " problem" | 3451 " question" | 0.1379 | +0.0625 (pre-round +0.0963) | f64's | +| items | prompt | 6 | 9402 " stone" | 24062 " brick" | 0.1680 | +1.0000 (pre-round +0.9824) | f64's | +| items | prompt | 8 | 229 " " | 267 " a" | 0.0823 | -0.1250 (pre-round -0.0781) | engine's | +| items | prompt | 314 | 27 "3" | 1035 "14" | 0.1292 | -0.3750 (pre-round -0.3704) | engine's | +| items | prompt | 338 | 25 "1" | 26 "2" | 0.3699 | +0.3750 (pre-round +0.4515) | f64's | +| items | prompt | 371 | 437 " are" | 296 " in" | 0.0565 | -0.6250 (pre-round -0.5979) | engine's | +| items | prompt | 375 | 342 " (" | 8187 " item" | 0.3365 | -1.1250 (pre-round -1.1544) | engine's | +| items | prompt | 377 | 342 " (" | 41501 " Item" | 0.6276 | -2.3750 (pre-round -2.3544) | engine's | +| items | reply | 387 | 41501 " Item" | 342 " (" | 0.7665 | +0.6250 (pre-round +0.5957) | f64's | +| items | reply | 396 | 2959 " red" | 6542 " blue" | 0.2816 | +0.5000 (pre-round +0.5314) | f64's | + +bf16-points misses where float64 agrees with the engine (margin = bf16 logit of its pick minus the engine's token; f64 margin = how far float64 prefers the engine's token): + +| case | row | t | engine | bf16 pick | bf16 margin | f64 margin | +|---|---|---|---|---|---|---| +| capitals | prompt | 10 | 10125 " Italy" | 6398 " Germany" | 0.0000 (pre-round -0.0632) | 0.0079 | +| code | reply | 44 | 1113 " add" | 1619 " main" | 0.6250 (pre-round +0.5908) | 0.1500 | +| code | reply | 47 | 71 "_" | 65296 "_sum" | 0.1250 (pre-round +0.1222) | 0.3751 | +| items | prompt | 7 | 22 "." | 20 "," | 0.0000 (pre-round +0.0882) | 0.0279 | +| items | prompt | 9 | 30 "6" | 27 "3" | 0.5000 (pre-round +0.4381) | 0.1054 | +| items | prompt | 36 | 440 " The" | 41501 " Item" | 0.1250 (pre-round +0.1183) | 0.1051 | +| items | prompt | 379 | 8187 " item" | 229 " " | 0.2500 (pre-round +0.2420) | 0.0420 | diff --git a/tools/lfm2_moe/ref64.py b/tools/lfm2_moe/ref64.py new file mode 100644 index 0000000000..bce94597dc --- /dev/null +++ b/tools/lfm2_moe/ref64.py @@ -0,0 +1,254 @@ +#!/usr/bin/env python3 +"""Float64 CPU reference of LFM2.5-24B-A2B (lfm2_moe) from its MLX affine checkpoint, checked against the CUDA +engine's greedy tokens. + +The stored weights are dequantized exactly (q * scale + bias) and every op runs in float64. The engine's greedy +tokens come from a running server (`collect`): + prompt rows the one-token greedy pick after each prefix [0, t] of a prompt (a fresh request per prefix) + reply rows a 32-token greedy reply, checked teacher-forced on the engine's own reply +Row t is the logits row after token t, so prompt row P - 1 and reply token 0 are the same position. + +--bf16 adds a second pass with the engine's bf16 storage points (float64 inside every op, round-to-nearest-even to +bf16 where the CUDA kernels store bf16): embed, residual, every RMSNorm (bf16(w * bf16(h * inv))), projection outputs, +q/k norm and RoPE, conv B*x / acc / C*y, SwiGLU, expert outputs, MoE combine and the logits. Prompt rows take the +prefill path (weights rounded once to bf16, attention P rounded before PV, expert down stored bf16), decode rows the +exact-weight GEMV path (fp32 attention P, fp32 expert down). The router stays unrounded. Softmax P is taken against +the row's final max (the kernels' running max over 64-key tiles is not replayed). Greedy over bf16 logits breaks ties +to the lower id, as the engine does. This pass shows which side of a miss rounding alone lands on; it is not bit-exact. + +usage: + ref64.py collect URL MODEL_DIR ENGINE_JSON the greedy rows of prompts() from a server at URL + ref64.py MODEL_DIR ENGINE_JSON [--bf16] prints the miss table (markdown) and the agreement totals +""" +import json +import math +import sys +import urllib.request + +import torch +from safetensors import safe_open + +torch.set_default_dtype(torch.float64) +F64 = torch.float64 + + +def complete(url, ids, n): + body = {"prompt": ids, "max_tokens": n, "temperature": 0, "ignore_eos": True, "draft": False, + "return_token_ids": True} + req = urllib.request.Request(url + "/v1/completions", data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=900) as r: + out = json.loads(r.read()) + return out.get("tensorfold", out)["token_ids"] + + +def prompts(n_items): + """The receipt's prompt set: texts, encoded with the checkpoint's tokenizer.json after BOS.""" + colors = "red green blue amber violet".split() + items = " ".join(f"Item {i}: the {colors[i % 5]} box holds {7 * i % 23} stones." for i in range(n_items)) + return [("capitals", "The capital of France is Paris. The capital of Germany is"), + ("code", "def add(a, b):\n return a + b\n\ndef mul(a, b):\n"), + ("chat", "<|im_start|>user\nExplain why the sky is blue in two sentences.<|im_end|>\n" + "<|im_start|>assistant\n"), + ("items", items + " Question: how many stones does item 3 hold? Answer: item 3 holds")] + + +def collect(url, model_dir, out_path, reply_len=32): + from tokenizers import Tokenizer + tok = Tokenizer.from_file(f"{model_dir}/tokenizer.json") + bos = json.load(open(f"{model_dir}/config.json"))["bos_token_id"] + cases = [] + for name, text in prompts(30): + ids = [bos] + tok.encode(text).ids + picks = [complete(url, ids[: t + 1], 1)[0] for t in range(len(ids))] + cases.append({"name": name, "prompt": ids, "prompt_picks": picks, "reply": complete(url, ids, reply_len)}) + print(f"{name}: {len(ids)} prompt rows, {len(cases[-1]['reply'])} reply tokens", file=sys.stderr) + json.dump({"cases": cases}, open(out_path, "w")) + + +def bf16(x): + """Nearest bf16 (8 significant bits, ties to even), kept as float64: exact, no float32 double rounding.""" + m, e = torch.frexp(x) + return torch.ldexp(torch.round(m * 256.0), (e - 8).to(torch.int32)).to(F64) + + +def same(x): + return x + + +class Model: + def __init__(self, d): + self.d = d + self.cfg = json.load(open(f"{d}/config.json")) + self.map = json.load(open(f"{d}/model.safetensors.index.json"))["weight_map"] + self.files = {} + q = self.cfg["quantization"] + self.embed_bits = q.get("model.embed_tokens", q)["bits"] + + def raw(self, k): + f = self.map[k] + if f not in self.files: + self.files[f] = safe_open(f"{self.d}/{f}", "pt") + return self.files[f].get_tensor(k) + + def deq(self, prefix, bits=4, rows=None): + """MLX affine: codes are a little-endian bitstream over uint32 words, value i at bits [i*bits, (i+1)*bits).""" + w, s, b = (self.raw(prefix + n) for n in (".weight", ".scales", ".biases")) + if rows is not None: + w, s, b = w[rows], s[rows], b[rows] + by = w.contiguous().view(torch.uint8).to(torch.int64) + if 32 % bits == 0: + words = by.reshape(*by.shape[:-1], -1, 4) << torch.tensor([0, 8, 16, 24]) + words = words.sum(-1, keepdim=True) + q = ((words >> (bits * torch.arange(32 // bits))) & ((1 << bits) - 1)).flatten(-2) + else: + nb = by.shape[-1] * 8 + bit = ((by.unsqueeze(-1) >> torch.arange(8)) & 1).reshape(*by.shape[:-1], nb) + q = (bit[..., : nb // bits * bits].reshape(*by.shape[:-1], -1, bits) << torch.arange(bits)).sum(-1) + q = q.to(F64).reshape(*q.shape[:-1], s.shape[-1], -1) + return (q * s.to(F64)[..., None] + b.to(F64)[..., None]).flatten(-2) + + def forward(self, ids, P, bf): + """Logits [T, V] (float64, not rounded). bf = same: plain float64; bf = bf16: the engine's storage points with + rows [0, P) on the prefill path and the rest on the decode path.""" + cfg, T = self.cfg, len(ids) + H, HK, hd = cfg["num_attention_heads"], cfg["num_key_value_heads"], 64 + silu = torch.nn.functional.silu + + def lin(x, W, p=P): # prefill rows multiply bf16(W), decode rows the exact W; fp64 accumulation + return torch.cat([x[:p] @ bf(W).T, x[p:] @ W.T]) + + def rms(h, w): # the engine's form: bf16(w * bf16(h * inv)), inv from the bf16 row + return bf(w.to(F64) * bf(h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + cfg["norm_eps"]))) + + inv = 1.0 / cfg["rope_parameters"]["rope_theta"] ** (torch.arange(0, hd, 2, dtype=F64) / hd) + ang = torch.arange(T, dtype=F64)[:, None] * inv[None] + cos, sin = torch.cat([ang.cos()] * 2, -1), torch.cat([ang.sin()] * 2, -1) + + def rot(x): + return x * cos + torch.cat([-x[..., hd // 2:], x[..., : hd // 2]], -1) * sin + + future = torch.triu(torch.ones(T, T, dtype=torch.bool), 1) + h = bf(self.deq("model.embed_tokens", self.embed_bits, rows=torch.tensor(ids))) + for L, kind in enumerate(cfg["layer_types"]): + p = f"model.layers.{L}." + x = rms(h, self.raw(p + "operator_norm.weight")) + if kind == "conv": + B, C, xx = bf(lin(x, self.deq(p + "conv.in_proj"))).chunk(3, -1) + bx = bf(B * xx) + w = self.raw(p + "conv.conv.weight").to(F64).reshape(-1, 3) # tap 2 is the current row + pad = torch.cat([torch.zeros(2, bx.shape[1]), bx]) + y = bf(C * bf(pad[0:T] * w[:, 0] + pad[1:T + 1] * w[:, 1] + pad[2:T + 2] * w[:, 2])) + o = bf(lin(y, self.deq(p + "conv.out_proj"))) + else: + q, k, v = (bf(lin(x, self.deq(p + f"self_attn.{n}_proj"))) for n in "qkv") + q = bf(rot(rms(q.reshape(T, H, hd), self.raw(p + "self_attn.q_layernorm.weight")).transpose(0, 1))) + k = bf(rot(rms(k.reshape(T, HK, hd), self.raw(p + "self_attn.k_layernorm.weight")).transpose(0, 1))) + k = k.repeat_interleave(H // HK, 0) + v = v.reshape(T, HK, hd).transpose(0, 1).repeat_interleave(H // HK, 0) + s = ((q @ k.transpose(1, 2)) / math.sqrt(hd)).masked_fill(future, -math.inf) + e = torch.exp(s - s.max(-1, keepdim=True).values) + pv = torch.cat([bf(e[:, :P]), e[:, P:]], 1) # prefill rounds P before the PV product, decode keeps it + a = bf((pv @ v) / e.sum(-1, keepdim=True)) + o = bf(lin(a.transpose(0, 1).reshape(T, H * hd), self.deq(p + "self_attn.out_proj"))) + h = bf(h + o) + x = rms(h, self.raw(p + "ffn_norm.weight")) + f = p + "feed_forward." + if L < cfg["num_dense_layers"]: + g, u = bf(lin(x, self.deq(f + "gate_proj"))), bf(lin(x, self.deq(f + "up_proj"))) + o = bf(lin(bf(bf(silu(g)) * u), self.deq(f + "down_proj"))) + else: + sc = torch.sigmoid(x @ self.raw(f + "gate.weight").to(F64).T) + pick = torch.topk(sc + self.raw(f + "expert_bias").to(F64), cfg["num_experts_per_tok"], -1).indices + wts = sc.gather(-1, pick) + wts = wts / (wts.sum(-1, keepdim=True) + 1e-6) * cfg["routed_scaling_factor"] + acc = torch.zeros_like(x) + m = f + "switch_mlp." + for ex in pick.unique().tolist(): + sel = pick == ex + rr = sel.any(-1).nonzero().squeeze(-1) # ascending: its first pe rows are prefill rows + pe = int((rr < P).sum()) + g = bf(lin(x[rr], self.deq(m + "gate_proj", rows=ex), pe)) + u = bf(lin(x[rr], self.deq(m + "up_proj", rows=ex), pe)) + y = lin(bf(bf(silu(g)) * u), self.deq(m + "down_proj", rows=ex), pe) + acc[rr] += (wts * sel).sum(-1)[rr, None] * torch.cat([bf(y[:pe]), y[pe:]]) + o = bf(acc) + h = bf(h + o) + print(f" layer {L} {kind}", file=sys.stderr, flush=True) + x = rms(h, self.raw("model.embedding_norm.weight")) + V = cfg["vocab_size"] + return torch.cat([x @ self.deq("model.embed_tokens", self.embed_bits, rows=torch.arange(a, min(a + 8000, V))).T + for a in range(0, V, 8000)], -1) + + +def greedy(lg): + """argmax in (value desc, id asc) order, as the engine picks over its bf16 logits.""" + return int((lg == lg.max()).nonzero()[0]) + + +def main(model_dir, engine_path, with_bf16): + mdl = Model(model_dir) + try: + from tokenizers import Tokenizer + tok = Tokenizer.from_file(f"{model_dir}/tokenizer.json") + + def show(t): + return f"{t} {json.dumps(tok.decode([t], skip_special_tokens=False))}".replace("|", "\\|") + except ImportError: + show = str + cases = json.load(open(engine_path))["cases"] + tot = {(m, k): [0, 0] for m in ("f64", "bf16") for k in ("prompt", "reply")} + f64_rows, bf16_rows = [], [] + for c in cases: + P = len(c["prompt"]) + ids = c["prompt"] + c["reply"][:-1] + print(f"{c['name']}: {len(ids)} rows", file=sys.stderr, flush=True) + lf = mdl.forward(ids, P, same) + lb_pre = mdl.forward(ids, P, bf16) if with_bf16 else None + lb = bf16(lb_pre) if with_bf16 else None + rows = [("prompt", t, tk) for t, tk in enumerate(c["prompt_picks"])] + rows += [("reply", P - 1 + j, tk) for j, tk in enumerate(c["reply"])] + for kind, t, eng in rows: + top = int(lf[t].argmax()) + tot["f64", kind][0] += top == eng + tot["f64", kind][1] += 1 + bcol = "" + if with_bf16: + bp = greedy(lb[t]) + tot["bf16", kind][0] += bp == eng + tot["bf16", kind][1] += 1 + if top != eng: + if with_bf16: + d, dp = lb[t, top] - lb[t, eng], lb_pre[t, top] - lb_pre[t, eng] + side = "engine's" if bp == eng else "f64's" if bp == top else f"other {show(bp)}" + bcol = f" {d:+.4f} (pre-round {dp:+.4f}) | {side} |" + f64_rows.append(f"| {c['name']} | {kind} | {t} | {show(eng)} | {show(top)} | " + f"{lf[t, top] - lf[t, eng]:.4f} |{bcol}") + elif with_bf16 and bp != eng: + bf16_rows.append(f"| {c['name']} | {kind} | {t} | {show(eng)} | {show(bp)} | " + f"{lb[t, bp] - lb[t, eng]:.4f} (pre-round {lb_pre[t, bp] - lb_pre[t, eng]:+.4f}) | " + f"{lf[t, eng] - lf[t, bp]:.4f} |") + mode = [m for m in ("f64", "bf16") if m == "f64" or with_bf16] + for m in mode: + print(f"{m}: prompt rows {tot[m, 'prompt'][0]}/{tot[m, 'prompt'][1]}, " + f"reply rows {tot[m, 'reply'][0]}/{tot[m, 'reply'][1]}") + print("\nFloat64 top-1 misses (margin = f64 logit of its top minus f64 logit of the engine's token" + + ("; bf16 margin = the same difference in the bf16-points pass, < 0 means it prefers the engine's token" + if with_bf16 else "") + "):\n") + print("| case | row | t | engine | float64 | f64 margin |" + (" bf16 margin | bf16 picks |" if with_bf16 else "")) + print("|---|---|---|---|---|---|" + ("---|---|" if with_bf16 else "")) + print("\n".join(f64_rows) or "(none)") + if with_bf16: + print("\nbf16-points misses where float64 agrees with the engine (margin = bf16 logit of its pick minus the " + "engine's token; f64 margin = how far float64 prefers the engine's token):\n") + print("| case | row | t | engine | bf16 pick | bf16 margin | f64 margin |\n|---|---|---|---|---|---|---|") + print("\n".join(bf16_rows) or "(none)") + + +if __name__ == "__main__": + if len(sys.argv) == 5 and sys.argv[1] == "collect": + collect(*sys.argv[2:]) + elif len(sys.argv) in (3, 4) and sys.argv[3:] in ([], ["--bf16"]): + main(sys.argv[1], sys.argv[2], len(sys.argv) == 4) + else: + sys.exit(__doc__) From 0a9de29092b6911152deb6a02e615a90128d753c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 28 Oct 2026 09:32:42 +0100 Subject: [PATCH 08/12] docs: LFM2.5-24B-A2B on CUDA, its receipt and the weights' licence The recipe gets a CUDA section (build, serve, revision, GB10 receipt, licence), the README row names one NVIDIA GPU beside the Mac, and the contract's coverage lists LFM2.5's CUDA kernels. --- README.md | 2 +- docs/kernel-contract.md | 7 ++++++ docs/recipes/lfm2.5-24b-a2b.md | 40 +++++++++++++++++++++++++++++++++- 3 files changed, 47 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 8925e38b8c..59c0d83c56 100644 --- a/README.md +++ b/README.md @@ -173,7 +173,7 @@ Conversions are explicitly marked: their upstream weights need the linked conver | GLM-5.3-Flash, FP4-Local | `tensorfold pull nvidia/GLM-5.3-Flash-FP4-Local@76348ddb2be2e35e39ecfc00e75dc07703f4f3d5` | Two GB10s at TP=2, or three at TP=3 over a ring of direct links: @Bizuayeu loaded it as stored and served it at both sizes. 165.3 GiB of weights, 13% less than NVFP4's 190.4 GiB. | [Native CUDA](docs/recipes/glm-5.3-flash-cuda.md), in 1.1.5 (#1050). NVIDIA's mixed-precision checkpoint: NVFP4 routed experts with compressed block scales, MXFP8 and 128×128-block FP8 for most of the rest. As close to BF16 as NVFP4 at its smaller size: at TP=2 its NLL is 1.11% above BF16's against NVFP4's 1.37% ([contributed receipt](docs/receipts/glm-fp4local-80517fc3/README.md) on #1050's head, not rerun on main), a difference the per-document intervals in [XSFP4's card](https://huggingface.co/Bizuayeu/GLM-5.3-Flash-XSFP4#measured) do not resolve. | | GLM-5.3-Flash, XSFP4 | `tensorfold pull Bizuayeu/GLM-5.3-Flash-XSFP4@88736f916f8939607fde8a4b881dcae749f7551d` | Two GB10s at TP=2, or three at TP=3 over a ring of direct links: on main `8086c205c` @Bizuayeu loaded it in 370.7 s at TP=2 (rank 0's model 96.97 GiB at a 524,288-token window) and in 313.8 s at TP=3 (63.60 GiB at 1,048,576), and both served. 165.7 GiB of weights, 0.4 GiB more than FP4-Local's. | [Native CUDA](docs/recipes/glm-5.3-flash-cuda.md), in 1.1.5 through FP4-Local's path. FP4-Local with its routed experts' NVFP4 blocks fitted again from the BF16 weights, a smoothing folded into them, so it runs on the same kernels. Closer to BF16 than FP4-Local and NVFP4 at TP=2 by all three measures in [its card](https://huggingface.co/Bizuayeu/GLM-5.3-Flash-XSFP4#measured), on #1050's code: top-1 agreement 93.26%, NLL 0.85% above and a KL bound of 0.0431, against FP4-Local's 92.82%, 1.11% and 0.0510. | | GLM-5.3, EXLR8 `home` | `tensorfold pull kindlingai/glm-5.3-exlr8-k2-k3-k4@f459f5d0272a529403dd9b3fd32f771f4250a0d4 --exlr8 home --node 0/2 --dir ~/models/glm-5.3-exlr8-home` | Two M5 Ultra Macs, 256 GB each; about 160 GB resident per Mac at 8k context / four streams. | [Two-Mac recipe](docs/glm-5.3-exlr8.md). Use `--node 1/2` on the other Mac; Thunderbolt 5 and MCDMA required. No CUDA backend. | -| LFM2.5-24B-A2B, MLX 4 / 6 / 8-bit | `tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-4bit@667408d494bde191c2cb9dd4f4549827816b9275`
`tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-6bit@1deec53eb49571e21cdcbf82fd5839248f07ed88`
`tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-8bit@67ab417dd02d0c9545802a726df668620eae049d` | One Mac; M4 Pro checked. | [Native recipe](docs/recipes/lfm2.5-24b-a2b.md), [release history](CHANGELOG.md). | +| LFM2.5-24B-A2B, MLX 4 / 6 / 8-bit | `tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-4bit@667408d494bde191c2cb9dd4f4549827816b9275`
`tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-6bit@1deec53eb49571e21cdcbf82fd5839248f07ed88`
`tensorfold pull LiquidAI/LFM2.5-24B-A2B-MLX-8bit@67ab417dd02d0c9545802a726df668620eae049d` | One Mac (M4 Pro checked), or one NVIDIA GPU with the 4-bit export (GB10 checked). | [Native recipe](docs/recipes/lfm2.5-24b-a2b.md) ([CUDA](docs/recipes/lfm2.5-24b-a2b.md#cuda)), [release history](CHANGELOG.md). Weights: LFM Open License v1.0. | | Qwen3.6-35B-A3B, MLX with MTP | `tensorfold pull TensorFold/Qwen3.6-35B-A3B-MLX-4bit-MTP@81169a9bc511a27c1b4eedb77a2cd98ced431847` | One GB10; 20.9 GB checkpoint. | [Qwen3.6](docs/recipes/qwen3.6-moe.md). Native CUDA; the recipe's Metal path belongs to Python. | | Gemma 4 31B, MLX 4-bit | `tensorfold pull mlx-community/gemma-4-31B-it-4bit@696d436c404745a59f30e4939a658162b0a9e57f`
`tensorfold pull z-lab/gemma-4-31B-it-DFlash@eabd648301ce28583cc14757912e5e0f84e152e1` | One M5 Mac; M5 Ultra checked. M1-M4 refused. | [Native 31B](docs/recipes/gemma-4-native.md). Name the drafter with `--drafter`; images with `--vision`. | | Gemma 4 12B, converted MLX 4-bit | `hf download google/gemma-4-12B-it --revision 707f0a3b8a3c7ad586ed01e27eafbad8a27dd0f7 --local-dir ~/models/gemma-4-12b-hf`
`hf download google/gemma-4-12B-it-assistant --revision 46d4c6f13f0ac0ad827b915669b8df9b81c64c51 --local-dir ~/models/gemma-4-12b-assistant` | One M5 Mac; converted text weights 6.3 GB, image weights 100 MB, plus assistant/caches. M1-M4 refused. | [Convert and serve](docs/recipes/gemma-4-12b-native.md). Upstream download needs `hf`; conversion needs Python, mlx and mlx-lm. Serving is native. | diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 1ed43ac146..972079f809 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -190,6 +190,13 @@ GLM's Metal window projections register narrow and wide rows at 2, 3, 4, 5, 6 an Independent FP64 linear cases accompany byte replacements against frozen pre-window-optimization kernels. The gathered SwiGLU keeps its predecessor's clipped BF16 output and nonfinite rules; gate/up taps qualify its dots separately. Controls dispatch real source mutations that couple rows or select the wrong expert; host tests validate their witnesses. +CUDA also checks LFM2.5-24B-A2B's kernels: the router's fp32 logits (`tf_lfm2m_route`) and its top-4 rule, Liquid's +3-tap short conv with its per-row state snapshots and the state commit (`liquid_conv.cuh`), each against the host's +replay of its fp32 and bf16 steps (three controls: ties to the higher expert, bx unrounded, every stream committing into +the first stream's slot). Their paged attention at head dim 64 runs over shuffled 512-key pages with several streams +a launch: the decode rows' chunks and merge (`tf_paged_attn_chunk_segs_64`, `tf_paged_attn_merge_segs_64`) and the +prompt rows (`pattn_kernel<64, 8, 4, 8>`), head by head within `attention_ref`'s float64 bound (two controls: the +decode reading its values from the neighbouring lane's dims, the prompt from the neighbouring 8 dims). CUDA also checks the affine-3 experts' split decode pair, bit-equal to the pair kernel it replaces (two controls, one dropping a last round of K groups short of the warps, one chaining the groups in reverse). CUDA also checks Flash Next's NVFP4 experts on prompt tiles with the decode form's sums (SwiGLU gate-up and fp32 diff --git a/docs/recipes/lfm2.5-24b-a2b.md b/docs/recipes/lfm2.5-24b-a2b.md index 7b8b2cdf4b..1df62b6c51 100644 --- a/docs/recipes/lfm2.5-24b-a2b.md +++ b/docs/recipes/lfm2.5-24b-a2b.md @@ -1,4 +1,4 @@ -# LFM2.5-24B-A2B, native Metal +# LFM2.5-24B-A2B, native Metal and CUDA This recipe covers [`LiquidAI/LFM2.5-24B-A2B-MLX-4bit`](https://huggingface.co/LiquidAI/LFM2.5-24B-A2B-MLX-4bit) at revision `667408d494bde191c2cb9dd4f4549827816b9275`. The numbers below are measured on this 4-bit export only. @@ -41,3 +41,41 @@ the shared-round kernels on it. The M4 Pro is not yet measured on a quiet machin Every reply was token for token equal before and after, and equal to its solo run and to `"draft": false`. Cold prefill at 2k, 8k and 32k tokens and peak memory did not change. + +## CUDA + +`lfm2_moe` also serves the 4-bit export above (same revision, `667408d494bde191c2cb9dd4f4549827816b9275`) on one +NVIDIA GPU, through the lane host. Checked on a GB10 (sm_121). + +```bash +zig build native -Doptimize=ReleaseFast -Dnvcc=/usr/local/cuda/bin/nvcc +zig-out/native/bin/tensorfold-native serve "$HOME/models/lfm2.5-24b-a2b-4bit" --port 8080 --parallel 8 +``` + +The default window is 32,768 tokens (`--context`), and `--parallel` holds up to 128 streams if the memory budget fits +them. Each stream has fixed pages for its keys and values. The family reads MLX affine 4-bit weights in groups of +64, with the embedding at 4, 6 or 8 bits. The 6- and 8-bit exports are refused by name and stay on Metal. Drafts +are the lane core's copy drafts from the prompt's own text; there is no draft head yet. + +Most launches are shared CUDA operators: the affine linears and experts, the norms, `qk_rows`, the MoE combine, the +pick and the draws. The family adds three kernels: the router (fp32 weights, sigmoid plus bias, top-4) in +`lfm2_moe.cu`, and the 3-tap short conv and its state commit in `liquid_conv.cuh`. It also adds two instances of shared +templates at head dim 64: `pattn_kernel<64, 8, 4, 8>` for prompt rows, and `tf_paged_attn_chunk_segs_64` with +`tf_paged_attn_merge_segs_64` (`paged_attn_segs.cuh`) for decode rows. All of them are under the kernel contract +(see [the contract](../kernel-contract.md#coverage)). + +Receipt (GB10, nvcc 13.0, `--parallel 8`, greedy): + +- drafted == plain and concurrent == solo on 8 of 8 prompts, and the token hashes are the same after a restart; +- a float64 forward of the same weights on the CPU (`tools/lfm2_moe/ref64.py`) picks the engine's greedy token on + 421 of 435 teacher-forced prompt rows and 123 of 128 reply rows; `tools/lfm2_moe/REF64.md` lists every miss with + both tokens and margins; +- decode, total tok/s: 87 at 1 stream, 194 at 4, 303 at 8; +- cold TTFT: 0.32 s at 2.8k prompt tokens, 1.32 s at 11.2k. + +## Licence + +The weights are under Liquid AI's LFM Open License v1.0 (`license: other`, `lfm1.0` on the model card; the text is the +`LICENSE` file in the checkpoint). Use, changes and redistribution are allowed; commercial use is licensed only to +an organisation under US$10 million in annual revenue, and a non-profit may use it for research or other +non-commercial purposes regardless. Read the licence before serving it commercially. From b4f4f49316147e3568ec79ebeaac927e2ec477cc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 27 Oct 2026 20:54:09 +0100 Subject: [PATCH 09/12] lfm2_moe: the attention caches in pages of the core pool, every stream leasing The kernels already read keys and values through a slot's page table, so the pool changes only what the table holds and how large the arena is; no kernel changes. A page is 512 positions of the ten attention layers' keys and values (10 MiB, 20 KiB a token); the conv states stay a fixed part a slot (240 KiB). - Loaded.layout with engine_paged, as Dory's: the plan sets a window of pages aside for the first stream; each stream leases its prompt, its reply and one verify window at its prompt pass and gives them back when it ends. A request that cannot lease waits (LaneMemory); fits is the host's admission. - Model.attach makes the arena and the tables and conv states of the slots the plan names; forward refuses rows past a slot's pages. The idle server logs the pool's census. - TENSORFOLD_LFM2_KV_PAGES=0 keeps every stream's own window of pages, made at open. Served on the GB10 at --parallel 8, 16k: drafted == plain and concurrent == solo, also with 8 long prompts on a 52-page pool where the last waits for pages; the replies equal the fixed pages' ones. --- zig/src/families/lfm2_moe/cuda_backend.zig | 94 +++++++++++-- zig/src/families/lfm2_moe/cuda_forward.zig | 1 + zig/src/families/lfm2_moe/cuda_model.zig | 146 +++++++++++++++++++-- zig/src/families/lfm2_moe/cuda_native.zig | 43 +++++- 4 files changed, 256 insertions(+), 28 deletions(-) diff --git a/zig/src/families/lfm2_moe/cuda_backend.zig b/zig/src/families/lfm2_moe/cuda_backend.zig index 2ebc676d1d..75aa1fc183 100644 --- a/zig/src/families/lfm2_moe/cuda_backend.zig +++ b/zig/src/families/lfm2_moe/cuda_backend.zig @@ -6,6 +6,7 @@ const cfg = @import("config.zig"); const fwd = @import("cuda_forward.zig"); const mdl = @import("cuda_model.zig"); const kern = @import("cuda_kernels.zig"); +const kv = @import("core").kv_pool; const be = lanes.backend; const Stream = lanes.Stream; @@ -30,8 +31,14 @@ const Lane = struct { rule: u64, // sample.cu's Rule, when the stream samples sampled: bool = false, fill: ?usize = null, // a prompt pass going a chunk at a time: where its next chunk starts + lease: kv.Lease = .{}, // a paged slot's pages }; +/// Pages a request holds: its prompt, its reply and one last verify window. +pub fn pagesFor(prompt: usize, max_new: usize) usize { + return (prompt + max_new + max_window + kern.page - 1) / kern.page; +} + pub const Lfm = struct { gpa: std.mem.Allocator, d: *const cuda.Driver, @@ -46,6 +53,7 @@ pub const Lfm = struct { meta: cuda.DeviceBuffer, // a draw's first position minus one, a window each slot_bufs: []cuda.DeviceBuffer, // a slot's last-row logits, then its rule: two a slot next: u64 = 0, + pool: ?kv.Pool = null, // the pages every slot leases (attach) /// The device bytes `init` allocates: the token and pick rings, metas, each slot's last row and rule. pub fn openBytes(slots: usize) usize { @@ -53,7 +61,7 @@ pub const Lfm = struct { } pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, m: *mdl.Model, k: *const kern.Kernels, s: cuda.Stream) !Lfm { - const free = try gpa.alloc(bool, m.slots); + const free = try gpa.alloc(bool, m.slotCap()); errdefer gpa.free(free); @memset(free, true); const greedy = try gpa.alloc(bool, ring + round_rows); @@ -65,7 +73,7 @@ pub const Lfm = struct { errdefer picks.free(); var meta = try cuda.DeviceBuffer.alloc(d, round_rows * 16); errdefer meta.free(); - const bufs = try gpa.alloc(cuda.DeviceBuffer, 2 * m.slots); + const bufs = try gpa.alloc(cuda.DeviceBuffer, 2 * m.slotCap()); errdefer gpa.free(bufs); var made: usize = 0; errdefer for (bufs[0..made]) |*b| b.free(); @@ -76,9 +84,19 @@ pub const Lfm = struct { return .{ .gpa = gpa, .d = d, .m = m, .k = k, .s = s, .free = free, .tok = tok, .picks = picks, .greedy = greedy, .meta = meta, .slot_bufs = bufs }; } + /// Makes a paged model's arena of `pages` pages for `extra + 1` slots, and the pool that leases them. + pub fn attach(x: *Lfm, extra: usize, pages: usize) !void { + try x.m.attach(x.s, extra, pages); + x.pool = try kv.Pool.init(x.gpa, pages, mdl.page_elems * 2); + } + pub fn deinit(x: *Lfm) void { var it = x.lanes.valueIterator(); - while (it.next()) |l| l.rows.deinit(x.gpa); + while (it.next()) |l| { + if (x.pool) |*p| p.finish(&l.lease); + l.rows.deinit(x.gpa); + } + if (x.pool) |*p| p.deinit(); x.lanes.deinit(x.gpa); for (x.slot_bufs) |*b| b.free(); x.gpa.free(x.slot_bufs); @@ -97,7 +115,7 @@ pub const Lfm = struct { /// The round loop's facts: no draft head; host copies verified in windows up to 16 rows, streams sharing a forward. pub fn facts(x: *const Lfm) lanes.Model { - return .{ .exact_width = max_window, .gpu_tokens = false, .mtp = false, .speculate = false, .hidden_rows = true, .batch_rows = round_rows, .max_streams = @intCast(@min(x.m.slots, round_rows)) }; + return .{ .exact_width = max_window, .gpu_tokens = false, .mtp = false, .speculate = false, .hidden_rows = true, .batch_rows = round_rows, .max_streams = @intCast(@min(x.m.slotCap(), round_rows)) }; } fn self(ptr: *anyopaque) *Lfm { @@ -133,7 +151,39 @@ pub const Lfm = struct { } fn freeSlot(x: *const Lfm) ?usize { - return std.mem.indexOfScalar(bool, x.free, true); + return std.mem.indexOfScalar(bool, x.free[0..x.m.slots], true); + } + + fn fitsContext(x: *const Lfm, prompt: usize, max_new: usize) bool { + return prompt + max_new + max_window <= x.m.context; + } + + /// Whether a request could start now: a free slot and (paged) its pages; one that never fits passes. + pub fn admits(x: *const Lfm, prompt: usize, max_new: usize) bool { + if (!x.fitsContext(prompt, max_new)) return true; // its prompt pass refuses it by name + const p = x.pool orelse return x.freeSlot() != null; + return x.freeSlot() != null and p.fits(pagesFor(prompt, max_new)); + } + + const Map = struct { + x: *Lfm, + slot: usize, + fn write(ctx: *anyopaque, ids: []const u32, from: usize) anyerror!void { + const m: *const Map = @ptrCast(@alignCast(ctx)); + try m.x.m.mapPages(m.x.s, m.slot, ids[from..], from); + } + }; + + /// A free slot for the request, its pages leased now when paged; LaneMemory: not yet. + fn takeSlot(x: *Lfm, lease: *kv.Lease, prompt: usize, max_new: usize) !usize { + const slot = x.freeSlot() orelse return error.NoFreeSlot; + const p = if (x.pool) |*p| p else return slot; + const n = pagesFor(prompt, max_new); + try p.promise(lease, n); + errdefer p.finish(lease); + var map: Map = .{ .x = x, .slot = slot }; + try p.cover(lease, n, .{ .ctx = &map, .write = Map.write }); + return slot; } /// A prompt pass begins: a slot, the stream's rule, a kept state restored; where its first chunk starts. @@ -141,15 +191,16 @@ pub const Lfm = struct { const p = s.prompt(); if (p.len == 0) return error.EmptyPrompt; if (s.images.len > 0) return error.LfmTextOnly; - if (p.len + @as(usize, s.max_new) + max_window > x.m.context) return error.PastContext; + if (!x.fitsContext(p.len, s.max_new)) return error.PastContext; const got = try x.lanes.getOrPut(x.gpa, s); if (!got.found_existing) { - const slot = x.freeSlot() orelse { + var lease: kv.Lease = .{}; + const slot = x.takeSlot(&lease, p.len, s.max_new) catch |e| { x.lanes.removeByPtr(got.key_ptr); - return error.NoFreeSlot; + return e; }; x.free[slot] = false; - got.value_ptr.* = .{ .slot = slot, .last = x.slot_bufs[2 * slot].ptr, .rule = x.slot_bufs[2 * slot + 1].ptr }; + got.value_ptr.* = .{ .slot = slot, .last = x.slot_bufs[2 * slot].ptr, .rule = x.slot_bufs[2 * slot + 1].ptr, .lease = lease }; } const l = got.value_ptr; l.sampled = false; @@ -305,14 +356,35 @@ pub const Lfm = struct { } fn dropLane(x: *Lfm, s: *const Stream) void { - if (x.lanes.fetchRemove(s)) |kv| { - var l = kv.value; + if (x.lanes.fetchRemove(s)) |entry| { + var l = entry.value; + if (x.pool) |*p| { + p.finish(&l.lease); + x.m.unmapPages(l.slot); + } x.free[l.slot] = true; l.rows.deinit(x.gpa); + if (x.lanes.count() == 0) x.report(); } } + /// The idle server's line: nothing promised, every page free (kept states are copies, not pages). + fn report(x: *const Lfm) void { + if (@import("builtin").is_test) return; + const p = x.pool orelse return; + const c = p.census(); + const ok = p.idle() and c.free == c.total; + std.log.info("kv pool idle{s}: promised {d}, free {d} + store {d} == total {d}", .{ if (ok) "" else " (UNEVEN)", c.promised, c.free, c.kept, c.total }); + } + fn release(ptr: *anyopaque, s: *Stream) void { self(ptr).dropLane(s); } }; + +test "a request holds its prompt, its reply and one verify window in whole pages" { + try std.testing.expectEqual(@as(usize, 1), pagesFor(1, 1)); + try std.testing.expectEqual(@as(usize, 1), pagesFor(400, 96)); + try std.testing.expectEqual(@as(usize, 2), pagesFor(400, 97)); + try std.testing.expectEqual(@as(usize, 65), pagesFor(32768, 0)); +} diff --git a/zig/src/families/lfm2_moe/cuda_forward.zig b/zig/src/families/lfm2_moe/cuda_forward.zig index 6c2b7ef357..7937f80d4a 100644 --- a/zig/src/families/lfm2_moe/cuda_forward.zig +++ b/zig/src/families/lfm2_moe/cuda_forward.zig @@ -148,6 +148,7 @@ pub fn forward(m: *const Model, k: *const Kernels, s: Stream, chains: []const Ch var rows: usize = 0; for (chains) |ch| { if (ch.tokens.len == 0 or ch.slot >= m.slots or ch.p0 + ch.tokens.len > m.context) return error.BadChain; + if (ch.p0 + ch.tokens.len > m.caps[ch.slot]) return error.PastPages; rows += ch.tokens.len; } if (rows > @as(usize, if (prompt) mdl.prompt_rows else mdl.round_rows) or (prompt and chains.len != 1)) return error.BadWindow; diff --git a/zig/src/families/lfm2_moe/cuda_model.zig b/zig/src/families/lfm2_moe/cuda_model.zig index 50f2c51277..7856e1bf53 100644 --- a/zig/src/families/lfm2_moe/cuda_model.zig +++ b/zig/src/families/lfm2_moe/cuda_model.zig @@ -1,6 +1,7 @@ //! LFM2.5-24B-A2B's device state on CUDA: KV in 512-position pages, each slot's conv states, the scratch. const std = @import("std"); const cuda = @import("cuda"); +const kv = @import("core").kv_pool; const c = @import("config.zig"); const wts = @import("cuda_weights.zig"); const kern = @import("cuda_kernels.zig"); @@ -60,9 +61,13 @@ pub const Model = struct { w: *const wts.Weights, context: usize, // positions a slot holds: whole pages slots: usize, + paged: bool = false, // slots lease pages of the core pool (attach makes the arena), else each owns a window + pages: usize = 0, // pages in the arena arena: u64 = 0, // [pages][attention layer][K, V][512][kv_width] bf16 - table: u64 = 0, // int32 [slots][pages a slot]: slot s's pages s * per .. (s + 1) * per + table: u64 = 0, // int32 [slots][pages a slot]: slot s's page ids in position order states: u64 = 0, // bf16 [slots][conv layer][2][hidden] + ids: []u32 = &.{}, // the tables on the host + caps: []usize = &.{}, // positions slot s's pages cover inv: u64 = 0, // fp32 [head_dim / 2]: RoPE's inverse frequencies s: Scratch = undefined, buffers: std.ArrayList(DeviceBuffer) = .empty, @@ -82,6 +87,16 @@ pub const Model = struct { return positions(context) / kern.page * page_elems * 2 + slot_elems * 2; } + /// What the host's pool plans: a stream's conv states and table row, a token its rows of the ten layers. + pub fn layout(m: *const Model) kv.Layout { + return .{ .fixed_bytes = slot_elems * 2 + m.perSlot() * 4, .token_bytes = page_elems * 2 / kern.page, .page_tokens = kern.page, .snap_bytes = 0 }; + } + + /// Slots a backend holds: the model's, or (paged) as many as a round takes. + pub fn slotCap(m: *const Model) usize { + return if (m.paged) round_rows else m.slots; + } + /// Decode attention's key chunks a row walks at most, and the blocks that walk them. pub fn chunks(m: *const Model) usize { return m.perSlot(); @@ -105,26 +120,89 @@ pub const Model = struct { return b.ptr; } - pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, w: *const wts.Weights, context: usize, slots: usize) !Model { - var m: Model = .{ .gpa = gpa, .d = d, .w = w, .context = positions(context), .slots = slots }; + /// `slots` streams, each owning a window of pages; paged: one slot's states now, `attach` makes the pages. + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, w: *const wts.Weights, context: usize, slots: usize, paged: bool) !Model { + var m: Model = .{ .gpa = gpa, .d = d, .w = w, .context = positions(context), .slots = if (paged) 1 else slots, .paged = paged }; errdefer m.deinit(); try m.build(s); return m; } /// The device bytes `init` allocates for these sizes, counted by the same path without allocating. - pub fn openBytes(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, w: *const wts.Weights, context: usize, slots: usize) !usize { - var m: Model = .{ .gpa = gpa, .d = d, .w = w, .context = positions(context), .slots = slots, .planned = 0 }; + pub fn openBytes(gpa: std.mem.Allocator, d: *const cuda.Driver, s: cuda.Stream, w: *const wts.Weights, context: usize, slots: usize, paged: bool) !usize { + var m: Model = .{ .gpa = gpa, .d = d, .w = w, .context = positions(context), .slots = if (paged) 1 else slots, .paged = paged, .planned = 0 }; defer m.deinit(); try m.build(s); return m.planned.?; } - fn build(m: *Model, s: cuda.Stream) !void { + /// Makes the arena of `pool` pages to lease and the states and tables of `extra + 1` slots (no pages mapped). + pub fn attach(m: *Model, s: cuda.Stream, extra: usize, pool: usize) !void { + if (!m.paged or m.pages != 0 or extra >= round_rows) return error.BadAttach; + m.forget(m.table); + m.forget(m.states); + m.gpa.free(m.ids); + m.gpa.free(m.caps); + m.ids = &.{}; + m.caps = &.{}; + m.slots = extra + 1; + m.pages = pool; + m.arena = try m.alloc(pool * page_elems * 2); + try m.d.check(m.d.api.cuMemsetD8Async(m.arena, 0, pool * page_elems * 2, s.handle), "cuMemsetD8Async"); + try m.makeSlots(s); + try s.synchronize(); + } + + /// Slot `slot`'s pages from table entry `from` on are `ids` (arena page numbers in position order). + pub fn mapPages(m: *Model, s: cuda.Stream, slot: usize, ids: []const u32, from: usize) !void { + const at = try m.note(slot, ids, from); + try m.d.check(m.d.api.cuMemcpyHtoDAsync_v2(m.table + at * 4, ids.ptr, ids.len * 4, s.handle), "cuMemcpyHtoDAsync"); + } + + /// mapPages on the host: the table's copy and the slot's reach; the entry its device copy starts at. + fn note(m: *Model, slot: usize, ids: []const u32, from: usize) !usize { + const per = m.perSlot(); + if (slot >= m.slots or from + ids.len > per) return error.BadChain; + for (ids) |id| if (id >= m.pages) return error.BadChain; + const at = slot * per + from; + @memcpy(m.ids[at..][0..ids.len], ids); + m.caps[slot] = (from + ids.len) * kern.page; + return at; + } + + /// Slot `slot` holds no pages: its table entries stay, no forward may reach them. + pub fn unmapPages(m: *Model, slot: usize) void { + m.caps[slot] = 0; + } + + /// Frees one buffer `alloc` made. + fn forget(m: *Model, ptr: u64) void { + for (m.buffers.items, 0..) |*b, i| if (b.ptr == ptr) { + b.free(); + _ = m.buffers.swapRemove(i); + return; + }; + } + + /// The slots' tables and conv states; a window of pages a slot unless paged (their pages come from leases). + fn makeSlots(m: *Model, s: cuda.Stream) !void { const per = m.perSlot(); - m.arena = try m.alloc(m.slots * per * page_elems * 2); m.table = try m.alloc(m.slots * per * 4); m.states = try m.alloc(m.slots * slot_elems * 2); + if (m.planned != null) return; + m.ids = try m.gpa.alloc(u32, m.slots * per); + m.caps = try m.gpa.alloc(usize, m.slots); + for (m.ids, 0..) |*v, i| v.* = if (m.paged) 0 else @intCast(i); + @memset(m.caps, if (m.paged) 0 else m.context); + try m.d.check(m.d.api.cuMemcpyHtoDAsync_v2(m.table, m.ids.ptr, m.ids.len * 4, s.handle), "cuMemcpyHtoDAsync"); + try (DeviceBuffer{ .d = m.d, .ptr = m.states, .len = m.slots * slot_elems * 2, .borrowed = true }).fill8(0, s.handle); + try s.synchronize(); + } + + fn build(m: *Model, s: cuda.Stream) !void { + if (!m.paged) m.pages = m.slots * m.perSlot(); + if (m.pages > 0) m.arena = try m.alloc(m.pages * page_elems * 2); + try m.makeSlots(s); m.inv = try m.alloc(c.head_dim / 2 * 4); const r = prompt_rows; const h = c.hidden; @@ -163,19 +241,15 @@ pub const Model = struct { .logits = try m.alloc(round_rows * c.vocab * 2), }; if (m.planned != null) return; - const ids = try m.gpa.alloc(i32, m.slots * per); - defer m.gpa.free(ids); - for (ids, 0..) |*v, i| v.* = @intCast(i); - try (DeviceBuffer{ .d = m.d, .ptr = m.table, .len = ids.len * 4, .borrowed = true }).upload(0, std.mem.sliceAsBytes(ids)); const inv = inverseFrequencies(m.w.config.theta); try (DeviceBuffer{ .d = m.d, .ptr = m.inv, .len = inv.len * 4, .borrowed = true }).upload(0, std.mem.sliceAsBytes(&inv)); - try (DeviceBuffer{ .d = m.d, .ptr = m.states, .len = m.slots * slot_elems * 2, .borrowed = true }).fill8(0, s.handle); - try s.synchronize(); } pub fn deinit(m: *Model) void { for (m.buffers.items) |*b| b.free(); m.buffers.deinit(m.gpa); + m.gpa.free(m.ids); + m.gpa.free(m.caps); } /// Attention layer a's key (or value) cache as the kernels take it: its part of page 0. @@ -213,3 +287,49 @@ test "a page holds 512 positions of the ten attention layers' keys and values: 1 try std.testing.expectEqual(@as(f32, 1), inv[0]); try std.testing.expect(@abs(inv[31] - 1.0 / std.math.pow(f64, 1e7, 62.0 / 64.0)) < 1e-12); } + +test "the pool's layout: 20 KiB a token, a stream's conv states and table row fixed, no snapshot slots" { + var m: Model = undefined; + m.context = Model.positions(32768); + const l = m.layout(); + try std.testing.expectEqual(@as(u64, 10 << 20), l.pageBytes()); + try std.testing.expectEqual(@as(u64, 20 << 10), l.token_bytes); + try std.testing.expectEqual(@as(u64, slot_elems * 2 + 64 * 4), l.fixed_bytes); + // 40 GiB past the weights, 8 streams at 32k, the default 16 GiB cache: a window for the first, the rest pages + const p = try kv.plan(l, .{ .room = 40 << 30, .engine = 0, .lanes = 8, .cache = 16 << 30, .window = 32768 + 16, .engine_paged = true }); + try std.testing.expectEqual(@as(u32, 8), p.lanes); + try std.testing.expectEqual(@as(usize, 0), p.snaps); + try std.testing.expect(p.pages * l.pageBytes() + p.copied + 7 * l.fixed_bytes <= 40 << 30); + try std.testing.expect(p.pages >= l.pages(32768 + 16)); +} + +test "leases map a slot's pages in position order; its reach ends at them; an id past the arena is refused" { + var ids: [8]u32 = @splat(0); + var caps = [_]usize{ 0, 0 }; + var m: Model = undefined; + m.context = 4 * kern.page; + m.slots = 2; + m.pages = 16; + m.ids = &ids; + m.caps = ∩︀ + var pool = try kv.Pool.init(std.testing.allocator, 16, page_elems * 2); + defer pool.deinit(); + var lease: kv.Lease = .{}; + try pool.promise(&lease, 3); + const Rec = struct { + m: *Model, + fn write(ctx: *anyopaque, got: []const u32, from: usize) anyerror!void { + const r: *@This() = @ptrCast(@alignCast(ctx)); + _ = try r.m.note(1, got[from..], from); + } + }; + var rec: Rec = .{ .m = &m }; + try pool.cover(&lease, 3, .{ .ctx = &rec, .write = Rec.write }); + try std.testing.expectEqualSlices(u32, lease.ids.items, ids[4..7]); + try std.testing.expectEqual(@as(usize, 3 * kern.page), caps[1]); + try std.testing.expectEqual(@as(usize, 13), pool.available()); + pool.finish(&lease); + try std.testing.expect(pool.idle() and pool.available() == 16); + try std.testing.expectError(error.BadChain, m.note(0, &.{16}, 0)); + try std.testing.expectError(error.BadChain, m.note(0, &.{ 1, 2, 3, 4, 5 }, 0)); +} diff --git a/zig/src/families/lfm2_moe/cuda_native.zig b/zig/src/families/lfm2_moe/cuda_native.zig index b5699b5795..73b0e410bb 100644 --- a/zig/src/families/lfm2_moe/cuda_native.zig +++ b/zig/src/families/lfm2_moe/cuda_native.zig @@ -7,6 +7,7 @@ const wts = @import("cuda_weights.zig"); const kern = @import("cuda_kernels.zig"); const mdl = @import("cuda_model.zig"); const be = @import("cuda_backend.zig"); +const kv = @import("core").kv_pool; pub const model_type = "lfm2_moe"; pub const formats: []const []const u8 = &.{"mlx-q4g64"}; @@ -14,6 +15,8 @@ pub const default_context: i64 = 32768; pub const max_segments: u32 = 1; pub const prompt_rows: u32 = be.prompt_rows; pub const snap_lookahead: u32 = 0; +/// A request's last verify window: its pages hold that many rows past its prompt and reply. +pub const window_rows: u32 = be.max_window; pub const explains_open = true; /// The server's reading of the room left after the weights; null if unreadable. @@ -30,6 +33,10 @@ pub const Loaded = struct { ctx: *anyopaque, deinit: *const fn (*anyopaque) void, target: *be.Lfm, + /// The attention caches as pages of the core pool (null: every stream owns a window, made at open). + layout: ?kv.Layout = null, + /// With a layout: the first stream's rows are pages too, a window of them set aside by the plan. + engine_paged: bool = true, }; const Owned = struct { gpa: std.mem.Allocator, w: wts.Weights, k: kern.Kernels, s: cuda.Stream, m: mdl.Model, lb: be.Lfm }; @@ -45,17 +52,37 @@ fn fit(room: u64, asked: u32, fixed: bool, plan: anytype) !usize { return n; } +/// TENSORFOLD_LFM2_KV_PAGES=0: every stream owns a window of pages, made at open; else they lease the pool's. +fn pagesOn() bool { + const v = std.c.getenv("TENSORFOLD_LFM2_KV_PAGES") orelse return true; + return !std.mem.eql(u8, std.mem.span(v), "0"); +} + +/// The device bytes `n` streams take (paged: the first's states and a round's slot buffers, no pages). const Plan = struct { own: *const Owned, d: *const cuda.Driver, context: usize, + paged: bool, fn bytes(p: Plan, n: usize) !u64 { const o = p.own; - return try mdl.Model.openBytes(o.gpa, p.d, o.s, &o.w, p.context, n) + be.Lfm.openBytes(n); + return try mdl.Model.openBytes(o.gpa, p.d, o.s, &o.w, p.context, n, p.paged) + be.Lfm.openBytes(if (p.paged) be.round_rows else n); } }; +/// The pool: `streams` more slots beside the first and an arena of `n` pages. +pub fn attach(p: *anyopaque, streams: u32, n: usize) !void { + const own: *Owned = @ptrCast(@alignCast(p)); + try own.lb.attach(streams, n); +} + +/// Whether a request's caches could start now (the lane host's admission, on its engine thread). +pub fn fits(p: *anyopaque, prompt: u32, max_tokens: u32) bool { + const own: *const Owned = @ptrCast(@alignCast(p)); + return own.lb.admits(prompt, max_tokens); +} + /// The attention layers the kernels' page arena holds: the admitted 10 of 40 (conv the rest). fn checkLayers(config: *const cfg.Config) !void { var n: usize = 0; @@ -77,16 +104,17 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ errdefer own.w.deinit(); try checkLayers(&own.w.config); const context = @min(mdl.Model.positions(o.context), own.w.config.context); - const plan: Plan = .{ .own = own, .d = ctx.d, .context = context }; + const paged = pagesOn(); + const plan: Plan = .{ .own = own, .d = ctx.d, .context = context, .paged = paged }; const room: u64 = if (o.room) |r| r.read(r.ctx) orelse 0 else std.math.maxInt(u64); - const n = fit(room, o.lanes, o.lanes_fixed, plan) catch |e| { + const n = if (paged) try fitFirst(room, plan) else fit(room, o.lanes, o.lanes_fixed, plan) catch |e| { if (e == error.ParallelPastMemory) { const most = fit(room, o.lanes, false, plan) catch 0; std.log.err("--parallel {d} needs {d:.2} GiB at a {d}-token window and the CUDA memory budget leaves {d:.2} GiB after the weights: {d} stream{s} fit", .{ o.lanes, gib(try plan.bytes(o.lanes)), context, gib(room), most, if (most == 1) "" else "s" }); } return e; }; - own.m = try mdl.Model.init(gpa, ctx.d, own.s, &own.w, context, n); + own.m = try mdl.Model.init(gpa, ctx.d, own.s, &own.w, context, n, paged); errdefer own.m.deinit(); own.lb = try be.Lfm.init(gpa, ctx.d, &own.m, &own.k, own.s); errdefer own.lb.deinit(); @@ -99,9 +127,16 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ .ctx = own, .deinit = release, .target = &own.lb, + .layout = if (paged) own.m.layout() else null, }; } +/// Paged: the first stream's states and the scratch must fit; the host plans the pool after. +fn fitFirst(room: u64, plan: anytype) !usize { + if (try plan.bytes(1) > room) return error.NoStreamFits; + return 1; +} + fn gib(bytes: u64) f64 { return @as(f64, @floatFromInt(bytes)) / (1 << 30); } From abae447b34fad2af6bb4aed27bcfb739efe11e0a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Tue, 27 Oct 2026 20:57:19 +0100 Subject: [PATCH 10/12] lfm2_moe: the prompt cache on CUDA, a kept state the conv states and the pages' rows through its position A state at `at` is the slot's conv states and its keys and values at positions [0, at) of the ten attention layers: one copy a run of consecutive whole pages, then the last page's rows a layer's keys and values. It is copied behind the chunk that ends at the mark and restored into the new stream's leased pages before its next chunk, so a later turn's pass starts at its history. Lookahead 0, as the other headless families. States live in 2 MiB snapshot pages allocated once at startup inside the prompt-cache budget (`snap_ready`), as GLM's: a keep or a restore leases or reads pages and allocates nothing, and a dropped state's pages go back to the pool at once (every copy that reads or writes them is queued on the backend's one stream). The store counts held and spare pages through `snap_charged` / `snap_spare` / `snap_reuses`. --- zig/src/families/lfm2_moe/cuda_backend.zig | 50 ++++- zig/src/families/lfm2_moe/cuda_family.zig | 2 + zig/src/families/lfm2_moe/cuda_native.zig | 53 +++++- zig/src/families/lfm2_moe/cuda_snap.zig | 209 +++++++++++++++++++++ 4 files changed, 311 insertions(+), 3 deletions(-) create mode 100644 zig/src/families/lfm2_moe/cuda_snap.zig diff --git a/zig/src/families/lfm2_moe/cuda_backend.zig b/zig/src/families/lfm2_moe/cuda_backend.zig index 75aa1fc183..7daf1e3de4 100644 --- a/zig/src/families/lfm2_moe/cuda_backend.zig +++ b/zig/src/families/lfm2_moe/cuda_backend.zig @@ -7,6 +7,7 @@ const fwd = @import("cuda_forward.zig"); const mdl = @import("cuda_model.zig"); const kern = @import("cuda_kernels.zig"); const kv = @import("core").kv_pool; +const snap = @import("cuda_snap.zig"); const be = lanes.backend; const Stream = lanes.Stream; @@ -54,6 +55,7 @@ pub const Lfm = struct { slot_bufs: []cuda.DeviceBuffer, // a slot's last-row logits, then its rule: two a slot next: u64 = 0, pool: ?kv.Pool = null, // the pages every slot leases (attach) + snaps: ?snap.Pages = null, // the prompt cache's pages, allocated once inside its budget (snap_ready) /// The device bytes `init` allocates: the token and pick rings, metas, each slot's last row and rule. pub fn openBytes(slots: usize) usize { @@ -97,6 +99,7 @@ pub const Lfm = struct { l.rows.deinit(x.gpa); } if (x.pool) |*p| p.deinit(); + if (x.snaps) |*p| p.deinit(); x.lanes.deinit(x.gpa); for (x.slot_bufs) |*b| b.free(); x.gpa.free(x.slot_bufs); @@ -150,6 +153,30 @@ pub const Lfm = struct { }; } + /// The stream's state at `at` for the prompt cache, copied behind the chunk that ends there. + pub fn saveState(x: *Lfm, owner: ?*anyopaque, at: u32) !*snap.Saved { + const s: *const Stream = @ptrCast(@alignCast(owner orelse return error.NoOwner)); + const l = x.lanes.getPtr(s) orelse return error.UnknownStream; + if ((l.fill orelse l.len) != at) return error.NotAtMark; // a state is copied only where the pass stands + return snap.save(x.gpa, try x.snapPages(), x.m, x.s, l.slot, at); + } + + /// The stream's slot made to hold `k`'s state (a restore outside its prompt pass). + pub fn restoreState(x: *Lfm, owner: ?*anyopaque, k: *const snap.Saved) !void { + const s: *const Stream = @ptrCast(@alignCast(owner orelse return error.NoOwner)); + const l = x.lanes.getPtr(s) orelse return error.UnknownStream; + try snap.restore(try x.snapPages(), x.m, x.s, l.slot, k); + l.len = k.at; + } + + pub fn dropState(x: *Lfm, k: *snap.Saved) void { + snap.drop(x.gpa, &x.snaps.?, k); + } + + fn snapPages(x: *Lfm) !*snap.Pages { + return if (x.snaps) |*p| p else error.NoSnapshotCache; + } + fn freeSlot(x: *const Lfm) ?usize { return std.mem.indexOfScalar(bool, x.free[0..x.m.slots], true); } @@ -210,13 +237,32 @@ pub const Lfm = struct { l.sampled = true; }; s.cached = 0; - if (s.reuse.saved != null) s.reuse_failed = true; // no prompt cache: every pass starts at 0 try x.m.clearSlot(x.s, l.slot); - l.fill = 0; + l.fill = x.resumeAt(s, l, p.len); s.fill_at.store(@intCast(l.fill.?), .monotonic); return l; } + /// The kept state's position, restored into the slot (0: none, or not this prompt's). + fn resumeAt(x: *Lfm, s: *Stream, l: *const Lane, len: usize) usize { + const saved = s.reuse.saved orelse return 0; + const k: *const snap.Saved = @ptrCast(@alignCast(saved)); + if (k.at != s.reuse.at or k.at == 0 or k.at >= len) { + s.reuse_failed = true; + return 0; + } + const p = x.snapPages() catch { + s.reuse_failed = true; + return 0; + }; + snap.restore(p, x.m, x.s, l.slot, k) catch { + s.reuse_failed = true; + return 0; + }; + s.cached = k.at; + return k.at; + } + /// One chunk from `l.fill` to the next grid line or mark; true once the prompt's last row's logits are written. fn chunk(x: *Lfm, s: *Stream, l: *Lane) !bool { const p = s.prompt(); diff --git a/zig/src/families/lfm2_moe/cuda_family.zig b/zig/src/families/lfm2_moe/cuda_family.zig index a938da0d0e..50263f9d8a 100644 --- a/zig/src/families/lfm2_moe/cuda_family.zig +++ b/zig/src/families/lfm2_moe/cuda_family.zig @@ -6,6 +6,7 @@ pub const model = @import("cuda_model.zig"); pub const forward = @import("cuda_forward.zig"); pub const backend = @import("cuda_backend.zig"); pub const native = @import("cuda_native.zig"); +pub const snap = @import("cuda_snap.zig"); test { _ = config; @@ -15,4 +16,5 @@ test { _ = forward; _ = backend; _ = native; + _ = snap; } diff --git a/zig/src/families/lfm2_moe/cuda_native.zig b/zig/src/families/lfm2_moe/cuda_native.zig index 73b0e410bb..b377869309 100644 --- a/zig/src/families/lfm2_moe/cuda_native.zig +++ b/zig/src/families/lfm2_moe/cuda_native.zig @@ -1,4 +1,4 @@ -//! LFM2.5-24B-A2B (MLX affine 4-bit) behind the native CUDA registry: plain and copy-drafted decoding, no prompt cache. +//! LFM2.5-24B-A2B (MLX affine 4-bit) behind the native CUDA registry: plain and copy-drafted decoding, prompt reuse. const std = @import("std"); const cuda = @import("cuda"); const lanes = @import("lanes"); @@ -7,6 +7,7 @@ const wts = @import("cuda_weights.zig"); const kern = @import("cuda_kernels.zig"); const mdl = @import("cuda_model.zig"); const be = @import("cuda_backend.zig"); +const snap = @import("cuda_snap.zig"); const kv = @import("core").kv_pool; pub const model_type = "lfm2_moe"; @@ -14,6 +15,7 @@ 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 = be.prompt_rows; +/// A kept state is the caches through its position and nothing past it (no draft head reads a token on). pub const snap_lookahead: u32 = 0; /// A request's last verify window: its pages hold that many rows past its prompt and reply. pub const window_rows: u32 = be.max_window; @@ -141,6 +143,49 @@ fn gib(bytes: u64) f64 { return @as(f64, @floatFromInt(bytes)) / (1 << 30); } +fn lfm(ptr: *anyopaque) *be.Lfm { + return @ptrCast(@alignCast(ptr)); +} + +/// Bytes a state at `at` prompt tokens takes: whole snapshot pages. +pub fn snap_bytes(_: *anyopaque, at: u32) u64 { + return snap.bytes(at); +} + +/// The cache's pages, allocated once before any request: the store counts held and spare pages through the hooks below. +pub fn snap_ready(ptr: *anyopaque, budget: u64) !u64 { + const x = lfm(ptr); + if (x.snaps != null) return error.SnapshotCacheMade; + x.snaps = try snap.Pages.init(x.gpa, x.d, budget); + return 0; +} + +pub fn snap_charged(_: *anyopaque, saved: *anyopaque) u64 { + const k: *const snap.Saved = @ptrCast(@alignCast(saved)); + return k.lease.ids.items.len * snap.page; +} + +pub fn snap_spare(ptr: *anyopaque) u64 { + return if (lfm(ptr).snaps) |*p| p.spare() else 0; +} + +pub fn snap_reuses(ptr: *anyopaque, at: u32) bool { + return snap_spare(ptr) >= snap.bytes(at); +} + +/// The prompt pass stands at `at` (a chunk end or a mark): copy the stream's caches there. +pub fn snap_save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { + return @ptrCast(try lfm(ptr).saveState(owner, at)); +} + +pub fn snap_restore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror!void { + return lfm(ptr).restoreState(owner, @ptrCast(@alignCast(saved))); +} + +pub fn snap_drop(ptr: *anyopaque, saved: *anyopaque) void { + lfm(ptr).dropState(@ptrCast(@alignCast(saved))); +} + pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { return switch (err) { error.PastContext => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", @@ -182,3 +227,9 @@ test "the admitted layer pattern has the arena's ten attention layers" { config.attention[0] = true; try std.testing.expectError(error.UnsupportedLfmGeometry, checkLayers(&config)); } + +test "the prompt cache's hooks compile against the store's vtable" { + const pc: struct { bytes: *const fn (*anyopaque, u32) u64, save: *const fn (*anyopaque, ?*anyopaque, u32) anyerror!*anyopaque, restore: *const fn (*anyopaque, ?*anyopaque, *anyopaque) anyerror!void, drop: *const fn (*anyopaque, *anyopaque) void } = .{ .bytes = snap_bytes, .save = snap_save, .restore = snap_restore, .drop = snap_drop }; + try std.testing.expect(@intFromPtr(pc.save) != 0); + try std.testing.expectEqual(snap.bytes(100), snap_bytes(undefined, 100)); +} diff --git a/zig/src/families/lfm2_moe/cuda_snap.zig b/zig/src/families/lfm2_moe/cuda_snap.zig new file mode 100644 index 0000000000..a3a9f41fa6 --- /dev/null +++ b/zig/src/families/lfm2_moe/cuda_snap.zig @@ -0,0 +1,209 @@ +//! A stream's kept state (conv states, K/V [0, at)) in pages allocated once in the prompt-cache budget, as GLM's. +const std = @import("std"); +const cuda = @import("cuda"); +const kv = @import("core").kv_pool; +const mdl = @import("cuda_model.zig"); +const kern = @import("cuda_kernels.zig"); + +/// A kept state: its position and the snapshot pages that hold it, in order. +pub const Saved = struct { at: u32, lease: kv.Lease }; + +const page_bytes = mdl.page_elems * 2; +const row_bytes = mdl.kv_width * 2; +const state_bytes = mdl.slot_elems * 2; + +/// A snapshot page: the unit the budget is cut into. +pub const page: usize = 2 << 20; + +/// Bytes a state at `at` holds: whole snapshot pages. +pub fn bytes(at: usize) u64 { + const n = state_bytes + at * (page_bytes / kern.page); + return (n + page - 1) / page * page; +} + +/// The budget's pages, allocated once; states lease them and give them back. +pub const Pages = struct { + pages: kv.Pool, + arena: ?cuda.DeviceBuffer, + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, budget: u64) !Pages { + const count = std.math.cast(u32, budget / page) orelse return error.SnapshotBudget; + var pages = try kv.Pool.init(gpa, count, page); + errdefer pages.deinit(); + return .{ .pages = pages, .arena = if (count > 0) try cuda.DeviceBuffer.alloc(d, @as(usize, count) * page) else null }; + } + + pub fn deinit(p: *Pages) void { + if (p.arena) |*a| a.free(); + p.pages.deinit(); + } + + /// Bytes of pages no state holds. + pub fn spare(p: *const Pages) u64 { + return p.pages.available() * page; + } + + fn take(p: *Pages, n: usize) !kv.Lease { + if (n == 0 or n > p.pages.available() or p.arena == null) return error.NoSnapshotRoom; + var lease: kv.Lease = .{}; + errdefer p.pages.finish(&lease); + try p.pages.promise(&lease, n); + try p.pages.cover(&lease, n, null); + return lease; + } +}; + +/// `slot`'s state at `at`, region by region: one copy a run of consecutive whole pages, a layer's rows of the last. +fn each(m: *const mdl.Model, slot: usize, at: usize, ctx: anytype, comptime f: fn (@TypeOf(ctx), dev: u64, off: usize, len: usize) anyerror!void) !void { + if (at > m.caps[slot]) return error.PastPages; + try f(ctx, m.states + slot * state_bytes, 0, state_bytes); + var off: usize = state_bytes; + const ids = m.ids[slot * m.perSlot() ..]; + const whole = at / kern.page; + var q: usize = 0; + while (q < whole) { + var n: usize = 1; + while (q + n < whole and ids[q + n] == ids[q] + n) n += 1; + try f(ctx, m.arena + @as(u64, ids[q]) * page_bytes, off, n * page_bytes); + off += n * page_bytes; + q += n; + } + const rows = at % kern.page; + if (rows == 0) return; + const last = m.arena + @as(u64, ids[whole]) * page_bytes; + for (0..2 * mdl.attention_layers) |i| { // layer a's keys at block 2a, its values at 2a + 1 + try f(ctx, last + i * kern.page * row_bytes, off, rows * row_bytes); + off += rows * row_bytes; + } +} + +/// Region bytes [off, off + len) over a state's pages: (arena offset, region offset, length), a run of consecutive ids. +fn pageRuns(ids: []const u32, off: usize, len: usize, ctx: anytype, comptime f: fn (@TypeOf(ctx), held: u64, done: usize, n: usize) anyerror!void) !void { + var done: usize = 0; + while (done < len) { + const at = off + done; + var q = at / page; + var n = @min(len - done, page - at % page); + while (done + n < len and ids[q + 1] == ids[q] + 1) : (q += 1) n = @min(len - done, n + page); + try f(ctx, @as(u64, ids[at / page]) * page + at % page, done, n); + done += n; + } +} + +/// Copies between a slot's region at `dev` and a lease's pages. +const Copy = struct { d: *const cuda.Driver, s: cuda.Stream, arena: u64, dev: u64, out: bool }; + +fn copyRun(c: Copy, held: u64, done: usize, n: usize) anyerror!void { + const src = if (c.out) c.dev + done else c.arena + held; + const dst = if (c.out) c.arena + held else c.dev + done; + try c.d.check(c.d.api.cuMemcpyDtoDAsync_v2(dst, src, n, c.s.handle), "cuMemcpyDtoDAsync"); +} + +const Leased = struct { d: *const cuda.Driver, s: cuda.Stream, arena: u64, ids: []const u32, out: bool }; + +fn copy(c: Leased, dev: u64, off: usize, len: usize) anyerror!void { + try pageRuns(c.ids, off, len, Copy{ .d = c.d, .s = c.s, .arena = c.arena, .dev = dev, .out = c.out }, copyRun); +} + +/// Copy `slot`'s state at `at` into leased pages, queued on `s` behind the chunk that ends there. +pub fn save(gpa: std.mem.Allocator, p: *Pages, m: *const mdl.Model, s: cuda.Stream, slot: usize, at: u32) !*Saved { + if (at == 0 or at > m.context) return error.NotAtMark; + const k = try gpa.create(Saved); + errdefer gpa.destroy(k); + k.* = .{ .at = at, .lease = try p.take(bytes(at) / page) }; + errdefer p.pages.finish(&k.lease); + try each(m, slot, at, Leased{ .d = m.d, .s = s, .arena = p.arena.?.ptr, .ids = k.lease.ids.items, .out = true }, copy); + return k; +} + +/// Make `slot`'s pages and conv states hold `k`'s state, queued on `s` before the next chunk. +pub fn restore(p: *const Pages, m: *const mdl.Model, s: cuda.Stream, slot: usize, k: *const Saved) !void { + if (k.at > m.context) return error.PastContext; + try each(m, slot, k.at, Leased{ .d = m.d, .s = s, .arena = p.arena.?.ptr, .ids = k.lease.ids.items, .out = false }, copy); +} + +/// Give a state's pages back: every copy that reads or writes them is queued on the one stream, ahead of any later one. +pub fn drop(gpa: std.mem.Allocator, p: *Pages, k: *Saved) void { + p.pages.finish(&k.lease); + gpa.destroy(k); +} + +const testing = std.testing; + +const Count = struct { + runs: *std.ArrayList([3]usize), + fn add(c: Count, dev: u64, off: usize, len: usize) anyerror!void { + try c.runs.append(testing.allocator, .{ dev, off, len }); + } +}; + +fn fake(ids: []u32, caps: []usize) mdl.Model { + var m: mdl.Model = undefined; + m.context = 4 * kern.page; + m.slots = caps.len; + m.arena = 1 << 40; + m.states = 1 << 36; + m.ids = ids; + m.caps = caps; + return m; +} + +test "a state: the conv states, consecutive whole pages as one copy, the last page's rows a layer's K and V" { + // slot 1 leases pages 7, 8, 3, 9: positions 0..1023 one run, 1024..1535 a page, 1536.. a partial page + var ids = [_]u32{ 0, 0, 0, 0, 7, 8, 3, 9 }; + var caps = [_]usize{ 0, 4 * kern.page }; + const m = fake(&ids, &caps); + const at = 3 * kern.page + 100; + var runs: std.ArrayList([3]usize) = .empty; + defer runs.deinit(testing.allocator); + try each(&m, 1, at, Count{ .runs = &runs }, Count.add); + try testing.expectEqual(@as(usize, 1 + 2 + 20), runs.items.len); + try testing.expectEqual([3]usize{ (1 << 36) + state_bytes, 0, state_bytes }, runs.items[0]); + try testing.expectEqual([3]usize{ (1 << 40) + 7 * page_bytes, state_bytes, 2 * page_bytes }, runs.items[1]); + try testing.expectEqual([3]usize{ (1 << 40) + 3 * page_bytes, state_bytes + 2 * page_bytes, page_bytes }, runs.items[2]); + const v9 = (1 << 40) + 9 * page_bytes + kern.page * row_bytes; // page 9, layer 0's values + try testing.expectEqual([3]usize{ v9, state_bytes + 3 * page_bytes + 100 * row_bytes, 100 * row_bytes }, runs.items[4]); + var total: usize = 0; + for (runs.items) |r| total += r[2]; + try testing.expectEqual(state_bytes + at * 20 * 1024, total); + try testing.expectEqual((total + page - 1) / page * page, bytes(at)); +} + +test "a state on whole pages has no partial copies; one past a slot's pages is refused" { + var ids = [_]u32{ 4, 5, 6, 7 }; + var caps = [_]usize{2 * kern.page}; + const m = fake(&ids, &caps); + var runs: std.ArrayList([3]usize) = .empty; + defer runs.deinit(testing.allocator); + try each(&m, 0, 2 * kern.page, Count{ .runs = &runs }, Count.add); + try testing.expectEqual(@as(usize, 2), runs.items.len); + try testing.expectEqual(@as(usize, 2 * page_bytes), runs.items[1][2]); + try testing.expectError(error.PastPages, each(&m, 0, 2 * kern.page + 1, Count{ .runs = &runs }, Count.add)); +} + +const Run = struct { + out: *std.ArrayList([3]usize), + fn add(r: Run, held: u64, done: usize, n: usize) anyerror!void { + try r.out.append(testing.allocator, .{ held, done, n }); + } +}; + +test "a region over snapshot pages: one copy a run of consecutive ids, split where the ids jump" { + const ids = [_]u32{ 4, 5, 9, 10, 11, 2 }; + var out: std.ArrayList([3]usize) = .empty; + defer out.deinit(testing.allocator); + // bytes [page / 2, 5 page + 10): pages 4-5 (from mid-page), 9-11, then 2 + try pageRuns(&ids, page / 2, 4 * page + page / 2 + 10, Run{ .out = &out }, Run.add); + try testing.expectEqual(@as(usize, 3), out.items.len); + try testing.expectEqual([3]usize{ 4 * page + page / 2, 0, page + page / 2 }, out.items[0]); + try testing.expectEqual([3]usize{ 9 * page, page + page / 2, 3 * page }, out.items[1]); + try testing.expectEqual([3]usize{ 2 * page, 4 * page + page / 2, 10 }, out.items[2]); + out.clearRetainingCapacity(); + try pageRuns(&ids, 3 * page + 7, 100, Run{ .out = &out }, Run.add); // inside one page + try testing.expectEqual([3]usize{ 10 * page + 7, 0, 100 }, out.items[0]); +} + +test "a state's pages: whole snapshot pages, the budget's pages leased and given back" { + try testing.expectEqual(@as(u64, page), bytes(1)); + try testing.expect(bytes(10_000) >= state_bytes + 10_000 * 20 * 1024 and bytes(10_000) % page == 0); +} From 928034a061a5b45aad0afe4d3402607eab6f6694 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 28 Oct 2026 11:35:40 +0100 Subject: [PATCH 11/12] lfm2_moe: keep a mark's state as the prompt pass goes by, so a keep adds no pass A chat mark inside a prompt chunk ended the chunk there, and every extra prompt pass reads the weights once more: a cold 10k chat turn, whose two marks sit in its last chunk, ran 12 passes instead of 10. Now up to four marks inside a chunk are tapped instead: the conv kernel runs once more on each mark's last row alone, writing the same y and that row's conv states into the model's tap buffer, and the keep copies those states with the K/V rows before the mark. The pass goes on to its grid line; a fifth mark still ends it. The kernel is unchanged. --- zig/src/families/lfm2_moe/cuda_backend.zig | 77 ++++++++++++++++++---- zig/src/families/lfm2_moe/cuda_forward.zig | 10 ++- zig/src/families/lfm2_moe/cuda_model.zig | 4 ++ zig/src/families/lfm2_moe/cuda_snap.zig | 20 +++--- 4 files changed, 88 insertions(+), 23 deletions(-) diff --git a/zig/src/families/lfm2_moe/cuda_backend.zig b/zig/src/families/lfm2_moe/cuda_backend.zig index 7daf1e3de4..681c7e64c2 100644 --- a/zig/src/families/lfm2_moe/cuda_backend.zig +++ b/zig/src/families/lfm2_moe/cuda_backend.zig @@ -55,6 +55,7 @@ pub const Lfm = struct { slot_bufs: []cuda.DeviceBuffer, // a slot's last-row logits, then its rule: two a slot next: u64 = 0, pool: ?kv.Pool = null, // the pages every slot leases (attach) + tapped: ?struct { stream: *const Stream, cut: Cut } = null, // the chunk's tapped marks while its hooks run snaps: ?snap.Pages = null, // the prompt cache's pages, allocated once inside its budget (snap_ready) /// The device bytes `init` allocates: the token and pick rings, metas, each slot's last row and rule. @@ -157,8 +158,10 @@ pub const Lfm = struct { pub fn saveState(x: *Lfm, owner: ?*anyopaque, at: u32) !*snap.Saved { const s: *const Stream = @ptrCast(@alignCast(owner orelse return error.NoOwner)); const l = x.lanes.getPtr(s) orelse return error.UnknownStream; - if ((l.fill orelse l.len) != at) return error.NotAtMark; // a state is copied only where the pass stands - return snap.save(x.gpa, try x.snapPages(), x.m, x.s, l.slot, at); + const tap: ?usize = if (x.tapped) |t| if (t.stream == s) std.mem.indexOfScalar(usize, t.cut.taps(), at) else null else null; + if (tap == null and (l.fill orelse l.len) != at) return error.NotAtMark; // where the pass stands, or a tap + const conv = if (tap) |j| x.m.s.tap + j * mdl.slot_elems * 2 else x.m.states + l.slot * mdl.slot_elems * 2; + return snap.save(x.gpa, try x.snapPages(), x.m, x.s, l.slot, at, conv); } /// The stream's slot made to hold `k`'s state (a restore outside its prompt pass). @@ -268,20 +271,55 @@ pub const Lfm = struct { const p = s.prompt(); const at = l.fill.?; if (s.isCancelled()) return error.Cancelled; - var end = @min(p.len, (at / prompt_rows + 1) * prompt_rows); - for (s.reuse.marks) |mk| if (mk > at and mk < end) { - end = mk; - }; - const done = end == p.len; - try fwd.forward(x.m, x.k, x.s, &.{.{ .slot = l.slot, .p0 = at, .tokens = p[at..end] }}, true, end - at - 1, @intFromBool(done)); + const c = cut(at, p.len, s.reuse.marks); + const done = c.end == p.len; + var rows: [mdl.max_taps]usize = undefined; + for (c.taps(), 0..) |t, j| rows[j] = t - at - 1; + try fwd.forward(x.m, x.k, x.s, &.{.{ .slot = l.slot, .p0 = at, .tokens = p[at..c.end], .taps = rows[0..c.n] }}, true, c.end - at - 1, @intFromBool(done)); if (done) try x.d.check(x.d.api.cuMemcpyDtoDAsync_v2(l.last, x.m.s.logits, cfg.vocab * 2, x.s.handle), "cuMemcpyDtoDAsync"); - s.fill_at.store(@intCast(end), .monotonic); - l.fill = if (done) null else end; + s.fill_at.store(@intCast(c.end), .monotonic); + l.fill = if (done) null else c.end; if (done) l.len = p.len; - if (s.reuse.hook) |hk| if (std.mem.indexOfScalar(u32, s.reuse.marks, @intCast(end)) != null) hk.at(hk.ptr, s, @intCast(end)); + if (s.reuse.hook) |hk| { + x.tapped = .{ .stream = s, .cut = c }; + defer x.tapped = null; + for (c.taps()) |t| hk.at(hk.ptr, s, @intCast(t)); + if (std.mem.indexOfScalar(u32, s.reuse.marks, @intCast(c.end)) != null) hk.at(hk.ptr, s, @intCast(c.end)); + } return done; } + /// A chunk's end (the grid line, the prompt's end or a mark) and the marks inside it, tapped in order. + const Cut = struct { + end: usize, + at: [mdl.max_taps]usize = undefined, + n: usize = 0, + fn taps(c: *const Cut) []const usize { + return c.at[0..c.n]; + } + }; + + /// The chunk from `at`: up to max_taps marks inside it are tapped, the next one ends it. + fn cut(at: usize, len: usize, marks: []const u32) Cut { + var c: Cut = .{ .end = @min(len, (at / prompt_rows + 1) * prompt_rows) }; + var from = at; + while (true) { // the marks inside in order, the first max_taps tapped + var next: ?usize = null; + for (marks) |mk| if (mk > from and mk < c.end and (next == null or mk < next.?)) { + next = mk; + }; + const mk = next orelse break; + if (c.n == mdl.max_taps) { + c.end = mk; + break; + } + c.at[c.n] = mk; + c.n += 1; + from = mk; + } + return c; + } + fn prefill(ptr: *anyopaque, s: *Stream) anyerror!void { const x = self(ptr); const l = try x.opening(s); @@ -434,3 +472,20 @@ test "a request holds its prompt, its reply and one verify window in whole pages try std.testing.expectEqual(@as(usize, 2), pagesFor(400, 97)); try std.testing.expectEqual(@as(usize, 65), pagesFor(32768, 0)); } + +test "a chunk taps the marks inside it, up to four, and ends at the next one, the grid line or the prompt's end" { + const cases = [_]struct { at: usize, len: usize, marks: []const u32, end: usize, taps: []const usize }{ + .{ .at = 0, .len = 3000, .marks = &.{}, .end = prompt_rows, .taps = &.{} }, + .{ .at = 0, .len = 3000, .marks = &.{700}, .end = prompt_rows, .taps = &.{700} }, + .{ .at = 0, .len = 3000, .marks = &.{ 900, 300, 2900 }, .end = prompt_rows, .taps = &.{ 300, 900 } }, + .{ .at = 0, .len = 3000, .marks = &.{ 5, 4, 3, 2, 1 }, .end = 5, .taps = &.{ 1, 2, 3, 4 } }, + .{ .at = 9216, .len = 10144, .marks = &.{ 10050, 10140 }, .end = 10144, .taps = &.{ 10050, 10140 } }, + .{ .at = 2048, .len = 2500, .marks = &.{ 2100, 2500 }, .end = 2500, .taps = &.{2100} }, + .{ .at = 700, .len = 3000, .marks = &.{700}, .end = prompt_rows, .taps = &.{} }, + }; + for (cases) |k| { + const c = Lfm.cut(k.at, k.len, k.marks); + try std.testing.expectEqual(k.end, c.end); + try std.testing.expectEqualSlices(usize, k.taps, c.taps()); + } +} diff --git a/zig/src/families/lfm2_moe/cuda_forward.zig b/zig/src/families/lfm2_moe/cuda_forward.zig index 7937f80d4a..f6723889f9 100644 --- a/zig/src/families/lfm2_moe/cuda_forward.zig +++ b/zig/src/families/lfm2_moe/cuda_forward.zig @@ -11,8 +11,8 @@ const Kernels = kern.Kernels; const Stream = cuda.Stream; const tail: cuda.norms.Tail = .{ .round = true }; // the norms' torch points: bf16(w * bf16(h * inv)) -/// A stream's rows in one forward: its slot, first position, tokens. -pub const Chain = struct { slot: usize, p0: usize, tokens: []const u32 }; +/// A stream's rows in one forward: its slot, first position, tokens; a prompt chunk's tap rows keep their conv states. +pub const Chain = struct { slot: usize, p0: usize, tokens: []const u32, taps: []const usize = &.{} }; const Pass = struct { m: *const Model, @@ -65,6 +65,10 @@ const Pass = struct { // a prompt chunk keeps its last row's state, a round each row's (a window keeps any prefix of its rows) const snap0: usize = if (p.prompt) p.rows - 1 else 0; try p.k.convRows(p.s, x.wide, cv.weight, state, mdl.slot_elems, x.slot, x.off, x.y, snaps, snap0, p.rows); + if (p.prompt) for (p.chains[0].taps, 0..) |r, j| { // the kernel again on row r: the same y, its states in tap j + const tap = x.tap + (j * mdl.slot_elems + index * 2 * c.hidden) * 2; + try p.k.convRows(p.s, x.wide + r * 3 * c.hidden * 2, cv.weight, state, mdl.slot_elems, x.slot + r * 4, x.off + r * 4, x.y + r * c.hidden * 2, tap, 0, 1); + }; try p.sums(x.y, c.hidden); try p.project(cv.out, x.y, x.xs, x.r, p.rows, p.prompt); } @@ -149,6 +153,8 @@ pub fn forward(m: *const Model, k: *const Kernels, s: Stream, chains: []const Ch for (chains) |ch| { if (ch.tokens.len == 0 or ch.slot >= m.slots or ch.p0 + ch.tokens.len > m.context) return error.BadChain; if (ch.p0 + ch.tokens.len > m.caps[ch.slot]) return error.PastPages; + if (ch.taps.len > 0 and (!prompt or ch.taps.len > mdl.max_taps)) return error.BadChain; + for (ch.taps) |r| if (r >= ch.tokens.len) return error.BadChain; rows += ch.tokens.len; } if (rows > @as(usize, if (prompt) mdl.prompt_rows else mdl.round_rows) or (prompt and chains.len != 1)) return error.BadWindow; diff --git a/zig/src/families/lfm2_moe/cuda_model.zig b/zig/src/families/lfm2_moe/cuda_model.zig index 7856e1bf53..6a70bf38b6 100644 --- a/zig/src/families/lfm2_moe/cuda_model.zig +++ b/zig/src/families/lfm2_moe/cuda_model.zig @@ -20,6 +20,8 @@ pub const kv_width = c.kv_heads * c.head_dim; pub const page_elems = attention_layers * 2 * kern.page * kv_width; /// Elements a slot's conv states hold: two bx rows a conv layer. pub const slot_elems = conv_layers * 2 * c.hidden; +/// Marks a prompt chunk taps (their conv states kept as the pass goes by); one more ends the chunk. +pub const max_taps = 4; /// Key chunks one attention launch walks in parallel at most (nemotron's chunk_blocks). const chunk_blocks = 64; @@ -52,6 +54,7 @@ pub const Scratch = struct { pm: u64, pl: u64, snaps: u64, // bf16 [conv_layers][round_rows][2][hidden]: each window row's conv state after it + tap: u64, // bf16 [max_taps][conv_layers][2][hidden]: the conv states after a chunk's tapped rows (a slot's layout) logits: u64, // bf16 [round_rows, vocab] }; @@ -238,6 +241,7 @@ pub const Model = struct { .pm = try m.alloc(parts * 4), .pl = try m.alloc(parts * 4), .snaps = try m.alloc(conv_layers * round_rows * 2 * h * 2), + .tap = try m.alloc(max_taps * slot_elems * 2), .logits = try m.alloc(round_rows * c.vocab * 2), }; if (m.planned != null) return; diff --git a/zig/src/families/lfm2_moe/cuda_snap.zig b/zig/src/families/lfm2_moe/cuda_snap.zig index a3a9f41fa6..965546d10d 100644 --- a/zig/src/families/lfm2_moe/cuda_snap.zig +++ b/zig/src/families/lfm2_moe/cuda_snap.zig @@ -53,10 +53,10 @@ pub const Pages = struct { } }; -/// `slot`'s state at `at`, region by region: one copy a run of consecutive whole pages, a layer's rows of the last. -fn each(m: *const mdl.Model, slot: usize, at: usize, ctx: anytype, comptime f: fn (@TypeOf(ctx), dev: u64, off: usize, len: usize) anyerror!void) !void { +/// `slot`'s state at `at`: conv states (from `conv`), a copy a run of consecutive whole pages, the last's rows a layer. +fn each(m: *const mdl.Model, slot: usize, conv: u64, at: usize, ctx: anytype, comptime f: fn (@TypeOf(ctx), dev: u64, off: usize, len: usize) anyerror!void) !void { if (at > m.caps[slot]) return error.PastPages; - try f(ctx, m.states + slot * state_bytes, 0, state_bytes); + try f(ctx, conv, 0, state_bytes); var off: usize = state_bytes; const ids = m.ids[slot * m.perSlot() ..]; const whole = at / kern.page; @@ -105,21 +105,21 @@ fn copy(c: Leased, dev: u64, off: usize, len: usize) anyerror!void { try pageRuns(c.ids, off, len, Copy{ .d = c.d, .s = c.s, .arena = c.arena, .dev = dev, .out = c.out }, copyRun); } -/// Copy `slot`'s state at `at` into leased pages, queued on `s` behind the chunk that ends there. -pub fn save(gpa: std.mem.Allocator, p: *Pages, m: *const mdl.Model, s: cuda.Stream, slot: usize, at: u32) !*Saved { +/// `slot`'s state at `at` copied into leased pages behind the chunk that ends there or taps it (conv states at `conv`). +pub fn save(gpa: std.mem.Allocator, p: *Pages, m: *const mdl.Model, s: cuda.Stream, slot: usize, at: u32, conv: u64) !*Saved { if (at == 0 or at > m.context) return error.NotAtMark; const k = try gpa.create(Saved); errdefer gpa.destroy(k); k.* = .{ .at = at, .lease = try p.take(bytes(at) / page) }; errdefer p.pages.finish(&k.lease); - try each(m, slot, at, Leased{ .d = m.d, .s = s, .arena = p.arena.?.ptr, .ids = k.lease.ids.items, .out = true }, copy); + try each(m, slot, conv, at, Leased{ .d = m.d, .s = s, .arena = p.arena.?.ptr, .ids = k.lease.ids.items, .out = true }, copy); return k; } /// Make `slot`'s pages and conv states hold `k`'s state, queued on `s` before the next chunk. pub fn restore(p: *const Pages, m: *const mdl.Model, s: cuda.Stream, slot: usize, k: *const Saved) !void { if (k.at > m.context) return error.PastContext; - try each(m, slot, k.at, Leased{ .d = m.d, .s = s, .arena = p.arena.?.ptr, .ids = k.lease.ids.items, .out = false }, copy); + try each(m, slot, m.states + slot * state_bytes, k.at, Leased{ .d = m.d, .s = s, .arena = p.arena.?.ptr, .ids = k.lease.ids.items, .out = false }, copy); } /// Give a state's pages back: every copy that reads or writes them is queued on the one stream, ahead of any later one. @@ -156,7 +156,7 @@ test "a state: the conv states, consecutive whole pages as one copy, the last pa const at = 3 * kern.page + 100; var runs: std.ArrayList([3]usize) = .empty; defer runs.deinit(testing.allocator); - try each(&m, 1, at, Count{ .runs = &runs }, Count.add); + try each(&m, 1, m.states + state_bytes, at, Count{ .runs = &runs }, Count.add); try testing.expectEqual(@as(usize, 1 + 2 + 20), runs.items.len); try testing.expectEqual([3]usize{ (1 << 36) + state_bytes, 0, state_bytes }, runs.items[0]); try testing.expectEqual([3]usize{ (1 << 40) + 7 * page_bytes, state_bytes, 2 * page_bytes }, runs.items[1]); @@ -175,10 +175,10 @@ test "a state on whole pages has no partial copies; one past a slot's pages is r const m = fake(&ids, &caps); var runs: std.ArrayList([3]usize) = .empty; defer runs.deinit(testing.allocator); - try each(&m, 0, 2 * kern.page, Count{ .runs = &runs }, Count.add); + try each(&m, 0, m.states, 2 * kern.page, Count{ .runs = &runs }, Count.add); try testing.expectEqual(@as(usize, 2), runs.items.len); try testing.expectEqual(@as(usize, 2 * page_bytes), runs.items[1][2]); - try testing.expectError(error.PastPages, each(&m, 0, 2 * kern.page + 1, Count{ .runs = &runs }, Count.add)); + try testing.expectError(error.PastPages, each(&m, 0, m.states, 2 * kern.page + 1, Count{ .runs = &runs }, Count.add)); } const Run = struct { From 10f32746c59ea17e858ff641a37647bacafeaf43 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 28 Oct 2026 11:35:40 +0100 Subject: [PATCH 12/12] docs: LFM2.5 on CUDA leases pool pages and keeps chat states in the prompt cache --- docs/recipes/lfm2.5-24b-a2b.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/docs/recipes/lfm2.5-24b-a2b.md b/docs/recipes/lfm2.5-24b-a2b.md index 1df62b6c51..2e770c0864 100644 --- a/docs/recipes/lfm2.5-24b-a2b.md +++ b/docs/recipes/lfm2.5-24b-a2b.md @@ -53,7 +53,10 @@ zig-out/native/bin/tensorfold-native serve "$HOME/models/lfm2.5-24b-a2b-4bit" -- ``` The default window is 32,768 tokens (`--context`), and `--parallel` holds up to 128 streams if the memory budget fits -them. Each stream has fixed pages for its keys and values. The family reads MLX affine 4-bit weights in groups of +them. Streams lease 512-position pages of keys and values from one pool as they need them +(`TENSORFOLD_LFM2_KV_PAGES=0` gives each stream fixed pages instead). The prompt cache keeps a chat's state at its +history marks, the conv states and the K/V rows before the mark, in pages allocated once inside its budget +(`--prompt-cache-gib`). The family reads MLX affine 4-bit weights in groups of 64, with the embedding at 4, 6 or 8 bits. The 6- and 8-bit exports are refused by name and stay on Metal. Drafts are the lane core's copy drafts from the prompt's own text; there is no draft head yet.