53 lines
2.1 KiB
C
53 lines
2.1 KiB
C
#pragma once
|
|
|
|
#include "ggml.h"
|
|
|
|
#include <stddef.h>
|
|
#include <stdbool.h>
|
|
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
// mapping registry; a registered range marks tensors as mmap-backed and
|
|
// eligible for prefetch. Nothing is opened or retained beyond (base, size).
|
|
void ggml_moe_prefetch_register_mapping (const void * base, size_t size);
|
|
void ggml_moe_prefetch_unregister_mapping(const void * base);
|
|
|
|
// worker pool control; n_threads > 0 (re)creates the pool, <= 0 shuts it down
|
|
void ggml_moe_prefetch_set_n_threads(int n_threads);
|
|
bool ggml_moe_prefetch_enabled(void);
|
|
|
|
// the epoch is bumped once per scheduler pass so per-tensor enqueues are
|
|
// idempotent within one graph execution but re-issued on the next.
|
|
void ggml_moe_prefetch_new_epoch(void);
|
|
|
|
// selective enqueue for one MoE node (GGML_OP_MUL_MAT_ID or
|
|
// GGML_OP_MOE_FUSED_UP_GATE); reads the ids tensor bytes (must be final and
|
|
// host-visible) and enqueues the selected expert ranges of src weights.
|
|
void ggml_moe_prefetch_node(const struct ggml_tensor * node);
|
|
|
|
// full-tensor lookahead enqueue (low priority), e.g. for prompt processing
|
|
// where the next layer's experts are needed in bulk.
|
|
void ggml_moe_prefetch_tensor(const struct ggml_tensor * w);
|
|
|
|
// block until all pending prefetch jobs for tensor w are complete.
|
|
// Returns immediately when nothing is pending.
|
|
void ggml_moe_prefetch_wait(const struct ggml_tensor * w);
|
|
|
|
// MADV_COLD the pages the streamer had to read from storage for tensor w,
|
|
// so prompt-processing streaming traffic is reclaimed ahead of the decode
|
|
// working set. Pages already resident before the sweep are left alone.
|
|
// Only applies to tensors enqueued via ggml_moe_prefetch_tensor.
|
|
void ggml_moe_prefetch_cold(const struct ggml_tensor * w);
|
|
|
|
// kernel-entry hook, called at MoE matmul start (gated by
|
|
// cplan->moe_expert_prefetch). Thread 0 fire-and-forget enqueues when the
|
|
// scheduler hook did not run for this epoch (pure-CPU graphs); falls back to
|
|
// madvise(MADV_WILLNEED) when the engine is off.
|
|
void ggml_moe_prefetch_kernel_hook(const struct ggml_tensor * node, int ith);
|
|
|
|
#ifdef __cplusplus
|
|
}
|
|
#endif
|