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;