From 1730b8760f2c7ba08a173a43aaf0ca3b4214be71 Mon Sep 17 00:00:00 2001 From: "s.anastasiou" Date: Thu, 8 Oct 2026 09:02:13 +0200 Subject: [PATCH] serve --batch: a layer split reads a prompt in STRATA_SPLIT_PIECE-token runs, decided at every boundary A layer split pipelines its stages across one sp.run's chunks, so a prompt read beside a decoding slot is cut into single chunks (a decode window between each) and loses that overlap. Until now the read kept that shape to its end, even after the slots finished their short turns a few seconds in: a long read ran at about half speed. read_part now asks at every chunk boundary how much to read next (prompt_read_piece): one chunk while a slot decodes, STRATA_SPLIT_PIECE tokens (rounded down to whole chunks) once none does, the rest of the prompt when the variable is unset. A piece bounds how long a shorter request's BYIELD waits. STRATA_SPLIT_PIECE defaults to 0 = the old behaviour (whole run when idle). Same change in sycl/ (not compiled here, no oneAPI). read_piece_test: 15 checks. --- CMakeLists.txt | 4 +++ include/strata/program/read_piece.hpp | 26 +++++++++++++++ src/program/generate.cpp | 18 +++++++++-- src/program/read_piece_test.cpp | 46 +++++++++++++++++++++++++++ sycl/src/program/generate.cpp | 18 +++++++++-- 5 files changed, 106 insertions(+), 6 deletions(-) create mode 100644 include/strata/program/read_piece.hpp create mode 100644 src/program/read_piece_test.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 995765af933..e88cb47427c 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/include/strata/program/read_piece.hpp b/include/strata/program/read_piece.hpp new file mode 100644 index 00000000000..07057b37906 --- /dev/null +++ b/include/strata/program/read_piece.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include +#include + +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(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 diff --git a/src/program/generate.cpp b/src/program/generate.cpp index ce8e0047443..fa50f0fa1b0 100644 --- a/src/program/generate.cpp +++ b/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 ` + 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(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(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; diff --git a/src/program/read_piece_test.cpp b/src/program/read_piece_test.cpp new file mode 100644 index 00000000000..f72a747e9f5 --- /dev/null +++ b/src/program/read_piece_test.cpp @@ -0,0 +1,46 @@ +#include "strata/program/read_piece.hpp" + +#include + +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); +} diff --git a/sycl/src/program/generate.cpp b/sycl/src/program/generate.cpp index d714f8c91aa..13817a92506 100644 --- a/sycl/src/program/generate.cpp +++ b/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 ` + 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(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(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;