The argmax and the draft's top probability over many blocks a row

The verify window's head ended with sampler_greedy_kernel, one 1024-thread block a token over the 248,320 logits,
and each draft pass with it and row_top_prob over the 40,525-token draft head, one block each.  argmax_rows scans a
row in up to 128 blocks (4,096 logits each), the row's last block merging their picks (the larger value, on equality
the lower index: the one-block kernel's order); row_top_prob_split computes row_top_prob's 32 warps' sums in 8 blocks,
the row's last block adding them in order.  Both keep their counters and partials in per-row slots of one scratch,
whatever a launch's row count.  Bitwise: verify_parity (ties, NaN, -inf, 1-8 rows on one scratch), and the window
hashes on fixed text equal the previous build's (IQ3_S, UD-Q4_K_XL; thinking and code text).

CUPTI, IQ3_S thinking text: the head's argmax 84-86 -> 9.7 us a window, a draft pass's argmax and top probability
15.1 + 10.1 -> 5.6 + 4.3 us; ~105 us a round (~0.7%).

Claude-Session: https://claude.ai/code/session_019XxRxQQNyk7Ycdr4ComSs1

include/strata/core/mtp.hpp

@@ -134,6 +134,8 @@ private:
     int64_t window_ = 0;        // attention over the last window_ cells (0 = every cell)
     int64_t prompt_len_ = 0;
     float* probs_ = nullptr;
+    uint8_t* arg_scratch_ = nullptr;   ///< argmax_rows' and row_top_prob_split's partials and counters
+    uint8_t* top_scratch_ = nullptr;
     // device
     int32_t *tok_ = nullptr, *step_ = nullptr, *pos_ = nullptr, *row_ = nullptr, *ident_ = nullptr;
     float *Rin_ = nullptr, *R_ = nullptr, *emb_ = nullptr, *en_ = nullptr, *e2_ = nullptr, *hn_ = nullptr, *h2_ = nullptr;

include/strata/core/verify.hpp

@@ -314,6 +314,7 @@ private:
     float* hit_xs_ = nullptr;
     void* hit_scratch_ = nullptr;
     float *head_mixed_ = nullptr, *head_inj_ = nullptr, *head_logits_ = nullptr;
+    uint8_t* arg_scratch_ = nullptr;   ///< argmax_rows' partials and counters
     float *sh_gate_ = nullptr, *sh_up_ = nullptr, *sh_g_ = nullptr;
     float* hist_snap_ = nullptr;                              // T * NG_HIST * NG_HC_DIM
     // the PLE block's projections (keys T * NG_HC_DIM, values T * n_embd) and the rest of the block's rows

include/strata/kernels/verify_kernels.hpp

@@ -86,6 +86,17 @@ void gather_rows(const uint8_t* src, int64_t row_bytes, const int32_t* ids, int6
 void map_ids(int32_t* ids, const int32_t* table, int n, void* stream);
 /// probs[t] = softmax(logits[t])[ids[t]] for n_rows rows of n_vocab (the probability of each row's argmax).
 void row_top_prob(const float* logits, int n_rows, int n_vocab, const int32_t* ids, float* probs, void* stream);
+/// out[t] = the argmax of row t (n floats a row; ties to the lowest index, NaN never picked, 0 when no value is above
+/// -inf): bitwise sample_tokens' greedy pick.  Up to 128 blocks a row scan slices of it, the row's last block merges
+/// their picks.  `scratch`: argmax_rows_scratch_bytes(n_rows), zero before the first launch (each launch leaves it
+/// so).  Graph-capturable.
+uint64_t argmax_rows_scratch_bytes(int n_rows);
+void argmax_rows(const float* logits, int n_rows, int n, void* scratch, int32_t* out, void* stream);
+/// row_top_prob over 8 blocks a row, each computing 4 of its 32 warps' sums, the row's last block adding the 32 in
+/// order: bitwise row_top_prob.  `scratch`: row_top_prob_scratch_bytes(n_rows), zero before the first launch.
+uint64_t row_top_prob_scratch_bytes(int n_rows);
+void row_top_prob_split(const float* logits, int n_rows, int n_vocab, const int32_t* ids, float* probs, void* scratch,
+                        void* stream);
 /// Dense-attention step records for `n` cells: [cell, cell+1, (cell+1)/4, cell+1] from cells[i] (device memory).
 void dense_steps(const int32_t* cells, int n, int32_t* steps, void* stream);
 /// A sliding attention window: for `n` step records (kStepCount ints each, n_kv at [1]) the selection becomes the

src/core/mtp.cpp

@@ -220,6 +220,8 @@ bool MtpDrafter::load(const std::string& rt_dir, const ModelGeometry& g, Session
         x_bf16_ = b.take<uint16_t>(N);
         out_ids_ = b.take<int32_t>(T + 4);
         probs_ = b.take<float>(T + 4);
+        arg_scratch_ = b.take<uint8_t>(strata::kernels::argmax_rows_scratch_bytes((int) T));
+        top_scratch_ = b.take<uint8_t>(strata::kernels::row_top_prob_scratch_bytes((int) T));
         dummy_inj_ = b.take<float>(HC);
     };
     Bump count;
@@ -417,11 +419,8 @@ bool MtpDrafter::record_forward(int T, int step_row0, cudaStream_t cs, std::stri
         const bool sub = dhead_ != nullptr;
         const int64_t nv = sub ? n_dvocab_ : n_vocab_;
         native_mmvq(head_->type(), sub ? dhead_ : head_->weights(), xq_, head_logits_, (int) N, (int) nv, T, cs);
-        SamplerParams sp;
-        sp.greedy = true;
-        sp.temperature = 0.0f;
-        sample_tokens(head_logits_, T, (int) nv, nullptr, 0, sp, out_ids_, cs);
-        row_top_prob(head_logits_, T, (int) nv, out_ids_, probs_, cs);
+        argmax_rows(head_logits_, T, (int) nv, arg_scratch_, out_ids_, cs);
+        row_top_prob_split(head_logits_, T, (int) nv, out_ids_, probs_, top_scratch_, cs);
         if (sub) map_ids(out_ids_, dvocab_, T, cs);
     } catch (const std::exception& e) {
         err = std::string("mtp: ") + e.what();

src/core/verify.cpp

@@ -273,6 +273,7 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState&
         sh_gate_ = b.take<float>(T * (uint64_t) g.n_ff);
         sh_up_ = b.take<float>(T * (uint64_t) g.n_ff); sh_g_ = b.take<float>(T + 4);
         head_logits_ = b.take<float>(T * (uint64_t) n_vocab_);
+        arg_scratch_ = b.take<uint8_t>(strata::kernels::argmax_rows_scratch_bytes((int) T));
         hist_snap_ = b.take<float>(T * HS);
         const uint64_t PD = (uint64_t) strata::kernels::NG_HC_DIM;
         ple_key_ = b.take<float>(T * PD); ple_val_ = b.take<float>(T * N); ple_nkey_ = b.take<float>(T * PD);
@@ -959,10 +960,7 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
                 return false;
             }
         }
-        SamplerParams sp;
-        sp.greedy = true;
-        sp.temperature = 0.0f;
-        sample_tokens(head_logits_, T, (int) n_vocab_, nullptr, 0, sp, m_out_, cs);
+        argmax_rows(head_logits_, T, (int) n_vocab_, arg_scratch_, m_out_, cs);
     }
     if (host_rows_) copy_from_mapped(m_rows_, R_, (int64_t) T * HC * N, cs);   // a plain copy kernel, here to host
     stamp(g.n_layers, 3, cs);

src/kernels/cuda/verify_kernels.cu

@@ -7,6 +7,7 @@
 
 #include <cuda_runtime.h>
 
+#include <algorithm>
 #include <cstdio>
 #include <cstdlib>
 
@@ -402,6 +403,142 @@ void row_top_prob(const float* logits, int n_rows, int n_vocab, const int32_t* i
     check("row_top_prob");
 }
 
+namespace {
+
+constexpr int kArgMaxBlocks = 128, kArgThreads = 256, kTopBlocks = 8;
+
+// the larger value, on equality the lower index: sampler_greedy_kernel's order
+__device__ __forceinline__ void arg_take(float ov, int oi, float& bv, int& bi) {
+    if (ov > bv || (ov == bv && oi < bi)) { bv = ov; bi = oi; }
+}
+
+// a block's (value, index) pick, in thread 0; `none` is the index of no candidate
+__device__ __forceinline__ void arg_block(float& bv, int& bi, int none) {
+    __shared__ float sv[32];
+    __shared__ int si[32];
+    for (int off = 16; off > 0; off >>= 1) {
+        const float ov = __shfl_down_sync(0xffffffffu, bv, off);
+        const int oi = __shfl_down_sync(0xffffffffu, bi, off);
+        arg_take(ov, oi, bv, bi);
+    }
+    const int warp = (int) (threadIdx.x >> 5), lane = (int) (threadIdx.x & 31);
+    if (lane == 0) { sv[warp] = bv; si[warp] = bi; }
+    __syncthreads();
+    if (warp == 0) {
+        const int nw = (int) (blockDim.x >> 5);
+        bv = lane < nw ? sv[lane] : __int_as_float(0xff800000);
+        bi = lane < nw ? si[lane] : none;
+        for (int off = 16; off > 0; off >>= 1) {
+            const float ov = __shfl_down_sync(0xffffffffu, bv, off);
+            const int oi = __shfl_down_sync(0xffffffffu, bi, off);
+            arg_take(ov, oi, bv, bi);
+        }
+    }
+}
+
+// A row's scratch, whatever the launch's row count (one scratch serves windows of every size): its counter, then
+// the blocks' values and indices.
+struct ArgRow {
+    unsigned counter;
+    unsigned pad[3];
+    float v[kArgMaxBlocks];
+    int i[kArgMaxBlocks];
+};
+
+// grid (blocks, rows): block b scans [b * per_block, +per_block) of its row
+__global__ void argmax_rows_kernel(const float* __restrict__ logits, int n, int per_block, ArgRow* __restrict__ rows,
+                                   int32_t* out) {
+    const int row = (int) blockIdx.y, b = (int) blockIdx.x, nb = (int) gridDim.x;
+    unsigned* counter = &rows[row].counter;
+    float* pv = rows[row].v;
+    int* pi = rows[row].i;
+    const float* l = logits + (size_t) row * n;
+    const int lo = b * per_block, hi = min(n, lo + per_block);
+    float bv = __int_as_float(0xff800000);   // -inf
+    int bi = n;
+    for (int v = lo + (int) threadIdx.x; v < hi; v += (int) blockDim.x) {
+        const float s = l[v];
+        if (s > bv) { bv = s; bi = v; }
+    }
+    arg_block(bv, bi, n);
+    __shared__ bool last;
+    if (threadIdx.x == 0) {
+        pv[b] = bv;
+        pi[b] = bi;
+        __threadfence();
+        last = atomicAdd(counter, 1u) == (unsigned) (nb - 1);
+    }
+    __syncthreads();
+    if (!last) return;
+    __threadfence();
+    bv = __int_as_float(0xff800000);
+    bi = n;
+    if ((int) threadIdx.x < nb) {
+        bv = ((volatile const float*) pv)[threadIdx.x];
+        bi = ((volatile const int*) pi)[threadIdx.x];
+    }
+    arg_block(bv, bi, n);
+    if (threadIdx.x == 0) {
+        out[row] = bi < n ? bi : 0;   // no value above -inf: 0, as the one-block kernel
+        *counter = 0;
+    }
+}
+
+// a row's top-probability scratch: its counter, then its 32 warps' sums
+struct TopRow {
+    unsigned counter;
+    unsigned pad[3];
+    float parts[32];
+};
+
+// grid (kTopBlocks, rows), 1024 / kTopBlocks threads: thread j of block b is row_top_prob's thread b * blockDim + j
+__global__ void row_top_prob_split_kernel(const float* __restrict__ logits, int n_vocab, const int32_t* __restrict__ ids,
+                                          float* __restrict__ probs, TopRow* __restrict__ rows) {
+    const int row = (int) blockIdx.y;
+    float* parts = rows[row].parts;
+    const float* l = logits + (size_t) row * n_vocab;
+    const float m = l[ids[row]];
+    const int vt = (int) (blockIdx.x * blockDim.x + threadIdx.x);
+    float s = 0.0f;
+    for (int i = vt; i < n_vocab; i += 1024) s += __expf(l[i] - m);
+    for (int o = 16; o > 0; o >>= 1) s += __shfl_xor_sync(0xffffffffu, s, o);
+    if ((threadIdx.x & 31) == 0) parts[vt >> 5] = s;
+    __threadfence();
+    __syncthreads();
+    __shared__ bool last;
+    if (threadIdx.x == 0) last = atomicAdd(&rows[row].counter, 1u) == (unsigned) (gridDim.x - 1);
+    __syncthreads();
+    if (!last || threadIdx.x != 0) return;
+    __threadfence();
+    float tot = 0.0f;
+    for (int w = 0; w < 32; ++w) tot += ((volatile const float*) parts)[w];
+    probs[row] = 1.0f / tot;
+    rows[row].counter = 0;
+}
+
+}  // namespace
+
+uint64_t argmax_rows_scratch_bytes(int n_rows) { return (uint64_t) n_rows * sizeof(ArgRow); }
+
+void argmax_rows(const float* logits, int n_rows, int n, void* scratch, int32_t* out, void* stream) {
+    if (n_rows <= 0) return;
+    const int nb = std::min(kArgMaxBlocks, std::max(1, (n + 4095) / 4096));
+    const int per_block = (n + nb - 1) / nb;
+    argmax_rows_kernel<<<dim3((unsigned) nb, (unsigned) n_rows), kArgThreads, 0, (cudaStream_t) stream>>>(
+        logits, n, per_block, (ArgRow*) scratch, out);
+    check("argmax_rows");
+}
+
+uint64_t row_top_prob_scratch_bytes(int n_rows) { return (uint64_t) n_rows * sizeof(TopRow); }
+
+void row_top_prob_split(const float* logits, int n_rows, int n_vocab, const int32_t* ids, float* probs, void* scratch,
+                        void* stream) {
+    if (n_rows <= 0) return;
+    row_top_prob_split_kernel<<<dim3(kTopBlocks, (unsigned) n_rows), 1024 / kTopBlocks, 0, (cudaStream_t) stream>>>(
+        logits, n_vocab, ids, probs, (TopRow*) scratch);
+    check("row_top_prob_split");
+}
+
 namespace {
 __global__ void window_ids_kernel(int32_t* steps, int window, int32_t* ids, long long stride) {
     const int q = blockIdx.y;

src/kernels/verify_parity.cpp

@@ -17,6 +17,7 @@
 #include "strata/kernels/native_qsa_indexer.hpp"
 #include "strata/kernels/native_rope.hpp"
 #include "strata/kernels/qsa_decode_attn.hpp"
+#include "strata/kernels/sampler.hpp"
 #include "strata/kernels/native_router.hpp"
 #include "strata/kernels/s2_expert_grouped.hpp"
 #include "strata/kernels/shared_expert.hpp"
@@ -1233,6 +1234,66 @@ int test_prefill_gdn(std::mt19937& rng, cudaStream_t s) {
     return bad;
 }
 
+// argmax_rows against sample_tokens' greedy pick and row_top_prob_split against row_top_prob, bitwise: rows of both
+// vocabularies with ties at the maximum (the lowest index wins), NaN, -inf, a row with no value above -inf, and
+// launches repeated on one scratch (each must leave its counters at zero)
+int test_argmax(std::mt19937& rng, cudaStream_t s) {
+    using namespace strata::kernels;
+    int bad = 0;
+    std::normal_distribution<float> nd(0.0f, 4.0f);
+    uint8_t* sa = dev<uint8_t>(argmax_rows_scratch_bytes(8));   // one scratch for every row count, as in the engine
+    uint8_t* st = dev<uint8_t>(row_top_prob_scratch_bytes(8));
+    check(cudaMemset(sa, 0, argmax_rows_scratch_bytes(8)), "memset");
+    check(cudaMemset(st, 0, row_top_prob_scratch_bytes(8)), "memset");
+    for (const int n : {248320, 40525, 4097, 33}) {
+        for (const int rows : {4, 8, 1, 3, 8, 2}) {
+            std::vector<float> h((size_t) rows * n);
+            for (float& v : h) v = nd(rng);
+            for (int r = 0; r < rows; ++r) {
+                float* row = h.data() + (size_t) r * n;
+                const int a = (int) (rng() % n), b = (int) (rng() % n);
+                row[a] = row[b] = 40.0f;                       // a tie at the maximum
+                row[(int) (rng() % n)] = std::nanf("");        // never picked
+                row[(int) (rng() % n)] = -INFINITY;
+                if (r == 2)                                    // no value above -inf: 0
+                    for (int i = 0; i < n; ++i) row[i] = i % 7 ? -INFINITY : std::nanf("");
+            }
+            float* d_l = dev<float>(h.size());
+            up(d_l, h);
+            int* d_ref = dev<int>(rows);
+            int32_t* d_out = dev<int32_t>(rows);
+            float* d_p1 = dev<float>(rows);
+            float* d_p2 = dev<float>(rows);
+            SamplerParams sp;
+            sp.greedy = true;
+            sp.temperature = 0.0f;
+            sample_tokens(d_l, rows, n, nullptr, 0, sp, d_ref, s);
+            row_top_prob(d_l, rows, n, d_ref, d_p1, s);
+            for (int rep = 0; rep < 3; ++rep) {
+                argmax_rows(d_l, rows, n, sa, d_out, s);
+                row_top_prob_split(d_l, rows, n, d_out, d_p2, st, s);
+                check(cudaStreamSynchronize(s), "argmax");
+                const std::vector<int> ref = down(d_ref, rows);
+                const std::vector<int32_t> out = down(d_out, rows);
+                const std::vector<float> p1 = down(d_p1, rows), p2 = down(d_p2, rows);
+                for (int r = 0; r < rows; ++r)
+                    if (ref[r] != out[r] || std::memcmp(&p1[r], &p2[r], 4) != 0) {
+                        if (bad < 5)
+                            std::printf("  argmax n %d rows %d row %d: %d / %d, p %.9g / %.9g\n", n, rows, r, ref[r],
+                                        out[r], p1[r], p2[r]);
+                        ++bad;
+                    }
+            }
+            cudaFree(d_l); cudaFree(d_ref); cudaFree(d_out); cudaFree(d_p1); cudaFree(d_p2);
+        }
+    }
+    cudaFree(sa);
+    cudaFree(st);
+    std::printf("argmax_rows, row_top_prob_split: %s\n",
+                bad ? "MISMATCH" : "bitwise equal to the one-block kernels (ties, NaN, -inf, 1-8 rows on one scratch)");
+    return bad;
+}
+
 }  // namespace
 
 int main(int argc, char** argv) {
@@ -1260,6 +1321,7 @@ int main(int argc, char** argv) {
     bad += test_ple_tokens(rng, s);
     bad += test_decode_attn(rng, s);
     bad += test_prefill_gdn(rng, s);
+    bad += test_argmax(rng, s);
     cudaStreamDestroy(s);
     std::printf("verify_parity: %s\n", bad ? "FAIL" : "PASS");
     return bad ? 1 : 0;