diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a381a1f72..ac2b0fc8ef 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -70,6 +70,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - A conversation state on disk that its model refuses while reading it back is forgotten after one recall, as a file that fails its own checks already is, so later turns stop reading it again (#1543, for #1539). Thanks to @wojo. +- Qwen3.6-35B-A3B on CUDA drafts every stream of a shared round together with its MTP layer: absorb passes, then passes of up to 16 rows a draft level for all streams, each stream's drafts unchanged, and the drafter's ring attention runs ahead of its loads (#1562, for #1560). On one GB10 with a 4.2K-token prompt, aggregate output at 2, 4 and 8 users is 137, 182 and 220 tok/s, against 116, 138 and 159 tok/s with `TF_MTP_BATCH=0` (one stream at a time) on the same build. Thanks to @laudney. + - GLM on CUDA: as the load packs each layer, 8 reader threads read that layer's weights into the page cache ahead of it, so GLM-5.3-Flash XSFP4 loads in about 155 s instead of 327 s on two GB10s (TP=2) and 123 s instead of 268 s on three (TP=3), with the same replies and the same memory free at serving (#1616). A rank now reads its layers whole from the disk, its unkept expert rows included, and still drops each layer's pages once it is packed. `TF_GLM_LOAD_PREFETCH` sets the count (`0` turns the readers off, at most 16); they run only on Linux. Thanks to @Bizuayeu. ## 1.1.9 diff --git a/docs/recipes/qwen3.6-moe.md b/docs/recipes/qwen3.6-moe.md index 28d7696b96..02665e0a9f 100644 --- a/docs/recipes/qwen3.6-moe.md +++ b/docs/recipes/qwen3.6-moe.md @@ -101,8 +101,9 @@ tensorfold serve TensorFold/Qwen3.6-35B-A3B-MLX-4bit-MTP --parallel 8 --name ben `--parallel N` decodes up to N requests in shared rounds, and every reply equals the same request served alone and its `"draft": false` run. A round verifies every stream's MTP chain or copied continuation in one forward; -the head absorbs every stream's kept rows in one call, then the chains advance a step at a time for all streams -by the one-stream rule, so a stream drafts what it drafts alone. A new prompt prefills 1,024 tokens a round +the head absorbs every stream's kept rows in passes of up to 16 rows, then each draft level runs every stream +still drafting in passes of up to 16 rows, with one vocabulary-head read and one synchronization a pass, so a +stream drafts exactly what it drafts alone (`TF_MTP_BATCH=0` drafts one stream at a time, for comparison). A new prompt prefills 1,024 tokens a round while the others decode, states are kept at message starts and prompt ends, and each stream's caches are sized once, at admission. Startup admits N full prompt/reply windows, three kept prompt ends and the graph buffers before loading; an explicit `--context` that does not fit is refused with the window that does. Rounds over @@ -124,6 +125,12 @@ Every reply's token SHA-256 is the same at each N and with `"draft": false`. The so copied continuations keep 9.3 tokens a round per stream; chats keep 2.9. One client at a time gets 912 tok/s on the label requests at `--parallel 8`, as a lone stream replays the graphs. +On one DGX Spark (GB10, sm_121, CUDA 13.0), checkpoint revision f84b054, the fleet benchmark's 4.2K-token prompt +(`bench.py` of spark-sglang-docker, 512 tokens, thinking on) gives 115.7 tok/s single-stream decode and 136.8 / +181.7 / 220.1 tok/s aggregate at 2 / 4 / 8 users, against 112.6 and 115.5 / 137.9 / 159.0 with `TF_MTP_BATCH=0` on +the same binary. The drafter's attention over its ring loads eight slots ahead of the fold, in the same order: a +4,200-slot ring takes 0.24 ms where it took 1.6 ms. + ## Measurements One DGX Spark (GB10) in NVIDIA's `pytorch:26.07-py3` container, checkpoint revision 81169a9, against vLLM serving diff --git a/zig/kernels/cuda/qwen36_mtp.cu b/zig/kernels/cuda/qwen36_mtp.cu index 196af0d088..642fe7538b 100644 --- a/zig/kernels/cuda/qwen36_mtp.cu +++ b/zig/kernels/cuda/qwen36_mtp.cu @@ -14,6 +14,62 @@ __device__ __forceinline__ float rbf(float x) { return __bfloat162float(__float2 using tf_common::warp_sum; +constexpr int DPL = HD / 32; // a lane's elements of a head row +constexpr int AHEAD = 8; // ring slots loaded a group, the next group's loads in flight while one folds + +// A key's score: the lane's products summed in order, then across the warp, scaled by 1/sqrt(HD). +__device__ __forceinline__ float score(const float (&q)[DPL], const __nv_bfloat16* k) { + float dot = 0.0f; +#pragma unroll + for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(k[i]), dot); + return warp_sum(dot) * (1.0f / 16.0f); +} + +// One key folded into the running softmax (mx, l) and the lane's weighted values o. +__device__ __forceinline__ void fold(float s, const __nv_bfloat16* v, float& mx, float& l, float (&o)[DPL]) { + const float mn = fmaxf(mx, s); + const float a = __expf(mx - mn), b = __expf(s - mn); + l = l * a + b; +#pragma unroll + for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); + mx = mn; +} + +// Ring slots [begin, end) of a ring of w folded in order: the same terms one at a time, with loads run ahead. +__device__ __forceinline__ void fold_ring(const __nv_bfloat16* rk, const __nv_bfloat16* rv, int w, int width, size_t off, long long begin, + long long end, const float (&q)[DPL], float& mx, float& l, float (&o)[DPL]) { + long long pos = begin; + int slot = (int)(begin % w); + uint4 kn[AHEAD], vn[AHEAD]; + const auto load = [&](uint4 (&k)[AHEAD], uint4 (&v)[AHEAD]) { +#pragma unroll + for (int u = 0; u < AHEAD; u++) { + k[u] = __ldg(reinterpret_cast(rk + (size_t)slot * width + off)); // read-only path, as a __restrict__ ring's + v[u] = __ldg(reinterpret_cast(rv + (size_t)slot * width + off)); + slot = slot + 1 == w ? 0 : slot + 1; + } + }; + if (end - pos >= AHEAD) load(kn, vn); + for (; end - pos >= AHEAD; pos += AHEAD) { + uint4 kr[AHEAD], vr[AHEAD]; +#pragma unroll + for (int u = 0; u < AHEAD; u++) { + kr[u] = kn[u]; + vr[u] = vn[u]; + } + if (end - pos >= 2 * AHEAD) load(kn, vn); // differences: no position near the type's end overflows + float s[AHEAD]; +#pragma unroll + for (int u = 0; u < AHEAD; u++) s[u] = score(q, reinterpret_cast(&kr[u])); +#pragma unroll + for (int u = 0; u < AHEAD; u++) fold(s[u], reinterpret_cast(&vr[u]), mx, l, o); + } + for (; pos < end; pos++) { + fold(score(q, rk + (size_t)slot * width + off), rv + (size_t)slot * width + off, mx, l, o); + slot = slot + 1 == w ? 0 : slot + 1; + } +} + } // namespace tf_q36m using namespace tf_q36m; @@ -22,7 +78,6 @@ using namespace tf_q36m; extern "C" __global__ void __launch_bounds__(32) tf_q36m_norm_rope(const __nv_bfloat16* __restrict__ X, const __nv_bfloat16* __restrict__ W, __nv_bfloat16* __restrict__ OUT, int heads, int stride, int in_head, long long pos0, int step, float eps, float theta) { - constexpr int DPL = HD / 32; const int head = blockIdx.x % heads; const size_t m = blockIdx.x / heads; const int t = threadIdx.x; @@ -62,12 +117,10 @@ extern "C" __global__ void __launch_bounds__(32) tf_q36m_attn(const __nv_bfloat1 const __nv_bfloat16* __restrict__ KB, const __nv_bfloat16* __restrict__ VB, const __nv_bfloat16* __restrict__ RK, const __nv_bfloat16* __restrict__ RV, int heads, int kv_heads, int W, long long begin, long long pos0, __nv_bfloat16* __restrict__ OUT) { - constexpr int DPL = HD / 32; const int lane = threadIdx.x; const int m = blockIdx.x, h = blockIdx.y; const int kvh = h / (heads / kv_heads); const int width = kv_heads * HD; - const float scale = 1.0f / 16.0f; float q[DPL], o[DPL]; #pragma unroll for (int i = 0; i < DPL; i++) { @@ -75,33 +128,50 @@ extern "C" __global__ void __launch_bounds__(32) tf_q36m_attn(const __nv_bfloat1 o[i] = 0.0f; } float mx = -INFINITY, l = 0.0f; - for (long long pos = begin; pos < pos0; pos++) { - const __nv_bfloat16* k = RK + (size_t)(pos % W) * width + kvh * HD + lane * DPL; - float dot = 0.0f; -#pragma unroll - for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(k[i]), dot); - const float s = warp_sum(dot) * scale; - const float mn = fmaxf(mx, s); - const float a = __expf(mx - mn), b = __expf(s - mn); - l = l * a + b; - const __nv_bfloat16* v = RV + (size_t)(pos % W) * width + kvh * HD + lane * DPL; -#pragma unroll - for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); - mx = mn; - } + fold_ring(RK, RV, W, width, (size_t)kvh * HD + lane * DPL, begin, pos0, q, mx, l, o); for (int r = 0; r <= m; r++) { - const __nv_bfloat16* k = KB + ((size_t)r * kv_heads + kvh) * HD + lane * DPL; - float dot = 0.0f; + const size_t at = ((size_t)r * kv_heads + kvh) * HD + lane * DPL; + fold(score(q, KB + at), VB + at, mx, l, o); + } + const __nv_bfloat16* g = G + (size_t)m * gstride + (size_t)h * (2 * HD) + HD + lane * DPL; // the gate half #pragma unroll - for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(k[i]), dot); - const float s = warp_sum(dot) * scale; - const float mn = fmaxf(mx, s); - const float a = __expf(mx - mn), b = __expf(s - mn); - l = l * a + b; - const __nv_bfloat16* v = VB + ((size_t)r * kv_heads + kvh) * HD + lane * DPL; + for (int i = 0; i < DPL; i++) { + const float gate = 1.0f / (1.0f + __expf(-bf(g[i]))); + OUT[((size_t)m * heads + h) * HD + lane * DPL + i] = __float2bfloat16_rn(rbf(o[i] / l) * gate); + } +} + +// A packed row's ring: its keys and values, its slots [begin, end) of a ring of w, and its segment's first row. +struct RingRow { + const __nv_bfloat16* rk; + const __nv_bfloat16* rv; + long long begin; + long long end; + int seg; + int w; +}; + +// tf_q36m_attn over many streams' rows, a warp a (row, head): row m's ring slots T[m], then the block rows T[m].seg..m. +extern "C" __global__ void __launch_bounds__(32) tf_q36m_attn_rows(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ G, int gstride, + const __nv_bfloat16* __restrict__ KB, const __nv_bfloat16* __restrict__ VB, + const RingRow* __restrict__ T, int heads, int kv_heads, + __nv_bfloat16* __restrict__ OUT) { + const int lane = threadIdx.x; + const int m = blockIdx.x, h = blockIdx.y; + const RingRow t = T[m]; + const int kvh = h / (heads / kv_heads); + const int width = kv_heads * HD; + float q[DPL], o[DPL]; #pragma unroll - for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); - mx = mn; + for (int i = 0; i < DPL; i++) { + q[i] = bf(Q[((size_t)m * heads + h) * HD + lane * DPL + i]); + o[i] = 0.0f; + } + float mx = -INFINITY, l = 0.0f; + fold_ring(t.rk, t.rv, t.w, width, (size_t)kvh * HD + lane * DPL, t.begin, t.end, q, mx, l, o); + for (int r = t.seg; r <= m; r++) { + const size_t at = ((size_t)r * kv_heads + kvh) * HD + lane * DPL; + fold(score(q, KB + at), VB + at, mx, l, o); } const __nv_bfloat16* g = G + (size_t)m * gstride + (size_t)h * (2 * HD) + HD + lane * DPL; // the gate half #pragma unroll @@ -118,12 +188,10 @@ extern "C" __global__ void __launch_bounds__(32) tf_q36m_attn_tree(const __nv_bf const int* __restrict__ ANC, int anc_stride, const __nv_bfloat16* __restrict__ RK, const __nv_bfloat16* __restrict__ RV, int heads, int kv_heads, int W, long long begin, long long end, __nv_bfloat16* __restrict__ OUT) { - constexpr int DPL = HD / 32; const int lane = threadIdx.x; const int m = blockIdx.x, h = blockIdx.y; const int kvh = h / (heads / kv_heads); const int width = kv_heads * HD; - const float scale = 1.0f / 16.0f; float q[DPL], o[DPL]; #pragma unroll for (int i = 0; i < DPL; i++) { @@ -131,36 +199,13 @@ extern "C" __global__ void __launch_bounds__(32) tf_q36m_attn_tree(const __nv_bf o[i] = 0.0f; } float mx = -INFINITY, l = 0.0f; - for (long long pos = begin; pos < end; pos++) { - const __nv_bfloat16* k = RK + (size_t)(pos % W) * width + kvh * HD + lane * DPL; - float dot = 0.0f; -#pragma unroll - for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(k[i]), dot); - const float s = warp_sum(dot) * scale; - const float mn = fmaxf(mx, s); - const float a = __expf(mx - mn), b = __expf(s - mn); - l = l * a + b; - const __nv_bfloat16* v = RV + (size_t)(pos % W) * width + kvh * HD + lane * DPL; -#pragma unroll - for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); - mx = mn; - } + const size_t off = (size_t)kvh * HD + lane * DPL; + fold_ring(RK, RV, W, width, off, begin, end, q, mx, l, o); const int* anc = ANC + (size_t)m * anc_stride; const int n = anc[0]; for (int j = 0; j <= n; j++) { - const bool self = j == n; - const __nv_bfloat16* k = (self ? KB + (size_t)m * width : LK + (size_t)anc[1 + j] * width) + kvh * HD + lane * DPL; - float dot = 0.0f; -#pragma unroll - for (int i = 0; i < DPL; i++) dot = fmaf(q[i], bf(k[i]), dot); - const float s = warp_sum(dot) * scale; - const float mn = fmaxf(mx, s); - const float a = __expf(mx - mn), b = __expf(s - mn); - l = l * a + b; - const __nv_bfloat16* v = (self ? VB + (size_t)m * width : LV + (size_t)anc[1 + j] * width) + kvh * HD + lane * DPL; -#pragma unroll - for (int i = 0; i < DPL; i++) o[i] = fmaf(b, bf(v[i]), o[i] * a); - mx = mn; + const size_t at = (j == n ? (size_t)m : (size_t)anc[1 + j]) * width + off; + fold(score(q, (j == n ? KB : LK) + at), (j == n ? VB : LV) + at, mx, l, o); } const __nv_bfloat16* g = G + (size_t)m * gstride + (size_t)h * (2 * HD) + HD + lane * DPL; #pragma unroll diff --git a/zig/qwen35_draft_absorb_host.zig b/zig/qwen35_draft_absorb_host.zig index 5e878a0faf..4e53a41a1d 100644 --- a/zig/qwen35_draft_absorb_host.zig +++ b/zig/qwen35_draft_absorb_host.zig @@ -83,7 +83,8 @@ test "Draft precision: MTP real absorb adapter forwards both row roles and resto d.ring_fns.gather_taps = function(&driver); d.ring_fns.gather_rows = function(&driver); d.taps_in = 1; - var ctx: q.mtp.Ctx = .{ .ring = .{ .d = &driver, .ptr = 1, .len = 1 }, .last = .{ .d = &driver, .ptr = 1, .len = 1 } }; + d.batch = false; // the serial absorb's roles and refusals; the batch has its own GPU check + var ctx: q.mtp.Ctx = .{ .ring = .{ .d = &driver, .ptr = 1, .len = 1 }, .last = .{ .d = &driver, .ptr = 1, .len = 1 }, .step = .{ .d = &driver, .ptr = 1, .len = 1 } }; var s = try lanes.Stream.init(T.allocator, .{ .id = "mtp", .prompt = &.{1}, .max_new = 8 }); defer s.deinit(T.allocator); try check(&d, &ctx, &s); diff --git a/zig/src/core/contract/qwen_served_records.zig b/zig/src/core/contract/qwen_served_records.zig index 2aa03bc528..40e3cc32ed 100644 --- a/zig/src/core/contract/qwen_served_records.zig +++ b/zig/src/core/contract/qwen_served_records.zig @@ -11,16 +11,15 @@ pub const records = [_]registry.CudaRecord{ .{ .launch = .{ .name = "experts plan_rank_rows", .symbol = "_ZN10tf_experts14plan_rank_rowsEPKiiiPiS2_S1_i", .image = "experts", .threads = 1024, .min_sm = 80 }, .usage = .{ .{ .regs = 32, .static = 4096, .bound = 1024 }, .{ .regs = 40, .static = 4096, .bound = 1024 }, .{ .regs = 40, .static = 4096, .bound = 1024 }, .{ .regs = 42, .static = 4096, .bound = 1024 }, .{ .regs = 42, .static = 4096, .bound = 1024 } } }, .{ .launch = .{ .name = "experts plan_scatter_rows", .symbol = "_ZN10tf_experts17plan_scatter_rowsEPKiiiS1_S1_PiS1_i", .image = "experts", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 14, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "tf_q36m_norm_rope", .symbol = "tf_q36m_norm_rope", .image = "qwen36_mtp", .threads = 32, .min_sm = 80 }, .usage = .{ .{ .regs = 36, .static = 0, .bound = 32 }, .{ .regs = 37, .static = 0, .bound = 32 }, .{ .regs = 39, .static = 0, .bound = 32 }, .{ .regs = 42, .static = 0, .bound = 32 }, .{ .regs = 42, .static = 0, .bound = 32 } } }, - .{ .launch = .{ .name = "tf_q36m_attn", .symbol = "tf_q36m_attn", .image = "qwen36_mtp", .threads = 32, .min_sm = 80 }, .usage = .{ .{ .regs = 50, .static = 0, .bound = 32 }, .{ .regs = 52, .static = 0, .bound = 32 }, .{ .regs = 51, .static = 0, .bound = 32 }, .{ .regs = 54, .static = 0, .bound = 32 }, .{ .regs = 54, .static = 0, .bound = 32 } } }, - .{ .launch = .{ .name = "tf_q36m_attn_tree", .symbol = "tf_q36m_attn_tree", .image = "qwen36_mtp", .threads = 32, .min_sm = 80 }, .usage = .{ .{ .regs = 53, .static = 0, .bound = 32 }, .{ .regs = 54, .static = 0, .bound = 32 }, .{ .regs = 52, .static = 0, .bound = 32 }, .{ .regs = 57, .static = 0, .bound = 32 }, .{ .regs = 57, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "tf_q36m_attn", .symbol = "tf_q36m_attn", .image = "qwen36_mtp", .threads = 32, .min_sm = 80 }, .usage = .{ .{ .regs = 164, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 162, .static = 0, .bound = 32 }, .{ .regs = 162, .static = 0, .bound = 32 } }, .later_pending = true }, + .{ .launch = .{ .name = "tf_q36m_attn_rows", .symbol = "tf_q36m_attn_rows", .image = "qwen36_mtp", .threads = 32, .min_sm = 80 }, .usage = .{ .{ .regs = 158, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 160, .static = 0, .bound = 32 }, .{ .regs = 160, .static = 0, .bound = 32 } }, .later_pending = true }, + .{ .launch = .{ .name = "tf_q36m_attn_tree", .symbol = "tf_q36m_attn_tree", .image = "qwen36_mtp", .threads = 32, .min_sm = 80 }, .usage = .{ .{ .regs = 163, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 168, .static = 0, .bound = 32 }, .{ .regs = 159, .static = 0, .bound = 32 }, .{ .regs = 159, .static = 0, .bound = 32 } }, .later_pending = true }, }; -/// nvcc 13.2's usage where it differs from 13.0's. -pub const cuda_132 = [_]registry.Override{ - .{ .name = "tf_q36m_attn", .usage = .{ .{ .regs = 50, .static = 0, .bound = 32 }, .{ .regs = 52, .static = 0, .bound = 32 }, .{ .regs = 51, .static = 0, .bound = 32 }, .{ .regs = 52, .static = 0, .bound = 32 }, .{ .regs = 52, .static = 0, .bound = 32 } } }, -}; +/// nvcc 13.2's usage where it differs from 13.0's (the MTP attention's own 13.2 receipts are pending). +pub const cuda_132 = [_]registry.Override{}; -test "the nine served MoE, plan and MTP launches are in the live registry and fit every SM by their records" { +test "the ten served MoE, plan and MTP launches are in the live registry and fit every SM by their records" { for (records) |record| { var found: usize = 0; for (registry.all_cuda) |active| if (std.mem.eql(u8, active.launch.symbol, record.launch.symbol) and std.mem.eql(u8, active.launch.image, record.launch.image)) { diff --git a/zig/src/families/qwen3_5/cuda.zig b/zig/src/families/qwen3_5/cuda.zig index 75f8c3598a..0230e0c199 100644 --- a/zig/src/families/qwen3_5/cuda.zig +++ b/zig/src/families/qwen3_5/cuda.zig @@ -24,6 +24,7 @@ test { _ = native; _ = native_moe; _ = mtp; + _ = @import("cuda_mtp_batch.zig"); _ = config; _ = geometry; _ = @import("draw.zig"); diff --git a/zig/src/families/qwen3_5/cuda_mtp.zig b/zig/src/families/qwen3_5/cuda_mtp.zig index e2166c3ea5..7dab9cf016 100644 --- a/zig/src/families/qwen3_5/cuda_mtp.zig +++ b/zig/src/families/qwen3_5/cuda_mtp.zig @@ -8,6 +8,7 @@ const geometry = @import("cuda_geometry.zig"); const weights = @import("cuda_mtp_weights.zig"); const timing = @import("cuda_draft.zig").timing; const precision = @import("cuda_draft_precision.zig"); +const batched = @import("cuda_mtp_batch.zig"); const E = engine.Engine; const DeviceBuffer = cuda.DeviceBuffer; @@ -25,15 +26,16 @@ const ranks = lanes.shape.ranks; /// The logprob row's words (sample.cu's LP_WORDS): the ids of a row's best tokens sit at words 2 onward. const lp_words = 3 + 2 * 20; -const Fns = struct { norm_rope: cuda.Function, attn: cuda.Function, attn_tree: cuda.Function }; +const Fns = struct { norm_rope: cuda.Function, attn: cuda.Function, attn_tree: cuda.Function, attn_rows: cuda.Function }; /// qwen35_draft.cu's ring kernels, shared with the DFlash2 drafter. const Ring = struct { gather_taps: cuda.Function, gather_rows: cuda.Function, ring_put: cuda.Function }; -const TapRows = extern struct { n: c_int, rows: [most_window]c_int }; +pub const TapRows = extern struct { n: c_int, rows: [most_window]c_int }; /// A stream's context: its ring of keys and values over [begin, end), the newest row's state, and its held drafts. pub const Ctx = struct { ring: DeviceBuffer, last: DeviceBuffer, // the newest absorbed row's normed output (bf16 hidden): the first draft row's state + step: DeviceBuffer, // a batched chain's previous draft level's normed output (cuda_mtp_batch.zig) begin: u64 = 0, end: u64 = 0, cands: [ranks]u32 = @splat(0), // the newest absorbed row's best tokens by rank (rank 0: its chain draft) @@ -89,6 +91,9 @@ pub const Mtp = struct { ids: u64, // u32 the rows' tokens tok: u64, // u32 the drawn token staged: cuda.HostBuffer, // pinned: the ids and ancestors up, the token and words down + tab: u64, // a batched pass's RingRow table (cuda_mtp_batch.zig) + batch: bool = true, // shared rounds draft every stream in one pass a level; TF_MTP_BATCH=0 keeps one at a time + batch_passes: u64 = 0, // batched passes run, for the GPU check streams: std.AutoHashMapUnmanaged(*Stream, *Ctx) = .empty, /// Contexts between streams, allocated and run once (`reserve`), as the driver stalls a buffer's first kernel. pool: std.ArrayList(*Ctx) = .empty, @@ -102,18 +107,21 @@ pub const Mtp = struct { const anc_stride = 1 + max_lanes; const ids_at = 0; - const tok_at = max_rows * 4; + pub const tok_at = max_rows * 4; const words_at = tok_at + 64; const anc_at = words_at + max_lanes * lp_words * 4; - const staged_len = anc_at + max_lanes * anc_stride * 4; + pub const bids_at = anc_at + max_lanes * anc_stride * 4; // a batched pass's ids, then its ring table + pub const btab_at = bids_at + batched.most * 4; + const staged_len = btab_at + batched.most * @sizeOf(batched.RingRow); pub fn init(gpa: std.mem.Allocator, io: std.Io, e: *E, dir: []const u8, capacity: usize) !*Mtp { if (!e.g.moe()) return error.NotRouted; const d = try gpa.create(Mtp); errdefer gpa.destroy(d); - d.* = .{ .gpa = gpa, .io = io, .e = e, .w = undefined, .mod = undefined, .dmod = undefined, .fns = undefined, .ring_fns = undefined, .scratch = undefined, .capacity = capacity, .depth = default_depth, .trees = platform.c.getenv("TF_MTP_TREES") != null, .tap = .{@intCast(e.g.layers - 1)}, .taps_in = 0, .hin = 0, .hn = 0, .emb = 0, .en = 0, .h = 0, .r = 0, .x = 0, .qraw = 0, .q = 0, .kraw = 0, .kb = 0, .vb = 0, .att = 0, .xs = 0, .logits = 0, .words = 0, .picks = 0, .lane_state = 0, .lane_k = 0, .lane_v = 0, .anc = 0, .ids = 0, .tok = 0, .staged = undefined }; + d.* = .{ .gpa = gpa, .io = io, .e = e, .w = undefined, .mod = undefined, .dmod = undefined, .fns = undefined, .ring_fns = undefined, .scratch = undefined, .capacity = capacity, .depth = default_depth, .trees = platform.c.getenv("TF_MTP_TREES") != null, .tap = .{@intCast(e.g.layers - 1)}, .taps_in = 0, .hin = 0, .hn = 0, .emb = 0, .en = 0, .h = 0, .r = 0, .x = 0, .qraw = 0, .q = 0, .kraw = 0, .kb = 0, .vb = 0, .att = 0, .xs = 0, .logits = 0, .words = 0, .picks = 0, .lane_state = 0, .lane_k = 0, .lane_v = 0, .anc = 0, .ids = 0, .tok = 0, .staged = undefined, .tab = 0 }; if (platform.c.getenv("TF_MTP_DRAFTS")) |v| d.depth = std.fmt.parseInt(u32, std.mem.span(v), 10) catch default_depth; d.depth = @min(@max(d.depth, 1), most_window - 1); + if (platform.c.getenv("TF_MTP_BATCH")) |v| d.batch = !std.mem.eql(u8, std.mem.span(v), "0"); d.mod = try cuda.Module.load(e.d, cuda.kernels.qwen36_mtp); errdefer d.mod.unload(); d.dmod = try cuda.Module.load(e.d, cuda.kernels.qwen35_draft); @@ -177,13 +185,17 @@ pub const Mtp = struct { return d.e.g.kv_heads * head_dim; } - fn ringBytes(d: *const Mtp) usize { + pub fn ringBytes(d: *const Mtp) usize { return 2 * d.capacity * d.kvWidth() * 2; } - /// A stream's buffers in a pool lane's run: its ring, then (256-aligned) its newest row's state. + /// A stream's buffers in a pool lane's run: ring, newest row's state, then chain step state, each 256-aligned. pub fn areaBytes(d: *const Mtp) usize { - return std.mem.alignForward(usize, d.ringBytes(), 256) + d.e.g.hidden * 2; + return d.stepOffset() + d.e.g.hidden * 2; + } + + fn stepOffset(d: *const Mtp) usize { + return std.mem.alignForward(usize, std.mem.alignForward(usize, d.ringBytes(), 256) + d.e.g.hidden * 2, 256); } fn scratchBytes(d: *Mtp) usize { @@ -199,12 +211,13 @@ pub const Mtp = struct { const aw = g.attnWidth(); const kw = d.kvWidth(); const sizes = .{ - .{ "taps_in", R * g.hidden * 2 }, .{ "hin", R * g.hidden * 2 }, .{ "hn", R * g.hidden * 2 }, .{ "emb", R * g.hidden * 2 }, - .{ "en", R * g.hidden * 2 }, .{ "h", R * g.hidden * 2 }, .{ "r", R * g.hidden * 2 }, .{ "x", R * g.hidden * 2 }, - .{ "qraw", R * 2 * aw * 2 }, .{ "q", R * aw * 2 }, .{ "kraw", R * kw * 2 }, .{ "kb", R * kw * 2 }, - .{ "vb", R * kw * 2 }, .{ "att", R * aw * 2 }, .{ "xs", R * (@max(g.hidden, aw) / 64) * 4 }, .{ "logits", L * g.vocab * 2 }, - .{ "words", L * lp_words * 4 }, .{ "picks", L * 4 }, .{ "lane_state", L * g.hidden * 2 }, .{ "lane_k", L * kw * 2 }, - .{ "lane_v", L * kw * 2 }, .{ "anc", L * anc_stride * 4 }, .{ "ids", R * 4 }, .{ "tok", 64 }, + .{ "taps_in", R * g.hidden * 2 }, .{ "hin", R * g.hidden * 2 }, .{ "hn", R * g.hidden * 2 }, .{ "emb", R * g.hidden * 2 }, + .{ "en", R * g.hidden * 2 }, .{ "h", R * g.hidden * 2 }, .{ "r", R * g.hidden * 2 }, .{ "x", R * g.hidden * 2 }, + .{ "qraw", R * 2 * aw * 2 }, .{ "q", R * aw * 2 }, .{ "kraw", R * kw * 2 }, .{ "kb", R * kw * 2 }, + .{ "vb", R * kw * 2 }, .{ "att", R * aw * 2 }, .{ "xs", R * (@max(g.hidden, aw) / 64) * 4 }, .{ "logits", L * g.vocab * 2 }, + .{ "words", L * lp_words * 4 }, .{ "picks", L * 4 }, .{ "lane_state", L * g.hidden * 2 }, .{ "lane_k", L * kw * 2 }, + .{ "lane_v", L * kw * 2 }, .{ "anc", L * anc_stride * 4 }, .{ "ids", R * 4 }, .{ "tok", 64 }, + .{ "tab", batched.most * @sizeOf(batched.RingRow) }, }; var total: usize = 0; inline for (sizes) |s| { @@ -216,6 +229,7 @@ pub const Mtp = struct { } pub fn drop(d: *Mtp, ctx: *Ctx) void { + ctx.step.free(); ctx.last.free(); ctx.ring.free(); d.gpa.destroy(ctx); @@ -223,17 +237,20 @@ pub const Mtp = struct { fn resetCtx(d: *Mtp, ctx: *Ctx) void { _ = d; - ctx.* = .{ .ring = ctx.ring, .last = ctx.last }; + ctx.* = .{ .ring = ctx.ring, .last = ctx.last, .step = ctx.step }; } /// A new context, touched once (an absorb, then a chain), so no request pays the driver's stall. pub fn create(d: *Mtp) !*Ctx { const ctx = try d.gpa.create(Ctx); errdefer d.gpa.destroy(ctx); - ctx.* = .{ .ring = try DeviceBuffer.alloc(d.e.d, d.ringBytes()), .last = undefined }; + ctx.* = .{ .ring = try DeviceBuffer.alloc(d.e.d, d.ringBytes()), .last = undefined, .step = undefined }; errdefer ctx.ring.free(); ctx.last = try DeviceBuffer.alloc(d.e.d, d.e.g.hidden * 2); errdefer ctx.last.free(); + ctx.step = try DeviceBuffer.alloc(d.e.d, d.e.g.hidden * 2); + errdefer ctx.step.free(); + try ctx.step.fill8(0, null); try ctx.ring.fill8(0, null); try ctx.last.fill8(0, null); try d.absorb(ctx, d.taps_in, 0, 1, null, &.{0}); // any tap rows do: the drafter's own zeroed staging @@ -273,7 +290,7 @@ pub const Mtp = struct { } /// out = norm(h + r) with w; h takes the sum (r 0: no residual). - fn norm(d: *Mtp, h: u64, r: u64, w: u64, out: u64, rows: usize) !void { + pub fn norm(d: *Mtp, h: u64, r: u64, w: u64, out: u64, rows: usize) !void { const e = d.e; try e.run(e.fns.norm, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ h, r, w, out, e.g.eps, @as(c_int, @intCast(e.g.hidden)), @as(u64, 0), @as(u64, 0) }); } @@ -346,7 +363,7 @@ pub const Mtp = struct { return d.absorbRole(ctx, tap_ring, start_, n_, rows_, follow, true); } - fn absorbRole(d: *Mtp, ctx: *Ctx, tap_ring: u64, start_: u64, n_: usize, rows_: ?[]const u32, follow: []const u32, prompt: bool) !void { + pub fn absorbRole(d: *Mtp, ctx: *Ctx, tap_ring: u64, start_: u64, n_: usize, rows_: ?[]const u32, follow: []const u32, prompt: bool) !void { const previous = d.projection_role; defer d.projection_role = previous; d.projection_role = if (prompt) .prompt else .live; @@ -391,7 +408,7 @@ pub const Mtp = struct { } /// `depth` drafts after `pending` at pos0 (the pending token's slot): a chain, or the planner's tree of lanes. - fn hold(d: *Mtp, ctx: *Ctx, pending: u32, pos0: u64, depth: usize, lanes_: ?*const lanes.shape.Shape) !void { + pub fn hold(d: *Mtp, ctx: *Ctx, pending: u32, pos0: u64, depth: usize, lanes_: ?*const lanes.shape.Shape) !void { _ = pending; // the newest absorbed row's follow token, whose candidates ctx.cands are if (ctx.end != pos0) return error.DraftContextBehind; if (lanes_) |l| return d.holdTree(ctx, pos0, l); @@ -499,7 +516,7 @@ pub const Mtp = struct { fn factsFn(p: *anyopaque) dr.Facts { const d = cast(p); - return .{ .depth = d.depth, .step_ms = d.step_ms, .prior = &prior, .plain_guard = true, .batched = false, .stream_bytes = d.ringBytes() + d.e.g.hidden * 2, .trees = d.trees, .window_costs = d.costs[0..d.n_costs], .lane_costs = d.costs[0..d.n_costs] }; + return .{ .depth = d.depth, .step_ms = d.step_ms, .prior = &prior, .plain_guard = true, .batched = !d.trees, .stream_bytes = d.ringBytes() + 2 * d.e.g.hidden * 2, .trees = d.trees, .window_costs = d.costs[0..d.n_costs], .lane_costs = d.costs[0..d.n_costs] }; } fn openFn(p: *anyopaque, s: *Stream) anyerror!void { @@ -509,7 +526,7 @@ pub const Mtp = struct { const ctx = try d.gpa.create(Ctx); errdefer d.gpa.destroy(ctx); const last = base + std.mem.alignForward(usize, d.ringBytes(), 256); - ctx.* = .{ .ring = .{ .d = d.e.d, .ptr = base, .len = d.ringBytes(), .borrowed = true }, .last = .{ .d = d.e.d, .ptr = last, .len = d.e.g.hidden * 2, .borrowed = true } }; + ctx.* = .{ .ring = .{ .d = d.e.d, .ptr = base, .len = d.ringBytes(), .borrowed = true }, .last = .{ .d = d.e.d, .ptr = last, .len = d.e.g.hidden * 2, .borrowed = true }, .step = .{ .d = d.e.d, .ptr = base + d.stepOffset(), .len = d.e.g.hidden * 2, .borrowed = true } }; return d.streams.put(d.gpa, s, ctx); }; const ctx = d.pool.pop() orelse try d.create(); @@ -522,20 +539,26 @@ pub const Mtp = struct { const d = cast(p); const t0 = std.Io.Clock.awake.now(d.io); defer timing("mtp absorb", d.io, t0); - for (items) |it| { - const ctx = d.streams.get(it.stream) orelse return error.UnknownStream; - if (it.features.space != .device or it.features.dtype != .bf16) return error.UnsupportedFeatures; - // the target's features: its ring, `offset` the first row's position (cuda_native.zig's featuresFn) - if (it.rows) |rs| { - try d.absorbRole(ctx, it.features.buffer, it.features.offset, rs.len, rs, it.follow, it.prompt); - } else try d.absorbRole(ctx, it.features.buffer, it.features.offset, it.features.rows, null, it.follow, it.prompt); - } + // a shared round's streams in passes; one stream alone runs the same bits without the batch's copies + if (d.batch and !d.trees and items.len > 1) return batched.absorb(d, items); + for (items) |it| try d.absorbOne(it); + } + + /// One stream's absorb, as the serial path runs it. + pub fn absorbOne(d: *Mtp, it: dr.Absorb) !void { + const ctx = d.streams.get(it.stream) orelse return error.UnknownStream; + if (it.features.space != .device or it.features.dtype != .bf16) return error.UnsupportedFeatures; + // the target's features: its ring, `offset` the first row's position (cuda_native.zig's featuresFn) + if (it.rows) |rs| { + try d.absorbRole(ctx, it.features.buffer, it.features.offset, rs.len, rs, it.follow, it.prompt); + } else try d.absorbRole(ctx, it.features.buffer, it.features.offset, it.features.rows, null, it.follow, it.prompt); } fn holdFn(p: *anyopaque, items: []const dr.Hold) anyerror!void { const d = cast(p); const t0 = std.Io.Clock.awake.now(d.io); defer timing("mtp hold", d.io, t0); + if (d.batch and !d.trees and items.len > 1) return batched.hold(d, items); for (items) |it| { const ctx = d.streams.get(it.stream) orelse return error.UnknownStream; try d.hold(ctx, it.pending, it.position - 1, it.depth, it.lanes); diff --git a/zig/src/families/qwen3_5/cuda_mtp_area_test.zig b/zig/src/families/qwen3_5/cuda_mtp_area_test.zig index 7af0626f13..4fe97c6d46 100644 --- a/zig/src/families/qwen3_5/cuda_mtp_area_test.zig +++ b/zig/src/families/qwen3_5/cuda_mtp_area_test.zig @@ -1,4 +1,4 @@ -//! MTP's ring and state in a pool lane's run: the area's layout, borrowed on open, given back on release. +//! MTP's ring and states in a pool lane's run: the area's layout, borrowed on open, given back on release. const std = @import("std"); const lanes = @import("lanes"); const mtp = @import("cuda_mtp.zig"); @@ -15,7 +15,7 @@ const Runs = struct { } }; -test "MTP on a pool lane borrows its run's area: ring, then 256-aligned state, given back with the lane" { +test "MTP on a pool lane borrows its run's area: ring, then 256-aligned states, given back with the lane" { var e: E = undefined; e.d = undefined; e.g.kv_heads = 2; @@ -31,7 +31,7 @@ test "MTP on a pool lane borrows its run's area: ring, then 256-aligned state, g defer d.streams.deinit(t.allocator); defer d.pool.deinit(t.allocator); // the reserved window's stream takes the drafter's own context - var own = mtp.Ctx{ .ring = .{ .d = undefined, .ptr = 0x1000, .len = 1 }, .last = .{ .d = undefined, .ptr = 0x2000, .len = 1 } }; + var own = mtp.Ctx{ .ring = .{ .d = undefined, .ptr = 0x1000, .len = 1 }, .last = .{ .d = undefined, .ptr = 0x2000, .len = 1 }, .step = .{ .d = undefined, .ptr = 0x3000, .len = 1 } }; try d.pool.append(t.allocator, &own); // with lane runs, the drafter keeps one context of its own however many streams are admitted try d.reserve(8); @@ -43,15 +43,19 @@ test "MTP on a pool lane borrows its run's area: ring, then 256-aligned state, g try drafter.open(&streams[0]); try t.expectError(error.StreamOpen, drafter.open(&streams[0])); const ctx = d.streams.get(&streams[0]).?; - try t.expect(ctx.ring.borrowed and ctx.last.borrowed); + try t.expect(ctx.ring.borrowed and ctx.last.borrowed and ctx.step.borrowed); try t.expectEqual(runs.base, ctx.ring.ptr); try t.expectEqual(ring, ctx.ring.len); try t.expectEqual(runs.base + std.mem.alignForward(usize, ring, 256), ctx.last.ptr); try t.expect(ctx.ring.ptr + ctx.ring.len <= ctx.last.ptr); try t.expectEqual(@as(u64, 0), ctx.last.ptr % 256); try t.expectEqual(e.g.hidden * 2, ctx.last.len); - // the area's last byte is the state's: areaBytes is the run's ring part, nothing past it - try t.expectEqual(runs.base + d.areaBytes(), ctx.last.ptr + ctx.last.len); + // the batched chain's step state follows, 256-aligned, the same size + try t.expect(ctx.last.ptr + ctx.last.len <= ctx.step.ptr); + try t.expectEqual(@as(u64, 0), ctx.step.ptr % 256); + try t.expectEqual(e.g.hidden * 2, ctx.step.len); + // the area's last byte is the step state's: areaBytes is the run's ring part, nothing past it + try t.expectEqual(runs.base + d.areaBytes(), ctx.step.ptr + ctx.step.len); try drafter.open(&streams[1]); try t.expectEqual(&own, d.streams.get(&streams[1]).?); try t.expectEqual(@as(usize, 0), d.pool.items.len); diff --git a/zig/src/families/qwen3_5/cuda_mtp_batch.zig b/zig/src/families/qwen3_5/cuda_mtp_batch.zig new file mode 100644 index 0000000000..cde5e42200 --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_mtp_batch.zig @@ -0,0 +1,254 @@ +//! The MTP drafter's shared rounds: every stream's absorbs and chained drafts in passes of up to 16 rows. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const engine = @import("cuda_engine.zig"); +const geometry = @import("cuda_geometry.zig"); +const mtp = @import("cuda_mtp.zig"); + +const Mtp = mtp.Mtp; +const Ctx = mtp.Ctx; +const dr = lanes.drafter; +const head_dim = geometry.head_dim; + +/// Rows a pass at most: the drafter's verify-tile and routed-chunk width, so every row keeps its solo bits. +pub const most = 16; + +/// tf_q36m_attn_rows' RingRow: a packed row's ring, its slots [begin, end) of a ring of w, its segment's first row. +pub const RingRow = extern struct { rk: u64, rv: u64, begin: i64, end: i64, seg: i32, w: i32 }; + +/// A stream's rows in a pass: its context, its first packed row, its rows, and their first slot. +const Seg = struct { ctx: *Ctx, first: usize, rows: usize, p0: u64 }; + +/// Consecutive items into passes of at most `most` rows (out[i]: its items [from, to)); an item over `most` alone. +pub fn plan(rows: []const usize, out: [][2]usize) usize { + var n: usize = 0; + var i: usize = 0; + while (i < rows.len) { + var j = i; + var sum: usize = 0; + while (j < rows.len and rows[j] <= most and sum + rows[j] <= most) : (j += 1) sum += rows[j]; + if (j == i) j = i + 1; + out[n] = .{ i, j }; + n += 1; + i = j; + } + return n; +} + +/// A row's attention over its context's ring: slots [begin, end), then its segment's rows from `seg`. +fn ringRow(d: *Mtp, ctx: *Ctx, end: u64, seg: usize) RingRow { + return .{ .rk = d.ringK(ctx), .rv = d.ringV(ctx), .begin = @intCast(ctx.begin), .end = @intCast(end), .seg = @intCast(seg), .w = @intCast(d.capacity) }; +} + +/// The pass's tokens and ring table up from the batch's own staging (the last pass's fetch has waited on it). +fn upload(d: *Mtp, ids: []const u32, tab: []const RingRow) !void { + const e = d.e; + const st = d.staged.bytes; + @memcpy(std.mem.bytesAsSlice(u32, st[Mtp.bids_at..][0 .. ids.len * 4]), ids); + @memcpy(st[Mtp.btab_at..][0 .. tab.len * @sizeOf(RingRow)], std.mem.sliceAsBytes(tab)); + try cuda.copy_source.enqueue(e.d, d.ids, st[Mtp.bids_at..][0 .. ids.len * 4], e.stream); + try cuda.copy_source.enqueue(e.d, d.tab, st[Mtp.btab_at..][0 .. tab.len * @sizeOf(RingRow)], e.stream); +} + +/// Mtp.layer over every segment's rows at once: d.ids their tokens, d.hin their states; x gets their normed outputs. +fn layerRows(d: *Mtp, segs: []const Seg, rows: usize) !void { + const e = d.e; + const g = e.g; + const w = &d.w; + const aw = g.attnWidth(); + const kw = d.kvWidth(); + try e.embed(d.ids, d.emb, rows, 0); + try d.norm(d.emb, 0, w.norm_e, d.en, rows); + try d.norm(d.hin, 0, w.norm_h, d.hn, rows); + try d.proj(.{ .a4 = w.fc_e }, d.en, d.h, rows); + try d.proj(.{ .a4 = w.fc_h }, d.hn, d.r, rows); + try d.norm(d.h, d.r, w.input_norm, d.x, rows); + try d.projQkv(rows); + for (segs) |s| { // rope positions run from each segment's first slot + const p0: c_longlong = @intCast(s.p0); + try e.run(d.fns.norm_rope, .{ s.rows * g.q_heads, 1, 1 }, .{ 32, 1, 1 }, .{ d.qraw + s.first * 2 * aw * 2, w.q_norm, d.q + s.first * aw * 2, @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(2 * aw)), @as(c_int, 2 * head_dim), p0, @as(c_int, 1), g.eps, geometry.theta }); + try e.run(d.fns.norm_rope, .{ s.rows * g.kv_heads, 1, 1 }, .{ 32, 1, 1 }, .{ d.kraw + s.first * kw * 2, w.k_norm, d.kb + s.first * kw * 2, @as(c_int, @intCast(g.kv_heads)), @as(c_int, @intCast(kw)), @as(c_int, head_dim), p0, @as(c_int, 1), g.eps, geometry.theta }); + } + try e.run(d.fns.attn_rows, .{ rows, g.q_heads, 1 }, .{ 32, 1, 1 }, .{ d.q, d.qraw, @as(c_int, @intCast(2 * aw)), d.kb, d.vb, d.tab, @as(c_int, @intCast(g.q_heads)), @as(c_int, @intCast(g.kv_heads)), d.att }); + try d.proj(w.o, d.att, d.r, rows); + try d.norm(d.h, d.r, w.post_norm, d.x, rows); + try d.routed(w.mlp.moe.router, w.mlp.moe.experts, d.x, d.r, rows); + try d.norm(d.h, d.r, w.norm, d.x, rows); +} + +/// Each segment's keys and values from the pass into its context's ring, at its slots. +fn ringPut(d: *Mtp, segs: []const Seg) !void { + const e = d.e; + const kw = d.kvWidth(); + for (segs) |s| try e.run(d.ring_fns.ring_put, .{ s.rows, 1, 1 }, .{ 256, 1, 1 }, .{ d.kb + s.first * kw * 2, d.vb + s.first * kw * 2, d.ringK(s.ctx), d.ringV(s.ctx), @as(c_int, @intCast(kw)), @as(c_longlong, @intCast(s.p0)), @as(c_int, @intCast(d.capacity)) }); +} + +/// Rows of `input` through the target's head, and each row's greedy token: one projection, one argmax launch. +fn head(d: *Mtp, input: u64, rows: usize) !void { + const e = d.e; + try d.proj(e.head, input, d.logits, rows); + try e.run(e.fns.argmax, .{ rows, 1, 1 }, .{ 1024, 1, 1 }, .{ d.logits, @as(c_int, @intCast(e.g.vocab)), d.tok }); +} + +/// The first `rows` drawn tokens, read back with one synchronization. +fn fetchRows(d: *Mtp, rows: usize) ![]align(1) const u32 { + const e = d.e; + const st = d.staged.bytes; + try e.d.check(e.d.api.cuMemcpyDtoHAsync_v2(st.ptr + Mtp.tok_at, d.tok, rows * 4, e.stream.handle), "cuMemcpyDtoHAsync"); + try e.stream.synchronize(); + return std.mem.bytesAsSlice(u32, st[Mtp.tok_at..][0 .. rows * 4]); +} + +/// The round's absorbs: up to `most` rows a pass, each stream's rows one segment; larger or empty ones as before. +pub fn absorb(d: *Mtp, items: []const dr.Absorb) !void { + const sizes = try d.gpa.alloc(usize, items.len); + defer d.gpa.free(sizes); + for (items, sizes) |it, *s| { + const n = if (it.rows) |rs| rs.len else it.features.rows; + s.* = if (n == 0) most + 1 else n; // an empty absorb keeps the serial path's handling + } + const passes = try d.gpa.alloc([2]usize, items.len); + defer d.gpa.free(passes); + for (passes[0..plan(sizes, passes)]) |p| { + if (sizes[p[0]] > most) { + try d.absorbOne(items[p[0]]); + continue; + } + try absorbPass(d, items[p[0]..p[1]]); + } +} + +fn absorbPass(d: *Mtp, items: []const dr.Absorb) !void { + const e = d.e; + const g = e.g; + const previous = d.projection_role; + defer d.projection_role = previous; + d.projection_role = .live; // up to `most` rows: prompt and live rows take the same arithmetic + var segs: [most]Seg = undefined; + var tab: [most]RingRow = undefined; + var ids: [most]u32 = undefined; + var at: usize = 0; + for (items, 0..) |it, si| { + const ctx = d.streams.get(it.stream) orelse return error.UnknownStream; + if (it.features.space != .device or it.features.dtype != .bf16) return error.UnsupportedFeatures; + const n = if (it.rows) |rs| rs.len else it.features.rows; + const start = it.features.offset; + if (it.follow.len < n) return error.BadDraftRows; + // absorbRole's bookkeeping; `most` rows never reach the ring's trim + if (start < ctx.end) ctx.end = start; + if (start > ctx.end) ctx.begin = start; + ctx.end = start; + const out = d.taps_in + at * g.hidden * 2; + if (it.rows) |rs| { + var tr: mtp.TapRows = .{ .n = @intCast(n), .rows = @splat(0) }; + for (rs, 0..) |r, i| tr.rows[i] = @intCast(r); + try e.run(d.ring_fns.gather_rows, .{ n, 1, 1 }, .{ 256, 1, 1 }, .{ it.features.buffer, @as(c_int, engine.tap_rows), @as(c_longlong, @intCast(start)), tr, @as(c_int, @intCast(g.hidden)), out }); + } else try e.run(d.ring_fns.gather_taps, .{ n, 1, 1 }, .{ 256, 1, 1 }, .{ it.features.buffer, @as(c_int, engine.tap_rows), @as(c_longlong, @intCast(start)), @as(c_int, @intCast(g.hidden)), out }); + for (0..n) |r| { + tab[at + r] = ringRow(d, ctx, start, at); + ids[at + r] = it.follow[r]; + } + segs[si] = .{ .ctx = ctx, .first = at, .rows = n, .p0 = start }; + at += n; + } + const s = segs[0..items.len]; + try d.norm(d.taps_in, 0, e.final_norm, d.hin, at); // the target's final norm: the rows the head read + try upload(d, ids[0..at], tab[0..at]); + try layerRows(d, s, at); + try ringPut(d, s); + for (s, 0..) |seg, j| { // each stream's newest row: its state kept for the drafts, its row of the head's input + const last = d.x + (seg.first + seg.rows - 1) * g.hidden * 2; + try seg.ctx.last.copyFrom(0, last, g.hidden * 2, e.stream.handle); + try d.scratch.copyFrom(d.hn - d.scratch.ptr + j * g.hidden * 2, last, g.hidden * 2, e.stream.handle); + seg.ctx.end = seg.p0 + seg.rows; + } + try head(d, d.hn, s.len); + const toks = try fetchRows(d, s.len); + for (s, toks) |seg, t| seg.ctx.cands[0] = t; + d.batch_passes += 1; +} + +/// A chain's stream and the drafts it asks for. +const Chain = struct { ctx: *Ctx, pos0: u64, depth: usize }; + +/// The round's holds: each draft level's chains in passes of up to `most` streams; trees keep the serial path. +pub fn hold(d: *Mtp, items: []const dr.Hold) !void { + const chains = try d.gpa.alloc(Chain, items.len); + defer d.gpa.free(chains); + var n: usize = 0; + var deepest: usize = 0; + for (items) |it| { + const ctx = d.streams.get(it.stream) orelse return error.UnknownStream; + if (it.lanes != null) { + try d.hold(ctx, it.pending, it.position - 1, it.depth, it.lanes); + continue; + } + const pos0 = it.position - 1; + if (ctx.end != pos0) return error.DraftContextBehind; + ctx.held[0] = ctx.cands[0]; + ctx.held_n = 1; + chains[n] = .{ .ctx = ctx, .pos0 = pos0, .depth = @min(it.depth, ctx.held.len) }; + deepest = @max(deepest, chains[n].depth); + n += 1; + } + var k: usize = 1; + while (k < deepest) : (k += 1) { + var group: [most]Chain = undefined; + var rows: usize = 0; + for (chains[0..n]) |c| if (c.depth > k) { + group[rows] = c; + rows += 1; + if (rows == most) { + try holdPass(d, group[0..rows], k); + rows = 0; + } + }; + if (rows > 0) try holdPass(d, group[0..rows], k); + } +} + +/// Draft level k of every chain in `group`, a row each: its previous state and draft at slot pos0 + k - 1. +fn holdPass(d: *Mtp, group: []const Chain, k: usize) !void { + const e = d.e; + const g = e.g; + const previous = d.projection_role; + defer d.projection_role = previous; + d.projection_role = .proposal; + var segs: [most]Seg = undefined; + var tab: [most]RingRow = undefined; + var ids: [most]u32 = undefined; + for (group, 0..) |c, r| { + const src = if (k == 1) c.ctx.last.ptr else c.ctx.step.ptr; + try d.scratch.copyFrom(d.hin - d.scratch.ptr + r * g.hidden * 2, src, g.hidden * 2, e.stream.handle); + const p0 = c.pos0 + k - 1; + ids[r] = c.ctx.held[k - 1]; + tab[r] = ringRow(d, c.ctx, p0, r); + segs[r] = .{ .ctx = c.ctx, .first = r, .rows = 1, .p0 = p0 }; + } + const s = segs[0..group.len]; + try upload(d, ids[0..group.len], tab[0..group.len]); + try layerRows(d, s, group.len); + try ringPut(d, s); + for (s) |seg| try seg.ctx.step.copyFrom(0, d.x + seg.first * g.hidden * 2, g.hidden * 2, e.stream.handle); + try head(d, d.x, group.len); + const toks = try fetchRows(d, group.len); + for (group, toks) |c, t| { + c.ctx.held[k] = t; + c.ctx.held_n = k + 1; + } + d.batch_passes += 1; +} + +test "passes pack consecutive items up to 16 rows; a longer item runs alone" { + var out: [8][2]usize = undefined; + const n = plan(&.{ 3, 4, 2, 8 }, &out); + try std.testing.expectEqualSlices([2]usize, &.{ .{ 0, 3 }, .{ 3, 4 } }, out[0..n]); + const m = plan(&.{ 2, 17, 1 }, &out); + try std.testing.expectEqualSlices([2]usize, &.{ .{ 0, 1 }, .{ 1, 2 }, .{ 2, 3 } }, out[0..m]); + var ones: [17]usize = @splat(1); + var wide: [17][2]usize = undefined; + const w = plan(&ones, &wide); + try std.testing.expectEqualSlices([2]usize, &.{ .{ 0, 16 }, .{ 16, 17 } }, wide[0..w]); + try std.testing.expectEqual(@as(usize, 40), @sizeOf(RingRow)); +} diff --git a/zig/src/families/qwen3_5/cuda_mtp_batch_check.zig b/zig/src/families/qwen3_5/cuda_mtp_batch_check.zig new file mode 100644 index 0000000000..4f4161e59b --- /dev/null +++ b/zig/src/families/qwen3_5/cuda_mtp_batch_check.zig @@ -0,0 +1,155 @@ +//! GPU check: batched MTP drafting leaves each stream as the serial path does (held tokens, ring, state), case by case. +const std = @import("std"); +const lanes = @import("lanes"); +const mtp = @import("cuda_mtp.zig"); + +const Mtp = mtp.Mtp; +const Ctx = mtp.Ctx; +const dr = lanes.drafter; + +/// Streams of one side: the serial set or the batched one, each its own contexts on the same drafter. +const Side = struct { + streams: []lanes.Stream, + ctxs: []*Ctx, +}; + +/// The two sides' contexts now: held tokens, candidates, ring slots, rings and newest states byte for byte. +fn same(gpa: std.mem.Allocator, d: *Mtp, a: Side, b: Side, n: usize) !bool { + const ring = d.ringBytes(); + const ra = try gpa.alloc(u8, ring + d.e.g.hidden * 2); + defer gpa.free(ra); + const rb = try gpa.alloc(u8, ring + d.e.g.hidden * 2); + defer gpa.free(rb); + for (a.ctxs[0..n], b.ctxs[0..n]) |x, y| { + if (x.held_n != y.held_n or !std.mem.eql(u32, x.held[0..x.held_n], y.held[0..y.held_n])) return false; + if (!std.mem.eql(u32, &x.cands, &y.cands) or x.begin != y.begin or x.end != y.end) return false; + try d.e.stream.synchronize(); + try x.ring.download(0, ra[0..ring]); + try x.last.download(0, ra[ring..]); + try y.ring.download(0, rb[0..ring]); + try y.last.download(0, rb[ring..]); + if (!std.mem.eql(u8, ra, rb)) return false; + } + return true; +} + +/// `n` streams a side, their contexts opened and zeroed (a context's warm-up leaves scratch-dependent bytes). +fn open(gpa: std.mem.Allocator, d: *Mtp, n: usize, tag: u8) !Side { + const side: Side = .{ .streams = try gpa.alloc(lanes.Stream, n), .ctxs = try gpa.alloc(*Ctx, n) }; + for (side.streams, side.ctxs, 0..) |*s, *c, i| { + var id: [16]u8 = undefined; + s.* = try lanes.Stream.init(gpa, .{ .id = try std.fmt.bufPrint(&id, "{c}{d}", .{ tag, i }), .prompt = &.{1}, .max_new = 8 }); + try d.drafter().open(s); + c.* = d.streams.get(s).?; + try c.*.ring.fill8(0, null); + try c.*.last.fill8(0, null); + try c.*.step.fill8(0, null); + } + return side; +} + +fn close(gpa: std.mem.Allocator, d: *Mtp, side: Side) void { + for (side.streams) |*s| { + d.drafter().release(s); + s.deinit(gpa); + } + gpa.free(side.streams); + gpa.free(side.ctxs); +} + +/// An absorb of `n` rows at `start` of the target's tap ring (`rows`: picked ones), followed by the ids after them. +fn absorbItem(s: *lanes.Stream, tap: u64, start: u64, n: u32, rows: ?[]const u32, ids: []const u32, prompt: bool) dr.Absorb { + return .{ .stream = s, .features = .{ .buffer = tap, .offset = start, .rows = n, .row_bytes = 0, .space = .device, .dtype = .bf16 }, .start = start, .follow = ids[start + 1 ..], .rows = rows, .prompt = prompt }; +} + +/// The same absorbs on both sides, serial then batched, compared; then the same holds, compared again. +fn step(gpa: std.mem.Allocator, d: *Mtp, a: Side, b: Side, n: usize, absorbs: []const dr.Absorb, depths: []const u32) !bool { + inline for (.{ a, b }, .{ false, true }) |side, batch| { + d.batch = batch; + const before = d.batch_passes; + const items = try gpa.alloc(dr.Absorb, absorbs.len); + defer gpa.free(items); + for (items, absorbs) |*x, y| { + x.* = y; + x.stream = &side.streams[streamIndex(a, y.stream)]; + } + try d.drafter().absorb(items); + if ((d.batch_passes > before) != batch) return error.WrongDraftPath; + } + if (!try same(gpa, d, a, b, n)) return false; + inline for (.{ a, b }, .{ false, true }) |side, batch| { + d.batch = batch; + const before = d.batch_passes; + const holds = try gpa.alloc(dr.Hold, depths.len); + defer gpa.free(holds); + for (holds, depths, 0..) |*h, depth, i| h.* = .{ .stream = &side.streams[i], .pending = 0, .position = side.ctxs[i].end + 1, .depth = depth }; + try d.drafter().hold(holds); + var deep = false; + for (depths) |x| deep = deep or x > 1; + if (deep and (d.batch_passes > before) != batch) return error.WrongDraftPath; + } + return same(gpa, d, a, b, n); +} + +fn streamIndex(a: Side, s: *lanes.Stream) usize { + for (a.streams, 0..) |*x, i| if (x == s) return i; + unreachable; +} + +/// The cases on the target's tap ring `tap` (filled by a prompt over `ids`); prints a line each. +pub fn run(gpa: std.mem.Allocator, d: *Mtp, tap: u64, ids: []const u32) !bool { + const trees = d.trees; + const batch = d.batch; + defer { + d.trees = trees; + d.batch = batch; + } + d.trees = false; + var ok = true; + { + const a = try open(gpa, d, 6, 'a'); + defer close(gpa, d, a); + const b = try open(gpa, d, 6, 'b'); + defer close(gpa, d, b); + const lens = [_]u32{ 1, 2, 3, 4, 1, 7 }; + var first: [6]dr.Absorb = undefined; + for (&first, lens, 0..) |*x, n, i| x.* = absorbItem(&a.streams[i], tap, 10 * i + 3, n, null, ids, i == 0); + const r1 = try step(gpa, d, a, b, 6, &first, &.{ 3, 1, 0, 3, 2, 3 }); + std.debug.print("{s} batched mtp: six streams, absorbs of 1 to 7 rows (one a prompt row), depths 0 to 3\n", .{if (r1) "PASS" else "FAIL"}); + // a second round: picked rows [0, 2] of a 3-row window at each end, and a gap past one stream's end + const picks = [_]u32{ 0, 2 }; + var second: [6]dr.Absorb = undefined; + for (&second, 0..) |*x, i| { + const end = a.ctxs[i].end; + x.* = if (i == 2) absorbItem(&a.streams[i], tap, end + 5, 2, null, ids, false) else absorbItem(&a.streams[i], tap, end, 2, &picks, ids, false); + } + const r2 = try step(gpa, d, a, b, 6, &second, &.{ 3, 3, 3, 3, 3, 3 }); + std.debug.print("{s} batched mtp: partial acceptance (picked rows), a gap past a stream's end, depth 3\n", .{if (r2) "PASS" else "FAIL"}); + ok = ok and r1 and r2; + } + { + // contexts released to the pool and reopened: a reused context drafts as a fresh one + const a = try open(gpa, d, 3, 'c'); + defer close(gpa, d, a); + const b = try open(gpa, d, 3, 'd'); + defer close(gpa, d, b); + var items: [3]dr.Absorb = undefined; + for (&items, 0..) |*x, i| x.* = absorbItem(&a.streams[i], tap, 40 + 2 * i, 3, null, ids, false); + const r = try step(gpa, d, a, b, 3, &items, &.{ 3, 2, 3 }); + std.debug.print("{s} batched mtp: reused contexts from the pool\n", .{if (r) "PASS" else "FAIL"}); + ok = ok and r; + } + { + const a = try open(gpa, d, 17, 'e'); + defer close(gpa, d, a); + const b = try open(gpa, d, 17, 'f'); + defer close(gpa, d, b); + var items: [17]dr.Absorb = undefined; + for (&items, 0..) |*x, i| x.* = absorbItem(&a.streams[i], tap, 5 + i, 1, null, ids, false); + const depths: [17]u32 = @splat(3); + const r = try step(gpa, d, a, b, 17, &items, &depths); + std.debug.print("{s} batched mtp: 17 streams, two passes a draft level\n", .{if (r) "PASS" else "FAIL"}); + ok = ok and r; + } + return ok; +} diff --git a/zig/src/families/qwen3_5/cuda_mtp_check.zig b/zig/src/families/qwen3_5/cuda_mtp_check.zig index 464fd8737e..5a6b99e48f 100644 --- a/zig/src/families/qwen3_5/cuda_mtp_check.zig +++ b/zig/src/families/qwen3_5/cuda_mtp_check.zig @@ -93,6 +93,8 @@ pub fn run(gpa: std.mem.Allocator, d: *Mtp, tap: u64, ids: []const u32) !bool { try download(d, chain.v[lane * kw ..][0..kw], d.vb); } ok = report("six siblings a depth, tree rows vs each lane's path chained", tree, chain, &internal) and ok; + // batched shared rounds against one stream at a time, on the same tap ring + ok = try @import("cuda_mtp_batch_check.zig").run(gpa, d, tap, ids) and ok; return ok; } diff --git a/zig/tests/cuda/contract_qwen_mtp_rows.zig b/zig/tests/cuda/contract_qwen_mtp_rows.zig new file mode 100644 index 0000000000..f7ea2da9cb --- /dev/null +++ b/zig/tests/cuda/contract_qwen_mtp_rows.zig @@ -0,0 +1,344 @@ +//! tf_q36m_attn_rows under the contract: several streams' rows over their own rings, each row tf_q36m_attn's bytes. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const tc = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; + +const heads = 16; +const kv_heads = 2; +const hd = 256; +const width = kv_heads * hd; // a ring slot's or a block row's keys (or values) +const gstride = heads * 2 * hd; // a row's q and gate halves +const nan_bits: u16 = 0x7fc0; +const rule = "as tf_q36m_attn gives them, launched once a segment over the same rows"; + +/// The packed row's ring, as the kernel's RingRow (40 bytes). +const RingRow = extern struct { rk: u64, rv: u64, begin: i64, end: i64, seg: i32, w: i32 }; + +/// One stream's part of a packing: its ring's capacity, first slot and slots, and its rows' packed range. +const Seg = struct { cap: usize, begin: usize, slots: usize, first: usize, rows: usize }; + +/// Each row's terms run over its ring's slots, then its segment's rows up to its own: two roundings a term. +fn depth(s: c.Shape) usize { + return 2 * (s.k + s.rows) + 16; +} + +pub fn spec() c.Spec { + return .{ + .name = "cuda qwen3.6 mtp attention over rings tf_q36m_attn_rows", + .accum = .f32, + .depth = depth, + .ops = 2, // o / l, then the gate's product + .mids = 1, // o / l rounded to bf16 before the gate + .out = .bf16, + .replaces = "tf_q36m_attn", + .nan = .{ .stated = rule }, + // the first: three rings, so the single-shape gates (repeats, guards, NaN) run across streams + .served = &.{ .{ .rows = 16, .n = heads * hd, .k = 6 }, .{ .rows = 1, .n = heads * hd, .k = 6 }, .{ .rows = 4, .n = heads * hd, .k = 6 }, .{ .rows = 4, .n = heads * hd, .k = 4200 }, .{ .rows = 16, .n = heads * hd, .k = 4200 } }, + }; +} + +/// Three streams at most: rows 1, up to 3, then the rest, each over its own ring (one of them wrapping). +fn pack(rows: usize, k: usize, out: *[3]Seg) []Seg { + const sizes = [3]usize{ 1, @min(3, rows -| 1), rows -| 4 }; + const caps = [3]usize{ k + 1, 2 * k + 64, 32768 }; + const begins = [3]usize{ 0, 5, 32768 - 3 }; + const slots = [3]usize{ k, k / 2 + 1, k -| 1 }; + var n: usize = 0; + var first: usize = 0; + for (sizes, caps, begins, slots) |size, cap, begin, slot| if (size > 0) { + out[n] = .{ .cap = cap, .begin = begin, .slots = slot, .first = first, .rows = size }; + first += size; + n += 1; + }; + return out[0..n]; +} + +pub const Rows = struct { + rig: rig_.Rig, + attn: cuda.Function, + rows_fn: cuda.Function, + arena: std.heap.ArenaAllocator, + shape: c.Shape = undefined, + fill: g.Fill = .random, + segs_buf: [3]Seg = undefined, + segs: []Seg = &.{}, + q: []u16 = &.{}, + gate: []u16 = &.{}, + kb: []u16 = &.{}, + vb: []u16 = &.{}, + rk: [3][]u16 = .{ &.{}, &.{}, &.{} }, + rv: [3][]u16 = .{ &.{}, &.{}, &.{} }, + dq: u64 = 0, + dg: u64 = 0, + dkb: u64 = 0, + dvb: u64 = 0, + drk: [3]u64 = @splat(0), + drv: [3]u64 = @splat(0), + table: u64 = 0, + out: u64 = 0, + base: u64 = 0, + + pub fn init(gpu: tc.Gpu, s: cuda.Stream, mod: cuda.Module) !Rows { + return .{ .rig = .{ .gpu = gpu, .s = s }, .attn = try mod.function("tf_q36m_attn"), .rows_fn = try mod.function("tf_q36m_attn_rows"), .arena = .init(gpu.gpa) }; + } + + pub fn deinit(self: *Rows) void { + self.rig.deinit(); + self.arena.deinit(); + } + + pub fn kernel(self: *Rows) g.Kernel { + return .{ .spec = spec(), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Rows); + + fn segOf(self: *const Rows, row: usize) usize { + for (self.segs, 0..) |s, i| if (row < s.first + s.rows) return i; + unreachable; + } + + /// Row `m`'s terms in the kernel's order: its ring's slots, then its segment's rows up to its own. + fn terms(self: *const Rows, m: usize) usize { + const s = self.segs[self.segOf(m)]; + return s.slots + (m - s.first) + 1; + } + + /// Key or value `j` of row `m`'s terms for kv head `kvh`. + fn term(self: *const Rows, m: usize, j: usize, kvh: usize, values: bool) []const u16 { + const si = self.segOf(m); + const s = self.segs[si]; + if (j < s.slots) { + const at = ((s.begin + j) % s.cap) * width + kvh * hd; + return (if (values) self.rv[si] else self.rk[si])[at..][0..hd]; + } + const r = s.first + (j - s.slots); + return (if (values) self.vb else self.kb)[r * width + kvh * hd ..][0..hd]; + } + + fn value(self: *const Rows, rng: std.Random, j: usize) u16 { + return switch (self.fill) { + .stress => c.toBf16(([3]f32{ 4096, 1, -4096 })[j % 3]), + else => c.toBf16((rng.float(f32) - 0.5) * 0.5), + }; + } + + pub fn prepare(self: *Rows, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.rows > 16 or shape.n != heads * hd or shape.k < 2) return false; + self.rig.reset(); + _ = self.arena.reset(.free_all); + const a = self.arena.allocator(); + self.shape = shape; + self.fill = fill; + self.segs = pack(shape.rows, shape.k, &self.segs_buf); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const rows = shape.rows; + self.q = try a.alloc(u16, rows * heads * hd); + self.gate = try a.alloc(u16, rows * gstride); + self.kb = try a.alloc(u16, rows * width); + self.vb = try a.alloc(u16, rows * width); + for (self.q) |*x| x.* = if (fill == .stress) 0 else c.toBf16((rng.float(f32) - 0.5) * 0.5); + for (self.gate) |*x| x.* = if (fill == .stress) 0 else c.toBf16((rng.float(f32) - 0.5) * 4); + for (self.segs, 0..) |s, si| { + // every physical slot outside [begin, begin + slots) of the ring holds NaN: a slot it should not read + self.rk[si] = try a.alloc(u16, s.cap * width); + self.rv[si] = try a.alloc(u16, s.cap * width); + @memset(self.rk[si], nan_bits); + @memset(self.rv[si], nan_bits); + for (0..s.slots) |j| { + const at = ((s.begin + j) % s.cap) * width; + for (0..width) |i| { + self.rk[si][at + i] = if (fill == .stress) 0 else c.toBf16((rng.float(f32) - 0.5) * 0.5); + self.rv[si][at + i] = self.value(rng, j); + } + } + for (0..s.rows) |r| for (0..width) |i| { + self.kb[(s.first + r) * width + i] = if (fill == .stress) 0 else c.toBf16((rng.float(f32) - 0.5) * 0.5); + self.vb[(s.first + r) * width + i] = self.value(rng, s.slots + r); + }; + } + if (fill == .nonfinite) { + self.q[1 * hd + 3] = nan_bits; // row 0, head 1 + const last = self.segs[self.segs.len - 1]; + const at = (last.begin % last.cap) * width + 7; + self.rk[self.segs.len - 1][at] = c.toBf16(std.math.inf(f32)); // the last stream's first slot + self.vb[(rows - 1) * width + 9] = c.toBf16(-std.math.inf(f32)); + } + self.dq = try self.rig.input(std.mem.sliceAsBytes(self.q)); + self.dg = try self.rig.input(std.mem.sliceAsBytes(self.gate)); + self.dkb = try self.rig.input(std.mem.sliceAsBytes(self.kb)); + self.dvb = try self.rig.input(std.mem.sliceAsBytes(self.vb)); + const tab = try a.alloc(RingRow, rows); + for (self.segs, 0..) |s, si| { + self.drk[si] = try self.rig.input(std.mem.sliceAsBytes(self.rk[si])); + self.drv[si] = try self.rig.input(std.mem.sliceAsBytes(self.rv[si])); + for (0..s.rows) |r| tab[s.first + r] = .{ .rk = self.drk[si], .rv = self.drv[si], .begin = @intCast(s.begin), .end = @intCast(s.begin + s.slots), .seg = @intCast(s.first), .w = @intCast(s.cap) }; + } + self.table = try self.rig.input(std.mem.sliceAsBytes(tab)); + self.out = try self.rig.output(rows * heads * hd * 2); + self.base = try self.rig.scratch(rows * heads * hd * 2); + return true; + } + + pub fn run(self: *Rows) !void { + var a: cuda.Args = .{}; + inline for (.{ self.dq, self.dg, @as(c_int, gstride), self.dkb, self.dvb, self.table, @as(c_int, heads), @as(c_int, kv_heads), self.out }) |v| a.add(v); + try cuda.launch.launch(self.rows_fn, .{ .grid = .{ .x = @intCast(self.shape.rows), .y = heads }, .block = .{ .x = 32 } }, self.rig.s, &a); + try self.rig.s.synchronize(); + } + + /// tf_q36m_attn once a segment, each on its own ring and its own rows' offsets. + pub fn replaced(self: *Rows) ![]const u8 { + for (self.segs, 0..) |s, si| { + const off = s.first; + var a: cuda.Args = .{}; + inline for (.{ self.dq + off * heads * hd * 2, self.dg + off * gstride * 2, @as(c_int, gstride), self.dkb + off * width * 2, self.dvb + off * width * 2, self.drk[si], self.drv[si], @as(c_int, heads), @as(c_int, kv_heads) }) |v| a.add(v); + a.add(@as(c_int, @intCast(s.cap))); + a.add(@as(c_longlong, @intCast(s.begin))); + a.add(@as(c_longlong, @intCast(s.begin + s.slots))); + a.add(self.base + off * heads * hd * 2); + try cuda.launch.launch(self.attn, .{ .grid = .{ .x = @intCast(s.rows), .y = heads }, .block = .{ .x = 32 } }, self.rig.s, &a); + } + try self.rig.s.synchronize(); + return self.rig.bytes(self.base, self.shape.rows * heads * hd * 2); + } + + /// The stated NaN rule: the outputs equal tf_q36m_attn's bytes on the same planted inputs. + pub fn check(self: *Rows, at: g.Gate) !?[]const u8 { + if (at != .nonfinite) return null; + const got = try self.arena.allocator().dupe(u8, try self.rig.outputBytes()); + const want = try self.replaced(); + return if (std.mem.eql(u8, got, want)) null else "its bytes differ from tf_q36m_attn's on the planted inputs"; + } + + pub fn count(self: *const Rows) usize { + return self.shape.rows * heads * hd; + } + + pub fn reference(self: *Rows, want: []f64, mag: []f64) !void { + const a = self.rig.gpu.gpa; + for (0..self.shape.rows) |m| for (0..heads) |h| { + const kvh = h / (heads / kv_heads); + const n = self.terms(m); + const p = try a.alloc(f64, n); + defer a.free(p); + for (p, 0..) |*x, j| { + var dot: f64 = 0; + for (self.q[(m * heads + h) * hd ..][0..hd], self.term(m, j, kvh, false)) |qv, kv| dot += c.fromBf16(qv) * c.fromBf16(kv); + x.* = dot / 16; + } + const top = std.mem.max(f64, p); + var total: f64 = 0; + for (p) |*x| { + x.* = @exp(x.* - top); + total += x.*; + } + for (0..hd) |d| { + var sum: f64 = 0; + var most: f64 = 0; + for (p, 0..) |w, j| { + const v = c.fromBf16(self.term(m, j, kvh, true)[d]); + sum += w * v; + most = @max(most, @abs(v)); + } + const gate = 1 / (1 + @exp(-c.fromBf16(self.gate[m * gstride + h * 2 * hd + hd + d]))); + const o = (m * heads + h) * hd + d; + want[o] = sum / total * gate; + mag[o] = most * gate; + } + }; + } + + /// Stress inputs: every weight is 1, so each output is its values' sum in `acc` and `order`, over their count. + pub fn simulate(self: *Rows, acc: c.Accum, order: g.Order, out: []f64) !void { + const a = self.rig.gpu.gpa; + for (0..self.shape.rows) |m| { + const n = self.terms(m); + const t = try a.alloc(f64, n); + defer a.free(t); + for (0..heads) |h| for (0..hd) |d| { + for (t, 0..) |*x, j| x.* = c.fromBf16(self.term(m, j, h / (heads / kv_heads), true)[d]); + const s = contract.rows.sumIn(acc, order, t); + const mean: f32 = @floatCast(s / @as(f64, @floatFromInt(n))); + out[(m * heads + h) * hd + d] = c.fromBf16(c.toBf16(c.fromBf16(c.toBf16(mean)) * 0.5)); + }; + } + } +}; + +/// The kernel and its module, registered in the contract's list; destroyed after the gates have run. +pub const Group = struct { + a: std.mem.Allocator, + mod: cuda.Module, + rows: Rows, + + pub fn create(a: std.mem.Allocator, gpu: tc.Gpu, s: cuda.Stream, kernels: *std.ArrayList(g.Kernel)) !*Group { + const self = try a.create(Group); + errdefer a.destroy(self); + self.a = a; + self.mod = try cuda.Module.load(gpu.d, cuda.kernels.qwen36_mtp); + errdefer self.mod.unload(); + self.rows = try Rows.init(gpu, s, self.mod); + errdefer self.rows.deinit(); + try kernels.append(a, self.rows.kernel()); + return self; + } + + pub fn destroy(self: *Group) void { + self.rows.deinit(); + self.mod.unload(); + self.a.destroy(self); + } +}; + +/// Gates 1-4 on this kernel alone, then its FP64 check: the kernel-contract run's share, for a quick look. +pub fn alone(gpu: tc.Gpu) !void { + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var kernels: std.ArrayList(g.Kernel) = .empty; + defer kernels.deinit(gpu.gpa); + const group = try Group.create(gpu.gpa, gpu, s, &kernels); + defer group.destroy(); + var t: g.Tally = .{}; + try g.runAll(gpu.gpa, kernels.items, 0x6b63, &t); + try fp64(gpu); + try tc.expect(t.problems() == 0, "tf_q36m_attn_rows: {d} gate failures", .{t.failed}); +} + +/// tf_q36m_attn_rows against float64 at the served packings (gate 3 compares a replacement by its bytes only). +pub fn fp64(gpu: tc.Gpu) !void { + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var mod = try cuda.Module.load(gpu.d, cuda.kernels.qwen36_mtp); + defer mod.unload(); + var x = try Rows.init(gpu, s, mod); + defer x.deinit(); + const sp = spec(); + for (sp.served, 0..) |shape, i| { + try tc.expect(try x.prepare(shape, .random, 0x5eed + i), "tf_q36m_attn_rows refused a served shape", .{}); + try x.run(); + const n = x.count(); + const want = try gpu.gpa.alloc(f64, 2 * n); + defer gpu.gpa.free(want); + try x.reference(want[0..n], want[n..]); + const out = try gpu.gpa.dupe(u8, try x.rig.outputBytes()); + defer gpu.gpa.free(out); + for (0..n) |o| { + const t: f64 = @floatFromInt(x.terms(o / (heads * hd))); + const bound = (2 * c.unit(.bf16) + c.gamma(256 + 8 * t, c.unit(.f32))) * want[n + o] + 0x1p-133; + const got = c.read(.bf16, out, o); + try tc.expect(std.math.isFinite(got) and @abs(got - want[o]) <= bound, "tf_q36m_attn_rows FP64 mismatch at {d}x{d}, output {d}: {d} vs {d}", .{ shape.rows, shape.k, o, got, want[o] }); + } + @memset(out[0..2], 0x40); // a planted finite corruption must leave the bound + const bad = c.read(.bf16, out, 0); + try tc.expect(@abs(bad - want[0]) > (2 * c.unit(.bf16) + c.gamma(256 + 8 * @as(f64, @floatFromInt(x.terms(0))), c.unit(.f32))) * want[n], "tf_q36m_attn_rows FP64 check missed a planted corruption", .{}); + } + std.debug.print("PASS tf_q36m_attn_rows: heterogeneous rings (one wrapping), NaN outside each ring's slots, FP64 at {d} served shapes; corruption RED\n", .{sp.served.len}); +} diff --git a/zig/tests/cuda/contract_qwen_served.zig b/zig/tests/cuda/contract_qwen_served.zig index 30230f5a1b..959e43f32e 100644 --- a/zig/tests/cuda/contract_qwen_served.zig +++ b/zig/tests/cuda/contract_qwen_served.zig @@ -56,6 +56,7 @@ pub fn run(gpu: check.Gpu) !void { std.debug.print("PASS {s}: served geometry, FP64, five repeats, guards, inputs; finite corruption RED\n", .{ref.symbol(mode)}); } try planRows(gpu, s); + try @import("contract_qwen_mtp_rows.zig").fp64(gpu); } fn served(fns: cuda.moe.Fns, s: cuda.Stream, mode: ref.Mode, p: []const u64, rows: usize, variant: bool, params: u64) !void { diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index fd53c75789..2c4b07526f 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -567,6 +567,8 @@ pub fn run(gpu: check.Gpu) !void { }; defer for (&contiguous) |*x| x.deinit(); for (&contiguous) |*x| try kernels.append(a, x.kernel()); + const mtp_rows = try @import("contract_qwen_mtp_rows.zig").Group.create(a, gpu, s, &kernels); + defer mtp_rows.destroy(); const weight_only = try @import("contract_weight_only_suite.zig").Suite.create(a, gpu, s, &kernels); defer weight_only.destroy(a); var t: g.Tally = .{}; diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 8f77f46dae..47a8e8c942 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -231,6 +231,7 @@ const usage = \\ qmm-prefill-many-contract the GB10 bit-exact multi-projection prompt launch against separate launches \\ qmm-prefill-tile22-contract the GB10 large-prompt 2x2 tile against the default tile \\ qmm-prefill-many-bench same-process A-B-A screening for the GB10 launch fusion + \\ mtp-rows tf_q36m_attn_rows: gates 1-4 (byte-equal to tf_q36m_attn a segment) and float64 alone \\ qwen-pages qwen35.cu's attention cache through a page table: paged == contiguous, float64, red controls \\ kernel-contract-limits DIR each registered launch against sm_80/86/89/120/121 from the fatbins in DIR \\ glm-exl3 GLM's EXL3 routed experts, decode and prompt, against float64 (synthetic) @@ -450,6 +451,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "qmm-prefill-many-contract")) return qmm_prefill_many_contract.run(gpu); if (std.mem.eql(u8, cmd, "qmm-prefill-tile22-contract")) return qmm_prefill_tile22_contract.run(gpu); if (std.mem.eql(u8, cmd, "qmm-prefill-many-bench")) return qmm_prefill_many_contract.bench(gpu); + if (std.mem.eql(u8, cmd, "mtp-rows")) return @import("contract_qwen_mtp_rows.zig").alone(gpu); if (std.mem.eql(u8, cmd, "qwen-pages")) return @import("contract_qwen_pages.zig").run(gpu); if (std.mem.eql(u8, cmd, "kernel-contract-limits")) return contract_tests.fatbins(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "tree")) return tree_tests.run(gpu);