nemotron: Super's MTP drafts on the draft vocabulary's FP8 rows while a stream's text is covered, same replies (stacked on #1006)

docs/kernel-contract.md

@@ -104,7 +104,8 @@ Metal: softmax, RMS norm, bf16 top-k, the row projection (3, 4, 6 and 8 bits), a
 formats on the M5's tensor units, where its reg kernel must equal the cooperative one bit for bit), and tile_qmm's
 prompt tiles and their decode twin (4, 6 and 8 bits; the twin bit-equal to the tiles). Gate 5 also holds the rest of
 tile_qmm's pipelines and GLM-5.3-Flash's hc_expand_10.
-CUDA: the RMS norm rows (16-byte form and the fp32-out tree), Nemotron's router logits and top-k, qmm_group's cluster
-tile, lane_gemv and its split-K form (both bit-equal to it), prompt16's FP4 tile with bf16 and fp32 out, qmmf's FP4
-tiles with fp32 out (K slices met in a cluster, fused, or by `reduce_kernel<true>`), and the reduce-scatter's
-`tf_rank_sum`.
+CUDA: the RMS norm rows (16-byte form and the fp32-out tree), Nemotron's router logits and top-k, the latent top-k's
+decode form (bit-equal to `tf_super_topk`), the Mamba scan of a verify window (its host twin's bytes), qmm_group's
+cluster tile, lane_gemv and its split-K form (both bit-equal to it), prompt16's FP4 tile with bf16 and fp32 out,
+qmmf's FP4 tiles with fp32 out (K slices met in a cluster, fused, or by `reduce_kernel<true>`), the checkpoint math's
+decode-sized NVFP4 experts (bit-equal to its prompt kernels, sm_120 up), and the reduce-scatter's `tf_rank_sum`.

docs/recipes/nemotron-3.5-super-vl.md

@@ -25,6 +25,15 @@ The weights take 70.8 GiB of the GB10's memory (73.2 GiB with the image tower an
 16,384-token window the server admits 8 streams at 0.35 GiB each and keeps a prompt cache of about 8 to 10 GiB, and
 about 24 GB of the GB10's memory stays available. Leave the other ranks' variables below unset.
 
+The MTP draft head can use the embedded draft vocabulary's FP8 output rows or the already-resident full output head.
+A stream drafts on the subset while the last 16 tokens of its context, its prompt's tail at first, are all draft ids,
+and on the full head otherwise; other streams do not change its choice. Each head has its own captured graphs and
+its own measured level cost beside the window costs both share, and shared rounds are priced at the full head's.
+The target still verifies proposals with its unchanged full-vocabulary head.
+For lone drafted streams, `--mtp-drafts` caps MTP depth (at most 31: Super verifies windows of up to 32 rows) and
+`--mtp-confidence` sets a cumulative-confidence threshold from 0 to 1 instead of the cost-based cutoff. A zero draft
+cap disables drafting. Copied continuations and concurrent streams otherwise retain their existing lane policies.
+
 ## Run on two GB10s
 
 Rank 0 runs the whole model and serves; rank 1 holds only its share of every MoE layer's experts (the MTP layer's
@@ -90,8 +99,7 @@ equal their solo runs, and three-turn chats with regenerate, edit and delete equ
 | Two GB10s | `checkpoint` (default) | 46.8 tok/s (35.9 to 63.6) | 20.2 tok/s | |
 | Two GB10s | `full` | 48.1 tok/s (36.4 to 65.6) | 21.4 tok/s | |
 
-Four concurrent greedy streams on two GB10s took 6.2 s together against 11.8 s one at a time. The MTP layer drafts
-with the target's own output head over the whole vocabulary.
+Four concurrent greedy streams on two GB10s took 6.2 s together against 11.8 s one at a time.
 
 ## Images
 

docs/recipes/nemotron-3.5.md

@@ -51,9 +51,10 @@ cuts, so cold and resumed prompts use the same chunks; templates without markers
 
 ## CUDA
 
-This section describes the Python 0.6.6 engine's CUDA family; its `--tp`, `--rank`, `--master`, `--mtp-drafts` and
-`--prefill-fp8` flags are that engine's, and the native binary refuses them. The native CUDA build serves Nemotron on
-one GPU, from the MLX 4-bit checkpoint or NVIDIA's NVFP4 export, with concurrent requests sharing each round; see the
+This section describes the Python 0.6.6 engine's CUDA family; its `--tp`, `--rank`, `--master` and `--prefill-fp8`
+flags are that engine's, and the native binary refuses them. The native CUDA build serves Nemotron on one GPU, from
+the MLX 4-bit checkpoint or NVIDIA's NVFP4 export, with concurrent requests sharing each round, and its own
+`--mtp-drafts` (at most 15 for this model) and `--mtp-confidence` set a lone stream's MTP drafts; see the
 [native CUDA runbook](../../RUNBOOK.md#linux-and-cuda).
 
 CUDA reads this model's MLX 4-bit checkpoint only; NVFP4 and EXL3 exports of it are not read yet. Its prompt matmuls

zig/kernels/cuda/ck_experts.cu

@@ -71,44 +71,66 @@ extern "C" __global__ void __launch_bounds__(256)
     dst[i] = src[(e * n + tile * 64 + c) * (ks / 16) + (k0 + g * 64) / 16 + q];
 }
 
-// A warp: a plan item's pairs (16 at most) by 64 columns; EPI 1 relu^2 of bf16(alpha_e sum) to bf16, EPI 0 fp32 sums.
-template <int EPI>
+// One 64-input step a warp reads: rows g and g + 8's codes and block scales, then each n8 tile's weights and scales.
+template <int TJ>
+struct Step {
+    uint32_t a[4], sa;
+    uint2 b[TJ];
+    uint32_t s[TJ];
+};
+
+// A warp: a plan item's pairs (16 at most) by TJ n8 tiles, U steps read ahead; EPI 1 relu^2 to bf16, EPI 0 fp32 sums.
+template <int EPI, int TJ, int U>
 __device__ __forceinline__ void grouped(const uint8_t* __restrict__ codes, const uint8_t* __restrict__ xs, int slots,
                                         const uint32_t* __restrict__ w, const uint8_t* __restrict__ ws,
                                         const float* __restrict__ alpha, int K, int N, const int* __restrict__ items,
                                         const int* __restrict__ counts, const int* __restrict__ members,
                                         void* __restrict__ out) {
-    const int lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3, NB = N / 64, KG = K / 64;
-    const int units = __ldg(counts) * NB;
+    const int lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3, NG = N / (8 * TJ), KG = K / 64;
+    const int units = __ldg(counts) * NG;
     for (int unit = blockIdx.x * 4 + (threadIdx.x >> 5); unit < units; unit += gridDim.x * 4) {
-        const int it = unit / NB, cb = unit - it * NB;
+        const int it = unit / NG, n0 = (unit - it * NG) * 8 * TJ;
         const int e = __ldg(items + 3 * it), first = __ldg(items + 3 * it + 1), cnt = __ldg(items + 3 * it + 2);
         const bool v0 = g < cnt, v1 = g + 8 < cnt;
         const int p0 = v0 ? __ldg(members + first + g) : 0, p1 = v1 ? __ldg(members + first + g + 8) : 0;
         const int x0 = slots ? p0 / slots : p0, x1 = slots ? p1 / slots : p1;
         const uint8_t* c0 = codes + static_cast<size_t>(x0) * (K / 2);
         const uint8_t* c1 = codes + static_cast<size_t>(x1) * (K / 2);
-        const uint8_t* s0 = xs + static_cast<size_t>(x0) * (K / 16);
-        const uint8_t* s1 = xs + static_cast<size_t>(x1) * (K / 16);
-        const size_t tile = (static_cast<size_t>(e) * NB + cb) * KG;
-        float acc[8][4];
+        const bool hi = t & 1;  // lanes t 0 and 1 give rows g and g + 8's four block scales; t 2, 3 are not read
+        const bool vr = hi ? v1 : v0;
+        const uint8_t* sr = xs + static_cast<size_t>(hi ? x1 : x0) * (K / 16);
+        const size_t tile = (static_cast<size_t>(e) * (N / 64) + n0 / 64) * KG;  // the 64-column tile holding n0
+        const uint2* wb = reinterpret_cast<const uint2*>(w) + tile * 256 + n0 % 64 / 8 * 32 + lane;
+        const uint32_t* sw = reinterpret_cast<const uint32_t*>(ws) + tile * 64 + n0 % 64 + g;
+        auto fetch = [&](int kg, Step<TJ>& f) {  // rows past cnt: zeros
+            f.a[0] = v0 ? __ldg(reinterpret_cast<const uint32_t*>(c0 + kg * 32 + 4 * t)) : 0u;
+            f.a[1] = v1 ? __ldg(reinterpret_cast<const uint32_t*>(c1 + kg * 32 + 4 * t)) : 0u;
+            f.a[2] = v0 ? __ldg(reinterpret_cast<const uint32_t*>(c0 + kg * 32 + 16 + 4 * t)) : 0u;
+            f.a[3] = v1 ? __ldg(reinterpret_cast<const uint32_t*>(c1 + kg * 32 + 16 + 4 * t)) : 0u;
+            f.sa = vr ? __ldg(reinterpret_cast<const uint32_t*>(sr + kg * 4)) : 0u;
+#pragma unroll
+            for (int j = 0; j < TJ; ++j) {
+                f.b[j] = __ldg(wb + kg * 256 + j * 32);
+                f.s[j] = __ldg(sw + kg * 64 + j * 8);
+            }
+        };
+        float acc[TJ][4];
+#pragma unroll
+        for (int j = 0; j < TJ; ++j) acc[j][0] = acc[j][1] = acc[j][2] = acc[j][3] = 0.0f;
+        Step<TJ> st[U];
+#pragma unroll
+        for (int u = 0; u < U - 1; ++u)
+            if (u < KG) fetch(u, st[u]);
+        for (int k0 = 0; k0 < KG; k0 += U) {  // a column's sum takes its steps in order whatever TJ and U are
+            // step k0 + u's slot is u: indices known at compile time keep st in registers
 #pragma unroll
-        for (int j = 0; j < 8; ++j) acc[j][0] = acc[j][1] = acc[j][2] = acc[j][3] = 0.0f;
-        for (int kg = 0; kg < KG; ++kg) {
-            uint32_t a[4];  // rows g and g + 8, inputs 8t .. and 32 + 8t .. of this step (rows past cnt: zeros)
-            a[0] = v0 ? __ldg(reinterpret_cast<const uint32_t*>(c0 + kg * 32 + 4 * t)) : 0u;
-            a[1] = v1 ? __ldg(reinterpret_cast<const uint32_t*>(c1 + kg * 32 + 4 * t)) : 0u;
-            a[2] = v0 ? __ldg(reinterpret_cast<const uint32_t*>(c0 + kg * 32 + 16 + 4 * t)) : 0u;
-            a[3] = v1 ? __ldg(reinterpret_cast<const uint32_t*>(c1 + kg * 32 + 16 + 4 * t)) : 0u;
-            const bool hi = t & 1;  // lanes t 0 and 1 give rows g and g + 8's four block scales; t 2, 3 are not read
-            const bool vr = hi ? v1 : v0;
-            const uint32_t sa = vr ? __ldg(reinterpret_cast<const uint32_t*>((hi ? s1 : s0) + kg * 4)) : 0u;
-            const uint2* wb = reinterpret_cast<const uint2*>(w) + (tile + kg) * 256;
-            const uint32_t* sw = reinterpret_cast<const uint32_t*>(ws) + (tile + kg) * 64;
+            for (int u = 0; u < U; ++u) {
+                const int kg = k0 + u;
+                if (kg + U - 1 < KG) fetch(kg + U - 1, st[(u + U - 1) % U]);
+                if (kg < KG)
 #pragma unroll
-            for (int j = 0; j < 8; ++j) {
-                const uint2 b = __ldg(wb + j * 32 + lane);
-                mma4::mma_fp4(acc[j], a, b.x, b.y, sa, __ldg(sw + j * 8 + g));
+                    for (int j = 0; j < TJ; ++j)
+                        mma4::mma_fp4(acc[j], st[u].a, st[u].b[j].x, st[u].b[j].y, st[u].sa, st[u].s[j]);
             }
         }
         const float al = __ldg(alpha + e);
@@ -117,8 +139,8 @@ __device__ __forceinline__ void grouped(const uint8_t* __restrict__ codes, const
             if (!(h ? v1 : v0)) continue;
             const size_t row = static_cast<size_t>(h ? p1 : p0) * N;
 #pragma unroll
-            for (int j = 0; j < 8; ++j) {
-                const int col = cb * 64 + j * 8 + 2 * t;
+            for (int j = 0; j < TJ; ++j) {
+                const int col = n0 + j * 8 + 2 * t;
                 const float a0 = acc[j][2 * h] * al, a1 = acc[j][2 * h + 1] * al;
                 if constexpr (EPI == 0) {
                     *reinterpret_cast<float2*>(reinterpret_cast<float*>(out) + row + col) = make_float2(a0, a1);
@@ -132,13 +154,16 @@ __device__ __forceinline__ void grouped(const uint8_t* __restrict__ codes, const
     }
 }
 
-#define TF_CK4(NAME, EPI)                                                                                              \
+#define TF_CK4(NAME, EPI, TJ, U)                                                                                       \
     extern "C" __global__ void __launch_bounds__(128)                                                                  \
         NAME(const uint8_t* __restrict__ codes, const uint8_t* __restrict__ xs, int slots,                            \
              const uint32_t* __restrict__ w, const uint8_t* __restrict__ ws, const float* __restrict__ alpha, int K,  \
              int N, const int* __restrict__ items, const int* __restrict__ counts, const int* __restrict__ members,   \
              void* __restrict__ out) {                                                                                 \
-        grouped<EPI>(codes, xs, slots, w, ws, alpha, K, N, items, counts, members, out);                              \
+        grouped<EPI, TJ, U>(codes, xs, slots, w, ws, alpha, K, N, items, counts, members, out);                       \
     }
-TF_CK4(tf_ck4_up_relu2, 1)
-TF_CK4(tf_ck4_down_f32, 0)
+// Prompt chunks: 64 columns a warp, a step at a time; decode-sized calls: 32 a warp, four steps in flight. Same bits.
+TF_CK4(tf_ck4_up_relu2, 1, 8, 1)
+TF_CK4(tf_ck4_down_f32, 0, 8, 1)
+TF_CK4(tf_ck4_up_relu2_few, 1, 4, 4)
+TF_CK4(tf_ck4_down_f32_few, 0, 4, 4)

zig/kernels/cuda/nemotron_latent.cu

@@ -11,6 +11,8 @@ using namespace tf_glue;
 namespace {
 
 constexpr int MAXC = 32;  // experts a lane holds: up to 1024 (picks a row: one a lane, up to 32)
+constexpr int MAXE = 32 * MAXC;  // experts tf_super_topk_few ranks: one a thread
+constexpr int MAXK = 32;         // picks a row
 
 // The larger of two (value, index) picks: a higher value, or the lower index on a tie (values are never NaN).
 __device__ __forceinline__ bool beats(float v, int i, float bv, int bi) { return v > bv || (v == bv && i < bi); }
@@ -64,6 +66,51 @@ extern "C" __global__ void __launch_bounds__(32)
     for (int s = 0; s < S; ++s) IDX[(r * TOPK + lane) * S + s] = s * E + mine_id;
 }
 
+// tf_super_topk's picks for decode-sized calls, a thread an expert: pick k is the expert k others beat.
+extern "C" __global__ void __launch_bounds__(MAXE)
+    tf_super_topk_few(const float* __restrict__ PART, const float* __restrict__ BIAS, int* __restrict__ IDX,
+                      float* __restrict__ WT, int R, float scaling, int E, int SK, int TOPK, int S, int norm) {
+    __shared__ __align__(16) float sel_s[MAXE];  // a block of whole warps: lanes past E hold -inf, never ahead
+    __shared__ int pick_s[MAXK];
+    __shared__ float prob_s[MAXK];
+    const int64_t r = blockIdx.x;
+    const int e = threadIdx.x;
+    float score = 0.0f, sel = -INFINITY;
+    if (e < E) {
+        float logit = 0.0f;
+        for (int s = 0; s < SK; ++s) logit = __fadd_rn(logit, PART[(s * R + r) * E + e]);
+        score = tf_sigmoid(logit);
+        const float v = __fadd_rn(score, BIAS[e]);
+        sel = v != v ? -INFINITY : v;  // a NaN logit is never picked ahead of a number
+    }
+    sel_s[e] = sel;
+    const int numbers = __syncthreads_count(e < E && sel != -INFINITY);
+    if (e < E && sel != -INFINITY) {
+        const float4* sel4 = reinterpret_cast<const float4*>(sel_s);
+        int rank = 0;  // how many experts beat this one: its pick, if under TOPK
+        for (int j = 0; j < blockDim.x && rank < TOPK; j += 16) {
+#pragma unroll
+            for (int q = 0; q < 4; ++q) {
+                const float4 v = sel4[j / 4 + q];
+                const int i = j + 4 * q;
+                rank += beats(v.x, i, sel, e) + beats(v.y, i + 1, sel, e) + beats(v.z, i + 2, sel, e) +
+                        beats(v.w, i + 3, sel, e);
+            }
+        }
+        if (rank < TOPK) pick_s[rank] = e, prob_s[rank] = score;
+    }
+    if (e == 0)  // past the numbers every pick is expert 0, the lowest id among equal -infs, as tf_super_topk finds
+        for (int k = numbers; k < TOPK; ++k) pick_s[k] = 0, prob_s[k] = score;
+    __syncthreads();
+    if (e >= TOPK) return;
+    float total = 0.0f;
+    for (int k = 0; k < TOPK; ++k) total = __fadd_rn(total, prob_s[k]);
+    const float mine = prob_s[e];
+    const float w = norm ? __fmul_rn(__fdiv_rn(mine, __fadd_rn(total, 1e-20f)), scaling) : __fmul_rn(mine, scaling);
+    WT[r * TOPK + e] = w;
+    for (int s = 0; s < S; ++s) IDX[(r * TOPK + e) * S + s] = s * E + pick_s[e];
+}
+
 // Row r's latent sum: each pick's shares added in share order, rounded to bf16, then weighted in pick order.
 extern "C" __global__ void __launch_bounds__(256)
     tf_super_combine(const float* __restrict__ P, long long share_stride, long long pair_stride, int S,