390 lines
18 KiB
C++
390 lines
18 KiB
C++
// [4] 段散写持久侧区缓存（zero-copy dual-source 默认路径）。
|
||
// 与持久 staging 缓存同构，但目标是“多个结构化 per-key 池数组”：每个池数组形如
|
||
// (cap+spec, ...)，行 [base_row, base_row+spec) 为侧区。命中缺口时 pread blob 整行，
|
||
// 再把该行内按固定顺序拼接的各段 memcpy 进对应 per-key 数组的同一物理行。
|
||
#include "side_region.h"
|
||
#include "../io/blob_io.h"
|
||
#include "../io/bg_reader.h"
|
||
|
||
#include <condition_variable>
|
||
#include <cstdio>
|
||
#include <cstdlib>
|
||
#include <functional>
|
||
#include <thread>
|
||
#include <unordered_set>
|
||
|
||
struct SideLayer {
|
||
std::map<int, int> e2r; // expert -> 物理侧区行 [base_row, base_row+spec)
|
||
std::vector<int> free_rows;
|
||
std::map<int, uint32_t> freq; // expert -> 预测频次(LFU 分数;仅 SIDEREGION_LFU 用)
|
||
bool inited = false;
|
||
int base = 0; // 侧区起始物理行 base_row
|
||
int spec = 0; // 侧区行数 spec_slots
|
||
};
|
||
static std::mutex g_side_mutex;
|
||
static std::map<std::pair<int, int>, SideLayer> g_side; // 键 (layer, gen)：双缓冲两代独立
|
||
|
||
// 侧区 fill 在途计数：eval_gpu 提交预取时 +1，后台 read_publish 写完字节后 -1。
|
||
// 消费方在前向开头 sideregion_drain() 排空上一前向的 fill，保证被消费的侧区行字节
|
||
// 已完全写好（含 GPU 完成回调滞后的情形）→ 消灭「GPU 消费 kernel 读到半写侧区行」竞态。
|
||
static std::atomic<long> g_side_inflight{0};
|
||
static std::mutex g_side_drain_mutex;
|
||
static std::condition_variable g_side_drain_cv;
|
||
static inline void side_inflight_done() {
|
||
{
|
||
std::lock_guard<std::mutex> lk(g_side_drain_mutex);
|
||
g_side_inflight.fetch_sub(1);
|
||
}
|
||
g_side_drain_cv.notify_all();
|
||
}
|
||
void sideregion_drain() {
|
||
std::unique_lock<std::mutex> lk(g_side_drain_mutex);
|
||
g_side_drain_cv.wait(lk, [] { return g_side_inflight.load() == 0; });
|
||
}
|
||
|
||
// 临时诊断：只追踪指定 (layer,row) 的所有账本/字节事件（SIDE_TRACE_LAYER/ROW）。
|
||
static inline bool side_trace_hit(int layer, int row) {
|
||
const char* le = std::getenv("SIDE_TRACE_LAYER");
|
||
const char* re = std::getenv("SIDE_TRACE_ROW");
|
||
if (!le || !re) return false;
|
||
return layer == atoi(le) && row == atoi(re);
|
||
}
|
||
|
||
// 临时诊断：跨线程全局事件序号 + 线程短 id，用于把 reserve(Metal 回调线程)/read_publish(bg 线程)/
|
||
// sideregion_kv(主线程) 的事件按发生顺序排出来（root-cause 取证，SIDE_TRACE_* 命中时才用）。
|
||
static std::atomic<uint64_t> g_side_ev{0};
|
||
static inline unsigned side_tid() {
|
||
return static_cast<unsigned>(
|
||
std::hash<std::thread::id>{}(std::this_thread::get_id()) & 0xffff);
|
||
}
|
||
|
||
class PrefetchPoolSideRegionPrimitive : public mx::Primitive {
|
||
public:
|
||
PrefetchPoolSideRegionPrimitive(mx::Stream s, std::vector<int> seg_nbytes, int layer, int gen,
|
||
std::string path, size_t stride, std::vector<int> resident,
|
||
int spec_slots, int base_row)
|
||
: Primitive(s), seg_(std::move(seg_nbytes)), layer_(layer), gen_(gen), path_(std::move(path)),
|
||
stride_(stride), resident_(std::move(resident)), spec_(spec_slots), base_(base_row) {}
|
||
const char* name() const override { return "PrefetchPoolSideRegionPrimitive"; }
|
||
|
||
void eval_cpu(const std::vector<mx::array>& in, std::vector<mx::array>& out) override {
|
||
out[0].set_data(mx::allocator::malloc(out[0].nbytes()));
|
||
run(in);
|
||
}
|
||
void eval_gpu(const std::vector<mx::array>& in, std::vector<mx::array>& out) override {
|
||
out[0].set_data(mx::allocator::malloc(out[0].nbytes()));
|
||
std::vector<uint8_t*> ptrs;
|
||
ptrs.reserve(in.size() - 1);
|
||
for (size_t i = 1; i < in.size(); ++i) {
|
||
mx::array a = in[i]; // 非 const 拷贝才能取可写指针
|
||
ptrs.push_back(a.data<uint8_t>());
|
||
}
|
||
mx::array ids = in[0];
|
||
const uint32_t* idp = ids.data<uint32_t>();
|
||
size_t n = ids.size();
|
||
std::vector<int> seg = seg_;
|
||
int layer = layer_;
|
||
int gen = gen_;
|
||
std::string path = path_;
|
||
size_t stride = stride_;
|
||
std::vector<int> resident = resident_;
|
||
int spec = spec_, base = base_;
|
||
auto& enc = mx::metal::get_command_encoder(stream());
|
||
MTL::CommandBuffer* cb = enc.get_command_buffer();
|
||
// 提交即计在途：在 eval_gpu（预取提交）时 +1，直到后台字节写完才 -1。这样消费方前向开头
|
||
// sideregion_drain() 能等到「即使 GPU 完成回调尚未触发」的 fill，闭合跨前向的写-读竞态。
|
||
g_side_inflight.fetch_add(1);
|
||
// in 按值捕获 → 保活 expert_ids 与所有池数组 buffer；idp/ptrs 在回调里指针有效。
|
||
cb->addCompletedHandler(
|
||
[in, ptrs, seg, idp, n, layer, gen, path, stride, resident, spec, base](MTL::CommandBuffer*) {
|
||
// 阶段1（回调线程、持锁极短）：读惰性 id（此刻已算完）、预留侧区行。
|
||
// 必须在回调里——id 只有 command buffer 完成后才有效。
|
||
auto to_read = reserve(idp, n, layer, gen, resident, spec, base);
|
||
if (to_read.empty()) {
|
||
side_inflight_done(); // 无缺口可读：立即消账，避免 drain 空等
|
||
return;
|
||
}
|
||
// 诊断门控 SIDEREGION_SYNC=1：回调内同步 pread+memcpy+publish（不派 bg），
|
||
// 消除「异步 bg fill 与下一前向 gather 竞态」这一变量，用于 systematic-debugging 取证。
|
||
static const bool kSync = []() {
|
||
const char* e = std::getenv("SIDEREGION_SYNC");
|
||
return e && e[0] == '1';
|
||
}();
|
||
if (kSync) {
|
||
read_publish(ptrs, seg, to_read, path, stride, layer, gen);
|
||
side_inflight_done();
|
||
return;
|
||
}
|
||
// 阶段2+3 派给自由后台线程：~40MB pread + memcpy + 发布 e2r 脱离 Metal 回调线程，
|
||
// 与主 stream 后续层计算、多层预取互相并发 → 真正恢复 I/O/计算重叠。
|
||
// 闭包再次按值捕获 in：保活池/ids buffer 到后台读完（ptrs 指向其内存）。
|
||
bg_submit_task([in, ptrs, seg, to_read, path, stride, layer, gen]() {
|
||
read_publish(ptrs, seg, to_read, path, stride, layer, gen);
|
||
side_inflight_done(); // 字节写完才消账 → drain 保证被消费行已就绪
|
||
});
|
||
});
|
||
}
|
||
|
||
private:
|
||
void run(const std::vector<mx::array>& in) {
|
||
std::vector<uint8_t*> ptrs;
|
||
ptrs.reserve(in.size() - 1);
|
||
for (size_t i = 1; i < in.size(); ++i) {
|
||
mx::array a = in[i];
|
||
a.eval(); // 与 PrefetchStagingCachedPrimitive::eval_cpu 一致：先物化输入
|
||
ptrs.push_back(a.data<uint8_t>());
|
||
}
|
||
mx::array ids = in[0];
|
||
ids.eval();
|
||
// CPU 路径同步执行（测试/无 GPU 时）：预留 + 读发布一气呵成。
|
||
auto to_read = reserve(ids.data<uint32_t>(), ids.size(), layer_, gen_, resident_, spec_, base_);
|
||
read_publish(ptrs, seg_, to_read, path_, stride_, layer_, gen_);
|
||
}
|
||
|
||
// 阶段1：过滤常驻/去重 → 淘汰 ∉P 的旧行 → 为缺口预留物理行（出 free，暂不入 e2r，
|
||
// 避免消费者在字节写好前看到 e2r 命中而 gather 到脏字节）。返回 (expert, 预留行)。
|
||
static std::vector<std::pair<int, int>> reserve(
|
||
const uint32_t* idp, size_t n, int layer, int gen, const std::vector<int>& resident,
|
||
int spec, int base) {
|
||
const char* lfu_env = std::getenv("SIDEREGION_LFU"); // 每次读,便于测试切换
|
||
// 默认开:持久 LFU 单缓冲=生产路径(cli 默认)。仅显式 SIDEREGION_LFU=0 回退旧 legacy 双缓冲。
|
||
bool lfu = !lfu_env || lfu_env[0] != '0';
|
||
std::unordered_set<int> res(resident.begin(), resident.end());
|
||
std::vector<int> P;
|
||
std::unordered_set<int> Pset, seen;
|
||
for (size_t i = 0; i < n; ++i) {
|
||
int e = static_cast<int>(idp[i]);
|
||
if (res.count(e) || !seen.insert(e).second) continue;
|
||
P.push_back(e);
|
||
Pset.insert(e);
|
||
}
|
||
std::vector<std::pair<int, int>> to_read;
|
||
std::lock_guard<std::mutex> lk(g_side_mutex);
|
||
SideLayer& c = g_side[{layer, gen}];
|
||
if (!c.inited) {
|
||
for (int r = 0; r < spec; ++r) c.free_rows.push_back(base + r);
|
||
c.base = base;
|
||
c.spec = spec;
|
||
c.inited = true;
|
||
}
|
||
if (!lfu) {
|
||
// 旧行为:∉P 全弃(一次性预取批)。
|
||
for (auto it = c.e2r.begin(); it != c.e2r.end();) {
|
||
if (!Pset.count(it->first)) {
|
||
c.free_rows.push_back(it->second);
|
||
it = c.e2r.erase(it);
|
||
} else {
|
||
++it;
|
||
}
|
||
}
|
||
for (int e : P) {
|
||
if (c.e2r.count(e) || c.free_rows.empty()) continue;
|
||
to_read.emplace_back(e, c.free_rows.back());
|
||
c.free_rows.pop_back();
|
||
}
|
||
return to_read;
|
||
}
|
||
// LFU 持久:∉P 不清;再预测命中已驻专家 freq+1(越常预测越热)。
|
||
for (int e : P) {
|
||
if (c.e2r.count(e)) c.freq[e] += 1;
|
||
}
|
||
for (int e : P) {
|
||
if (c.e2r.count(e)) continue; // 已驻,跳过(不重读)
|
||
int row;
|
||
if (!c.free_rows.empty()) {
|
||
row = c.free_rows.back();
|
||
c.free_rows.pop_back();
|
||
if (side_trace_hit(layer, row))
|
||
fprintf(stderr, "[SIDE_TRACE ev=%llu tid=%u] L%d gen%d RESERVE_FROM_FREE row=%d expert=%d\n",
|
||
(unsigned long long)g_side_ev.fetch_add(1), side_tid(), layer, gen, row, e);
|
||
} else {
|
||
// free 空:LFU 淘汰 e2r 中 freq 最小且 ∉P 者(tie-break:最小 expert id)。
|
||
int victim = -1;
|
||
uint32_t best = 0;
|
||
for (auto& kv : c.e2r) {
|
||
if (Pset.count(kv.first)) continue; // 不淘本步要用的
|
||
uint32_t f = c.freq.count(kv.first) ? c.freq[kv.first] : 0;
|
||
if (victim < 0 || f < best || (f == best && kv.first < victim)) {
|
||
victim = kv.first;
|
||
best = f;
|
||
}
|
||
}
|
||
if (victim < 0) continue; // 全是 P 热,无可淘 → 本步不读
|
||
row = c.e2r[victim];
|
||
c.e2r.erase(victim);
|
||
c.freq.erase(victim);
|
||
if (side_trace_hit(layer, row))
|
||
fprintf(stderr, "[SIDE_TRACE ev=%llu tid=%u] L%d gen%d EVICT_REUSE row=%d victim=%d newExpert=%d\n",
|
||
(unsigned long long)g_side_ev.fetch_add(1), side_tid(), layer, gen, row, victim, e);
|
||
}
|
||
to_read.emplace_back(e, row);
|
||
}
|
||
// ===== 临时诊断（SIDE_AUDIT=1）：reserve 结束时审计侧区行账本自洽性 =====
|
||
if (std::getenv("SIDE_AUDIT")) side_audit(c, layer, gen, "reserve", to_read);
|
||
return to_read;
|
||
}
|
||
|
||
// 审计：检测「一行被两专家占用」「free 与 e2r 行重叠」「to_read 行仍被 e2r 占用」。
|
||
static void side_audit(SideLayer& c, int layer, int gen, const char* where,
|
||
const std::vector<std::pair<int, int>>& to_read) {
|
||
std::map<int, int> row_owner; // row -> expert
|
||
for (auto& p : c.e2r) {
|
||
auto it = row_owner.find(p.second);
|
||
if (it != row_owner.end())
|
||
fprintf(stderr, "[SIDE_AUDIT] %s L%d gen%d DOUBLE_OWN row=%d experts=%d,%d\n",
|
||
where, layer, gen, p.second, it->second, p.first);
|
||
row_owner[p.second] = p.first;
|
||
}
|
||
std::unordered_set<int> free_set(c.free_rows.begin(), c.free_rows.end());
|
||
for (auto& p : c.e2r)
|
||
if (free_set.count(p.second))
|
||
fprintf(stderr, "[SIDE_AUDIT] %s L%d gen%d FREE_E2R_OVERLAP row=%d owned_by=%d\n",
|
||
where, layer, gen, p.second, p.first);
|
||
if (c.free_rows.size() != free_set.size())
|
||
fprintf(stderr, "[SIDE_AUDIT] %s L%d gen%d FREE_DUP free_n=%zu uniq=%zu\n",
|
||
where, layer, gen, c.free_rows.size(), free_set.size());
|
||
for (auto& pr : to_read)
|
||
if (row_owner.count(pr.second))
|
||
fprintf(stderr, "[SIDE_AUDIT] %s L%d gen%d TOREAD_LIVE_ROW row=%d assigned_to=%d still_owned_by=%d\n",
|
||
where, layer, gen, pr.second, pr.first, row_owner[pr.second]);
|
||
}
|
||
|
||
// 阶段2+3：pread blob 整行 + 各段 memcpy 直写进对应 per-key 池数组的物理侧区行（不持锁），
|
||
// 写完后持锁发布 e2r。ptrs[i] 为第 i 个池 key 数组的 buffer 指针（C++ 拥有、地址恒定不迁移，
|
||
// 由 Route 3 底座保证——见 pool_owned_zeros），故后台异步直写安全、消费侧读同一 buffer。
|
||
// 在后台线程跑：与计算并发，且不阻塞消费侧的 sideregion_kv。
|
||
static void read_publish(const std::vector<uint8_t*>& ptrs, const std::vector<int>& seg,
|
||
const std::vector<std::pair<int, int>>& to_read,
|
||
const std::string& path, size_t stride, int layer, int gen) {
|
||
// 段在 blob 记录内的偏移（段顺序 = seg 顺序 = ptrs/池 key 顺序）。
|
||
std::vector<size_t> seg_off(seg.size(), 0);
|
||
size_t acc = 0;
|
||
for (size_t i = 0; i < seg.size(); ++i) { seg_off[i] = acc; acc += static_cast<size_t>(seg[i]); }
|
||
int fd = open_blob_nocache(path.c_str());
|
||
if (fd < 0) { // 失败则把预留行还回 free
|
||
std::lock_guard<std::mutex> lk(g_side_mutex);
|
||
SideLayer& c = g_side[{layer, gen}];
|
||
for (auto& pr : to_read) c.free_rows.push_back(pr.second);
|
||
return;
|
||
}
|
||
std::vector<uint8_t> rec(stride); // 整条 blob 记录临时缓冲
|
||
std::vector<std::pair<int, int>> done;
|
||
for (auto& pr : to_read) {
|
||
int e = pr.first, row = pr.second;
|
||
if (::pread(fd, rec.data(), stride, static_cast<off_t>(static_cast<size_t>(e) * stride)) !=
|
||
static_cast<ssize_t>(stride)) {
|
||
std::lock_guard<std::mutex> lk(g_side_mutex); // 读失败：行还回 free
|
||
g_side[{layer, gen}].free_rows.push_back(row);
|
||
continue;
|
||
}
|
||
// 各段 memcpy 进对应 per-key 池数组的物理行 row（(n_slots,*shape) 布局 → 行偏移 = row*seg[i]）。
|
||
for (size_t i = 0; i < seg.size(); ++i)
|
||
std::memcpy(ptrs[i] + static_cast<size_t>(row) * static_cast<size_t>(seg[i]),
|
||
rec.data() + seg_off[i], static_cast<size_t>(seg[i]));
|
||
if (side_trace_hit(layer, row))
|
||
fprintf(stderr, "[SIDE_TRACE ev=%llu tid=%u] L%d gen%d WRITEPOOL row=%d expert=%d\n",
|
||
(unsigned long long)g_side_ev.fetch_add(1), side_tid(), layer, gen, row, e);
|
||
done.emplace_back(e, row);
|
||
}
|
||
::close(fd);
|
||
{
|
||
std::lock_guard<std::mutex> lk(g_side_mutex);
|
||
SideLayer& c = g_side[{layer, gen}];
|
||
for (auto& pr : done) { // 字节就绪后才发布 e2r
|
||
c.e2r[pr.first] = pr.second;
|
||
if (!c.freq.count(pr.first)) c.freq[pr.first] = 1; // 新专家初始 freq
|
||
if (side_trace_hit(layer, pr.second))
|
||
fprintf(stderr, "[SIDE_TRACE ev=%llu tid=%u] L%d gen%d PUBLISH row=%d expert=%d\n",
|
||
(unsigned long long)g_side_ev.fetch_add(1), side_tid(), layer, gen, pr.second,
|
||
pr.first);
|
||
}
|
||
if (std::getenv("SIDE_AUDIT")) side_audit(c, layer, gen, "publish", {});
|
||
}
|
||
}
|
||
std::vector<int> seg_;
|
||
int layer_;
|
||
int gen_;
|
||
std::string path_;
|
||
size_t stride_;
|
||
std::vector<int> resident_;
|
||
int spec_;
|
||
int base_;
|
||
};
|
||
|
||
mx::array prefetch_pool_sideregion(
|
||
const std::vector<mx::array>& pool_list, const std::vector<int>& seg_nbytes,
|
||
const mx::array& expert_ids, int layer, const std::string& path, int stride,
|
||
const std::vector<int>& resident, int spec_slots, int base_row, int gen,
|
||
mx::StreamOrDevice s = {}) {
|
||
// 防御：段数必须与池数组个数一一对应（每段映射唯一一个 per-key 数组）。
|
||
if (seg_nbytes.size() != pool_list.size()) {
|
||
throw std::invalid_argument(
|
||
"prefetch_pool_sideregion: seg_nbytes.size()=" + std::to_string(seg_nbytes.size()) +
|
||
" != pool_list.size()=" + std::to_string(pool_list.size()));
|
||
}
|
||
// 防御：各段字节之和必须等于 stride（blob 记录恰好是各段的拼接）。
|
||
size_t seg_sum = 0;
|
||
for (int b : seg_nbytes) seg_sum += static_cast<size_t>(b);
|
||
if (seg_sum != static_cast<size_t>(stride)) {
|
||
throw std::invalid_argument(
|
||
"prefetch_pool_sideregion: sum(seg_nbytes)=" + std::to_string(seg_sum) +
|
||
" != stride=" + std::to_string(stride));
|
||
}
|
||
std::vector<mx::array> inputs;
|
||
inputs.push_back(expert_ids);
|
||
for (auto& a : pool_list) inputs.push_back(a);
|
||
return mx::array(
|
||
mx::Shape{1}, mx::uint8,
|
||
std::make_shared<PrefetchPoolSideRegionPrimitive>(
|
||
mx::to_stream(s), seg_nbytes, layer, gen, path, static_cast<size_t>(stride), resident,
|
||
spec_slots, base_row),
|
||
inputs);
|
||
}
|
||
|
||
std::vector<int> sideregion_contents(int layer, int gen) {
|
||
std::lock_guard<std::mutex> lk(g_side_mutex);
|
||
std::vector<int> out;
|
||
auto it = g_side.find({layer, gen});
|
||
if (it != g_side.end())
|
||
for (auto& p : it->second.e2r) { out.push_back(p.first); out.push_back(p.second); }
|
||
return out;
|
||
}
|
||
|
||
// 侧区 e2r → 两个 device mx.array (keys uint32, vals int32),直接在 C++ 从 map 建连续 buffer。
|
||
// 消掉 Python 侧 dict 构建 + list(...)→mx.array 的每层 host 胶水。
|
||
std::pair<mx::array, mx::array> sideregion_kv(int layer, int gen) {
|
||
std::vector<uint32_t> keys;
|
||
std::vector<int32_t> vals;
|
||
{
|
||
std::lock_guard<std::mutex> lk(g_side_mutex);
|
||
auto it = g_side.find({layer, gen});
|
||
if (it != g_side.end()) {
|
||
keys.reserve(it->second.e2r.size());
|
||
vals.reserve(it->second.e2r.size());
|
||
for (auto& p : it->second.e2r) {
|
||
keys.push_back(static_cast<uint32_t>(p.first));
|
||
vals.push_back(static_cast<int32_t>(p.second));
|
||
if (side_trace_hit(layer, p.second))
|
||
fprintf(stderr, "[SIDE_TRACE ev=%llu tid=%u] L%d gen%d CONSUME_KV row=%d expert=%d\n",
|
||
(unsigned long long)g_side_ev.fetch_add(1), side_tid(), layer, gen, p.second, p.first);
|
||
}
|
||
}
|
||
}
|
||
int n = static_cast<int>(keys.size());
|
||
return {mx::array(keys.data(), mx::Shape{n}, mx::uint32),
|
||
mx::array(vals.data(), mx::Shape{n}, mx::int32)};
|
||
}
|
||
|
||
void sideregion_reset() {
|
||
std::lock_guard<std::mutex> lk(g_side_mutex);
|
||
g_side.clear();
|
||
}
|
||
|
||
std::unordered_map<int, int> sideregion_snapshot(int layer, int gen) {
|
||
std::unordered_map<int, int> side;
|
||
std::lock_guard<std::mutex> lk(g_side_mutex);
|
||
auto it = g_side.find({layer, gen});
|
||
if (it != g_side.end()) for (auto& p : it->second.e2r) side[p.first] = p.second;
|
||
return side;
|
||
}
|