diff --git a/CMakeLists.txt b/CMakeLists.txt index 4efc29766f8..59c46772d4a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -350,6 +350,14 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) add_test(NAME file_expert_source_test COMMAND file_expert_source_test) set_tests_properties(file_expert_source_test PROPERTIES TIMEOUT 60) endif() + # P1: the verify host pin picks a core no pool worker uses, and refuses under a layer split. + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/tests/core/verify_host_pin_test.cpp") + enable_testing() + add_executable(verify_host_pin_test tests/core/verify_host_pin_test.cpp) + target_link_libraries(verify_host_pin_test PRIVATE strata_engine) + add_test(NAME verify_host_pin_test COMMAND verify_host_pin_test) + set_tests_properties(verify_host_pin_test PROPERTIES TIMEOUT 60) + endif() if(STRATA_BUILD_TESTS AND STRATA_NATIVE_EXPERTS) # native_experts.txt v2-v4 (the real packs' copies and the per-role shard column) and the arena filled from # synthetic GGUF shards at the real expert dimensions diff --git a/include/strata/core/verify.hpp b/include/strata/core/verify.hpp index 5717eb4fa9c..42b537f6f40 100644 --- a/include/strata/core/verify.hpp +++ b/include/strata/core/verify.hpp @@ -33,6 +33,7 @@ #include #include #include +#include #include namespace strata::core { @@ -67,6 +68,11 @@ class Verifier { /// the flags guarded, so this verifier refuses every later window. True when the streams finished. bool release_gpu_waits(int timeout_ms); + /// P1: the core `ExpertPool` reserves for the host loop (the first physical core `physical_cores(true)` + /// drops), or -1 to refuse. A layer split returns -1: every stage's host would otherwise fight for the one + /// reserved core and serialize the stages. Pure topology, no side effects - `maybe_pin_host` calls it. + static int host_pin_core(bool split); + /// `max_t` <= kVerifyMaxT. `head` may be null (the canonical head is then run per token). bool init(const WeightTable& wt, const ModelGeometry& g, SessionState& ss, const VerifyHits& hits, const NativeHead* head, int max_t, std::string& err); @@ -174,6 +180,15 @@ class Verifier { int device_ = -1; ///< the device `init` ran on: run/commit switch to it (layer split) std::atomic released_{false}; ///< #267: release_gpu_waits ran (maybe on the watchdog thread): no more windows bool device_plan_ = false; ///< E-6: resident-only layers planned on the device (STRATA_VERIFY_DEVICE_PLAN) + // P1: the serve/spec host thread pins itself to the core ExpertPool reserves for the host - the other half of + // the reservation SessionLoopScratch::init honours on the generate path - once per verifier. Layer splits do + // not pin (see host_pin_core). Restored in the destructor when the same thread destroys the verifier. + bool pin_attempted_ = false; + bool pinned_ = false; + long long pin_prev_ = -1; + int pin_core_ = -1; + std::thread::id pin_thread_; + void maybe_pin_host(); uint32_t* skip_ = nullptr; ///< E-6: per group, the ring whose plan the device built (0: the host's) unsigned long long* slot_off_d_ = nullptr; ///< E-6: the slot offsets on the device int64_t lb_ = 0, le_ = -1; ///< set_stage: the layers this verifier runs (-1: to the last) diff --git a/src/core/verify.cpp b/src/core/verify.cpp index 999f85d8601..0fb6562b75b 100644 --- a/src/core/verify.cpp +++ b/src/core/verify.cpp @@ -8,6 +8,7 @@ #include "strata/core/on_device.hpp" #include "strata/kernels/iq_kernels.hpp" #include "strata/kernels/cpu/expert_layout.hpp" +#include "strata/kernels/cpu/pool.hpp" #include "strata/kernels/bf16_gemv.hpp" #include "strata/kernels/native_router.hpp" #include "strata/kernels/native_moe.hpp" @@ -171,6 +172,20 @@ Verifier::~Verifier() { Verifier* me = this; slot.compare_exchange_strong(me, nullptr); } + // P1: give the caller's thread its affinity back - but only the thread that took it. A verifier that landed + // on another thread cannot restore this thread's mask, and nailing the wrong thread to one core is worse than + // the pin leaking; log once and leave it. + if (pinned_) { + pinned_ = false; + if (std::this_thread::get_id() == pin_thread_) { + strata::kernels::cpu::restore_thread_affinity(pin_prev_); + } else { + static std::atomic warned{false}; + if (!warned.exchange(true)) + std::fprintf(stderr, "strata verify: the host pin (core %d) was taken on another thread; leaving " + "this thread's affinity untouched\n", pin_core_); + } + } if (cs_) cudaStreamSynchronize(cs_); for (auto& e : exec_) if (e) cudaGraphExecDestroy(e); @@ -1013,6 +1028,34 @@ bool Verifier::capture_commit(std::string& err) { return true; } +int Verifier::host_pin_core(bool split) { + if (split) return -1; // one reserved core, one stage's host: pinning every stage serializes them + const std::vector workers = strata::kernels::cpu::physical_cores(true); + for (int c : strata::kernels::cpu::physical_cores(false)) + if (std::find(workers.begin(), workers.end(), c) == workers.end()) return c; + return -1; +} + +void Verifier::maybe_pin_host() { + if (pin_attempted_) return; + pin_attempted_ = true; + if (std::getenv("STRATA_NO_VERIFY_PIN") != nullptr) return; // the A/B off arm + const bool split = next_ != nullptr || lb_ > 0 || (g_ != nullptr && le_ < g_->n_layers); + const int core = host_pin_core(split); + if (core < 0 || core >= 64) return; + const long long prev = strata::kernels::cpu::pin_current_thread(core); + // Refuse when the platform failed, or when the calling thread's affinity did not include the reserved core (a + // worker's thread, or a cpuset that excludes it): undo and leave the thread exactly as it was. + if (prev <= 0 || ((prev >> core) & 1) == 0) { + if (prev > 0) strata::kernels::cpu::restore_thread_affinity(prev); + return; + } + pin_prev_ = prev; + pin_core_ = core; + pinned_ = true; + pin_thread_ = std::this_thread::get_id(); +} + bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, void* user, int32_t* out, std::string& err) { using namespace strata::kernels; @@ -1022,6 +1065,7 @@ bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, const ModelGeometry& g = *g_; SessionState& ss = *ss_; if (pos0 + T > ss.qsa_states[ss.qsa_primary()].max_cells) { err = "verify: the window runs past the context"; return false; } + maybe_pin_host(); if (!capture(T, err) || !capture_commit(err)) return false; VDBG("captured; staging\n"); const Clock::time_point t0 = Clock::now(); diff --git a/tests/core/verify_host_pin_test.cpp b/tests/core/verify_host_pin_test.cpp new file mode 100644 index 00000000000..c664df5c2d1 --- /dev/null +++ b/tests/core/verify_host_pin_test.cpp @@ -0,0 +1,30 @@ +// tests/core/verify_host_pin_test.cpp - P1: the verify host thread pins to a core no pool worker uses, and never +// under a layer split (one reserved core, one stage's host - otherwise the stages serialize). +#include "strata/core/verify.hpp" +#include "strata/kernels/cpu/pool.hpp" + +#include +#include +#include + +int main() { + using strata::core::Verifier; + const std::vector workers = strata::kernels::cpu::physical_cores(true); + const std::vector all = strata::kernels::cpu::physical_cores(false); + const int core = Verifier::host_pin_core(false); + if (!all.empty() && core < 0) { + std::fprintf(stderr, "verify_host_pin: no reserved host core among %zu physical cores\n", all.size()); + return 1; + } + if (core >= 0 && std::find(workers.begin(), workers.end(), core) != workers.end()) { + std::fprintf(stderr, "verify_host_pin: chosen host core %d is a pool worker's\n", core); + return 1; + } + if (Verifier::host_pin_core(true) != -1) { + std::fprintf(stderr, "verify_host_pin: a layer split must not pin\n"); + return 1; + } + std::printf("verify_host_pin: ok (host core %d; %zu worker cores, %zu physical cores)\n", core, workers.size(), + all.size()); + return 0; +}