serve --batch: a layer split reads a prompt in STRATA_SPLIT_PIECE-token runs, decided at every boundary

CMakeLists.txt

@@ -1060,6 +1060,10 @@ if(STRATA_BUILD_TESTS OR STRATA_BUILD_CONVERSATION_TESTS)
   add_executable(message_boundary_test src/program/message_boundary_test.cpp)
   target_include_directories(message_boundary_test PRIVATE include)
   add_test(NAME message_boundary_test COMMAND message_boundary_test)
+  # serve --batch: how many tokens one prompt-read run takes (header-only, no library to link)
+  add_executable(read_piece_test src/program/read_piece_test.cpp)
+  target_include_directories(read_piece_test PRIVATE include)
+  add_test(NAME read_piece_test COMMAND read_piece_test)
   # the serve loop's conversation-cache retention policy (header-only, no library to link)
   add_executable(conv_cache_test src/program/conv_cache_test.cpp)
   target_include_directories(conv_cache_test PRIVATE include)

include/strata/program/read_piece.hpp

@@ -0,0 +1,26 @@
+#pragma once
+
+#include <algorithm>
+#include <cstdint>
+
+namespace strata::program {
+
+// How many prompt tokens one run of a --batch prompt read covers.  Beside decoding slots (and on one GPU) it is one
+// prompt chunk, so the slots decode between chunks.  A layer split with no slot decoding reads the rest as ONE run,
+// which pipelines its stages over the chunks (returned as 0) -- unless `piece` (STRATA_SPLIT_PIECE) is set: then it
+// reads that many tokens (a whole number of chunks) per run, so a waiting shorter request's BYIELD is seen within one
+// piece instead of after the whole read.
+inline int64_t prompt_read_piece(int64_t chunk, bool split_idle, int64_t piece) {
+    chunk = std::max<int64_t>(chunk, 1);
+    if (!split_idle) return chunk;
+    if (piece <= 0) return 0;
+    return std::max(chunk, piece / chunk * chunk);
+}
+
+// Where one run of that read ends: `piece` from prompt_read_piece, decided again at every boundary, so a read that
+// started beside a decoding slot goes back to whole pieces (or the whole rest) once no slot decodes.
+inline int64_t prompt_read_run_end(int64_t q, int64_t end, int64_t piece) {
+    return piece <= 0 ? end : std::min(end, q + piece);
+}
+
+} // namespace strata::program

src/program/generate.cpp

@@ -60,6 +60,7 @@
 #include "strata/program/logits_selection.hpp"
 #include "strata/program/conv_cache.hpp"
 #include "strata/program/message_boundary.hpp"
+#include "strata/program/read_piece.hpp"
 #include "strata/spec/draft_policy.hpp"
 #include "strata/spec/draft_source.hpp"
 #include "strata/spec/suffix_drafter.hpp"
@@ -9508,11 +9509,22 @@ int main(int argc, char** argv) {
             // the request ends with `YIELDED <slot> <tokens>` + DONE cancel, and the server sends it again later: it
             // continues from the slot with the same chunks.  #656's cooperative preemption, with a slot as the park.
             auto read_part = [&](int64_t a0, int64_t b0, std::string& e) -> bool {
-                // (a layer split reads its stages as a pipeline over one run's chunks: in pieces only beside slots)
-                if (o.batch <= 0 || piped || (!stages.empty() && !batch_on())) return sp.run(ids.data() + a0, b0 - a0, a0, e);
+                // (a layer split reads its stages as a pipeline over one run's chunks: in pieces only beside slots,
+                // or in STRATA_SPLIT_PIECE-token runs so a shorter request's BYIELD is seen within one piece)
+                static const int64_t split_piece = [] {
+                    const char* v = std::getenv("STRATA_SPLIT_PIECE");
+                    return v != nullptr ? std::max<int64_t>(0, std::atoll(v)) : int64_t(0);
+                }();
+                if (o.batch <= 0 || piped) return sp.run(ids.data() + a0, b0 - a0, a0, e);
+                // decided again at every boundary: a read that began beside a decoding slot goes back to whole
+                // pieces once no slot decodes (they finish their short turns within seconds of a long read)
+                const auto piece_now = [&] {
+                    return strata::program::prompt_read_piece(sp.chunk(), !stages.empty() && !batch_on(), split_piece);
+                };
+                if (piece_now() == 0) return sp.run(ids.data() + a0, b0 - a0, a0, e);
                 const int64_t C = std::max<int64_t>(sp.chunk(), 1);
                 for (int64_t q = a0; q < b0;) {
-                    const int64_t r = std::min(b0, q + C);
+                    const int64_t r = strata::program::prompt_read_run_end(q, b0, piece_now());
                     const auto tq = Clock::now();
                     if (!sp.run(ids.data() + q, r - q, q, e)) return false;
                     q = r;

src/program/read_piece_test.cpp

@@ -0,0 +1,46 @@
+#include "strata/program/read_piece.hpp"
+
+#include <cstdio>
+
+int main() {
+    using strata::program::prompt_read_piece;
+    int checks = 0;
+    auto check = [&](int64_t chunk, bool split_idle, int64_t piece, int64_t expected) {
+        ++checks;
+        const int64_t got = prompt_read_piece(chunk, split_idle, piece);
+        if (got != expected)
+            std::printf("FAIL prompt_read_piece(%lld, %d, %lld) = %lld, expected %lld\n", (long long) chunk,
+                        (int) split_idle, (long long) piece, (long long) got, (long long) expected);
+        return got == expected;
+    };
+    bool ok = true;
+    // beside decoding slots (or one GPU): one prompt chunk at a time, whatever the piece
+    ok &= check(8192, false, 0, 8192);
+    ok &= check(8192, false, 65536, 8192);
+    ok &= check(0, false, 65536, 1);          // a chunk of 0 still advances
+    // a layer split with no slot decoding: one run over the rest unless a piece is set
+    ok &= check(8192, true, 0, 0);
+    ok &= check(8192, true, -5, 0);
+    // a piece is a whole number of chunks (the runs keep the same chunk grid), at least one
+    ok &= check(8192, true, 65536, 65536);
+    ok &= check(8192, true, 70000, 65536);
+    ok &= check(8192, true, 100, 8192);
+    ok &= check(32768, true, 65536, 65536);
+    ok &= check(0, true, 3, 3);
+    // a run ends one piece on, at the end of the read, or (piece 0: no slot decoding, no piece set) at the end
+    auto end_is = [&](int64_t q, int64_t end, int64_t piece, int64_t expected) {
+        ++checks;
+        const int64_t got = strata::program::prompt_read_run_end(q, end, piece);
+        if (got != expected)
+            std::printf("FAIL prompt_read_run_end(%lld, %lld, %lld) = %lld, expected %lld\n", (long long) q,
+                        (long long) end, (long long) piece, (long long) got, (long long) expected);
+        return got == expected;
+    };
+    ok &= end_is(0, 249251, 8192, 8192);
+    ok &= end_is(245760, 249251, 8192, 249251);
+    ok &= end_is(8192, 249251, 65536, 73728);
+    ok &= end_is(8192, 249251, 0, 249251);
+    ok &= end_is(8192, 249251, -1, 249251);
+    if (!ok) return 1;
+    std::printf("read_piece_test OK (%d checks)\n", checks);
+}

sycl/src/program/generate.cpp

@@ -73,6 +73,7 @@ namespace strata::prefill { void set_nonresident_share(double share); }   // SYC
 #include "strata/program/logits_selection.hpp"
 #include "strata/program/conv_cache.hpp"
 #include "strata/program/message_boundary.hpp"
+#include "strata/program/read_piece.hpp"
 #include "strata/spec/draft_policy.hpp"
 #include "strata/spec/draft_source.hpp"
 #include "strata/spec/suffix_drafter.hpp"
@@ -10036,11 +10037,22 @@ int main(int argc, char **argv) try {
             // the request ends with `YIELDED <slot> <tokens>` + DONE cancel, and the server sends it again later: it
             // continues from the slot with the same chunks.  #656's cooperative preemption, with a slot as the park.
             auto read_part = [&](int64_t a0, int64_t b0, std::string& e) -> bool {
-                // (a layer split reads its stages as a pipeline over one run's chunks: in pieces only beside slots)
-                if (o.batch <= 0 || piped || (!stages.empty() && !batch_on())) return sp.run(ids.data() + a0, b0 - a0, a0, e);
+                // (a layer split reads its stages as a pipeline over one run's chunks: in pieces only beside slots,
+                // or in STRATA_SPLIT_PIECE-token runs so a shorter request's BYIELD is seen within one piece)
+                static const int64_t split_piece = [] {
+                    const char* v = std::getenv("STRATA_SPLIT_PIECE");
+                    return v != nullptr ? std::max<int64_t>(0, std::atoll(v)) : int64_t(0);
+                }();
+                if (o.batch <= 0 || piped) return sp.run(ids.data() + a0, b0 - a0, a0, e);
+                // decided again at every boundary: a read that began beside a decoding slot goes back to whole
+                // pieces once no slot decodes (they finish their short turns within seconds of a long read)
+                const auto piece_now = [&] {
+                    return strata::program::prompt_read_piece(sp.chunk(), !stages.empty() && !batch_on(), split_piece);
+                };
+                if (piece_now() == 0) return sp.run(ids.data() + a0, b0 - a0, a0, e);
                 const int64_t C = std::max<int64_t>(sp.chunk(), 1);
                 for (int64_t q = a0; q < b0;) {
-                    const int64_t r = std::min(b0, q + C);
+                    const int64_t r = strata::program::prompt_read_run_end(q, b0, piece_now());
                     const auto tq = Clock::now();
                     if (!sp.run(ids.data() + q, r - q, q, e)) return false;
                     q = r;