batch-mtp: work with a layer split (slot drafters on the last stage's GPU, grouped slot rows through every stage)

- generate.cpp: --batch-mtp accepted with a layer split on separate GPUs (refused with --batch-groups); slot drafters
  load from the last stage's slot sessions under the last stage's device guard and bind to its weights/head; their
  residual buffers, row copy (async on the drafter's stream) and draft-KV copies run on the drafter's device; every
  stage verifier gets the 8-row window size and the batch graph limit; slot-drafter logits are not charged to CUDA0.
- verify.cpp: stage_batch's layout rules move to include/strata/core/batch_rows.hpp (CPU test batch_rows_test); grouped
  rows are allowed on a stage that hands rows on, the hand-off bound is kVerifyMaxT rows.
- docs/BATCHING.md, --batch-mtp help text.

CMakeLists.txt

@@ -1130,6 +1130,10 @@ if(STRATA_BUILD_TESTS OR STRATA_BUILD_CONVERSATION_TESTS)
   add_executable(bf16_bits_test src/kernels/bf16_bits_test.cpp)
   target_include_directories(bf16_bits_test PRIVATE include)
   add_test(NAME bf16_bits_test COMMAND bf16_bits_test)
+  # the batch window's row layout (a slot's group of proposed rows, the layer split's hand-off capacity; header-only, CPU)
+  add_executable(batch_rows_test tests/core/batch_rows_test.cpp)
+  target_include_directories(batch_rows_test PRIVATE include)
+  add_test(NAME batch_rows_test COMMAND batch_rows_test)
 endif()
 
 add_library(strata_kernels_cpu STATIC src/kernels/cpu/expert.cpp src/kernels/cpu/pool.cpp

docs/BATCHING.md

@@ -19,12 +19,22 @@ asked, with a note when it is more than setup would recommend.
 "parallel": 2
 ```
 
-On one GPU with MTP (`--mtp` and `--spec`), `--batch-mtp` (in the config's `args`, or `STRATA_BATCH_MTP=1` in the
-server's environment) lets each batch slot verify one MTP proposal per window. It is opt-in; without it the batch
-behaviour described below is exactly the one without MTP. It needs VRAM per slot for the draft state and buffers, so
-check the engine's free-memory log before using it on a smaller card. If it cannot run (one slot, no `--mtp`, a layer
-split or helper GPU) the engine says so and batches as usual. RTX PRO 5000 owners measured +31% to +39% total
-throughput with 2 to 4 clients (a RX R9700 run too); it has not been validated with a layer split.
+With MTP (`--mtp` and `--spec`), `--batch-mtp` (in the config's `args`, or `STRATA_BATCH_MTP=1` in the server's
+environment) lets each batch slot verify one MTP proposal per window. It is opt-in; without it the batch behaviour
+described below is exactly the one without MTP. It needs VRAM per slot for the draft state and buffers, so check the
+engine's free-memory log before using it on a smaller card. If it cannot run (one slot, no `--mtp`, `--batch-groups`
+with a layer split, a layer split on one GPU) the engine says so and batches as usual. RTX PRO 5000 owners measured
++31% to +39% total throughput with 2 to 4 clients on one GPU (a RX R9700 run too).
+
+**With a layer split** (e.g. `"layer_split": "24"` on two GPUs, `"parallel": 2`, `--spec 4 --mtp ...`) it works the
+same way, with the slot drafters on the **last stage's GPU** (where the solo drafter and the head are): each slot gets a
+drafter that shares the solo drafter's weights and owns only its K/V ring and buffers (the engine logs `--batch-mtp: N
+slot drafters on CUDAk (X MiB ... each)` at start; the K/V ring of each also keeps a pinned-RAM copy, as the solo
+drafter's does). A window carries two rows per active slot (its token and its proposal) through every stage, so at most
+four slots are in one window; more slots (`--batch` up to what fits) rotate through. The split keeps `--kv-pool-tokens`
+and `--kv-resident` as without `--batch-mtp`. Not supported with a split: `--batch-groups` (pipelined slot groups; the
+engine turns `--batch-mtp` off and says so). The code does not depend on the number of stages, but the two-stage split
+is the one it was written for.
 
 With a layer split, the engine options go into the config's `args`:
 
@@ -124,7 +134,8 @@ counter-based draw (Philox(seed, position)).
 
 - By default, batch windows carry no MTP drafts: a conversation in a slot decodes one token per window (the solo
   path keeps its drafts, which is why a request alone is not put in a slot, and goes back to it when left alone).
-- Grouped MTP currently uses one proposal per slot and requires one GPU; it does not support a layer split.
+- `--batch-mtp` uses one proposal per slot. With a layer split it needs each stage on its own GPU and does not combine
+  with `--batch-groups` (pipelined slot groups).
 - Repetition / frequency / presence penalties are not applied in batch windows.
 - A prompt shorter than one chunk is read in one piece (the slots wait for it); a read gives way only at a chunk
   boundary, and not for pictures.

include/strata/core/batch_rows.hpp

@@ -0,0 +1,30 @@
+#pragma once
+// The row layout of a batch window (Verifier::stage_batch), as pure logic so a CPU test can pin it down.
+//
+// A batch window carries S rows, each a token of one slot.  A slot owns one row (plain --batch) or a contiguous group
+// of rows at consecutive positions (--batch-mtp: the slot's current token and its MTP proposal; a layer split's every
+// stage sees the same layout).  The rows travel between the stages of a layer split in hand-off buffers of
+// `handoff_rows` rows, the window's group starting at row `hbase`.
+
+#include <cstdint>
+
+namespace strata::core {
+
+/// nullptr when the layout is valid, else why not.  `rows[t]` is row t's slot, `pos[t]` its position.  `n_slots`: the
+/// slots the verifier holds; `max_t`: its window capacity; `via_handoff`: the stage hands its rows on or takes them from
+/// another stage (a layer split), so [hbase, hbase + S) must fit the `handoff_rows` rows of the hand-off buffers.
+inline const char* batch_rows_error(const int* rows, int S, const int64_t* pos, int n_slots, int max_t, int hbase,
+                                    bool via_handoff, int handoff_rows) {
+    if (S < 1 || S > max_t || hbase < 0 || (via_handoff && hbase + S > handoff_rows))
+        return "verify: batch rows out of range (init_slots)";
+    for (int t = 0; t < S; ++t) {
+        if (rows[t] < 0 || rows[t] >= n_slots) return "verify: a batch row's slot is out of range";
+        for (int u = 0; u < t - 1; ++u)
+            if (rows[u] == rows[t] && rows[t - 1] != rows[t]) return "verify: a slot's proposed rows must be contiguous";
+        if (t > 0 && rows[t] == rows[t - 1] && pos[t] != pos[t - 1] + 1)
+            return "verify: proposed rows must have consecutive positions";
+    }
+    return nullptr;
+}
+
+}  // namespace strata::core

src/core/verify.cpp

@@ -1,4 +1,5 @@
 // src/core/verify.cpp - see include/strata/core/verify.hpp.
+#include "strata/core/batch_rows.hpp"
 #include "strata/core/verify.hpp"
 #include "strata/core/remote_expert_opt.hpp"
 #include "strata/core/dma_batch.hpp"
@@ -2552,30 +2553,13 @@ bool Verifier::instantiate_evicting(cudaGraphExec_t& ex, cudaGraph_t graph, cons
 bool Verifier::stage_batch(const int* rows, int S, int hbase, const int32_t* tokens, const int64_t* pos,
                            std::string& err) {
     using namespace strata::kernels;
-    if (S < 1 || S > max_t_ || hbase < 0 ||
-        (next_ != nullptr && hbase + S > (int) slots_.size())) {
-        err = "verify: batch rows out of range (init_slots)";
+    // the row layout: a slot may own a contiguous group of rows (--batch-mtp), also across the stages of a layer split,
+    // whose hand-off buffers hold kVerifyMaxT rows (a stage that hands rows on or takes them)
+    if (const char* why = batch_rows_error(rows, S, pos, (int) slots_.size(), max_t_, hbase,
+                                           next_ != nullptr || lb_ > 0, kVerifyMaxT)) {
+        err = why;
         return false;
     }
-    for (int t = 0; t < S; ++t) {
-        if (rows[t] < 0 || rows[t] >= (int) slots_.size()) {
-            err = "verify: a batch row's slot is out of range";
-            return false;
-        }
-        for (int u = 0; u < t - 1; ++u)
-            if (rows[u] == rows[t] && rows[t - 1] != rows[t]) {
-                err = "verify: a slot's proposed rows must be contiguous";
-                return false;
-            }
-        if (t > 0 && rows[t] == rows[t - 1] && pos[t] != pos[t - 1] + 1) {
-            err = "verify: proposed rows must have consecutive positions";
-            return false;
-        }
-        if (t > 0 && rows[t] == rows[t - 1] && next_ != nullptr) {
-            err = "verify: grouped slot rows do not support a layer split yet";
-            return false;
-        }
-    }
     if (released_.load()) { err = "verify: an earlier window never finished on the GPU (#267); restart the engine"; return false; }
     const ModelGeometry& g = *g_;
     for (int t = 0; t < S; ++t)

src/program/generate.cpp

@@ -770,8 +770,9 @@ void usage() {
                  "  --batch N / --slots N  --serve (opt-in): up to N requests decode together in batch slots (2..8\n"
                  "                       normally; --batch-mtp waves more through eight-row windows),\n"
                  "                       each slot with its own session (VRAM like the main one); docs/BATCHING.md\n"
-                 "  --batch-mtp          --batch (opt-in, one GPU, needs --mtp and --spec): each slot also verifies one MTP\n"
-                 "                       proposal per window (STRATA_BATCH_MTP=1 does the same); needs VRAM per slot\n"
+                 "  --batch-mtp          --batch (opt-in, needs --mtp and --spec): each slot also verifies one MTP proposal per\n"
+                 "                       window (STRATA_BATCH_MTP=1 does the same); works on one GPU and on a layer split (the\n"
+                 "                       slot drafters sit on the last stage's GPU; not with --batch-groups); needs VRAM per slot\n"
                  "  --batch-groups G     --batch with a layer split: the slots in G groups pipelined through the GPUs\n"
                  "  --pipeline-windows N --serve with a layer split on two GPUs (opt-in): one conversation's verify windows\n"
                  "                       with the GPUs overlapped - 1 = the prompt's short reads, 2 = decode as well\n"
@@ -4114,12 +4115,17 @@ int main(int argc, char** argv) {
     // engine still starts and serves one request at a time.
     // --batch-mtp / STRATA_BATCH_MTP=1 (opt-in): a slot also verifies one MTP proposal per window.  Recommend, never
     // force: when it cannot run it is said and left off (plain batching still starts); default off = exactly 0.1.39.
+    // With a layer split the slots' drafters live on the LAST stage's GPU (where the solo drafter, the head and the final
+    // residual rows are); the windows are the ones --batch runs, two rows per slot, through every stage.
     const char* batch_mtp_env = std::getenv("STRATA_BATCH_MTP");
     bool batch_mtp = o.batch_mtp || (batch_mtp_env != nullptr && batch_mtp_env[0] != '\0' && batch_mtp_env[0] != '0');
     if (batch_mtp) {
+        // (--pipeline-windows is already off with --batch slots)
         const char* why = o.batch < 2 ? "it needs --batch 2 or more" : o.mtp.empty() ? "it needs --mtp"
-                        : o.spec < 2 ? "it needs --spec T (T >= 2)" : (multi_gpu || split_same || !stages.empty())
-                        ? "it is for one GPU (no layer split or helper) for now" : !o.serve ? "it needs --serve" : nullptr;
+                        : o.spec < 2 ? "it needs --spec T (T >= 2)" : split_same
+                        ? "it needs each stage of a layer split on its own GPU" : (o.batch_groups > 1 && !stages.empty())
+                        ? "it does not combine with --batch-groups (pipelined slot groups) yet"
+                        : !o.serve ? "it needs --serve" : nullptr;
         if (why != nullptr) {
             std::fprintf(stderr, "strata generate: WARNING: --batch-mtp is off: %s\n", why);
             batch_mtp = false;
@@ -4237,9 +4243,13 @@ int main(int argc, char** argv) {
         if (!o.mtp.empty() && !mtp.load(o.mtp, draft_geometry, last_st ? last_st->ss : ss, mtp_t, err, o.mtp_window)) { std::fprintf(stderr, "strata generate: %s%s\n", err.c_str(), vram_free_note().c_str()); return 1; }
         mtp.set_ple_session(&ss);
         if (batch_mtp) {
+            // Slot b's drafter lives where the solo drafter does (the guard above: the last stage's GPU), carved from the
+            // LAST stage's slot-b session (bslot_ss.back(): stage 0's on one GPU): it borrows that session's context
+            // length and RoPE table, which are on that device.  Its PLE history is the first stage's slot session: PLE
+            // runs on the first stage, as for the solo drafter (`mtp.set_ple_session(&ss)` above).
             for (int b = 0; b < o.batch; ++b) {
                 auto d = std::make_unique<strata::core::MtpDrafter>();
-                if (!d->load(o.mtp, draft_geometry, *bslot_ss[0][(size_t) b], o.spec, err, o.mtp_window, &mtp)) {
+                if (!d->load(o.mtp, draft_geometry, *bslot_ss.back()[(size_t) b], o.spec, err, o.mtp_window, &mtp)) {
                     std::fprintf(stderr, "strata generate: batch MTP slot %d: %s%s\n", b, err.c_str(),
                                  vram_free_note().c_str());
                     return 1;
@@ -4248,6 +4258,14 @@ int main(int argc, char** argv) {
                 d->set_ple_session(bslot_ss[0][(size_t) b].get());
                 slot_mtp.push_back(std::move(d));
             }
+            if (!slot_mtp.empty()) {
+                size_t fb = 0, tb = 0;
+                cudaMemGetInfo(&fb, &tb);
+                std::fprintf(stderr, "strata generate: --batch-mtp: %d slot drafters on CUDA%d (%.0f MiB of private K/V "
+                                     "state and buffers each, sharing the solo drafter's weights); %.2f GiB free there\n",
+                             (int) slot_mtp.size(), mtp.device(), (double) slot_mtp.front()->vram_bytes() / 1048576.0,
+                             (double) fb / 1073741824.0);
+            }
         }
     }
     // THE HEAD BEFORE THE CACHE, AND BEFORE THE ARENA.  The expert cache takes what is free minus the reserve, so
@@ -4589,8 +4607,10 @@ int main(int argc, char** argv) {
         // ~110-180 MiB larger, and out of the reserve they left 16 GB cards below the stall line (#199)
         int64_t mtp_bind = (!o.mtp.empty() && native_head.loaded())
                                ? (int64_t) mtp.bind_bytes(native_head.row_bytes(), n_vocab) : 0;
-        for (const auto& d : slot_mtp)
-            mtp_bind += (int64_t) d->bind_bytes(native_head.row_bytes(), n_vocab);
+        // (with a layer split the slot drafters bind on the last stage, whose cache is sized from what is free there)
+        if (stages.empty())
+            for (const auto& d : slot_mtp)
+                mtp_bind += (int64_t) d->bind_bytes(native_head.row_bytes(), n_vocab);
         const int64_t reserve = (((int64_t) o.vram_reserve_mib + prefill_mib) << 20) + mtp_bind + pipe_first;
         const int64_t blob = (int64_t) strata::kernels::cpu::expert_layout().max_blob;
         int64_t slots = ((int64_t) free_b - reserve) / blob;
@@ -7143,7 +7163,11 @@ int main(int argc, char** argv) {
                     vs.slot_off = gs.cache.slot_offsets();
                     vs.n_slots = gs.cache.slots();
                     gs.ver.set_remote_expert_opt(remote_opt.get());
-                    ok_s = gs.ver.init(gs.wt, g, gs.ss, vs, gs.head.loaded() ? &gs.head : nullptr, std::max(o.spec, o.batch), err);
+                    // --batch-mtp: every stage runs the same 2-rows-per-slot windows as the first (up to kVerifyMaxT
+                    // rows, the hand-off's size) and bounds its captured batch layouts like the first (below)
+                    if (batch_mtp) gs.ver.set_batch_graph_limit(64);
+                    ok_s = gs.ver.init(gs.wt, g, gs.ss, vs, gs.head.loaded() ? &gs.head : nullptr,
+                                       batch_mtp ? strata::kernels::kVerifyMaxT : std::max(o.spec, o.batch), err);
                     split_drive.cache_base[st] = gs.cache.device_slot(0);
                     split_drive.cache_slot_off[st] = gs.cache.slot_offsets();
                     split_drive.pcie_num[st] = pcie_num_of(gs.pcie_frac);
@@ -7223,6 +7247,9 @@ int main(int argc, char** argv) {
         auto free_slot_mtp_rows = [](float* p) { if (p != nullptr) (void) cudaFree(p); };
         std::vector<std::unique_ptr<float, decltype(free_slot_mtp_rows)>> slot_mtp_rows;
         if (batch_mtp) {
+            // like the solo drafter's `bind` above: the last stage's weights (embedding) and head, and the residual
+            // buffers on the drafter's device (the last stage's GPU with a layer split)
+            const strata::core::OnDevice on_slot_mtp(mtp.device());
             for (int b = 0; b < o.batch; ++b) {
                 float* r = nullptr;
                 const size_t bytes = (size_t) o.spec * (size_t) g.hc * (size_t) g.n_embd * sizeof(float);
@@ -7232,7 +7259,8 @@ int main(int argc, char** argv) {
                     return 1;
                 }
                 slot_mtp_rows.emplace_back(r, free_slot_mtp_rows);
-                if (!slot_mtp[(size_t) b]->bind(wt, &native_head, r, err, &mtp)) {
+                if (!slot_mtp[(size_t) b]->bind(last_st ? last_st->wt : wt, last_st ? &last_st->head : &native_head, r, err,
+                                                &mtp)) {
                     std::fprintf(stderr, "strata serve: batch MTP slot %d: %s%s\n", b, err.c_str(),
                                  vram_free_note().c_str());
                     return 1;
@@ -8799,7 +8827,10 @@ int main(int argc, char** argv) {
                 if (cudaDeviceSynchronize() != cudaSuccess) { e = "batch admission: device sync failed"; return false; }
             }
             if (batch_mtp) {
-                // Admission first builds the solo draft KV; copy it into the slot before drafting.
+                // Admission first builds the solo draft KV; copy it into the slot before drafting.  Both drafters' K/V
+                // are on the drafter's device (the last stage's GPU with a layer split): the copy's ring refill is a
+                // kernel launch there.
+                const strata::core::OnDevice on_mtp(mtp.device());
                 strata::core::ConversationKv image;
                 if (!mtp.idle(e) || !slot_mtp[(size_t) b]->idle(e) ||
                     !strata::core::conversation_kv_save(image, mtp.kv_state(), draft_geometry, upto, false, e) ||
@@ -8843,7 +8874,8 @@ int main(int argc, char** argv) {
                 if (cudaDeviceSynchronize() != cudaSuccess) { e = "batch slot restore: device sync failed"; return false; }
             }
             if (batch_mtp) {
-                // A returning solo request resumes from its slot's draft KV.
+                // A returning solo request resumes from its slot's draft KV (on the drafter's device, as above).
+                const strata::core::OnDevice on_mtp(mtp.device());
                 strata::core::ConversationKv image;
                 if (!slot_mtp[(size_t) b]->idle(e) || !mtp.idle(e) ||
                     !strata::core::conversation_kv_save(image, slot_mtp[(size_t) b]->kv_state(),
@@ -8947,10 +8979,15 @@ int main(int argc, char** argv) {
                     sl.p += 1;
                 }
                 if (batch_mtp && sl.active) {
+                    // The slot's two residual rows (the final residuals of the last stage: ver.final_R_all() follows the
+                    // chain, so with a layer split they are on the last stage's GPU, where the drafters are).  The copy is
+                    // queued on the drafter's own stream after the window's host sync, so the draft graph reads it.
+                    const strata::core::OnDevice on_slot_mtp(slot_mtp[(size_t) b]->device());
                     const size_t stride = (size_t) g.hc * (size_t) g.n_embd;
-                    if (cudaMemcpy(slot_mtp_rows[(size_t) b].get(),
-                                   ver.final_R_all() + (size_t) first[t] * stride,
-                                   2 * stride * sizeof(float), cudaMemcpyDeviceToDevice) != cudaSuccess ||
+                    if (cudaMemcpyAsync(slot_mtp_rows[(size_t) b].get(),
+                                        ver.final_R_all() + (size_t) first[t] * stride,
+                                        2 * stride * sizeof(float), cudaMemcpyDeviceToDevice,
+                                        slot_mtp[(size_t) b]->stream()) != cudaSuccess ||
                         !slot_mtp[(size_t) b]->draft(2, outb + first[t], pos[first[t]], keep[b] - 1,
                                                     sl.draft.data(), err)) {
                         std::printf("ERR batch MTP slot %d: %s\n", b, err.c_str());

tests/core/batch_rows_test.cpp

@@ -0,0 +1,78 @@
+// The batch window's row layout (include/strata/core/batch_rows.hpp): plain one-row slots, a slot's contiguous group of
+// proposed rows (--batch-mtp), and the hand-off capacity of a layer split's stages.  Header-only, CPU.
+#include "strata/core/batch_rows.hpp"
+
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+
+namespace {
+int failures = 0;
+void check(bool ok, const char* what) {
+    if (!ok) {
+        std::fprintf(stderr, "FAIL: %s\n", what);
+        ++failures;
+    }
+}
+const char* run(const int* rows, const long long* p, int S, int n_slots, int max_t, int hbase, bool handoff) {
+    int64_t pos[8];
+    for (int t = 0; t < S; ++t) pos[t] = p[t];
+    return strata::core::batch_rows_error(rows, S, pos, n_slots, max_t, hbase, handoff, 8);
+}
+}  // namespace
+
+int main() {
+    using strata::core::batch_rows_error;
+    // plain batch: one row per slot, any positions
+    {
+        const int rows[] = {0, 1, 2};
+        const long long pos[] = {10, 500, 3};
+        check(run(rows, pos, 3, 3, 8, 0, false) == nullptr, "one row per slot");
+        check(run(rows, pos, 3, 3, 8, 0, true) == nullptr, "one row per slot through the hand-off");
+    }
+    // --batch-mtp: a slot's two rows at consecutive positions; also on a layer split (S > slots)
+    {
+        const int rows[] = {1, 1, 0, 0};
+        const long long pos[] = {40, 41, 7, 8};
+        check(run(rows, pos, 4, 2, 8, 0, false) == nullptr, "two rows per slot, one stage");
+        check(run(rows, pos, 4, 2, 8, 0, true) == nullptr, "two rows per slot across a layer split (S > slots)");
+    }
+    // eight rows = four slots of two: the whole hand-off
+    {
+        const int rows[] = {0, 0, 1, 1, 2, 2, 3, 3};
+        const long long pos[] = {1, 2, 5, 6, 9, 10, 20, 21};
+        check(run(rows, pos, 8, 4, 8, 0, true) == nullptr, "eight rows fill the hand-off");
+        check(run(rows, pos, 8, 4, 8, 1, true) != nullptr, "a group past the hand-off's end");
+        check(run(rows, pos, 8, 4, 8, 1, false) == nullptr, "no hand-off: the base does not bound the rows");
+    }
+    // refused layouts
+    {
+        const int rows[] = {0, 1, 0};
+        const long long pos[] = {5, 9, 6};
+        check(run(rows, pos, 3, 2, 8, 0, false) != nullptr, "a slot's rows must be contiguous");
+    }
+    {
+        const int rows[] = {0, 0};
+        const long long pos[] = {5, 7};
+        check(run(rows, pos, 2, 1, 8, 0, false) != nullptr, "a slot's rows need consecutive positions");
+    }
+    {
+        const int rows[] = {0, 2};
+        const long long pos[] = {5, 7};
+        check(run(rows, pos, 2, 2, 8, 0, false) != nullptr, "a slot past the slots held");
+        const int neg[] = {-1};
+        const long long p1[] = {1};
+        check(run(neg, p1, 1, 2, 8, 0, false) != nullptr, "a negative slot");
+    }
+    {
+        const int rows[] = {0, 0, 0};
+        const long long pos[] = {1, 2, 3};
+        check(run(rows, pos, 3, 1, 2, 0, false) != nullptr, "more rows than the window holds");
+        check(run(rows, pos, 0, 1, 8, 0, false) != nullptr, "an empty window");
+        check(run(rows, pos, 3, 1, 8, -1, false) != nullptr, "a negative hand-off base");
+        check(run(rows, pos, 3, 1, 8, 0, false) == nullptr, "three rows of one slot are a group");
+    }
+    if (failures != 0) return EXIT_FAILURE;
+    std::printf("batch_rows_test: ok\n");
+    return EXIT_SUCCESS;
+}