diff --git a/README.md b/README.md index 08cd62581c4..9b01777a9ad 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,133 @@ Supported: - Qwen3Next - Mellum MoE models +## Cached-experts configuration examples + +The examples below use `llama-server` and keep the normal model path on the primary CUDA device. Replace `MODEL.gguf` and `/path/to/moe-perf-data` with your model and the directory or JSON data produced for the hot-cache planner. + +### Default cached-experts, one CUDA card + +Use this when one CUDA device should hold the normal graph/KV path and as many cached experts as the automatic budget allows. Remaining experts stay on CPU via `--cpu-moe`. + +```sh +LLAMA_MOE_HOT_CACHE_CPU_DECODE_ROUTING=1 \ +LLAMA_MOE_HOT_CACHE_PARALLEL=1 \ +./build/bin/llama-server \ + --model MODEL.gguf \ + --device CUDA0 \ + --split-mode none \ + --main-gpu 0 \ + --n-gpu-layers 99 \ + --cpu-moe \ + --ctx-size 4096 \ + --ubatch-size 32 \ + --flash-attn on \ + --cache-type-k q8_0 \ + --cache-type-v q8_0 \ + --moe-hot-cache /path/to/moe-perf-data \ + --moe-hot-cache-max-mib -1 \ + --moe-hot-cache-auto-reserve-mib 1024 \ + --moe-hot-cache-pp-reduce-merge on +``` + +### Two CUDA cards + +Use this when `CUDA0` should remain the primary card for graph/KV/router/final merge and `CUDA1` should act as an additional expert lane. `warm` fills the primary lane first, then the second lane. For similar cards, try `hot-even`. + +```sh +GGML_CUDA_P2P=1 \ +LLAMA_MOE_HOT_CACHE_CPU_DECODE_ROUTING=1 \ +LLAMA_MOE_HOT_CACHE_PARALLEL=1 \ +./build/bin/llama-server \ + --model MODEL.gguf \ + --device CUDA0 \ + --split-mode none \ + --main-gpu 0 \ + --n-gpu-layers 99 \ + --cpu-moe \ + --ctx-size 4096 \ + --ubatch-size 32 \ + --flash-attn on \ + --cache-type-k q8_0 \ + --cache-type-v q8_0 \ + --moe-hot-cache /path/to/moe-perf-data \ + --moe-hot-cache-max-mib -1 \ + --moe-hot-cache-auto-reserve-mib 1024 \ + --moe-hot-cache-second-device CUDA1 \ + --moe-hot-cache-second-max-mib -1 \ + --moe-hot-cache-second-auto-reserve-mib 512 \ + --moe-hot-cache-device-strategy warm \ + --moe-hot-cache-pp-reduce-merge on +``` + +For a small primary GPU that should only run graph/KV/router/final merge, disable the primary expert cache and place experts on the second device: + +```sh +GGML_CUDA_P2P=1 \ +LLAMA_MOE_HOT_CACHE_CPU_DECODE_ROUTING=1 \ +LLAMA_MOE_HOT_CACHE_PARALLEL=1 \ +./build/bin/llama-server \ + --model MODEL.gguf \ + --device CUDA0 \ + --split-mode none \ + --main-gpu 0 \ + --n-gpu-layers 99 \ + --cpu-moe \ + --ctx-size 4096 \ + --ubatch-size 32 \ + --flash-attn on \ + --cache-type-k q8_0 \ + --cache-type-v q8_0 \ + --moe-hot-cache /path/to/moe-perf-data \ + --moe-hot-cache-max-mib 0 \ + --moe-hot-cache-second-device CUDA1 \ + --moe-hot-cache-second-max-mib -1 \ + --moe-hot-cache-second-auto-reserve-mib 512 \ + --moe-hot-cache-device-strategy warm \ + --moe-hot-cache-pp-reduce-merge on +``` + +### Three CUDA cards + +Use this when `CUDA0` is the primary card and `CUDA1`/`CUDA2` are additional expert lanes. This is the intended shape for one primary GPU plus two expert-only GPUs. Keep per-device reserve high enough for temporary buffers; reduce `--ctx-size` or `--ubatch-size` first if CUDA allocation fails. + +```sh +GGML_CUDA_P2P=1 \ +LLAMA_MOE_HOT_CACHE_CPU_DECODE_ROUTING=1 \ +LLAMA_MOE_HOT_CACHE_PARALLEL=1 \ +./build/bin/llama-server \ + --model MODEL.gguf \ + --device CUDA0 \ + --split-mode none \ + --main-gpu 0 \ + --n-gpu-layers 99 \ + --cpu-moe \ + --ctx-size 4096 \ + --ubatch-size 32 \ + --flash-attn on \ + --cache-type-k q8_0 \ + --cache-type-v q8_0 \ + --moe-hot-cache /path/to/moe-perf-data \ + --moe-hot-cache-max-mib -1 \ + --moe-hot-cache-auto-reserve-mib 1024 \ + --moe-hot-cache-second-device CUDA1 \ + --moe-hot-cache-second-max-mib -1 \ + --moe-hot-cache-second-auto-reserve-mib 512 \ + --moe-hot-cache-third-device CUDA2 \ + --moe-hot-cache-third-max-mib -1 \ + --moe-hot-cache-third-auto-reserve-mib 512 \ + --moe-hot-cache-device-strategy hot-even \ + --moe-hot-cache-pp-reduce-merge on +``` + +Notes: + +- `--moe-hot-cache-max-mib -1` auto-sizes a lane from currently free VRAM minus its reserve. +- `--moe-hot-cache-max-mib 0` disables the primary expert lane while keeping secondary or tertiary expert lanes available. +- `GGML_CUDA_P2P=1` enables CUDA peer-copy when the cards and driver support it; unsupported pairs fall back internally. +- `LLAMA_MOE_HOT_CACHE_PARALLEL=force` is a debugging mode for valid parallel regions. Use `1`/`auto` for normal runs. +- A general speedup claim for two GPUs could not be validated on the available test hardware because the cards are very asymmetric. Treat the two-GPU examples as configuration starting points, not as benchmark guidance. + These changes will probably never reach upstream llama because I broke the contribution rules hardly. I am a Java developer and the last time I wrote anything in C is, I even don't remember when it was, therefore, the bit of knowlegde of C that I had is gone. Secondly, this is a tool for me, I want it to function, I want it to be easy and I used other tools to create it faster. And lastly, I saw some discussions in the PRs and the tone is not what I would expect. I know especially big PRs are hard to overlook, but great features often create big PRs. I also hate big PRs. But, sometimes they are necessary. Anyway, I don't want to have such discussions, it's just a waste of time. diff --git a/common/arg.cpp b/common/arg.cpp index 190ac2b0198..985e69058cb 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -73,6 +73,10 @@ static bool llama_moe_hot_cache_weighting_valid(const std::string & value) { value == "flat"; } +static bool llama_moe_hot_cache_device_strategy_valid(const std::string & value) { + return value == "warm" || value == "hot-even"; +} + static void llama_moe_hot_cache_set_weighting_env(const std::string & value) { #if defined(_WIN32) _putenv_s("LLAMA_MOE_HOT_CACHE_WEIGHTING", value.c_str()); @@ -999,17 +1003,35 @@ bool common_params_parse(int argc, char ** argv, common_params & params, llama_e common_params_print_completion(ctx_arg); exit(0); } - if (ctx_arg.params.moe_hot_cache_max_mib != 0 && ctx_arg.params.moe_hot_cache.empty()) { + const bool moe_hot_cache_any_lane = + ctx_arg.params.moe_hot_cache_max_mib != 0 || + ctx_arg.params.moe_hot_cache_second_max_mib != 0 || + ctx_arg.params.moe_hot_cache_third_max_mib != 0; + if (moe_hot_cache_any_lane && ctx_arg.params.moe_hot_cache.empty()) { throw std::invalid_argument("--moe-hot-cache is required when --moe-hot-cache-max-mib is not 0"); } if (ctx_arg.params.moe_hot_cache_max_mib < -1) { throw std::invalid_argument("--moe-hot-cache-max-mib must be -1 or greater"); } - if (ctx_arg.params.moe_hot_cache_max_mib == -1 && ctx_arg.params.n_ctx <= 0) { - throw std::invalid_argument("--moe-hot-cache-max-mib -1 requires an explicit --ctx-size"); + if (ctx_arg.params.moe_hot_cache_second_max_mib < -1) { + throw std::invalid_argument("--moe-hot-cache-second-max-mib must be -1 or greater"); + } + if (ctx_arg.params.moe_hot_cache_third_max_mib < -1) { + throw std::invalid_argument("--moe-hot-cache-third-max-mib must be -1 or greater"); + } + if ((ctx_arg.params.moe_hot_cache_max_mib == -1 || + ctx_arg.params.moe_hot_cache_second_max_mib == -1 || + ctx_arg.params.moe_hot_cache_third_max_mib == -1) && ctx_arg.params.n_ctx <= 0) { + throw std::invalid_argument("--moe-hot-cache-*-max-mib -1 requires an explicit --ctx-size"); + } + if (ctx_arg.params.moe_hot_cache_second_max_mib != 0 && ctx_arg.params.moe_hot_cache_second_device.empty()) { + throw std::invalid_argument("--moe-hot-cache-second-device is required when --moe-hot-cache-second-max-mib is not 0"); + } + if (ctx_arg.params.moe_hot_cache_third_max_mib != 0 && ctx_arg.params.moe_hot_cache_third_device.empty()) { + throw std::invalid_argument("--moe-hot-cache-third-device is required when --moe-hot-cache-third-max-mib is not 0"); } - if (ctx_arg.params.moe_hot_cache_update_rate > 0.0f && ctx_arg.params.moe_hot_cache_max_mib == 0) { - throw std::invalid_argument("--moe-hot-cache-update-rate requires --moe-hot-cache-max-mib"); + if (ctx_arg.params.moe_hot_cache_update_rate > 0.0f && !moe_hot_cache_any_lane) { + throw std::invalid_argument("--moe-hot-cache-update-rate requires a MoE hot-cache lane budget"); } if (!ctx_arg.params.moe_layer_perf_out.empty()) { ctx_arg.params.no_perf = false; @@ -2417,6 +2439,13 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.moe_hot_cache_max_mib = value; } ).set_env("LLAMA_ARG_MOE_HOT_CACHE_MAX_MIB")); + add_opt(common_arg( + {"--moe-hot-cache-device"}, "DEV", + "experimental: backend device for the primary MoE hot-cache expert lane (default: first model GPU/iGPU)", + [](common_params & params, const std::string & value) { + params.moe_hot_cache_device = value; + } + ).set_env("LLAMA_ARG_MOE_HOT_CACHE_DEVICE")); add_opt(common_arg( {"--moe-hot-cache-auto-reserve-mib"}, "N", string_format("experimental: MiB to keep free when --moe-hot-cache-max-mib -1 auto-sizes the hot cache (default: %zu)", (size_t) params.moe_hot_cache_auto_reserve_mib), @@ -2428,6 +2457,74 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.moe_hot_cache_auto_reserve_mib = uint64_t(value); } ).set_env("LLAMA_ARG_MOE_HOT_CACHE_AUTO_RESERVE_MIB")); + add_opt(common_arg( + {"--moe-hot-cache-second-device"}, "DEV", + "experimental: backend device for the optional second MoE hot-cache expert lane", + [](common_params & params, const std::string & value) { + params.moe_hot_cache_second_device = value; + } + ).set_env("LLAMA_ARG_MOE_HOT_CACHE_SECOND_DEVICE")); + add_opt(common_arg( + {"--moe-hot-cache-second-max-mib"}, "N", + "experimental: max MiB for the optional second MoE hot-cache expert lane (0 = disabled, -1 = auto)", + [](common_params & params, const std::string & value_str) { + const int64_t value = std::stoll(value_str); + if (value < -1) { + throw std::invalid_argument("invalid value"); + } + params.moe_hot_cache_second_max_mib = value; + } + ).set_env("LLAMA_ARG_MOE_HOT_CACHE_SECOND_MAX_MIB")); + add_opt(common_arg( + {"--moe-hot-cache-second-auto-reserve-mib"}, "N", + string_format("experimental: MiB to keep free when --moe-hot-cache-second-max-mib -1 auto-sizes the second lane (default: %zu)", (size_t) params.moe_hot_cache_second_auto_reserve_mib), + [](common_params & params, const std::string & value_str) { + const int64_t value = std::stoll(value_str); + if (value < 0) { + throw std::invalid_argument("invalid value"); + } + params.moe_hot_cache_second_auto_reserve_mib = uint64_t(value); + } + ).set_env("LLAMA_ARG_MOE_HOT_CACHE_SECOND_AUTO_RESERVE_MIB")); + add_opt(common_arg( + {"--moe-hot-cache-third-device"}, "DEV", + "experimental: backend device for the optional third MoE hot-cache expert lane", + [](common_params & params, const std::string & value) { + params.moe_hot_cache_third_device = value; + } + ).set_env("LLAMA_ARG_MOE_HOT_CACHE_THIRD_DEVICE")); + add_opt(common_arg( + {"--moe-hot-cache-third-max-mib"}, "N", + "experimental: max MiB for the optional third MoE hot-cache expert lane (0 = disabled, -1 = auto)", + [](common_params & params, const std::string & value_str) { + const int64_t value = std::stoll(value_str); + if (value < -1) { + throw std::invalid_argument("invalid value"); + } + params.moe_hot_cache_third_max_mib = value; + } + ).set_env("LLAMA_ARG_MOE_HOT_CACHE_THIRD_MAX_MIB")); + add_opt(common_arg( + {"--moe-hot-cache-third-auto-reserve-mib"}, "N", + string_format("experimental: MiB to keep free when --moe-hot-cache-third-max-mib -1 auto-sizes the third lane (default: %zu)", (size_t) params.moe_hot_cache_third_auto_reserve_mib), + [](common_params & params, const std::string & value_str) { + const int64_t value = std::stoll(value_str); + if (value < 0) { + throw std::invalid_argument("invalid value"); + } + params.moe_hot_cache_third_auto_reserve_mib = uint64_t(value); + } + ).set_env("LLAMA_ARG_MOE_HOT_CACHE_THIRD_AUTO_RESERVE_MIB")); + add_opt(common_arg( + {"--moe-hot-cache-device-strategy"}, "{warm,hot-even}", + "experimental: distribute MoE hot-cache experts as warm lanes or per-layer hot-even lanes (default: warm)", + [](common_params & params, const std::string & value) { + if (!llama_moe_hot_cache_device_strategy_valid(value)) { + throw std::invalid_argument("--moe-hot-cache-device-strategy must be one of: warm, hot-even"); + } + params.moe_hot_cache_device_strategy = value; + } + ).set_env("LLAMA_ARG_MOE_HOT_CACHE_DEVICE_STRATEGY")); add_opt(common_arg( {"--moe-hot-cache"}, "FNAME", "experimental: path to /moe-layer-perf JSON used by --moe-hot-cache-max-mib", diff --git a/common/common.cpp b/common/common.cpp index 470e1605852..edceaf587ae 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1534,6 +1534,14 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.no_host = params.no_host; mparams.moe_hot_cache_max_mib = params.moe_hot_cache_max_mib; mparams.moe_hot_cache_path = params.moe_hot_cache.empty() ? nullptr : params.moe_hot_cache.c_str(); + mparams.moe_hot_cache_device = params.moe_hot_cache_device.empty() ? nullptr : params.moe_hot_cache_device.c_str(); + mparams.moe_hot_cache_second_max_mib = params.moe_hot_cache_second_max_mib; + mparams.moe_hot_cache_second_device = params.moe_hot_cache_second_device.empty() ? nullptr : params.moe_hot_cache_second_device.c_str(); + mparams.moe_hot_cache_second_auto_reserve_mib = params.moe_hot_cache_second_auto_reserve_mib; + mparams.moe_hot_cache_third_max_mib = params.moe_hot_cache_third_max_mib; + mparams.moe_hot_cache_third_device = params.moe_hot_cache_third_device.empty() ? nullptr : params.moe_hot_cache_third_device.c_str(); + mparams.moe_hot_cache_third_auto_reserve_mib = params.moe_hot_cache_third_auto_reserve_mib; + mparams.moe_hot_cache_device_strategy = params.moe_hot_cache_device_strategy.empty() ? "warm" : params.moe_hot_cache_device_strategy.c_str(); mparams.moe_hot_cache_auto_n_ctx = params.n_ctx > 0 ? uint32_t(params.n_ctx) : 0; mparams.moe_hot_cache_auto_n_seq_max = params.n_parallel > 0 ? uint32_t(params.n_parallel) : 1; mparams.moe_hot_cache_auto_n_ubatch = params.n_ubatch > 0 ? uint32_t(params.n_ubatch) : uint32_t(params.n_batch); diff --git a/common/common.h b/common/common.h index 2f8385336f2..d78df5d00e4 100644 --- a/common/common.h +++ b/common/common.h @@ -562,6 +562,14 @@ struct common_params { int64_t moe_hot_cache_max_mib = 0; // max MiB for experimental MoE hot expert cache, 0 = disabled, -1 = auto uint64_t moe_hot_cache_auto_reserve_mib = 1024; // MiB kept free when auto-sizing the MoE hot expert cache + std::string moe_hot_cache_device; // optional backend device for primary MoE hot-cache expert lane + int64_t moe_hot_cache_second_max_mib = 0; // max MiB for optional second MoE hot-cache expert lane + uint64_t moe_hot_cache_second_auto_reserve_mib = 512; // MiB kept free when auto-sizing the second lane + std::string moe_hot_cache_second_device; // backend device for optional second expert lane + int64_t moe_hot_cache_third_max_mib = 0; // max MiB for optional third MoE hot-cache expert lane + uint64_t moe_hot_cache_third_auto_reserve_mib = 512; // MiB kept free when auto-sizing the third lane + std::string moe_hot_cache_third_device; // backend device for optional third expert lane + std::string moe_hot_cache_device_strategy = "warm"; // MoE hot-cache device strategy: warm or hot-even std::string moe_hot_cache; // path to /moe-layer-perf JSON float moe_hot_cache_update_rate = 0.0f; // fraction of hot-cache entries to update after each completed server run float moe_hot_cache_layer_curve = 0.5f; // MoE hot-cache layer-pressure weighting curve, 0 = flat, 1 = aggressive diff --git a/docs/moe-hot-cache/moe-hot-cache-architecture-explainer.html b/docs/moe-hot-cache/moe-hot-cache-architecture-explainer.html index 4a7b7ba8a1f..301e6693b70 100644 --- a/docs/moe-hot-cache/moe-hot-cache-architecture-explainer.html +++ b/docs/moe-hot-cache/moe-hot-cache-architecture-explainer.html @@ -1462,8 +1462,10 @@

@@ -918,11 +936,35 @@

Hot-cache workflow arguments

Integer >= 0, default 1024 Used only with --moe-hot-cache-max-mib -1. Leaves this many MiB free for warmup, compute buffers, and transient CUDA allocations. + + --moe-hot-cache-device DEV + LLAMA_ARG_MOE_HOT_CACHE_DEVICE + Backend device name, default unset + Optional device for the primary expert lane. If omitted, the hot cache keeps the current behavior and uses the first GPU/iGPU seen by the loaded model. + + + --moe-hot-cache-second-device DEV
--moe-hot-cache-second-max-mib N
--moe-hot-cache-second-auto-reserve-mib N + LLAMA_ARG_MOE_HOT_CACHE_SECOND_DEVICE
LLAMA_ARG_MOE_HOT_CACHE_SECOND_MAX_MIB
LLAMA_ARG_MOE_HOT_CACHE_SECOND_AUTO_RESERVE_MIB + Optional second lane; max MiB uses the same 0, fixed, and -1 semantics. Auto reserve default 512 MiB. + Creates an expert-only worker lane resolved independently from normal --device, so the GPU can hold cached experts without receiving normal layers or KV cache. + + + --moe-hot-cache-third-device DEV
--moe-hot-cache-third-max-mib N
--moe-hot-cache-third-auto-reserve-mib N + LLAMA_ARG_MOE_HOT_CACHE_THIRD_DEVICE
LLAMA_ARG_MOE_HOT_CACHE_THIRD_MAX_MIB
LLAMA_ARG_MOE_HOT_CACHE_THIRD_AUTO_RESERVE_MIB + Optional third lane; max MiB uses the same 0, fixed, and -1 semantics. Auto reserve default 512 MiB. + Intended for primary-GPU graph/KV setups where two additional worker GPUs hold and compute routed expert slices. + + + --moe-hot-cache-device-strategy warm|hot-even + LLAMA_ARG_MOE_HOT_CACHE_DEVICE_STRATEGY + warm default, or hot-even + warm fills the first configured lane, then later lanes. hot-even distributes selected experts per layer across configured expert lanes while respecting each lane budget. + --moe-hot-cache-update-rate N moe-hot-cache-update-rate = N
LLAMA_ARG_MOE_HOT_CACHE_UPDATE_RATE Float 0.0 to 1.0, default 0.0 - After a completed server request, replaces up to this fraction of hot-cache entries with better observed candidates. Needs active perf counters. + After a completed server request, replaces up to this fraction of hot-cache entries with better observed candidates. Needs active perf counters. Multi-device expert lanes currently skip runtime replacement to avoid inconsistent per-lane maps. --moe-hot-cache-weighting MODE
--moe-hot-cache-qwen-weighting MODE @@ -1014,8 +1056,8 @@

Related llama.cpp arguments that affect hot-cache runs

--device CUDA0 - Controls which backend devices are used by the model. The hot cache picks the first GPU/IGPU device seen by the loaded model. - Use the intended fast GPU first when multiple devices are present. + Controls which backend devices are used by the normal model graph, including layer offload, KV placement, router logits, and the final multi-lane merge. + For expert-only CUDA1/CUDA2 lanes, keep normal --device to the primary graph GPU and name the other GPUs only with --moe-hot-cache-second-device and --moe-hot-cache-third-device. --n-gpu-layers N diff --git a/include/llama.h b/include/llama.h index 27c8d3ca075..ebc24ce0245 100644 --- a/include/llama.h +++ b/include/llama.h @@ -307,6 +307,14 @@ extern "C" { // Experimental Qwen3.5 MoE hot expert cache. 0 disables it, -1 uses remaining VRAM after model load and context reservation. int64_t moe_hot_cache_max_mib; const char * moe_hot_cache_path; + const char * moe_hot_cache_device; + int64_t moe_hot_cache_second_max_mib; + const char * moe_hot_cache_second_device; + uint64_t moe_hot_cache_second_auto_reserve_mib; + int64_t moe_hot_cache_third_max_mib; + const char * moe_hot_cache_third_device; + uint64_t moe_hot_cache_third_auto_reserve_mib; + const char * moe_hot_cache_device_strategy; uint32_t moe_hot_cache_auto_n_ctx; uint32_t moe_hot_cache_auto_n_seq_max; uint32_t moe_hot_cache_auto_n_ubatch; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 4ba36aafa56..a95f13e2cf1 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -65,6 +65,47 @@ static llm_graph_phase infer_decode_graph_phase( return LLM_GRAPH_PHASE_DECODE; } +// Opens expert-only MoE hot-cache devices so tensors allocated on those devices +// are visible to the scheduler. Keep GPU backends before ACCEL/CPU backends, +// because the scheduler assumes the last backend is CPU for graph inputs. +static void llama_context_add_moe_hot_cache_backends( + const llama_model & model, + std::vector & backends) { + if (model.moe_hot_cache == nullptr) { + return; + } + + for (ggml_backend_dev_t dev : model.moe_hot_cache->devices) { + const bool already_open = std::any_of(backends.begin(), backends.end(), [&](const auto & backend) { + return ggml_backend_get_device(backend.get()) == dev; + }); + if (already_open) { + continue; + } + + ggml_backend_t backend = ggml_backend_dev_init(dev, nullptr); + if (backend == nullptr) { + throw std::runtime_error(format("failed to initialize MoE hot-cache expert backend %s", ggml_backend_dev_name(dev))); + } + + const auto insert_pos = std::find_if(backends.begin(), backends.end(), [](const auto & existing) { + ggml_backend_dev_t existing_dev = ggml_backend_get_device(existing.get()); + if (existing_dev == nullptr) { + return true; + } + + const auto type = ggml_backend_dev_type(existing_dev); + return type != GGML_BACKEND_DEVICE_TYPE_GPU && + type != GGML_BACKEND_DEVICE_TYPE_IGPU; + }); + + LLAMA_LOG_INFO("%s: adding MoE hot-cache expert-only backend %s\n", + __func__, + ggml_backend_dev_name(dev)); + backends.insert(insert_pos, ggml_backend_ptr(backend)); + } +} + llama_context::llama_context( const llama_model & model, llama_context_params params) : @@ -287,6 +328,8 @@ llama_context::llama_context( backends.emplace_back(backend); } + llama_context_add_moe_hot_cache_backends(model, backends); + // add ACCEL backends (such as BLAS) for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { ggml_backend_dev_t dev = ggml_backend_dev_get(i); @@ -345,6 +388,7 @@ llama_context::llama_context( } llama_moe_hot_cache_init_after_context_memory(model); + llama_context_add_moe_hot_cache_backends(model, backends); // init backends if (!hparams.vocab_only) { diff --git a/src/llama-model.cpp b/src/llama-model.cpp index d800f3a69a6..40d319679b6 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2223,6 +2223,14 @@ llama_model_params llama_model_default_params() { /*.tensor_split =*/ nullptr, /*.moe_hot_cache_max_mib =*/ 0, /*.moe_hot_cache_path =*/ nullptr, + /*.moe_hot_cache_device =*/ nullptr, + /*.moe_hot_cache_second_max_mib=*/ 0, + /*.moe_hot_cache_second_device =*/ nullptr, + /*.moe_hot_cache_second_auto_reserve_mib =*/ 512, + /*.moe_hot_cache_third_max_mib =*/ 0, + /*.moe_hot_cache_third_device =*/ nullptr, + /*.moe_hot_cache_third_auto_reserve_mib =*/ 512, + /*.moe_hot_cache_device_strategy =*/ "warm", /*.moe_hot_cache_auto_n_ctx =*/ 0, /*.moe_hot_cache_auto_n_seq_max=*/ 1, /*.moe_hot_cache_auto_n_ubatch =*/ 512, diff --git a/src/models/gemma4.cpp b/src/models/gemma4.cpp index 2002fdaac92..226b25ee339 100644 --- a/src/models/gemma4.cpp +++ b/src/models/gemma4.cpp @@ -317,7 +317,11 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para ggml_tensor * logits = build_lora_mm(model.layers[il].ffn_gate_inp, tmp); // [n_expert, n_tokens] cb(logits, "ffn_moe_logits", il); - if (llama_moe_hot_cache_layer_active_for_graph(model, il, llama_moe_hot_cache_graph_kind::logits)) { + const bool hot_cache_layer_active = + llama_moe_hot_cache_layer_active_for_graph(model, il, llama_moe_hot_cache_graph_kind::logits); + const bool hot_cache_multi_lane = + hot_cache_layer_active && !model.moe_hot_cache->layers[il].lanes.empty(); + if (hot_cache_layer_active && (!hot_cache_multi_lane || (!cparams.warmup && cur_moe->ne[1] == 1))) { cur_moe = build_layer_moe_hot(cur_moe, logits, il); } else { cur_moe = build_moe_ffn(cur_moe, diff --git a/src/models/mellum.cpp b/src/models/mellum.cpp index 1a201a777ea..b7023c78699 100644 --- a/src/models/mellum.cpp +++ b/src/models/mellum.cpp @@ -180,7 +180,11 @@ llama_model_mellum::graph::graph(const llama_model & model, const llm_grap cb(cur, "ffn_norm", il); ggml_tensor * moe_out = nullptr; - if (llama_moe_hot_cache_layer_active_for_graph(model, il, llama_moe_hot_cache_graph_kind::logits)) { + const bool hot_cache_layer_active = + llama_moe_hot_cache_layer_active_for_graph(model, il, llama_moe_hot_cache_graph_kind::logits); + const bool hot_cache_multi_lane = + hot_cache_layer_active && !model.moe_hot_cache->layers[il].lanes.empty(); + if (hot_cache_layer_active && (!hot_cache_multi_lane || (!cparams.warmup && cur->ne[1] == 1))) { ggml_tensor * logits = build_lora_mm(model.layers[il].ffn_gate_inp, cur); cb(logits, "ffn_moe_logits", il); diff --git a/src/models/qwen35moe.cpp b/src/models/qwen35moe.cpp index e4be9f444be..7c47c5144d2 100644 --- a/src/models/qwen35moe.cpp +++ b/src/models/qwen35moe.cpp @@ -501,14 +501,23 @@ ggml_tensor * llama_model_qwen35moe::graph::build_layer_ffn(ggml_tensor * cur, c llama_moe_hot_cache_layer_active_for_graph(model, il, llama_moe_hot_cache_graph_kind::qwen35_ffn); const llama_moe_hot_cache_layer * hot_cache_layer = hot_cache_layer_active ? &model.moe_hot_cache->layers[il] : nullptr; + const bool hot_cache_multi_lane = hot_cache_layer != nullptr && !hot_cache_layer->lanes.empty(); + uint32_t hot_cache_n_hot = hot_cache_layer != nullptr ? hot_cache_layer->n_hot : 0; + if (hot_cache_multi_lane) { + hot_cache_n_hot = 0; + for (const auto & lane : hot_cache_layer->lanes) { + hot_cache_n_hot += lane.n_hot; + } + } const bool hot_cache_active = hot_cache_layer_active && + (!hot_cache_multi_lane || (!cparams.warmup && cur->ne[1] == 1)) && !llama_moe_hot_cache_pp_policy::bypass_hot_cache_for_prompt_processing( gphase, cparams.warmup, cur->ne[1], 0, - hot_cache_layer->n_hot, + hot_cache_n_hot, hot_cache_layer->n_expert, 0.0); if (hot_cache_active) { diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index ed7839e0751..6c7dc791ff0 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -540,8 +540,13 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_ffn(ggml_tensor * cur, c if (model.layers[il].ffn_gate_inp != nullptr) { // MoE branch ggml_tensor * moe_out = nullptr; + const bool hot_cache_layer_active = + llama_moe_hot_cache_layer_active_for_graph(model, il, llama_moe_hot_cache_graph_kind::logits); + const bool hot_cache_multi_lane = + hot_cache_layer_active && !model.moe_hot_cache->layers[il].lanes.empty(); const bool hot_cache_active = - llama_moe_hot_cache_layer_active_for_graph(model, il, llama_moe_hot_cache_graph_kind::logits) && + hot_cache_layer_active && + (!hot_cache_multi_lane || (!cparams.warmup && cur->ne[1] == 1)) && !llama_moe_hot_cache_pp_policy::bypass_hot_cache_for_prompt_processing(gphase, cparams.warmup, cur->ne[1], 1); if (hot_cache_active) { ggml_tensor * logits = build_lora_mm(model.layers[il].ffn_gate_inp, cur); diff --git a/src/moe-hot-cache/llama-moe-hot-cache-budget.cpp b/src/moe-hot-cache/llama-moe-hot-cache-budget.cpp index 340e35b7239..18b0b642410 100644 --- a/src/moe-hot-cache/llama-moe-hot-cache-budget.cpp +++ b/src/moe-hot-cache/llama-moe-hot-cache-budget.cpp @@ -9,6 +9,7 @@ #include #include #include +#include namespace { @@ -121,22 +122,56 @@ ggml_backend_dev_t llama_moe_hot_cache_select_gpu_dev(const llama_model * model) return dev; } +ggml_backend_dev_t llama_moe_hot_cache_resolve_gpu_dev( + const llama_model * model, + const char * name) { + if (name == nullptr || name[0] == '\0') { + return llama_moe_hot_cache_select_gpu_dev(model); + } + + ggml_backend_dev_t dev = ggml_backend_dev_by_name(name); + if (dev == nullptr) { + throw std::runtime_error(std::string("unknown --moe-hot-cache device: ") + name); + } + + const auto type = ggml_backend_dev_type(dev); + if (type != GGML_BACKEND_DEVICE_TYPE_GPU && type != GGML_BACKEND_DEVICE_TYPE_IGPU) { + throw std::runtime_error(std::string("--moe-hot-cache device is not a GPU/iGPU backend: ") + name); + } + + return dev; +} + size_t llama_moe_hot_cache_auto_budget_bytes( const llama_model & model, const llama_model_params & params, ggml_backend_dev_t dev, bool reserve_kv_cache) { + return llama_moe_hot_cache_auto_budget_bytes( + model, + params, + dev, + reserve_kv_cache, + params.moe_hot_cache_auto_reserve_mib); +} + +size_t llama_moe_hot_cache_auto_budget_bytes( + const llama_model & model, + const llama_model_params & params, + ggml_backend_dev_t dev, + bool reserve_kv_cache, + uint64_t reserve_mib) { size_t free = 0; size_t total = 0; ggml_backend_dev_memory(dev, &free, &total); GGML_UNUSED(total); const size_t kv_reserve = reserve_kv_cache ? estimate_kv_cache_bytes_on_device(model, params, dev) : 0; - const size_t safety_reserve = mul_mib_saturating(params.moe_hot_cache_auto_reserve_mib); + const size_t safety_reserve = mul_mib_saturating(reserve_mib); const size_t budget = llama_moe_hot_cache_compute_auto_budget_bytes( free, kv_reserve, - params.moe_hot_cache_auto_reserve_mib); + reserve_mib); if (budget == 0) { LLAMA_LOG_WARN("%s: auto hot-cache budget on %s is 0 MiB: free before hot-cache = %zu MiB, %s KV reserve = %zu MiB, safety reserve = %zu MiB\n", diff --git a/src/moe-hot-cache/llama-moe-hot-cache-budget.h b/src/moe-hot-cache/llama-moe-hot-cache-budget.h index f1580a4a073..45c26160f7c 100644 --- a/src/moe-hot-cache/llama-moe-hot-cache-budget.h +++ b/src/moe-hot-cache/llama-moe-hot-cache-budget.h @@ -6,6 +6,10 @@ ggml_backend_dev_t llama_moe_hot_cache_select_gpu_dev(const llama_model * model = nullptr); +ggml_backend_dev_t llama_moe_hot_cache_resolve_gpu_dev( + const llama_model * model, + const char * name); + size_t llama_moe_hot_cache_compute_auto_budget_bytes( size_t free_bytes, size_t kv_reserve_bytes, @@ -16,3 +20,10 @@ size_t llama_moe_hot_cache_auto_budget_bytes( const llama_model_params & params, ggml_backend_dev_t dev, bool reserve_kv_cache); + +size_t llama_moe_hot_cache_auto_budget_bytes( + const llama_model & model, + const llama_model_params & params, + ggml_backend_dev_t dev, + bool reserve_kv_cache, + uint64_t reserve_mib); diff --git a/src/moe-hot-cache/llama-moe-hot-cache-builder.cpp b/src/moe-hot-cache/llama-moe-hot-cache-builder.cpp index 2c068754ef4..2e039f999c6 100644 --- a/src/moe-hot-cache/llama-moe-hot-cache-builder.cpp +++ b/src/moe-hot-cache/llama-moe-hot-cache-builder.cpp @@ -125,6 +125,7 @@ std::unique_ptr llama_moe_hot_cache_build( ggml_backend_dev_t cache_dev) { auto cache = std::make_unique(); cache->layers.resize(model.hparams.n_layer()); + cache->devices.push_back(cache_dev); const auto selected_by_layer = llama_moe_hot_cache_group_selected_by_layer(plan, model.hparams.n_layer()); const auto stats = llama_moe_hot_cache_summarize_selected_layers(selected_by_layer); @@ -264,3 +265,194 @@ std::unique_ptr llama_moe_hot_cache_build( cache->ctxs.emplace_back(std::move(ctx)); return cache; } + +std::unique_ptr llama_moe_hot_cache_build_multi( + const llama_model & model, + const llama_moe_hot_cache_multi_plan & plan, + const std::vector & cache_devs) { + if (plan.lanes.empty() || plan.lanes.size() > LLAMA_MOE_HOT_CACHE_MAX_EXPERT_LANES) { + throw std::runtime_error("MoE hot-cache multi-device builder requires 1..3 lanes"); + } + if (cache_devs.size() != plan.lanes.size()) { + throw std::runtime_error("MoE hot-cache multi-device builder got mismatched plan/device lanes"); + } + if (plan.lanes.size() == 1) { + return llama_moe_hot_cache_build(model, plan.lanes[0], cache_devs[0]); + } + + auto cache = std::make_unique(); + cache->layers.resize(model.hparams.n_layer()); + cache->devices = cache_devs; + for (auto & layer : cache->layers) { + layer.lanes.resize(plan.lanes.size()); + } + + for (size_t lane_index = 0; lane_index < plan.lanes.size(); ++lane_index) { + const auto selected_by_layer = llama_moe_hot_cache_group_selected_by_layer( + plan.lanes[lane_index], + model.hparams.n_layer()); + const auto stats = llama_moe_hot_cache_summarize_selected_layers(selected_by_layer); + + LLAMA_LOG_INFO("%s: hot-cache lane %zu on %s active layers = %zu/%zu, hot experts per active layer min/avg/max = %zu/%.1f/%zu\n", + __func__, + lane_index, + ggml_backend_dev_name(cache_devs[lane_index]), + stats.active_layers, + selected_by_layer.size(), + stats.min_hot, + stats.avg_hot(), + stats.max_hot); + + size_t n_tensors = 0; + for (uint32_t il = 0; il < selected_by_layer.size(); ++il) { + if (selected_by_layer[il].empty()) { + continue; + } + + const auto & src = model.layers[il]; + auto & cache_layer = cache->layers[il]; + auto & dst_lane = cache_layer.lanes[lane_index]; + + n_tensors += 4; // map + hot mask + cold mask + down + n_tensors += src.ffn_gate_up_exps != nullptr ? 1 : 2; + n_tensors += src.ffn_gate_exps_s != nullptr ? 1 : 0; + n_tensors += src.ffn_up_exps_s != nullptr ? 1 : 0; + n_tensors += src.ffn_down_exps_s != nullptr ? 1 : 0; + + dst_lane.n_hot = selected_by_layer[il].size(); + dst_lane.n_expert = src.ffn_down_exps ? src.ffn_down_exps->ne[2] : 0; + dst_lane.expert_weights_scale = model.hparams.expert_weights_scale; + + cache_layer.n_expert = dst_lane.n_expert; + cache_layer.expert_weights_scale = model.hparams.expert_weights_scale; + if (cache_layer.expert_lane_map_host.empty() && dst_lane.n_expert > 0) { + cache_layer.expert_lane_map_host.assign(dst_lane.n_expert, -1); + } + } + + if (n_tensors == 0) { + continue; + } + + static constexpr size_t EXTRA_PER_TENSOR = 64; + const size_t n_extra_tensors = std::max(64, n_tensors / 4); + const size_t ctx_mem_size = ggml_tensor_overhead() * (n_tensors + n_extra_tensors) + + EXTRA_PER_TENSOR * n_tensors; + + ggml_init_params ctx_params = { + /*.mem_size =*/ ctx_mem_size, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + + ggml_context_ptr ctx { ggml_init(ctx_params) }; + if (!ctx) { + throw std::runtime_error("failed to create MoE hot-cache lane ggml context"); + } + + for (uint32_t il = 0; il < selected_by_layer.size(); ++il) { + const auto & experts = selected_by_layer[il]; + if (experts.empty()) { + continue; + } + + const auto & src = model.layers[il]; + auto & dst_lane = cache->layers[il].lanes[lane_index]; + const int64_t n_cache = int64_t(experts.size()) + 1; + const int64_t n_expert = src.ffn_down_exps->ne[2]; + + dst_lane.ffn_gate_up_exps = new_tensor_like_experts(ctx.get(), src.ffn_gate_up_exps, n_cache, format("blk.%u.ffn_gate_up_exps.hot_cache.%zu", il, lane_index).c_str()); + dst_lane.ffn_gate_exps = new_tensor_like_experts(ctx.get(), src.ffn_gate_exps, n_cache, format("blk.%u.ffn_gate_exps.hot_cache.%zu", il, lane_index).c_str()); + dst_lane.ffn_up_exps = new_tensor_like_experts(ctx.get(), src.ffn_up_exps, n_cache, format("blk.%u.ffn_up_exps.hot_cache.%zu", il, lane_index).c_str()); + dst_lane.ffn_down_exps = new_tensor_like_experts(ctx.get(), src.ffn_down_exps, n_cache, format("blk.%u.ffn_down_exps.hot_cache.%zu", il, lane_index).c_str()); + dst_lane.ffn_gate_exps_s = new_tensor_like_scale (ctx.get(), src.ffn_gate_exps_s, n_cache, format("blk.%u.ffn_gate_exps_s.hot_cache.%zu", il, lane_index).c_str()); + dst_lane.ffn_up_exps_s = new_tensor_like_scale (ctx.get(), src.ffn_up_exps_s, n_cache, format("blk.%u.ffn_up_exps_s.hot_cache.%zu", il, lane_index).c_str()); + dst_lane.ffn_down_exps_s = new_tensor_like_scale (ctx.get(), src.ffn_down_exps_s, n_cache, format("blk.%u.ffn_down_exps_s.hot_cache.%zu", il, lane_index).c_str()); + + dst_lane.hot_id_map = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_I32, 1, n_expert); + ggml_format_name(dst_lane.hot_id_map, "blk.%u.moe_hot_id_map.%zu", il, lane_index); + dst_lane.hot_mask = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_F32, 1, n_expert); + ggml_format_name(dst_lane.hot_mask, "blk.%u.moe_hot_mask.%zu", il, lane_index); + dst_lane.cold_mask = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_F32, 1, n_expert); + ggml_format_name(dst_lane.cold_mask, "blk.%u.moe_cold_mask.%zu", il, lane_index); + } + + ggml_backend_buffer_type_t buft = ggml_backend_dev_buffer_type(cache_devs[lane_index]); + ggml_backend_buffer_ptr buf { ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft) }; + if (!buf) { + throw std::runtime_error("failed to allocate MoE hot-cache lane buffer"); + } + ggml_backend_buffer_set_usage(buf.get(), GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + + for (uint32_t il = 0; il < selected_by_layer.size(); ++il) { + const auto & experts = selected_by_layer[il]; + if (experts.empty()) { + continue; + } + + const auto & src = model.layers[il]; + auto & cache_layer = cache->layers[il]; + auto & dst_lane = cache_layer.lanes[lane_index]; + const uint32_t dummy_id = experts.size(); + const uint32_t n_expert = dst_lane.n_expert; + + if (dst_lane.ffn_gate_up_exps) { zero_tensor(dst_lane.ffn_gate_up_exps); } + if (dst_lane.ffn_gate_exps) { zero_tensor(dst_lane.ffn_gate_exps); } + if (dst_lane.ffn_up_exps) { zero_tensor(dst_lane.ffn_up_exps); } + if (dst_lane.ffn_down_exps) { zero_tensor(dst_lane.ffn_down_exps); } + if (dst_lane.ffn_gate_exps_s) { zero_tensor(dst_lane.ffn_gate_exps_s); } + if (dst_lane.ffn_up_exps_s) { zero_tensor(dst_lane.ffn_up_exps_s); } + if (dst_lane.ffn_down_exps_s) { zero_tensor(dst_lane.ffn_down_exps_s); } + + std::vector hot_id_map(n_expert, int32_t(dummy_id)); + std::vector hot_mask(n_expert, 0.0f); + std::vector cold_mask(n_expert, 1.0f); + dst_lane.hot_id_map_host.assign(n_expert, -1); + + for (uint32_t cache_id = 0; cache_id < experts.size(); ++cache_id) { + const uint32_t expert = experts[cache_id]; + if (expert >= n_expert) { + continue; + } + + hot_id_map[expert] = int32_t(cache_id); + hot_mask[expert] = 1.0f; + cold_mask[expert] = 0.0f; + dst_lane.hot_id_map_host[expert] = int32_t(cache_id); + cache_layer.expert_lane_map_host[expert] = int32_t(lane_index); + + llama_moe_hot_cache_copy_expert_slice(src.ffn_gate_up_exps, dst_lane.ffn_gate_up_exps, expert, cache_id); + llama_moe_hot_cache_copy_expert_slice(src.ffn_gate_exps, dst_lane.ffn_gate_exps, expert, cache_id); + llama_moe_hot_cache_copy_expert_slice(src.ffn_up_exps, dst_lane.ffn_up_exps, expert, cache_id); + llama_moe_hot_cache_copy_expert_slice(src.ffn_down_exps, dst_lane.ffn_down_exps, expert, cache_id); + llama_moe_hot_cache_copy_scale_slice(src.ffn_gate_exps_s, dst_lane.ffn_gate_exps_s, expert, cache_id); + llama_moe_hot_cache_copy_scale_slice(src.ffn_up_exps_s, dst_lane.ffn_up_exps_s, expert, cache_id); + llama_moe_hot_cache_copy_scale_slice(src.ffn_down_exps_s, dst_lane.ffn_down_exps_s, expert, cache_id); + } + + ggml_backend_tensor_set(dst_lane.hot_id_map, hot_id_map.data(), 0, hot_id_map.size()*sizeof(hot_id_map[0])); + ggml_backend_tensor_set(dst_lane.hot_mask, hot_mask.data(), 0, hot_mask.size()*sizeof(hot_mask[0])); + ggml_backend_tensor_set(dst_lane.cold_mask, cold_mask.data(), 0, cold_mask.size()*sizeof(cold_mask[0])); + } + + LLAMA_LOG_WARN("%s: %12s hot-cache lane %zu buffer size = %8.2f MiB\n", + __func__, + ggml_backend_buffer_name(buf.get()), + lane_index, + ggml_backend_buffer_get_size(buf.get())/1024.0/1024.0); + + cache->bufs.emplace_back(std::move(buf)); + cache->ctxs.emplace_back(std::move(ctx)); + } + + for (auto & layer : cache->layers) { + for (const auto & lane : layer.lanes) { + if (lane.active()) { + static_cast(layer) = lane; + break; + } + } + } + + return cache; +} diff --git a/src/moe-hot-cache/llama-moe-hot-cache-builder.h b/src/moe-hot-cache/llama-moe-hot-cache-builder.h index fdddb6c6c90..1d80532f375 100644 --- a/src/moe-hot-cache/llama-moe-hot-cache-builder.h +++ b/src/moe-hot-cache/llama-moe-hot-cache-builder.h @@ -46,3 +46,8 @@ std::unique_ptr llama_moe_hot_cache_build( const llama_model & model, const llama_moe_hot_cache_plan & plan, ggml_backend_dev_t cache_dev); + +std::unique_ptr llama_moe_hot_cache_build_multi( + const llama_model & model, + const llama_moe_hot_cache_multi_plan & plan, + const std::vector & cache_devs); diff --git a/src/moe-hot-cache/llama-moe-hot-cache-graph.cpp b/src/moe-hot-cache/llama-moe-hot-cache-graph.cpp index b501a7ebb31..8c50e4952a6 100644 --- a/src/moe-hot-cache/llama-moe-hot-cache-graph.cpp +++ b/src/moe-hot-cache/llama-moe-hot-cache-graph.cpp @@ -473,6 +473,316 @@ static ggml_tensor * llama_moe_hot_cache_build_moe_ffn_with_ids( return moe_out; } +static ggml_backend_t llama_moe_hot_cache_backend_for_dev( + ggml_backend_sched_t sched, + ggml_backend_dev_t dev) { + if (sched == nullptr || dev == nullptr) { + return nullptr; + } + + const int n_backends = ggml_backend_sched_get_n_backends(sched); + for (int i = 0; i < n_backends; ++i) { + ggml_backend_t backend = ggml_backend_sched_get_backend(sched, i); + if (ggml_backend_get_device(backend) == dev) { + return backend; + } + } + + return nullptr; +} + +static ggml_backend_t llama_moe_hot_cache_primary_merge_backend( + ggml_backend_sched_t sched, + const llama_model & model, + int il) { + // Worker lanes are expert-only. The final lane join follows the normal + // layer placement, which is the primary graph device for split-mode none. + return llama_moe_hot_cache_backend_for_dev(sched, model.dev_layer(il)); +} + +static int32_t llama_moe_hot_cache_lane_id_field(size_t lane) { + switch (lane) { + case 0: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT_ID; + case 1: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_ID; + case 2: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT2_ID; + } + GGML_ABORT("invalid MoE hot-cache lane"); +} + +static int32_t llama_moe_hot_cache_lane_weight_field(size_t lane) { + switch (lane) { + case 0: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT_WEIGHT; + case 1: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_WEIGHT; + case 2: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT2_WEIGHT; + } + GGML_ABORT("invalid MoE hot-cache lane"); +} + +static int32_t llama_moe_hot_cache_lane_count_field(size_t lane) { + switch (lane) { + case 0: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT_COUNT; + case 1: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_COUNT; + case 2: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT2_COUNT; + } + GGML_ABORT("invalid MoE hot-cache lane"); +} + +static bool llama_moe_hot_cache_layer_has_cold_experts( + const llama_moe_hot_cache_layer & cache) { + if (cache.expert_lane_map_host.empty()) { + return cache.n_hot != cache.n_expert; + } + + for (int32_t lane : cache.expert_lane_map_host) { + if (lane < 0) { + return true; + } + } + + return false; +} + +static ggml_tensor * llama_moe_hot_cache_build_moe_hot_multi_from_logits( + const llm_graph_context & graph, + const llama_model & model, + ggml_tensor * cur, + ggml_tensor * logits, + int il, + const llama_moe_hot_cache_model_adapter & adapter) { + ggml_context * ctx0 = graph.ctx0; + ggml_cgraph * gf = graph.gf; + ggml_backend_sched_t sched = graph.sched; + const llama_hparams & hparams = graph.hparams; + const llama_cparams & cparams = graph.cparams; + const int64_t n_embd = cur->ne[0]; + const int64_t n_tokens = cur->ne[1]; + const int64_t n_expert = graph.n_expert; + const int64_t n_moe_slots = cparams.warmup ? hparams.n_expert_used : graph.n_expert_used; + const auto & layer = model.layers[il]; + const auto & cache = model.moe_hot_cache->layers[il]; + const auto profile = adapter.profile(); + + GGML_ASSERT(adapter.graph_kind != llama_moe_hot_cache_graph_kind::none); + GGML_ASSERT(!cache.lanes.empty()); + GGML_ASSERT(cache.lanes.size() <= LLAMA_MOE_HOT_CACHE_MAX_EXPERT_LANES); + GGML_ASSERT(n_tokens == 1); + GGML_ASSERT(n_moe_slots > 0); + GGML_ASSERT(n_moe_slots <= LLAMA_MAX_EXPERTS); + + const int64_t capacity = n_moe_slots*n_tokens; + ggml_tensor * worklist_shape = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, capacity, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COUNT); + ggml_tensor * worklist = nullptr; + + const bool cpu_decode_routing = + profile.cpu_decode_routing && + n_tokens >= 1 && n_tokens <= profile.cpu_decode_routing_max_tokens; + if (cpu_decode_routing) { + worklist = ggml_map_custom2( + ctx0, + worklist_shape, + logits, + llama_moe_hot_cache_select_worklist_from_logits_op(llama_moe_hot_cache_worklist_order::token_major), + 1, + const_cast(&cache)); + } else { + ggml_tensor * selected_experts = ggml_argsort_top_k(ctx0, logits, n_moe_slots); + graph.cb(selected_experts->src[0], "ffn_moe_argsort", il); + graph.cb(selected_experts, "ffn_moe_topk", il); + + ggml_tensor * logits_rows = ggml_reshape_3d(ctx0, logits, 1, n_expert, n_tokens); + ggml_tensor * weights = ggml_get_rows(ctx0, logits_rows, selected_experts); + graph.cb(weights, "ffn_moe_weights", il); + + weights = ggml_reshape_2d(ctx0, weights, n_moe_slots, n_tokens); + weights = ggml_soft_max(ctx0, weights); + graph.cb(weights, "ffn_moe_weights_norm", il); + + weights = ggml_reshape_3d(ctx0, weights, 1, n_moe_slots, n_tokens); + if (hparams.expert_weights_scale != 0.0f && hparams.expert_weights_scale != 1.0f) { + weights = ggml_scale(ctx0, weights, hparams.expert_weights_scale); + graph.cb(weights, "ffn_moe_weights_scaled", il); + } + + worklist = ggml_map_custom3( + ctx0, + worklist_shape, + selected_experts, + weights, + llama_moe_hot_cache_select_worklist_op(llama_moe_hot_cache_worklist_order::token_major), + 1, + const_cast(&cache)); + } + graph.cb(worklist, "ffn_moe_worklist", il); + + const auto view_worklist_field = [&](int32_t field) { + return ggml_view_1d(ctx0, worklist, capacity, field*worklist->nb[1]); + }; + const auto view_worklist_count = [&](int32_t field) { + return ggml_view_1d(ctx0, worklist, 1, field*worklist->nb[1]); + }; + + const auto merge_compact_slots = [&]( + ggml_tensor * branch_out, + ggml_backend_t branch_backend, + const char * name) { + ggml_tensor * merged = ggml_reshape_3d(ctx0, branch_out, n_embd, capacity, 1); + if (branch_backend != nullptr) { + ggml_backend_sched_set_tensor_backend(sched, merged, branch_backend); + } + merged = ggml_permute(ctx0, merged, 1, 0, 2, 3); + if (branch_backend != nullptr) { + ggml_backend_sched_set_tensor_backend(sched, merged, branch_backend); + } + if (!profile.decode_strided_sum_rows) { + merged = ggml_cont(ctx0, merged); + if (branch_backend != nullptr) { + ggml_backend_sched_set_tensor_backend(sched, merged, branch_backend); + } + } + merged = ggml_sum_rows(ctx0, merged); + if (branch_backend != nullptr) { + ggml_backend_sched_set_tensor_backend(sched, merged, branch_backend); + } + merged = ggml_reshape_2d(ctx0, merged, n_embd, n_tokens); + if (branch_backend != nullptr) { + ggml_backend_sched_set_tensor_backend(sched, merged, branch_backend); + } + graph.cb(merged, name, il); + ggml_build_forward_expand(gf, merged); + return merged; + }; + + const uint32_t hot_mul_mat_id_flags = llama_moe_hot_cache_graph_tweaks::hot_dummy_padding() + ? LLAMA_MOE_HOT_CACHE_MUL_MAT_ID_FLAG_NONE + : LLAMA_MOE_HOT_CACHE_MUL_MAT_ID_FLAG_ALLOW_NEGATIVE_IDS; + + std::vector branch_outputs; + branch_outputs.reserve(cache.lanes.size() + 1); + + for (size_t lane_index = 0; lane_index < cache.lanes.size(); ++lane_index) { + const auto & lane = cache.lanes[lane_index]; + if (!lane.active()) { + continue; + } + + ggml_tensor * lane_count = view_worklist_count(llama_moe_hot_cache_lane_count_field(lane_index)); + graph.cb(lane_count, format("ffn_moe_hot%zu_count", lane_index).c_str(), il); + ggml_build_forward_expand(gf, lane_count); + + ggml_tensor * lane_ids = ggml_cast(ctx0, view_worklist_field(llama_moe_hot_cache_lane_id_field(lane_index)), GGML_TYPE_I32); + lane_ids = ggml_reshape_2d(ctx0, lane_ids, 1, capacity); + graph.cb(lane_ids, format("ffn_moe_hot%zu_ids_compact", lane_index).c_str(), il); + ggml_build_forward_expand(gf, lane_ids); + + ggml_tensor * lane_weights = ggml_reshape_3d(ctx0, view_worklist_field(llama_moe_hot_cache_lane_weight_field(lane_index)), 1, 1, capacity); + graph.cb(lane_weights, format("ffn_moe_hot%zu_weights_compact", lane_index).c_str(), il); + ggml_build_forward_expand(gf, lane_weights); + + ggml_backend_t lane_backend = nullptr; + if (lane_index < model.moe_hot_cache->devices.size()) { + lane_backend = llama_moe_hot_cache_backend_for_dev(sched, model.moe_hot_cache->devices[lane_index]); + } + + ggml_tensor * lane_inputs = ggml_repeat_4d(ctx0, cur, n_embd, capacity, 1, 1); + if (lane_backend != nullptr) { + ggml_backend_sched_set_tensor_backend(sched, lane_inputs, lane_backend); + } + graph.cb(lane_inputs, format("ffn_moe_hot%zu_inputs", lane_index).c_str(), il); + + const std::string branch_name = format("hot%zu", lane_index); + ggml_tensor * lane_out = llama_moe_hot_cache_build_moe_ffn_with_ids( + graph, + lane_inputs, + lane_ids, + lane_weights, + lane.ffn_up_exps, + lane.ffn_gate_exps, + lane.ffn_down_exps, + lane.n_hot + 1, + 1, + adapter.ffn_op, + il, + lane.ffn_gate_up_exps, + lane.ffn_up_exps_s, + lane.ffn_gate_exps_s, + lane.ffn_down_exps_s, + hot_mul_mat_id_flags, + branch_name.c_str(), + lane_backend); + graph.cb(lane_out, format("ffn_moe_hot%zu_out", lane_index).c_str(), il); + + branch_outputs.push_back(merge_compact_slots( + lane_out, + lane_backend, + format("ffn_moe_hot%zu_slots", lane_index).c_str())); + } + + if (llama_moe_hot_cache_layer_has_cold_experts(cache)) { + ggml_backend_t cold_backend = cparams.warmup ? nullptr : graph.backend_cpu; + ggml_tensor * cold_ids = ggml_cast(ctx0, view_worklist_field(LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COLD_ID), GGML_TYPE_I32); + cold_ids = ggml_reshape_2d(ctx0, cold_ids, 1, capacity); + graph.cb(cold_ids, "ffn_moe_cold_ids_compact", il); + ggml_build_forward_expand(gf, cold_ids); + + ggml_tensor * cold_weights = ggml_reshape_3d(ctx0, view_worklist_field(LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COLD_WEIGHT), 1, 1, capacity); + graph.cb(cold_weights, "ffn_moe_cold_weights_compact", il); + ggml_build_forward_expand(gf, cold_weights); + + ggml_tensor * cold_count = view_worklist_count(LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COLD_COUNT); + graph.cb(cold_count, "ffn_moe_cold_count", il); + ggml_build_forward_expand(gf, cold_count); + + ggml_tensor * cold_inputs = ggml_repeat_4d(ctx0, cur, n_embd, capacity, 1, 1); + if (cold_backend != nullptr) { + ggml_backend_sched_set_tensor_backend(sched, cold_inputs, cold_backend); + } + graph.cb(cold_inputs, "ffn_moe_cold_inputs", il); + + const uint32_t cold_mul_mat_id_flags = + LLAMA_MOE_HOT_CACHE_MUL_MAT_ID_FLAG_ALLOW_NEGATIVE_IDS | + LLAMA_MOE_HOT_CACHE_MUL_MAT_ID_FLAG_SKIP_NEGATIVE_ID_OUTPUT_ZERO; + ggml_tensor * cold_out = llama_moe_hot_cache_build_moe_ffn_with_ids( + graph, + cold_inputs, + cold_ids, + cold_weights, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + n_expert, + 1, + adapter.ffn_op, + il, + layer.ffn_gate_up_exps, + layer.ffn_up_exps_s, + layer.ffn_gate_exps_s, + layer.ffn_down_exps_s, + cold_mul_mat_id_flags, + "cold", + cold_backend); + graph.cb(cold_out, "ffn_moe_cold_out", il); + branch_outputs.push_back(merge_compact_slots(cold_out, cold_backend, "ffn_moe_cold_slots")); + } + + GGML_ASSERT(!branch_outputs.empty()); + + ggml_backend_t merge_backend = llama_moe_hot_cache_primary_merge_backend(sched, model, il); + ggml_tensor * out = branch_outputs[0]; + for (size_t i = 1; i < branch_outputs.size(); ++i) { + out = ggml_add(ctx0, out, branch_outputs[i]); + if (merge_backend != nullptr) { + ggml_backend_sched_set_tensor_backend(sched, out, merge_backend); + } + ggml_build_forward_expand(gf, out); + } + + if (merge_backend != nullptr) { + ggml_backend_sched_set_tensor_backend(sched, out, merge_backend); + } + graph.cb(out, "ffn_moe_out", il); + return out; +} + static ggml_tensor * llama_moe_hot_cache_build_moe_hot_from_logits( const llm_graph_context & graph, const llama_model & model, @@ -500,6 +810,9 @@ static ggml_tensor * llama_moe_hot_cache_build_moe_hot_from_logits( const bool is_warmup_phase = graph_phase == llama_moe_hot_cache_graph_phase::warmup; GGML_ASSERT(adapter.graph_kind == llama_moe_hot_cache_graph_kind::logits); + if (!cache.lanes.empty()) { + return llama_moe_hot_cache_build_moe_hot_multi_from_logits(graph, model, cur, logits, il, adapter); + } GGML_ASSERT(!cache.hot_id_map_host.empty()); GGML_ASSERT(n_moe_slots > 0); GGML_ASSERT(n_moe_slots <= LLAMA_MAX_EXPERTS); @@ -998,13 +1311,18 @@ ggml_tensor * llama_model_qwen35moe::graph::build_layer_ffn_hot(ggml_tensor * cu const bool is_decode_phase = graph_phase == llama_moe_hot_cache_graph_phase::decode; const bool is_warmup_phase = graph_phase == llama_moe_hot_cache_graph_phase::warmup; - GGML_ASSERT(!cache.hot_id_map_host.empty()); GGML_ASSERT(n_moe_slots > 0); GGML_ASSERT(n_moe_slots <= LLAMA_MAX_EXPERTS); ggml_tensor * logits = build_lora_mm(layer.ffn_gate_inp, cur); cb(logits, "ffn_moe_logits", il); + if (!cache.lanes.empty()) { + return llama_moe_hot_cache_build_moe_hot_multi_from_logits(*this, model, cur, logits, il, adapter); + } + + GGML_ASSERT(!cache.hot_id_map_host.empty()); + const int64_t capacity = n_moe_slots*n_tokens; const llama_moe_hot_cache_pp_execution_plan pp_plan = llama_moe_hot_cache_pp_policy::build( graph_phase, diff --git a/src/moe-hot-cache/llama-moe-hot-cache-planner.cpp b/src/moe-hot-cache/llama-moe-hot-cache-planner.cpp index fb25c39ccec..6e10884176c 100644 --- a/src/moe-hot-cache/llama-moe-hot-cache-planner.cpp +++ b/src/moe-hot-cache/llama-moe-hot-cache-planner.cpp @@ -2,8 +2,10 @@ #include "llama-model.h" +#include #include #include +#include #include #include @@ -13,8 +15,112 @@ static uint64_t key(uint32_t layer, uint32_t expert) { return (uint64_t(layer) << 32) | uint64_t(expert); } +struct lane_state { + std::unordered_set active_layers; + std::unordered_map selected_by_layer; +}; + +static bool select_entry_into_lane( + llama_moe_hot_cache_plan & lane, + lane_state & state, + const llama_moe_hot_cache_entry & entry, + size_t bytes) { + size_t cost = bytes; + if (state.active_layers.find(entry.layer) == state.active_layers.end()) { + if (cost > std::numeric_limits::max() - bytes) { + return false; + } + cost += bytes; + } + + if (cost > lane.budget_bytes || lane.used_bytes > lane.budget_bytes - cost) { + return false; + } + + lane.selected.push_back({ entry.layer, entry.expert, bytes }); + lane.used_bytes += cost; + state.active_layers.insert(entry.layer); + state.selected_by_layer[entry.layer]++; + return true; +} + +static size_t best_hot_even_lane( + const std::vector & lanes, + const std::vector & states, + const llama_moe_hot_cache_entry & entry, + const std::vector & candidates) { + size_t best = std::numeric_limits::max(); + for (size_t lane : candidates) { + if (lane >= lanes.size()) { + continue; + } + + if (best == std::numeric_limits::max()) { + best = lane; + continue; + } + + const size_t lane_layer_count = states[lane].selected_by_layer.count(entry.layer) + ? states[lane].selected_by_layer.at(entry.layer) + : 0; + const size_t best_layer_count = states[best].selected_by_layer.count(entry.layer) + ? states[best].selected_by_layer.at(entry.layer) + : 0; + + if (lane_layer_count != best_layer_count) { + if (lane_layer_count < best_layer_count) { + best = lane; + } + continue; + } + + if (lanes[lane].used_bytes != lanes[best].used_bytes) { + if (lanes[lane].used_bytes < lanes[best].used_bytes) { + best = lane; + } + continue; + } + + if (lane < best) { + best = lane; + } + } + + return best; +} + } // namespace +size_t llama_moe_hot_cache_multi_plan::selected_count() const { + size_t result = 0; + for (const auto & lane : lanes) { + result += lane.selected.size(); + } + return result; +} + +size_t llama_moe_hot_cache_multi_plan::used_bytes() const { + size_t result = 0; + for (const auto & lane : lanes) { + if (result > std::numeric_limits::max() - lane.used_bytes) { + return std::numeric_limits::max(); + } + result += lane.used_bytes; + } + return result; +} + +size_t llama_moe_hot_cache_multi_plan::budget_bytes() const { + size_t result = 0; + for (const auto & lane : lanes) { + if (result > std::numeric_limits::max() - lane.budget_bytes) { + return std::numeric_limits::max(); + } + result += lane.budget_bytes; + } + return result; +} + size_t llama_moe_hot_cache_tensor_expert_bytes(const ggml_tensor * t) { if (t == nullptr) { return 0; @@ -97,3 +203,92 @@ llama_moe_hot_cache_plan llama_moe_hot_cache_select( return plan; } + +llama_moe_hot_cache_device_strategy llama_moe_hot_cache_parse_device_strategy( + const char * name) { + if (name == nullptr || name[0] == '\0' || std::string(name) == "warm") { + return llama_moe_hot_cache_device_strategy::warm; + } + if (std::string(name) == "hot-even") { + return llama_moe_hot_cache_device_strategy::hot_even; + } + + throw std::runtime_error("--moe-hot-cache-device-strategy must be one of: warm, hot-even"); +} + +llama_moe_hot_cache_multi_plan llama_moe_hot_cache_select_multi( + const std::vector & observed, + const std::vector & sizes, + const std::vector & lane_budget_bytes, + llama_moe_hot_cache_device_strategy strategy) { + if (lane_budget_bytes.empty() || lane_budget_bytes.size() > LLAMA_MOE_HOT_CACHE_MAX_EXPERT_LANES) { + throw std::runtime_error("MoE hot-cache multi-device planner requires 1..3 expert lanes"); + } + + llama_moe_hot_cache_multi_plan plan; + plan.observed = observed; + plan.lanes.resize(lane_budget_bytes.size()); + for (size_t i = 0; i < plan.lanes.size(); ++i) { + plan.lanes[i].observed = observed; + plan.lanes[i].budget_bytes = lane_budget_bytes[i]; + } + + std::unordered_map size_by_expert; + size_by_expert.reserve(sizes.size()); + for (const auto & size : sizes) { + size_by_expert[key(size.layer, size.expert)] = size.bytes; + } + + std::unordered_set selected; + std::vector states(plan.lanes.size()); + + for (const auto & entry : observed) { + const uint64_t entry_key = key(entry.layer, entry.expert); + if (selected.find(entry_key) != selected.end()) { + continue; + } + + const auto size_it = size_by_expert.find(entry_key); + if (size_it == size_by_expert.end()) { + continue; + } + + std::vector lane_order(plan.lanes.size()); + for (size_t i = 0; i < lane_order.size(); ++i) { + lane_order[i] = i; + } + + if (strategy == llama_moe_hot_cache_device_strategy::hot_even) { + std::sort(lane_order.begin(), lane_order.end(), [&](size_t a, size_t b) { + const size_t a_count = states[a].selected_by_layer.count(entry.layer) + ? states[a].selected_by_layer.at(entry.layer) + : 0; + const size_t b_count = states[b].selected_by_layer.count(entry.layer) + ? states[b].selected_by_layer.at(entry.layer) + : 0; + if (a_count != b_count) { + return a_count < b_count; + } + if (plan.lanes[a].used_bytes != plan.lanes[b].used_bytes) { + return plan.lanes[a].used_bytes < plan.lanes[b].used_bytes; + } + return a < b; + }); + + const size_t preferred = best_hot_even_lane(plan.lanes, states, entry, lane_order); + if (preferred != std::numeric_limits::max()) { + lane_order.erase(std::remove(lane_order.begin(), lane_order.end(), preferred), lane_order.end()); + lane_order.insert(lane_order.begin(), preferred); + } + } + + for (size_t lane : lane_order) { + if (select_entry_into_lane(plan.lanes[lane], states[lane], entry, size_it->second)) { + selected.insert(entry_key); + break; + } + } + } + + return plan; +} diff --git a/src/moe-hot-cache/llama-moe-hot-cache-planner.h b/src/moe-hot-cache/llama-moe-hot-cache-planner.h index 23da969fee8..83fe7333a30 100644 --- a/src/moe-hot-cache/llama-moe-hot-cache-planner.h +++ b/src/moe-hot-cache/llama-moe-hot-cache-planner.h @@ -8,3 +8,12 @@ size_t llama_moe_hot_cache_tensor_expert_bytes(const ggml_tensor * t); std::vector llama_moe_hot_cache_collect_expert_sizes( const llama_model & model); + +llama_moe_hot_cache_device_strategy llama_moe_hot_cache_parse_device_strategy( + const char * name); + +llama_moe_hot_cache_multi_plan llama_moe_hot_cache_select_multi( + const std::vector & observed, + const std::vector & sizes, + const std::vector & lane_budget_bytes, + llama_moe_hot_cache_device_strategy strategy); diff --git a/src/moe-hot-cache/llama-moe-hot-cache-updater.cpp b/src/moe-hot-cache/llama-moe-hot-cache-updater.cpp index 1ab4f7587f4..bfe291770f6 100644 --- a/src/moe-hot-cache/llama-moe-hot-cache-updater.cpp +++ b/src/moe-hot-cache/llama-moe-hot-cache-updater.cpp @@ -2,6 +2,7 @@ #include "llama-moe-hot-cache-builder.h" +#include "llama-impl.h" #include "llama-model.h" #include @@ -154,6 +155,22 @@ llama_moe_hot_cache_update_stats llama_moe_hot_cache_update_from_scored_observat return stats; } + const bool has_multi_lane_cache = std::any_of( + model.moe_hot_cache->layers.begin(), + model.moe_hot_cache->layers.end(), + [](const llama_moe_hot_cache_layer & layer) { + return !layer.lanes.empty(); + }); + if (has_multi_lane_cache) { + static bool logged = false; + if (!logged) { + LLAMA_LOG_WARN("%s: runtime hot-cache replacement is not yet supported for multi-device expert lanes\n", __func__); + logged = true; + } + stats.active = true; + return stats; + } + stats.active = true; struct layer_counts { diff --git a/src/moe-hot-cache/llama-moe-hot-cache-worklist.cpp b/src/moe-hot-cache/llama-moe-hot-cache-worklist.cpp index cbf66afbe7e..30584c7c2bd 100644 --- a/src/moe-hot-cache/llama-moe-hot-cache-worklist.cpp +++ b/src/moe-hot-cache/llama-moe-hot-cache-worklist.cpp @@ -21,6 +21,338 @@ static bool llama_moe_hot_cache_hot_dummy_padding() { return enabled; } +static int32_t lane_id_field(size_t lane) { + switch (lane) { + case 0: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT_ID; + case 1: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_ID; + case 2: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT2_ID; + } + GGML_ABORT("invalid MoE hot-cache lane"); +} + +static int32_t lane_src_slot_field(size_t lane) { + switch (lane) { + case 0: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT_SRC_SLOT; + case 1: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_SRC_SLOT; + case 2: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT2_SRC_SLOT; + } + GGML_ABORT("invalid MoE hot-cache lane"); +} + +static int32_t lane_token_id_field(size_t lane) { + switch (lane) { + case 0: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT_TOKEN_ID; + case 1: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_TOKEN_ID; + case 2: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT2_TOKEN_ID; + } + GGML_ABORT("invalid MoE hot-cache lane"); +} + +static int32_t lane_weight_field(size_t lane) { + switch (lane) { + case 0: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT_WEIGHT; + case 1: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_WEIGHT; + case 2: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT2_WEIGHT; + } + GGML_ABORT("invalid MoE hot-cache lane"); +} + +static int32_t lane_expert_id_field(size_t lane) { + switch (lane) { + case 0: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT_EXPERT_ID; + case 1: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_EXPERT_ID; + case 2: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT2_EXPERT_ID; + } + GGML_ABORT("invalid MoE hot-cache lane"); +} + +static int32_t lane_count_field(size_t lane) { + switch (lane) { + case 0: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT_COUNT; + case 1: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_COUNT; + case 2: return LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT2_COUNT; + } + GGML_ABORT("invalid MoE hot-cache lane"); +} + +static int32_t lane_hot_id_for_expert( + const llama_moe_hot_cache_layer & layer, + int32_t expert, + size_t & lane_index) { + if (expert < 0 || expert >= int32_t(layer.expert_lane_map_host.size())) { + return -1; + } + + const int32_t mapped_lane = layer.expert_lane_map_host[expert]; + if (mapped_lane < 0 || size_t(mapped_lane) >= layer.lanes.size() || + size_t(mapped_lane) >= LLAMA_MOE_HOT_CACHE_MAX_EXPERT_LANES) { + return -1; + } + + const auto & lane = layer.lanes[size_t(mapped_lane)]; + if (expert >= int32_t(lane.hot_id_map_host.size())) { + return -1; + } + + const int32_t hot_id = lane.hot_id_map_host[expert]; + if (hot_id < 0 || uint32_t(hot_id) >= lane.n_hot) { + return -1; + } + + lane_index = size_t(mapped_lane); + return hot_id; +} + +static void build_worklist_multi_from_selected( + ggml_tensor * dst, + const ggml_tensor * selected_experts, + const ggml_tensor * weights, + const llama_moe_hot_cache_layer & layer, + llama_moe_hot_cache_worklist_order order) { + GGML_ASSERT(dst != nullptr); + GGML_ASSERT(selected_experts != nullptr); + GGML_ASSERT(weights != nullptr); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(selected_experts->type == GGML_TYPE_I32); + GGML_ASSERT(weights->type == GGML_TYPE_F32); + GGML_ASSERT(dst->nb[0] == sizeof(float)); + GGML_ASSERT(dst->ne[1] == LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COUNT); + GGML_ASSERT(weights->ne[0] == 1); + GGML_ASSERT(selected_experts->ne[0] == weights->ne[1]); + GGML_ASSERT(selected_experts->ne[1] == weights->ne[2]); + GGML_ASSERT(layer.lanes.size() <= LLAMA_MOE_HOT_CACHE_MAX_EXPERT_LANES); + GGML_ASSERT(int64_t(layer.expert_lane_map_host.size()) == layer.n_expert); + + const int32_t capacity = dst->ne[0]; + const int32_t n_expert_used = selected_experts->ne[0]; + const int32_t n_tokens = selected_experts->ne[1]; + const int32_t total_slots = n_expert_used * n_tokens; + const int32_t dummy_src_slot = total_slots; + GGML_ASSERT(capacity == total_slots); + + auto set_field = [&](int32_t field, int32_t slot, float value) { + char * row = (char *) dst->data + field*dst->nb[1]; + *(float *)(row + slot*dst->nb[0]) = value; + }; + + auto fill_field = [&](int32_t field, float value) { + float * row = (float *) ((char *) dst->data + field*dst->nb[1]); + std::fill(row, row + capacity, value); + }; + + for (size_t lane = 0; lane < LLAMA_MOE_HOT_CACHE_MAX_EXPERT_LANES; ++lane) { + const bool lane_active = lane < layer.lanes.size() && layer.lanes[lane].n_hot > 0; + const float hot_padding_id = + llama_moe_hot_cache_hot_dummy_padding() && lane_active ? float(layer.lanes[lane].n_hot) : -1.0f; + fill_field(lane_id_field(lane), hot_padding_id); + fill_field(lane_src_slot_field(lane), float(dummy_src_slot)); + fill_field(lane_token_id_field(lane), 0.0f); + fill_field(lane_weight_field(lane), 0.0f); + fill_field(lane_expert_id_field(lane), -1.0f); + fill_field(lane_count_field(lane), 0.0f); + } + fill_field(LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COLD_ID, -1.0f); + fill_field(LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COLD_SRC_SLOT, float(dummy_src_slot)); + fill_field(LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COLD_TOKEN_ID, 0.0f); + fill_field(LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COLD_WEIGHT, 0.0f); + fill_field(LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COLD_COUNT, 0.0f); + + int32_t hot_slots[LLAMA_MOE_HOT_CACHE_MAX_EXPERT_LANES] = {}; + int32_t cold_slot = 0; + + const auto selected_expert_at = [&](int32_t iex, int32_t token) { + return *(const int32_t *) ((const char *) selected_experts->data + iex*selected_experts->nb[0] + token*selected_experts->nb[1]); + }; + + const auto weight_at = [&](int32_t iex, int32_t token) { + return *(const float *) ((const char *) weights->data + iex*weights->nb[1] + token*weights->nb[2]); + }; + + const auto write_hot = [&](size_t lane, int32_t slot, int32_t token, int32_t iex, int32_t expert, int32_t hot_id) { + const float weight = weight_at(iex, token); + const int32_t src_slot = token*n_expert_used + iex; + set_field(lane_id_field(lane), slot, float(hot_id)); + set_field(lane_src_slot_field(lane), slot, float(src_slot)); + set_field(lane_token_id_field(lane), slot, float(token)); + set_field(lane_weight_field(lane), slot, weight); + set_field(lane_expert_id_field(lane), slot, float(expert)); + }; + + const auto write_cold = [&](int32_t slot, int32_t token, int32_t iex, int32_t expert) { + const float weight = weight_at(iex, token); + const int32_t src_slot = token*n_expert_used + iex; + set_field(LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COLD_ID, slot, float(expert)); + set_field(LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COLD_SRC_SLOT, slot, float(src_slot)); + set_field(LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COLD_TOKEN_ID, slot, float(token)); + set_field(LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COLD_WEIGHT, slot, weight); + }; + + if (order == llama_moe_hot_cache_worklist_order::expert_major) { + int32_t hot_offsets[LLAMA_MOE_HOT_CACHE_MAX_EXPERT_LANES][LLAMA_MAX_EXPERTS] = {}; + int32_t cold_offsets[LLAMA_MAX_EXPERTS] = {}; + + for (size_t lane = 0; lane < layer.lanes.size(); ++lane) { + GGML_ASSERT(layer.lanes[lane].n_hot <= LLAMA_MAX_EXPERTS); + } + GGML_ASSERT(layer.n_expert <= LLAMA_MAX_EXPERTS); + + for (int32_t token = 0; token < n_tokens; ++token) { + for (int32_t iex = 0; iex < n_expert_used; ++iex) { + const int32_t expert = selected_expert_at(iex, token); + GGML_ASSERT(expert >= 0); + GGML_ASSERT(expert < int32_t(layer.expert_lane_map_host.size())); + + size_t lane = 0; + const int32_t hot_id = lane_hot_id_for_expert(layer, expert, lane); + if (hot_id >= 0) { + ++hot_offsets[lane][hot_id]; + } else { + ++cold_offsets[expert]; + } + } + } + + for (size_t lane = 0; lane < layer.lanes.size(); ++lane) { + int32_t running = 0; + for (uint32_t ih = 0; ih < layer.lanes[lane].n_hot; ++ih) { + const int32_t start = running; + running += hot_offsets[lane][ih]; + hot_offsets[lane][ih] = start; + } + hot_slots[lane] = running; + } + + int32_t running = 0; + for (uint32_t expert = 0; expert < layer.n_expert; ++expert) { + const int32_t start = running; + running += cold_offsets[expert]; + cold_offsets[expert] = start; + } + cold_slot = running; + + for (int32_t token = 0; token < n_tokens; ++token) { + for (int32_t iex = 0; iex < n_expert_used; ++iex) { + const int32_t expert = selected_expert_at(iex, token); + size_t lane = 0; + const int32_t hot_id = lane_hot_id_for_expert(layer, expert, lane); + if (hot_id >= 0) { + const int32_t slot = hot_offsets[lane][hot_id]++; + write_hot(lane, slot, token, iex, expert, hot_id); + } else { + const int32_t slot = cold_offsets[expert]++; + write_cold(slot, token, iex, expert); + } + } + } + } else { + for (int32_t token = 0; token < n_tokens; ++token) { + for (int32_t iex = 0; iex < n_expert_used; ++iex) { + const int32_t expert = selected_expert_at(iex, token); + GGML_ASSERT(expert >= 0); + GGML_ASSERT(expert < int32_t(layer.expert_lane_map_host.size())); + + size_t lane = 0; + const int32_t hot_id = lane_hot_id_for_expert(layer, expert, lane); + if (hot_id >= 0) { + write_hot(lane, hot_slots[lane], token, iex, expert, hot_id); + ++hot_slots[lane]; + } else { + write_cold(cold_slot, token, iex, expert); + ++cold_slot; + } + } + } + } + + if (capacity > 0) { + for (size_t lane = 0; lane < LLAMA_MOE_HOT_CACHE_MAX_EXPERT_LANES; ++lane) { + set_field(lane_count_field(lane), 0, float(hot_slots[lane])); + } + set_field(LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COLD_COUNT, 0, float(cold_slot)); + } +} + +static void build_worklist_multi_from_logits( + ggml_tensor * dst, + const ggml_tensor * logits, + const llama_moe_hot_cache_layer & layer, + llama_moe_hot_cache_worklist_order order) { + GGML_ASSERT(dst != nullptr); + GGML_ASSERT(logits != nullptr); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(logits->type == GGML_TYPE_F32); + GGML_ASSERT(dst->nb[0] == sizeof(float)); + GGML_ASSERT(dst->ne[1] == LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COUNT); + GGML_ASSERT(layer.lanes.size() <= LLAMA_MOE_HOT_CACHE_MAX_EXPERT_LANES); + GGML_ASSERT(int64_t(layer.expert_lane_map_host.size()) == layer.n_expert); + GGML_ASSERT(logits->ne[0] == layer.n_expert); + + const int32_t capacity = dst->ne[0]; + const int32_t n_expert = logits->ne[0]; + const int32_t n_tokens = logits->ne[1]; + GGML_ASSERT(n_tokens > 0); + GGML_ASSERT(capacity % n_tokens == 0); + + const int32_t n_expert_used = capacity / n_tokens; + GGML_ASSERT(n_expert_used > 0); + GGML_ASSERT(n_expert_used <= n_expert); + GGML_ASSERT(n_expert_used <= LLAMA_MAX_EXPERTS); + + ggml_init_params params = { + /*.mem_size =*/ size_t(n_expert_used * n_tokens) * (sizeof(int32_t) + sizeof(float)) + 16*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ false, + }; + ggml_context_ptr tmp(ggml_init(params)); + GGML_ASSERT(tmp != nullptr); + + ggml_tensor * selected = ggml_new_tensor_2d(tmp.get(), GGML_TYPE_I32, n_expert_used, n_tokens); + ggml_tensor * weights = ggml_new_tensor_3d(tmp.get(), GGML_TYPE_F32, 1, n_expert_used, n_tokens); + + int32_t top_experts[LLAMA_MAX_EXPERTS]; + float top_logits[LLAMA_MAX_EXPERTS]; + for (int32_t token = 0; token < n_tokens; ++token) { + for (int32_t i = 0; i < n_expert_used; ++i) { + top_experts[i] = -1; + top_logits[i] = -std::numeric_limits::infinity(); + } + + for (int32_t expert = 0; expert < n_expert; ++expert) { + const float logit = *(const float *) ((const char *) logits->data + expert*logits->nb[0] + token*logits->nb[1]); + for (int32_t pos = 0; pos < n_expert_used; ++pos) { + if (logit <= top_logits[pos]) { + continue; + } + for (int32_t move = n_expert_used - 1; move > pos; --move) { + top_logits[move] = top_logits[move - 1]; + top_experts[move] = top_experts[move - 1]; + } + top_logits[pos] = logit; + top_experts[pos] = expert; + break; + } + } + + const float max_logit = top_logits[0]; + float weight_sum = 0.0f; + for (int32_t iex = 0; iex < n_expert_used; ++iex) { + top_logits[iex] = std::exp(top_logits[iex] - max_logit); + weight_sum += top_logits[iex]; + } + + const float weight_scale = + layer.expert_weights_scale != 0.0f && layer.expert_weights_scale != 1.0f ? layer.expert_weights_scale : 1.0f; + + for (int32_t iex = 0; iex < n_expert_used; ++iex) { + *(int32_t *) ((char *) selected->data + iex*selected->nb[0] + token*selected->nb[1]) = top_experts[iex]; + *(float *) ((char *) weights->data + iex*weights->nb[1] + token*weights->nb[2]) = + top_logits[iex] / weight_sum * weight_scale; + } + } + + build_worklist_multi_from_selected(dst, selected, weights, layer, order); +} + } // namespace const char * llama_moe_hot_cache_worklist_order_name(llama_moe_hot_cache_worklist_order order) { @@ -48,6 +380,11 @@ void llama_moe_hot_cache_build_worklist( return; } + if (!layer.lanes.empty()) { + build_worklist_multi_from_selected(dst, selected_experts, weights, layer, order); + return; + } + GGML_ASSERT(dst != nullptr); GGML_ASSERT(selected_experts != nullptr); GGML_ASSERT(weights != nullptr); @@ -222,6 +559,11 @@ void llama_moe_hot_cache_build_worklist_from_logits( return; } + if (!layer.lanes.empty()) { + build_worklist_multi_from_logits(dst, logits, layer, order); + return; + } + GGML_ASSERT(dst != nullptr); GGML_ASSERT(logits != nullptr); GGML_ASSERT(dst->type == GGML_TYPE_F32); diff --git a/src/moe-hot-cache/llama-moe-hot-cache.cpp b/src/moe-hot-cache/llama-moe-hot-cache.cpp index 7dbd486dfc6..7ed288b9b5e 100644 --- a/src/moe-hot-cache/llama-moe-hot-cache.cpp +++ b/src/moe-hot-cache/llama-moe-hot-cache.cpp @@ -15,6 +15,7 @@ #include #include #include +#include namespace { @@ -109,6 +110,48 @@ static std::vector score_observations_for_arch( } } +static bool hot_cache_lane_enabled(int64_t max_mib) { + return max_mib != 0; +} + +static bool hot_cache_any_lane_enabled(const llama_model_params & params) { + return hot_cache_lane_enabled(params.moe_hot_cache_max_mib) || + hot_cache_lane_enabled(params.moe_hot_cache_second_max_mib) || + hot_cache_lane_enabled(params.moe_hot_cache_third_max_mib); +} + +static bool hot_cache_any_static_lane_enabled(const llama_model_params & params) { + return params.moe_hot_cache_max_mib > 0 || + params.moe_hot_cache_second_max_mib > 0 || + params.moe_hot_cache_third_max_mib > 0; +} + +static bool hot_cache_any_auto_lane_enabled(const llama_model_params & params) { + return params.moe_hot_cache_max_mib < 0 || + params.moe_hot_cache_second_max_mib < 0 || + params.moe_hot_cache_third_max_mib < 0; +} + +struct hot_cache_lane_request { + const char * name = nullptr; + int64_t max_mib = 0; + uint64_t reserve_mib = 0; +}; + +static std::vector hot_cache_lane_requests(const llama_model_params & params) { + std::vector lanes; + if (hot_cache_lane_enabled(params.moe_hot_cache_max_mib)) { + lanes.push_back({ params.moe_hot_cache_device, params.moe_hot_cache_max_mib, params.moe_hot_cache_auto_reserve_mib }); + } + if (hot_cache_lane_enabled(params.moe_hot_cache_second_max_mib)) { + lanes.push_back({ params.moe_hot_cache_second_device, params.moe_hot_cache_second_max_mib, params.moe_hot_cache_second_auto_reserve_mib }); + } + if (hot_cache_lane_enabled(params.moe_hot_cache_third_max_mib)) { + lanes.push_back({ params.moe_hot_cache_third_device, params.moe_hot_cache_third_max_mib, params.moe_hot_cache_third_auto_reserve_mib }); + } + return lanes; +} + } // namespace std::vector llama_moe_hot_cache_parse_perf_json_observations( @@ -122,7 +165,7 @@ std::vector llama_moe_hot_cache_parse_perf_json(const } void llama_moe_hot_cache_init(llama_model & model, const llama_model_params & params, bool reserve_kv_cache) { - if (params.moe_hot_cache_max_mib == 0) { + if (!hot_cache_any_lane_enabled(params)) { return; } @@ -139,9 +182,11 @@ void llama_moe_hot_cache_init(llama_model & model, const llama_model_params & pa throw std::runtime_error("--moe-hot-cache is required when --moe-hot-cache-max-mib is not 0"); } - LLAMA_LOG_WARN("%s: building hot-cache: max_mib = %lld, n_ctx = %u, n_seq_max = %u, n_ubatch = %u, swa_full = %d, kv_unified = %d, offload_kqv = %d, reserve_kv_cache = %d, path = %s\n", + LLAMA_LOG_WARN("%s: building hot-cache: max_mib = %lld, second_max_mib = %lld, third_max_mib = %lld, n_ctx = %u, n_seq_max = %u, n_ubatch = %u, swa_full = %d, kv_unified = %d, offload_kqv = %d, reserve_kv_cache = %d, path = %s\n", __func__, (long long) params.moe_hot_cache_max_mib, + (long long) params.moe_hot_cache_second_max_mib, + (long long) params.moe_hot_cache_third_max_mib, params.moe_hot_cache_auto_n_ctx, params.moe_hot_cache_auto_n_seq_max, params.moe_hot_cache_auto_n_ubatch, @@ -165,41 +210,78 @@ void llama_moe_hot_cache_init(llama_model & model, const llama_model_params & pa config.layer_curve); const auto observed = score_observations_for_arch(model.arch, observations, ¶ms); const auto sizes = llama_moe_hot_cache_collect_expert_sizes(model); - ggml_backend_dev_t cache_dev = llama_moe_hot_cache_select_gpu_dev(&model); - const size_t budget_bytes = params.moe_hot_cache_max_mib < 0 - ? llama_moe_hot_cache_auto_budget_bytes(model, params, cache_dev, reserve_kv_cache) - : size_t(params.moe_hot_cache_max_mib)*LLAMA_MOE_HOT_CACHE_MIB; - if (budget_bytes == 0) { - LLAMA_LOG_WARN("%s: hot-cache budget is 0 MiB; disabling hot-cache\n", __func__); + const auto lane_requests = hot_cache_lane_requests(params); + if (lane_requests.empty() || lane_requests.size() > LLAMA_MOE_HOT_CACHE_MAX_EXPERT_LANES) { + LLAMA_LOG_WARN("%s: no hot-cache lanes configured; disabling hot-cache\n", __func__); return; } - const auto plan = llama_moe_hot_cache_select(observed, sizes, budget_bytes); + std::vector cache_devs; + std::vector lane_budgets; + cache_devs.reserve(lane_requests.size()); + lane_budgets.reserve(lane_requests.size()); + + std::unordered_set seen_devs; + for (size_t lane = 0; lane < lane_requests.size(); ++lane) { + const auto & request = lane_requests[lane]; + ggml_backend_dev_t dev = llama_moe_hot_cache_resolve_gpu_dev(&model, request.name); + if (!seen_devs.insert(dev).second) { + throw std::runtime_error(std::string("duplicate MoE hot-cache expert device: ") + ggml_backend_dev_name(dev)); + } + + const size_t budget_bytes = request.max_mib < 0 + ? llama_moe_hot_cache_auto_budget_bytes(model, params, dev, reserve_kv_cache, request.reserve_mib) + : size_t(request.max_mib)*LLAMA_MOE_HOT_CACHE_MIB; + + LLAMA_LOG_WARN("%s: expert lane %zu device = %s, max_mib = %lld, reserve_mib = %zu, budget = %zu MiB\n", + __func__, + lane, + ggml_backend_dev_name(dev), + (long long) request.max_mib, + (size_t) request.reserve_mib, + budget_bytes/LLAMA_MOE_HOT_CACHE_MIB); + + cache_devs.push_back(dev); + lane_budgets.push_back(budget_bytes); + } + + const auto strategy = llama_moe_hot_cache_parse_device_strategy(params.moe_hot_cache_device_strategy); + const auto plan = llama_moe_hot_cache_select_multi(observed, sizes, lane_budgets, strategy); const uint32_t n_expert_per_layer = model.hparams.n_expert; if (n_expert_per_layer > 0) { - const double cpu_moe_layer_equiv = (double) plan.selected.size() / (double) n_expert_per_layer; - LLAMA_LOG_WARN("%s: selected %zu/%zu observed experts for hot-cache (n-cpu-moe equivalent = %.1f layers @ %u experts/layer, %zu/%zu MiB)\n", - __func__, plan.selected.size(), plan.observed.size(), + const double cpu_moe_layer_equiv = (double) plan.selected_count() / (double) n_expert_per_layer; + LLAMA_LOG_WARN("%s: selected %zu/%zu observed experts for hot-cache across %zu lanes (n-cpu-moe equivalent = %.1f layers @ %u experts/layer, %zu/%zu MiB)\n", + __func__, plan.selected_count(), plan.observed.size(), plan.lanes.size(), cpu_moe_layer_equiv, n_expert_per_layer, - plan.used_bytes/LLAMA_MOE_HOT_CACHE_MIB, plan.budget_bytes/LLAMA_MOE_HOT_CACHE_MIB); + plan.used_bytes()/LLAMA_MOE_HOT_CACHE_MIB, plan.budget_bytes()/LLAMA_MOE_HOT_CACHE_MIB); } else { - LLAMA_LOG_WARN("%s: selected %zu/%zu observed experts for hot-cache (%zu/%zu MiB)\n", - __func__, plan.selected.size(), plan.observed.size(), - plan.used_bytes/LLAMA_MOE_HOT_CACHE_MIB, plan.budget_bytes/LLAMA_MOE_HOT_CACHE_MIB); + LLAMA_LOG_WARN("%s: selected %zu/%zu observed experts for hot-cache across %zu lanes (%zu/%zu MiB)\n", + __func__, plan.selected_count(), plan.observed.size(), plan.lanes.size(), + plan.used_bytes()/LLAMA_MOE_HOT_CACHE_MIB, plan.budget_bytes()/LLAMA_MOE_HOT_CACHE_MIB); + } + for (size_t lane = 0; lane < plan.lanes.size(); ++lane) { + const auto & lane_plan = plan.lanes[lane]; + LLAMA_LOG_WARN("%s: expert lane %zu plan on %s: selected = %zu, used = %zu/%zu MiB\n", + __func__, + lane, + ggml_backend_dev_name(cache_devs.at(lane)), + lane_plan.selected.size(), + lane_plan.used_bytes/LLAMA_MOE_HOT_CACHE_MIB, + lane_plan.budget_bytes/LLAMA_MOE_HOT_CACHE_MIB); } - if (plan.selected.empty()) { + if (plan.selected_count() == 0) { LLAMA_LOG_WARN("%s: no experts selected; disabling hot-cache\n", __func__); return; } - model.moe_hot_cache = llama_moe_hot_cache_build(model, plan, cache_dev); + model.moe_hot_cache = llama_moe_hot_cache_build_multi(model, plan, cache_devs); } void llama_moe_hot_cache_init_after_model_load(llama_model & model, const llama_model_params & params) { - if (params.moe_hot_cache_max_mib <= 0) { + if (!hot_cache_any_static_lane_enabled(params)) { return; } @@ -208,7 +290,7 @@ void llama_moe_hot_cache_init_after_model_load(llama_model & model, const llama_ void llama_moe_hot_cache_init_after_context_memory(const llama_model & model) { const auto & params = model.get_params(); - if (model.hparams.vocab_only || params.moe_hot_cache_max_mib >= 0 || model.moe_hot_cache != nullptr) { + if (model.hparams.vocab_only || !hot_cache_any_auto_lane_enabled(params) || model.moe_hot_cache != nullptr) { return; } diff --git a/src/moe-hot-cache/llama-moe-hot-cache.h b/src/moe-hot-cache/llama-moe-hot-cache.h index cfd9fc0676c..84bde410299 100644 --- a/src/moe-hot-cache/llama-moe-hot-cache.h +++ b/src/moe-hot-cache/llama-moe-hot-cache.h @@ -11,6 +11,8 @@ #include #include +static constexpr size_t LLAMA_MOE_HOT_CACHE_MAX_EXPERT_LANES = 3; + struct llama_model; struct llama_model_params; @@ -52,6 +54,20 @@ struct llama_moe_hot_cache_plan { size_t used_bytes = 0; }; +enum class llama_moe_hot_cache_device_strategy { + warm, + hot_even, +}; + +struct llama_moe_hot_cache_multi_plan { + std::vector observed; + std::vector lanes; + + size_t selected_count() const; + size_t used_bytes() const; + size_t budget_bytes() const; +}; + struct llama_moe_hot_cache_update_stats { bool active = false; double update_rate = 0.0; @@ -81,7 +97,7 @@ struct llama_moe_hot_cache_weighting_config { using llama_moe_hot_cache_qwen35moe_weighting_mode = llama_moe_hot_cache_weighting_mode; using llama_moe_hot_cache_qwen35moe_weighting_config = llama_moe_hot_cache_weighting_config; -struct llama_moe_hot_cache_layer { +struct llama_moe_hot_cache_layer_lane { ggml_tensor * ffn_gate_up_exps = nullptr; ggml_tensor * ffn_gate_exps = nullptr; ggml_tensor * ffn_up_exps = nullptr; @@ -105,6 +121,31 @@ struct llama_moe_hot_cache_layer { } }; +struct llama_moe_hot_cache_layer : llama_moe_hot_cache_layer_lane { + std::vector lanes; + std::vector expert_lane_map_host; + + bool multi_lane_active() const { + size_t active_lanes = 0; + for (const auto & lane : lanes) { + active_lanes += lane.active() ? 1 : 0; + } + return active_lanes > 1; + } + + bool active() const { + if (llama_moe_hot_cache_layer_lane::active()) { + return true; + } + for (const auto & lane : lanes) { + if (lane.active()) { + return true; + } + } + return false; + } +}; + enum llama_moe_hot_cache_worklist_field : int32_t { LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT_ID = 0, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT_SRC_SLOT, @@ -117,11 +158,24 @@ enum llama_moe_hot_cache_worklist_field : int32_t { LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT_EXPERT_ID, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT_COUNT, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COLD_COUNT, + LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_ID, + LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_SRC_SLOT, + LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_TOKEN_ID, + LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_WEIGHT, + LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_EXPERT_ID, + LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_COUNT, + LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT2_ID, + LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT2_SRC_SLOT, + LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT2_TOKEN_ID, + LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT2_WEIGHT, + LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT2_EXPERT_ID, + LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT2_COUNT, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COUNT, }; struct llama_moe_hot_cache { std::vector layers; + std::vector devices; std::vector ctxs; std::vector bufs; diff --git a/tests/test-moe-hot-cache-adapter.cpp b/tests/test-moe-hot-cache-adapter.cpp index 9a0ecf1d2d6..495bb2327ce 100644 --- a/tests/test-moe-hot-cache-adapter.cpp +++ b/tests/test-moe-hot-cache-adapter.cpp @@ -1,4 +1,5 @@ #include "../src/moe-hot-cache/llama-moe-hot-cache-adapter.h" +#include "../include/llama.h" #include #include @@ -140,11 +141,19 @@ static void test_parallel_mode_is_runtime_switchable() { set_env_var("LLAMA_MOE_HOT_CACHE_PARALLEL", nullptr); } +static void test_default_worker_lane_reserve() { + const auto params = llama_model_default_params(); + require(params.moe_hot_cache_auto_reserve_mib == 1024); + require(params.moe_hot_cache_second_auto_reserve_mib == 512); + require(params.moe_hot_cache_third_auto_reserve_mib == 512); +} + int main() { test_find_supported_adapters(); test_rejects_unsupported_arch(); test_graph_kind_capability_checks(); test_profile_defaults_are_arch_specific(); test_parallel_mode_is_runtime_switchable(); + test_default_worker_lane_reserve(); return 0; } diff --git a/tests/test-moe-hot-cache-planner.cpp b/tests/test-moe-hot-cache-planner.cpp index 657573608ed..91ce4339fbf 100644 --- a/tests/test-moe-hot-cache-planner.cpp +++ b/tests/test-moe-hot-cache-planner.cpp @@ -105,11 +105,92 @@ static void test_select_rejects_overflowing_layer_dummy_cost() { require(plan.selected.empty()); } +static void test_select_multi_warm_fills_lanes_without_duplicates() { + const std::vector observed = { + { 0, 0, 100 }, + { 0, 1, 90 }, + { 0, 2, 80 }, + { 0, 3, 70 }, + }; + + const std::vector sizes = { + { 0, 0, 10 }, + { 0, 1, 10 }, + { 0, 2, 10 }, + { 0, 3, 10 }, + }; + + const auto plan = llama_moe_hot_cache_select_multi( + observed, + sizes, + { 30, 20 }, + llama_moe_hot_cache_device_strategy::warm); + + require(plan.lanes.size() == 2); + require(plan.selected_count() == 3); + require(plan.used_bytes() == 50); + require(plan.budget_bytes() == 50); + + require(plan.lanes[0].selected.size() == 2); + require(plan.lanes[0].selected[0].expert == 0); + require(plan.lanes[0].selected[1].expert == 1); + require(plan.lanes[1].selected.size() == 1); + require(plan.lanes[1].selected[0].expert == 2); +} + +static void test_select_multi_hot_even_balances_same_layer() { + const std::vector observed = { + { 2, 0, 100 }, + { 2, 1, 90 }, + { 2, 2, 80 }, + { 2, 3, 70 }, + }; + + const std::vector sizes = { + { 2, 0, 10 }, + { 2, 1, 10 }, + { 2, 2, 10 }, + { 2, 3, 10 }, + }; + + const auto plan = llama_moe_hot_cache_select_multi( + observed, + sizes, + { 40, 40 }, + llama_moe_hot_cache_device_strategy::hot_even); + + require(plan.lanes.size() == 2); + require(plan.selected_count() == 4); + require(plan.lanes[0].selected.size() == 2); + require(plan.lanes[1].selected.size() == 2); + require(plan.lanes[0].selected[0].expert == 0); + require(plan.lanes[1].selected[0].expert == 1); + require(plan.lanes[0].selected[1].expert == 2); + require(plan.lanes[1].selected[1].expert == 3); +} + +static void test_parse_device_strategy() { + require(llama_moe_hot_cache_parse_device_strategy(nullptr) == llama_moe_hot_cache_device_strategy::warm); + require(llama_moe_hot_cache_parse_device_strategy("warm") == llama_moe_hot_cache_device_strategy::warm); + require(llama_moe_hot_cache_parse_device_strategy("hot-even") == llama_moe_hot_cache_device_strategy::hot_even); + + bool threw = false; + try { + (void) llama_moe_hot_cache_parse_device_strategy("invalid"); + } catch (const std::runtime_error &) { + threw = true; + } + require(threw); +} + int main() { test_tensor_expert_bytes_splits_by_expert_dimension(); test_select_accounts_for_one_dummy_expert_per_active_layer(); test_select_ignores_observed_entries_without_size(); test_select_skips_too_expensive_entries_and_continues(); test_select_rejects_overflowing_layer_dummy_cost(); + test_select_multi_warm_fills_lanes_without_duplicates(); + test_select_multi_hot_even_balances_same_layer(); + test_parse_device_strategy(); return 0; } diff --git a/tests/test-moe-hot-cache-worklist.cpp b/tests/test-moe-hot-cache-worklist.cpp index 22d84d0e78e..5902467c7f4 100644 --- a/tests/test-moe-hot-cache-worklist.cpp +++ b/tests/test-moe-hot-cache-worklist.cpp @@ -165,6 +165,67 @@ static void test_build_worklist_all_hot_or_cold() { } } +static void test_build_worklist_multi_lane_routes_by_lane_map() { + auto ctx = make_ctx(); + require(ctx != nullptr); + + const int32_t n_expert_used = 2; + const int32_t n_tokens = 2; + const int32_t capacity = n_expert_used*n_tokens; + + ggml_tensor * selected = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_I32, n_expert_used, n_tokens); + ggml_tensor * weights = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, 1, n_expert_used, n_tokens); + ggml_tensor * packed = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_F32, capacity, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COUNT); + + set_selected(selected, 0, 0, 0); + set_selected(selected, 1, 0, 1); + set_selected(selected, 0, 1, 2); + set_selected(selected, 1, 1, 3); + + set_weight(weights, 0, 0, 0.10f); + set_weight(weights, 1, 0, 0.20f); + set_weight(weights, 0, 1, 0.30f); + set_weight(weights, 1, 1, 0.40f); + + llama_moe_hot_cache_layer layer; + layer.n_expert = 4; + layer.expert_weights_scale = 1.0f; + layer.expert_lane_map_host = { 0, 1, -1, 1 }; + layer.lanes.resize(2); + layer.lanes[0].n_expert = 4; + layer.lanes[0].n_hot = 1; + layer.lanes[0].hot_id_map_host = { 0, -1, -1, -1 }; + layer.lanes[1].n_expert = 4; + layer.lanes[1].n_hot = 2; + layer.lanes[1].hot_id_map_host = { -1, 0, -1, 1 }; + + llama_moe_hot_cache_build_worklist(packed, selected, weights, layer, 0, 1); + + require(get_worklist_field(packed, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT_ID, 0) == 0.0f); + require(get_worklist_field(packed, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT_SRC_SLOT, 0) == 0.0f); + require_close(get_worklist_field(packed, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT_WEIGHT, 0), 0.10f); + require(get_worklist_field(packed, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT_EXPERT_ID, 0) == 0.0f); + + require(get_worklist_field(packed, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_ID, 0) == 0.0f); + require(get_worklist_field(packed, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_SRC_SLOT, 0) == 1.0f); + require_close(get_worklist_field(packed, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_WEIGHT, 0), 0.20f); + require(get_worklist_field(packed, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_EXPERT_ID, 0) == 1.0f); + + require(get_worklist_field(packed, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_ID, 1) == 1.0f); + require(get_worklist_field(packed, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_SRC_SLOT, 1) == 3.0f); + require_close(get_worklist_field(packed, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_WEIGHT, 1), 0.40f); + require(get_worklist_field(packed, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_EXPERT_ID, 1) == 3.0f); + + require(get_worklist_field(packed, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COLD_ID, 0) == 2.0f); + require(get_worklist_field(packed, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COLD_SRC_SLOT, 0) == 2.0f); + require_close(get_worklist_field(packed, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COLD_WEIGHT, 0), 0.30f); + + require(get_worklist_field(packed, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT_COUNT, 0) == 1.0f); + require(get_worklist_field(packed, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT1_COUNT, 0) == 2.0f); + require(get_worklist_field(packed, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_HOT2_COUNT, 0) == 0.0f); + require(get_worklist_field(packed, LLAMA_MOE_HOT_CACHE_WORKLIST_FIELD_COLD_COUNT, 0) == 1.0f); +} + static void test_build_worklist_from_logits() { auto ctx = make_ctx(); require(ctx != nullptr); @@ -207,6 +268,7 @@ static void test_build_worklist_from_logits() { int main() { test_build_worklist_mixed(); test_build_worklist_all_hot_or_cold(); + test_build_worklist_multi_lane_routes_by_lane_map(); test_build_worklist_from_logits(); return 0; }