cpu: preserve full thread affinity on Windows and Linux

CMakeLists.txt

@@ -1043,6 +1043,15 @@ if(STRATA_ENABLE_CUDA AND STRATA_BUILD_TESTS AND EXISTS "${CMAKE_CURRENT_SOURCE_
 endif()
 
 if(STRATA_BUILD_TESTS)
+  if(WIN32)
+    add_executable(pool_affinity_test src/kernels/cpu/pool_affinity_test.cpp)
+    target_link_libraries(pool_affinity_test PRIVATE strata_kernels_cpu)
+    add_test(NAME pool_affinity_test COMMAND pool_affinity_test)
+  elseif(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+    add_executable(pool_affinity_test src/kernels/cpu/pool_affinity_linux_test.cpp)
+    target_link_libraries(pool_affinity_test PRIVATE strata_kernels_cpu)
+    add_test(NAME pool_affinity_test COMMAND pool_affinity_test)
+  endif()
   # Plan v0.3 P6: one expert applied to several tokens, bitwise vs the single-token kernel (synthetic, CPU only).
   # plan v0.3 P5: the device dequantizers vs the artifact's CPU dequantizers on the real shards (not a ctest: needs
   # the 38 GB model; run it by hand or from tools/gpu_session.py)

include/strata/core/session.hpp

@@ -16,6 +16,7 @@
 
 #include "strata/core/hit_hook.hpp"
 #include "strata/core/layer.hpp"
+#include "strata/kernels/cpu/pool.hpp"
 
 #include <cuda_runtime.h>
 
@@ -369,7 +370,7 @@ struct SessionLoopScratch {
     float* y_miss = nullptr;        ///< pinned host staging for the pool's answer, `parts_bytes` long
     size_t parts_bytes = 0;
     cudaEvent_t probe = nullptr;
-    long long pinned_core = -1;     ///< the affinity to restore, or -1 if the host was never pinned
+    strata::kernels::cpu::ThreadAffinity pinned_core{};  ///< affinity to restore, or invalid if the host was not pinned
     bool pinned = false;
 
     /// Allocates the buffers and pins the host thread.  Call ONCE, at session setup.

include/strata/kernels/cpu/pool.hpp

@@ -76,7 +76,7 @@ struct CpuTopology {
     int p_cores = 0;                ///< Physical performance cores
     int p_threads = 0;              ///< Total logical threads on performance cores
     int e_cores = 0;                ///< Efficient cores
-    std::vector<int> worker_cores;  ///< Ordered logical core IDs for workers (excluding host if skip_first)
+    std::vector<int> worker_cores;  ///< Ordered CPU IDs; on Windows, group * 64 + processor within the group
     int host_core = -1;             ///< Logical core reserved for host thread
 };
 
@@ -100,9 +100,23 @@ std::vector<int> physical_cores(bool skip_first, PoolAffinity affinity = PoolAff
 /// - exactly 5/6 of L9's 44.14 on 6 - and at **26.9 GB/s inside the host loop**, where the unpinned spinning
 /// host is free to land on a worker's core or its SMT sibling.  That 1.35x is not the kernel.
 ///
-/// Returns the PREVIOUS affinity mask, or -1 if the platform refused; pass it to `restore_thread_affinity`.
-long long pin_current_thread(int core);
-void restore_thread_affinity(long long previous);
+/// The previous host placement. `valid` is false when querying or setting placement failed.
+/// Windows uses a reversible CPU Set selection, leaving hard affinity (including Windows 11's implicit
+/// all-group eligibility) untouched. An empty selection restores inheritance from the process defaults.
+/// Linux stores the complete, dynamically sized native CPU mask, including CPU IDs above 63.
+struct ThreadAffinity {
+#if defined(_WIN32)
+    std::vector<unsigned long> cpu_sets;
+#else
+    std::vector<unsigned long> mask;
+#endif
+    bool valid = false;
+};
+
+/// Selects one encoded processor for the caller; Windows CPU Sets respect existing hard-affinity limits.
+/// Restore on the same calling thread. Pool-owned workers use hard group affinity instead.
+ThreadAffinity pin_current_thread(int core);
+void restore_thread_affinity(const ThreadAffinity& previous);
 
 class ExpertPool {
 public:

src/core/session.cpp

@@ -547,7 +547,7 @@ bool SessionLoopScratch::init(size_t parts_bytes_in, std::string& err) {
     const std::vector<int> cores = strata::kernels::cpu::physical_cores(false);
     if (!cores.empty()) {
         pinned_core = strata::kernels::cpu::pin_current_thread(cores[0]);
-        pinned = true;
+        pinned = pinned_core.valid;
     }
     return true;
 }
@@ -558,7 +558,7 @@ void SessionLoopScratch::free() {
     if (pinned) {
         strata::kernels::cpu::restore_thread_affinity(pinned_core);
         pinned = false;
-        pinned_core = -1;
+        pinned_core = {};
     }
     if (probe != nullptr) { cudaEventDestroy(probe); probe = nullptr; }
     if (y_miss != nullptr) { cudaFreeHost(y_miss); y_miss = nullptr; }

src/kernels/cpu/pool.cpp

@@ -14,9 +14,11 @@
 #if defined(_WIN32)
 #define WIN32_LEAN_AND_MEAN
 #include <windows.h>
+#include "pool_affinity_win.hpp"
 #else
 #include <pthread.h>
 #include <sched.h>
+#include "pool_affinity_linux.hpp"
 #endif
 
 namespace strata::kernels::cpu {
@@ -36,7 +38,13 @@ CpuTopology detect_cpu_topology(bool skip_first, PoolAffinity affinity) {
     DWORD len = 0;
     GetLogicalProcessorInformationEx(RelationProcessorCore, nullptr, &len);
     if (len == 0) {
-        for (unsigned i = 0; i < std::thread::hardware_concurrency(); ++i) topo.worker_cores.push_back((int) i);
+        // Preserve group identity even when detailed core topology is unavailable.
+        const WORD groups = GetActiveProcessorGroupCount();
+        for (WORD group = 0; group < groups; ++group) {
+            const DWORD count = GetActiveProcessorCount(group);
+            if (count == 0 || count == (DWORD) -1 || count > 64) continue;
+            for (DWORD i = 0; i < count; ++i) topo.worker_cores.push_back((int) group * 64 + (int) i);
+        }
         if (skip_first && !topo.worker_cores.empty()) {
             topo.host_core = topo.worker_cores.front();
             topo.worker_cores.erase(topo.worker_cores.begin());
@@ -60,10 +68,12 @@ CpuTopology detect_cpu_topology(bool skip_first, PoolAffinity affinity) {
                 CoreDesc cd;
                 cd.efficiency = e->Processor.EfficiencyClass;
                 cd.has_smt = (e->Processor.Flags & LTP_PC_SMT) != 0;
-                const GROUP_AFFINITY& g = e->Processor.GroupMask[0];
-                for (int bit = 0; bit < 64; ++bit) {
-                    if (g.Mask & (1ull << bit)) {
-                        cd.lps.push_back((int) (g.Group * 64 + bit));
+                for (WORD group = 0; group < e->Processor.GroupCount; ++group) {
+                    const GROUP_AFFINITY& g = e->Processor.GroupMask[group];
+                    for (int bit = 0; bit < 64; ++bit) {
+                        if (g.Mask & (KAFFINITY(1) << bit)) {
+                            cd.lps.push_back((int) (g.Group * 64 + bit));
+                        }
                     }
                 }
                 if (!cd.lps.empty()) {
@@ -162,12 +172,8 @@ CpuTopology detect_cpu_topology(bool skip_first, PoolAffinity affinity) {
     };
 
     std::vector<int> allowed;
-    cpu_set_t set;
-    CPU_ZERO(&set);
-    if (sched_getaffinity(0, sizeof set, &set) == 0) {
-        for (int i = 0; i < CPU_SETSIZE; ++i)
-            if (CPU_ISSET(i, &set)) allowed.push_back(i);
-    } else {
+    std::vector<unsigned long> allowed_mask;
+    if (detail::get_thread_affinity(allowed_mask, &allowed) != 0) {
         for (unsigned i = 0; i < std::thread::hardware_concurrency(); ++i) allowed.push_back((int) i);
     }
 
@@ -264,49 +270,60 @@ std::vector<int> physical_cores(bool skip_first, PoolAffinity affinity) {
 
 namespace {
 
-void pin_this_thread(int core) {
-    if (core < 0) return;
+bool pin_this_thread(int core, [[maybe_unused]] int worker = -1) {
+    if (core < 0) return false;
 #if defined(_WIN32)
-    SetThreadAffinityMask(GetCurrentThread(), (DWORD_PTR) 1 << (core & 63));
+    return detail::set_thread_group_affinity(core, worker);
 #else
-    cpu_set_t set;
-    CPU_ZERO(&set);
-    CPU_SET(core, &set);
-    pthread_setaffinity_np(pthread_self(), sizeof set, &set);
+    const int error = detail::pin_thread_to_cpu(core);
+    if (error != 0)
+        std::fprintf(stderr, "strata cpu pool: affinity for worker %d (CPU %d) failed: %d; previous affinity kept\n",
+                     worker, core, error);
+    return error == 0;
 #endif
 }
 
 }  // namespace
 
-long long pin_current_thread(int core) {
-    if (core < 0) return -1;
+ThreadAffinity pin_current_thread(int core) {
+    if (core < 0) return {};
 #if defined(_WIN32)
-    // `SetThreadAffinityMask` RETURNS the previous mask, or 0 on failure - so 0 doubles as the error, which is
-    // why the caller must not treat it as a restorable value.
-    const DWORD_PTR prev = SetThreadAffinityMask(GetCurrentThread(), (DWORD_PTR) 1 << (core & 63));
-    return prev == 0 ? -1 : (long long) prev;
+    ThreadAffinity previous;
+    ULONG target = 0;
+    if (!detail::get_thread_cpu_sets(previous.cpu_sets) || !detail::cpu_set_for_core(core, target) ||
+        !SetThreadSelectedCpuSets(GetCurrentThread(), &target, 1)) {
+        std::fprintf(stderr, "strata cpu pool: host CPU Set selection for processor %d failed: %lu; previous placement kept\n",
+                     core, (unsigned long) GetLastError());
+        return {};
+    }
+    previous.valid = true;
+    return previous;
 #else
-    cpu_set_t prev;
-    CPU_ZERO(&prev);
-    if (pthread_getaffinity_np(pthread_self(), sizeof prev, &prev) != 0) return -1;
-    unsigned long mask = 0;
-    for (int i = 0; i < CPU_SETSIZE && i < 64; ++i)
-        if (CPU_ISSET(i, &prev)) mask |= 1ul << i;
-    pin_this_thread(core);
-    return (long long) mask;
+    ThreadAffinity previous;
+    int error = detail::get_thread_affinity(previous.mask);
+    if (error == 0) error = detail::pin_thread_to_cpu(core);
+    if (error != 0) {
+        std::fprintf(stderr, "strata cpu pool: host affinity for CPU %d failed: %d; previous affinity kept\n", core, error);
+        return {};
+    }
+    previous.valid = true;
+    return previous;
 #endif
 }
 
-void restore_thread_affinity(long long previous) {
-    if (previous <= 0) return;
+void restore_thread_affinity(const ThreadAffinity& previous) {
+    if (!previous.valid) return;
 #if defined(_WIN32)
-    SetThreadAffinityMask(GetCurrentThread(), (DWORD_PTR) previous);
+    // Clearing an originally empty selection restores process-default/all-group eligibility without
+    // turning the caller's implicit Windows 11 affinity into an explicit single-group hard mask.
+    if (!SetThreadSelectedCpuSets(GetCurrentThread(), previous.cpu_sets.empty() ? nullptr : previous.cpu_sets.data(),
+                                 (ULONG) previous.cpu_sets.size()))
+        std::fprintf(stderr, "strata cpu pool: host CPU Set restoration failed: %lu\n",
+                     (unsigned long) GetLastError());
 #else
-    cpu_set_t set;
-    CPU_ZERO(&set);
-    for (int i = 0; i < 64; ++i)
-        if ((previous >> i) & 1) CPU_SET(i, &set);
-    pthread_setaffinity_np(pthread_self(), sizeof set, &set);
+    const int error = detail::set_thread_affinity(previous.mask);
+    if (error != 0)
+        std::fprintf(stderr, "strata cpu pool: host affinity restoration failed: %d\n", error);
 #endif
 }
 
@@ -366,7 +383,7 @@ ExpertPool::ExpertPool(int n_workers, bool pin, bool host_works, PoolAffinity af
     for (int i = 0; i < n_; ++i) {
         const int core = pin ? (i < (int) topo_.worker_cores.size() ? topo_.worker_cores[(size_t) i] : -1) : -1;
         threads_.emplace_back([this, i, core] {
-            pin_this_thread(core);
+            pin_this_thread(core, i);
             worker(i);
         });
     }

src/kernels/cpu/pool_affinity_linux.hpp

@@ -0,0 +1,69 @@
+#pragma once
+
+#if defined(__linux__)
+#include <pthread.h>
+#include <sched.h>
+
+#include <cerrno>
+#include <climits>
+#include <cstring>
+#include <vector>
+
+namespace strata::kernels::cpu::detail {
+
+// CPU_ALLOC and the sized macros also handle kernels whose CPU mask exceeds CPU_SETSIZE.
+// Keep the native mask opaque: copy its bytes for storage and use CPU_*_S to inspect it.
+class LinuxCpuSet {
+public:
+    explicit LinuxCpuSet(int cpus) : bytes(CPU_ALLOC_SIZE(cpus)), set(CPU_ALLOC(cpus)) {
+        if (set) CPU_ZERO_S(bytes, set);
+    }
+    ~LinuxCpuSet() { CPU_FREE(set); }
+    LinuxCpuSet(const LinuxCpuSet&) = delete;
+    LinuxCpuSet& operator=(const LinuxCpuSet&) = delete;
+
+    const size_t bytes;
+    cpu_set_t* const set;
+};
+
+// pthread affinity functions return an error number directly; they do not set errno.
+// A too-small query buffer returns EINVAL, even when the number of online CPUs is small.
+inline int get_thread_affinity(std::vector<unsigned long>& mask, std::vector<int>* cpus = nullptr) {
+    for (int count = CPU_SETSIZE;;) {
+        LinuxCpuSet current(count);
+        if (!current.set) return ENOMEM;
+        const int error = pthread_getaffinity_np(pthread_self(), current.bytes, current.set);
+        if (error == 0) {
+            mask.resize(current.bytes / sizeof(unsigned long));
+            std::memcpy(mask.data(), current.set, current.bytes);
+            if (cpus) {
+                cpus->clear();
+                for (int cpu = 0; cpu < count; ++cpu)
+                    if (CPU_ISSET_S(cpu, current.bytes, current.set)) cpus->push_back(cpu);
+            }
+            return 0;
+        }
+        if (error != EINVAL || count > INT_MAX / 2) return error;
+        count *= 2;
+    }
+}
+
+inline int set_thread_affinity(const std::vector<unsigned long>& mask) {
+    constexpr size_t word_bits = sizeof(unsigned long) * CHAR_BIT;
+    if (mask.empty() || mask.size() > INT_MAX / word_bits) return EINVAL;
+    LinuxCpuSet target(static_cast<int>(mask.size() * word_bits));
+    if (!target.set) return ENOMEM;
+    std::memcpy(target.set, mask.data(), target.bytes);
+    return pthread_setaffinity_np(pthread_self(), target.bytes, target.set);
+}
+
+inline int pin_thread_to_cpu(int core) {
+    if (core < 0 || core == INT_MAX) return EINVAL;
+    LinuxCpuSet target(core < CPU_SETSIZE ? CPU_SETSIZE : core + 1);
+    if (!target.set) return ENOMEM;
+    CPU_SET_S(core, target.bytes, target.set);
+    return pthread_setaffinity_np(pthread_self(), target.bytes, target.set);
+}
+
+}  // namespace strata::kernels::cpu::detail
+#endif