Qwen3.6 MTP drafter on CUDA: draft a shared round's streams together

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

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

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<const uint4*>(rk + (size_t)slot * width + off));  // read-only path, as a __restrict__ ring's
+            v[u] = __ldg(reinterpret_cast<const uint4*>(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<const __nv_bfloat16*>(&kr[u]));
+#pragma unroll
+        for (int u = 0; u < AHEAD; u++) fold(s[u], reinterpret_cast<const __nv_bfloat16*>(&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,46 +117,61 @@ 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++) {
         q[i] = bf(Q[((size_t)m * heads + h) * HD + lane * DPL + i]);
         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,49 +188,24 @@ 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++) {
         q[i] = bf(Q[((size_t)m * heads + h) * HD + lane * DPL + i]);
         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

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);

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)) {

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");