EPQ expert-paging port for Laguna-S-2.1 Q4_K_M (source + results)

37c1f39 verified 3 days ago

8.82 kB

// SDQ expert buffer type
//
// 這是一個「假」buffer type：它宣稱接受所有張量，但實際上
//   * 不配置任何 RAM（alloc_buffer 回傳 0 大小的 buffer）
//   * set_tensor 忽略傳入的資料（資料來自 llama.cpp 的 mmap，我們不碰）
//   * 真正的資料由 sdq::pager 從檔案 pread 到自己的 arena
//
// 用途：讓 llama.cpp 把 18.2 GiB 的 MoE 專家權重「掛在模型上」卻不放進 RAM，
// 改由分頁器管理 SSD → RAM → CPU 的三層資料流。

#include "sdq_pager.h"

#include "ggml-backend-impl.h"
#include "llama-mmap.h"

#include <cstdio>
#include <string>
#include <cstdlib>
#include <cstring>
#include <algorithm>
#include <mutex>
#include <unordered_map>
#include <vector>

namespace sdq {

// 給 tensor->data 的假位址。每個 expert 權重張量分到一小塊（預設 4 KB），
// 指標必須有效（ggml 會做指標運算），但內容永遠不會被讀。
//
// 這塊記憶體必須自己對齊到 get_alignment()：ggml_tallocr_new() 的起始 offset 是
// aligned_offset(base, 0, align)，base 若沒有對齊就會被往上推，最多吃掉
// (align - align(base)) 位元組 —— 117 個張量時足以讓最後一個剛好配置失敗
// （症狀：not enough space in the buffer ... needed 4096, available 880）。
struct expert_buffer {
    uint8_t * scratch = nullptr;
    size_t    size    = 0;
};

static bool g_enabled = false;

void    sdq_set_enabled(bool on) { g_enabled = on; }
bool    sdq_enabled()           { return g_enabled; }

static config g_cfg;
void        sdq_set_config(const config & c) { g_cfg = c; }
const config & sdq_get_config() { return g_cfg; }

// ---------------------------------------------------------------- buffer iface

// ggml_backend_buffer_i 的欄位依賴下面的實作，先給宣告
struct ggml_backend_buffer_i sdq_buffer_iface();

static const char * sdq_buft_get_name(ggml_backend_buffer_type_t buft) {
    (void) buft;
    return "EPQ_ExpertSSD";
}

static size_t sdq_buft_get_alignment(ggml_backend_buffer_type_t buft) {
    (void) buft;
    return 4096;
}

static size_t sdq_buft_get_max_size(ggml_backend_buffer_type_t buft) {
    (void) buft;
    return SIZE_MAX;
}

// 每個 expert 權重張量只回報 4 KB：18.2 GiB 的權重「不會」被載入 RAM 的關鍵。
static size_t sdq_buft_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) {
    (void) buft;
    (void) tensor;
    return 4096;
}

static bool sdq_buft_is_host(ggml_backend_buffer_type_t buft) {
    (void) buft;
    return false;
}

static ggml_backend_buffer_t sdq_buft_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) {
    // llama.cpp 會用 get_alloc_size() 計算總大小後呼叫本函式；
    // 我們只保留真正需要的假位址空間（每個張量 4 KB）。
    // 多配一個 slot：base 對齊後 aligned_offset 仍可能吃掉一點空間。
    const size_t align = sdq_buft_get_alignment(buft);
    const size_t bytes = size < 2 * align ? 2 * align : size + align;
    auto * buf = new expert_buffer();
    if (posix_memalign((void **) &buf->scratch, align, bytes) != 0 || !buf->scratch) {
        delete buf;
        return nullptr;
    }
    buf->size = bytes;
    memset(buf->scratch, 0, bytes);
    auto * b = new ggml_backend_buffer{
        /* .iface   = */ sdq_buffer_iface(),
        /* .buft    = */ buft,
        /* .context = */ buf,
        /* .size    = */ bytes,
        /* .usage   = */ GGML_BACKEND_BUFFER_USAGE_WEIGHTS,
    };
    return b;
}

static void sdq_buft_free_buffer(ggml_backend_buffer_t buffer) {
    // 只釋放自己的資料；ggml_backend_buffer_free() 會負責 delete buffer 本體，
    // 這裡再 delete 一次會造成 double free。
    auto * ctx = (expert_buffer *) buffer->context;
    if (ctx) {
        free(ctx->scratch);
        ctx->scratch = nullptr;
        delete ctx;
    }
    buffer->context = nullptr;
}

static void * sdq_buft_get_base(ggml_backend_buffer_t buffer) {
    return ((expert_buffer *) buffer->context)->scratch;
}

static ggml_status sdq_buft_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) {
    (void) buffer;
    (void) tensor;
    return GGML_STATUS_SUCCESS;
}

static void sdq_buft_memset_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor,
                                   uint8_t value, size_t offset, size_t size) {
    (void) buffer; (void) tensor; (void) value; (void) offset; (void) size;
}

static void sdq_buft_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor,
                                const void * data, size_t offset, size_t size) {
    // 關鍵：忽略。llama.cpp 會把 mmap 的資料交給我們，這裡一個位元組都不落地，
    // 專家權利永遠不會佔用 RAM。
    (void) buffer; (void) tensor; (void) data; (void) offset; (void) size;
}

static void sdq_buft_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor,
                                void * data, size_t offset, size_t size) {
    (void) buffer; (void) tensor; (void) offset;
    memset(data, 0, size);
}

static bool sdq_buft_cpy_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * src, ggml_tensor * dst) {
    (void) buffer; (void) src; (void) dst;
    return true;
}

static void sdq_buft_clear(ggml_backend_buffer_t buffer, uint8_t value) {
    (void) buffer; (void) value;
}

struct ggml_backend_buffer_i sdq_buffer_iface() {
    static ggml_backend_buffer_i i = [] {
        ggml_backend_buffer_i x {};
        x.free_buffer   = sdq_buft_free_buffer;
        x.get_base      = sdq_buft_get_base;
        x.init_tensor   = sdq_buft_init_tensor;
        x.memset_tensor = sdq_buft_memset_tensor;
        x.set_tensor    = sdq_buft_set_tensor;
        x.get_tensor    = sdq_buft_get_tensor;
        x.cpy_tensor    = sdq_buft_cpy_tensor;
        x.clear         = sdq_buft_clear;
        return x;
    }();
    return i;
}

ggml_backend_buffer_type_t sdq_expert_buft() {
    static ggml_backend_buffer_type_t buft = [] {
        return new ggml_backend_buffer_type{
            /* .iface   = */ {
                /* .get_name         = */ sdq_buft_get_name,
                /* .alloc_buffer     = */ sdq_buft_alloc_buffer,
                /* .get_alignment    = */ sdq_buft_get_alignment,
                /* .get_max_size     = */ sdq_buft_get_max_size,
                /* .get_alloc_size   = */ sdq_buft_get_alloc_size,
                /* .is_host          = */ sdq_buft_is_host,
            },
            /* .device  = */ nullptr,
            /* .context = */ nullptr,
        };
    }();
    return buft;
}

// llama.cpp 端需要的介面宣告
void sdq_install_buft_overrides(llama_model_params * params);

// EPQ_BUFT_KEEP 保留開關：std::regex 的負向先行斷言在 llama.cpp 的比對方式下
// 不可靠（會讓整個 override 失效、權重被載進 RAM），所以先不開放這個 knob。
static std::string sdq_expert_pattern() {
    return "blk\\.[0-9]+\\.ffn_(gate|up|down|gate_up)_exps\\.weight";
}

// 這兩個 static 只能用「初始化一次」的方式寫：libstdc++ 的 std::string move-assign
// 會把舊的 buffer 交給右邊的 temporary，temporary 析構時就 free 掉。若改成每次呼叫
// pattern = sdq_expert_pattern()，先前存進 override 陣列的 .pattern 會變成懸空指標；
// 讀到空字串時 std::regex("") 會匹配「任何」張量名，於是整個模型（含 token_embd）
// 都被塞進假 buffer，server 直接在 sched_reserve 裡 ggml_abort。
static const llama_model_tensor_buft_override * sdq_overrides() {
    static const std::string pattern = sdq_expert_pattern();
    static const llama_model_tensor_buft_override ov[2] = {
        { pattern.c_str(), sdq_expert_buft() },
        { nullptr, nullptr },
    };
    return ov;
}

void sdq_install_buft_overrides(llama_model_params * params) {
    if (!params) {
        return;
    }
    if (params->tensor_buft_overrides == nullptr) {
        params->tensor_buft_overrides = sdq_overrides();
    } else {
        // 使用者已有 override：把手動指定的順序原樣保留，附加在尾端
        size_t n = 0;
        while (params->tensor_buft_overrides[n].pattern != nullptr) {
            n++;
        }
        static std::vector<llama_model_tensor_buft_override> merged;
        merged.clear();
        for (size_t i = 0; i < n; i++) {
            merged.push_back(params->tensor_buft_overrides[i]);
        }
        const llama_model_tensor_buft_override * sdq_ov = sdq_overrides();
        for (size_t i = 0; sdq_ov[i].pattern != nullptr; i++) {
            merged.push_back(sdq_ov[i]);
        }
        merged.push_back({ nullptr, nullptr });
        params->tensor_buft_overrides = merged.data();
    }
}

} // namespace sdq

	// SDQ expert buffer type
	//
	// 這是一個「假」buffer type：它宣稱接受所有張量，但實際上
	// * 不配置任何 RAM（alloc_buffer 回傳 0 大小的 buffer）
	// * set_tensor 忽略傳入的資料（資料來自 llama.cpp 的 mmap，我們不碰）
	// * 真正的資料由 sdq::pager 從檔案 pread 到自己的 arena
	//
	// 用途：讓 llama.cpp 把 18.2 GiB 的 MoE 專家權重「掛在模型上」卻不放進 RAM，
	// 改由分頁器管理 SSD → RAM → CPU 的三層資料流。

	#include "sdq_pager.h"

	#include "ggml-backend-impl.h"
	#include "llama-mmap.h"

	#include <cstdio>
	#include <string>
	#include <cstdlib>
	#include <cstring>
	#include <algorithm>
	#include <mutex>
	#include <unordered_map>
	#include <vector>

	namespace sdq {

	// 給 tensor->data 的假位址。每個 expert 權重張量分到一小塊（預設 4 KB），
	// 指標必須有效（ggml 會做指標運算），但內容永遠不會被讀。
	//
	// 這塊記憶體必須自己對齊到 get_alignment()：ggml_tallocr_new() 的起始 offset 是
	// aligned_offset(base, 0, align)，base 若沒有對齊就會被往上推，最多吃掉
	// (align - align(base)) 位元組 —— 117 個張量時足以讓最後一個剛好配置失敗
	// （症狀：not enough space in the buffer ... needed 4096, available 880）。
	struct expert_buffer {
	uint8_t * scratch = nullptr;
	size_t size = 0;
	};

	static bool g_enabled = false;

	void sdq_set_enabled(bool on) { g_enabled = on; }
	bool sdq_enabled() { return g_enabled; }

	static config g_cfg;
	void sdq_set_config(const config & c) { g_cfg = c; }
	const config & sdq_get_config() { return g_cfg; }

	// ---------------------------------------------------------------- buffer iface

	// ggml_backend_buffer_i 的欄位依賴下面的實作，先給宣告
	struct ggml_backend_buffer_i sdq_buffer_iface();

	static const char * sdq_buft_get_name(ggml_backend_buffer_type_t buft) {
	(void) buft;
	return "EPQ_ExpertSSD";
	}

	static size_t sdq_buft_get_alignment(ggml_backend_buffer_type_t buft) {
	(void) buft;
	return 4096;
	}

	static size_t sdq_buft_get_max_size(ggml_backend_buffer_type_t buft) {
	(void) buft;
	return SIZE_MAX;
	}

	// 每個 expert 權重張量只回報 4 KB：18.2 GiB 的權重「不會」被載入 RAM 的關鍵。
	static size_t sdq_buft_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) {
	(void) buft;
	(void) tensor;
	return 4096;
	}

	static bool sdq_buft_is_host(ggml_backend_buffer_type_t buft) {
	(void) buft;
	return false;
	}

	static ggml_backend_buffer_t sdq_buft_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) {
	// llama.cpp 會用 get_alloc_size() 計算總大小後呼叫本函式；
	// 我們只保留真正需要的假位址空間（每個張量 4 KB）。
	// 多配一個 slot：base 對齊後 aligned_offset 仍可能吃掉一點空間。
	const size_t align = sdq_buft_get_alignment(buft);
	const size_t bytes = size < 2 * align ? 2 * align : size + align;
	auto * buf = new expert_buffer();
	if (posix_memalign((void **) &buf->scratch, align, bytes) != 0 \|\| !buf->scratch) {
	delete buf;
	return nullptr;
	}
	buf->size = bytes;
	memset(buf->scratch, 0, bytes);
	auto * b = new ggml_backend_buffer{
	/* .iface = */ sdq_buffer_iface(),
	/* .buft = */ buft,
	/* .context = */ buf,
	/* .size = */ bytes,
	/* .usage = */ GGML_BACKEND_BUFFER_USAGE_WEIGHTS,
	};
	return b;
	}

	static void sdq_buft_free_buffer(ggml_backend_buffer_t buffer) {
	// 只釋放自己的資料；ggml_backend_buffer_free() 會負責 delete buffer 本體，
	// 這裡再 delete 一次會造成 double free。
	auto * ctx = (expert_buffer *) buffer->context;
	if (ctx) {
	free(ctx->scratch);
	ctx->scratch = nullptr;
	delete ctx;
	}
	buffer->context = nullptr;
	}

	static void * sdq_buft_get_base(ggml_backend_buffer_t buffer) {
	return ((expert_buffer *) buffer->context)->scratch;
	}

	static ggml_status sdq_buft_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) {
	(void) buffer;
	(void) tensor;
	return GGML_STATUS_SUCCESS;
	}

	static void sdq_buft_memset_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor,
	uint8_t value, size_t offset, size_t size) {
	(void) buffer; (void) tensor; (void) value; (void) offset; (void) size;
	}

	static void sdq_buft_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor,
	const void * data, size_t offset, size_t size) {
	// 關鍵：忽略。llama.cpp 會把 mmap 的資料交給我們，這裡一個位元組都不落地，
	// 專家權利永遠不會佔用 RAM。
	(void) buffer; (void) tensor; (void) data; (void) offset; (void) size;
	}

	static void sdq_buft_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor,
	void * data, size_t offset, size_t size) {
	(void) buffer; (void) tensor; (void) offset;
	memset(data, 0, size);
	}

	static bool sdq_buft_cpy_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * src, ggml_tensor * dst) {
	(void) buffer; (void) src; (void) dst;
	return true;
	}

	static void sdq_buft_clear(ggml_backend_buffer_t buffer, uint8_t value) {
	(void) buffer; (void) value;
	}

	struct ggml_backend_buffer_i sdq_buffer_iface() {
	static ggml_backend_buffer_i i = [] {
	ggml_backend_buffer_i x {};
	x.free_buffer = sdq_buft_free_buffer;
	x.get_base = sdq_buft_get_base;
	x.init_tensor = sdq_buft_init_tensor;
	x.memset_tensor = sdq_buft_memset_tensor;
	x.set_tensor = sdq_buft_set_tensor;
	x.get_tensor = sdq_buft_get_tensor;
	x.cpy_tensor = sdq_buft_cpy_tensor;
	x.clear = sdq_buft_clear;
	return x;
	}();
	return i;
	}

	ggml_backend_buffer_type_t sdq_expert_buft() {
	static ggml_backend_buffer_type_t buft = [] {
	return new ggml_backend_buffer_type{
	/* .iface = */ {
	/* .get_name = */ sdq_buft_get_name,
	/* .alloc_buffer = */ sdq_buft_alloc_buffer,
	/* .get_alignment = */ sdq_buft_get_alignment,
	/* .get_max_size = */ sdq_buft_get_max_size,
	/* .get_alloc_size = */ sdq_buft_get_alloc_size,
	/* .is_host = */ sdq_buft_is_host,
	},
	/* .device = */ nullptr,
	/* .context = */ nullptr,
	};
	}();
	return buft;
	}

	// llama.cpp 端需要的介面宣告
	void sdq_install_buft_overrides(llama_model_params * params);

	// EPQ_BUFT_KEEP 保留開關：std::regex 的負向先行斷言在 llama.cpp 的比對方式下
	// 不可靠（會讓整個 override 失效、權重被載進 RAM），所以先不開放這個 knob。
	static std::string sdq_expert_pattern() {
	return "blk\\.[0-9]+\\.ffn_(gate\|up\|down\|gate_up)_exps\\.weight";
	}

	// 這兩個 static 只能用「初始化一次」的方式寫：libstdc++ 的 std::string move-assign
	// 會把舊的 buffer 交給右邊的 temporary，temporary 析構時就 free 掉。若改成每次呼叫
	// pattern = sdq_expert_pattern()，先前存進 override 陣列的 .pattern 會變成懸空指標；
	// 讀到空字串時 std::regex("") 會匹配「任何」張量名，於是整個模型（含 token_embd）
	// 都被塞進假 buffer，server 直接在 sched_reserve 裡 ggml_abort。
	static const llama_model_tensor_buft_override * sdq_overrides() {
	static const std::string pattern = sdq_expert_pattern();
	static const llama_model_tensor_buft_override ov[2] = {
	{ pattern.c_str(), sdq_expert_buft() },
	{ nullptr, nullptr },
	};
	return ov;
	}

	void sdq_install_buft_overrides(llama_model_params * params) {
	if (!params) {
	return;
	}
	if (params->tensor_buft_overrides == nullptr) {
	params->tensor_buft_overrides = sdq_overrides();
	} else {
	// 使用者已有 override：把手動指定的順序原樣保留，附加在尾端
	size_t n = 0;
	while (params->tensor_buft_overrides[n].pattern != nullptr) {
	n++;
	}
	static std::vector<llama_model_tensor_buft_override> merged;
	merged.clear();
	for (size_t i = 0; i < n; i++) {
	merged.push_back(params->tensor_buft_overrides[i]);
	}
	const llama_model_tensor_buft_override * sdq_ov = sdq_overrides();
	for (size_t i = 0; sdq_ov[i].pattern != nullptr; i++) {
	merged.push_back(sdq_ov[i]);
	}
	merged.push_back({ nullptr, nullptr });
	params->tensor_buft_overrides = merged.data();
	}
	}

	} // namespace sdq