owen/longhaul.cpp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf: accelerate longhaul prompt loading

Browse Source
This commit is contained in:
Owen Qwen
2026-07-29 12:44:50 -05:00
parent c2baa66534
commit fdd9fcf85e
7 changed files with 525 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/llama-bench/llama-bench.cpp
+39 -4
View File
@@ -341,6 +341,7 @@ struct cmd_params {
std::vector<int> n_cpu_moe;
std::vector<llama_split_mode> split_mode;
std::vector<llama_load_mode> load_mode;
uint64_t longhaul_cache_bytes;
std::vector<int> main_gpu;
std::vector<bool> no_kv_offload;
std::vector<llama_flash_attn_type> flash_attn;
@@ -385,6 +386,7 @@ static const cmd_params cmd_params_defaults = {
/* n_cpu_moe */ { 0 },
/* split_mode */ { LLAMA_SPLIT_MODE_LAYER },
/* load_mode */ { LLAMA_LOAD_MODE_MMAP },
/* longhaul_cache_bytes */ 0,
/* main_gpu */ { 0 },
/* no_kv_offload */ { false },
/* flash_attn */ { LLAMA_FLASH_ATTN_TYPE_AUTO },
@@ -459,7 +461,8 @@ static void print_usage(int /* argc */, char ** argv) {
printf(" -nkvo, --no-kv-offload <0|1> (default: %s)\n", join(cmd_params_defaults.no_kv_offload, ",").c_str());
printf(" -fa, --flash-attn <on|off|auto> (default: %s)\n", join(transform_to_str(cmd_params_defaults.flash_attn, llama_flash_attn_type_name), ",").c_str());
printf(" -dev, --device <dev0/dev1/...> (default: auto)\n");
printf(" -lm, --load-mode <none|mmap|mlock|mmap+mlock|dio> (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str());
printf(" -lm, --load-mode <none|mmap|mlock|mmap+mlock|dio|longhaul> (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str());
printf(" --longhaul-cache <GiB> expert cache size for longhaul mode\n");
printf(" -mmp, --mmap <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n");
printf(" -dio, --direct-io <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n");
printf(" -embd, --embeddings <0|1> (default: %s)\n", join(cmd_params_defaults.embeddings, ",").c_str());
@@ -521,6 +524,7 @@ static cmd_params parse_cmd_params(int argc, char ** argv) {
params.progress = cmd_params_defaults.progress;
params.no_warmup = cmd_params_defaults.no_warmup;
params.offline = cmd_params_defaults.offline;
params.longhaul_cache_bytes = cmd_params_defaults.longhaul_cache_bytes;
if (const char * env = getenv("HF_TOKEN")) {
params.hf_token = env;
@@ -774,6 +778,8 @@ static cmd_params parse_cmd_params(int argc, char ** argv) {
mode = LLAMA_LOAD_MODE_MMAP_MLOCK;
} else if (m == "dio") {
mode = LLAMA_LOAD_MODE_DIRECT_IO;
} else if (m == "longhaul") {
mode = LLAMA_LOAD_MODE_LONGHAUL;
} else {
invalid_param = true;
break;
@@ -784,6 +790,17 @@ static cmd_params parse_cmd_params(int argc, char ** argv) {
break;
}
params.load_mode.insert(params.load_mode.end(), modes.begin(), modes.end());
} else if (arg == "--longhaul-cache") {
if (++i >= argc) {
invalid_param = true;
break;
}
const uint64_t gib = std::stoull(argv[i]);
if (gib == 0 || gib > UINT64_MAX / 1024 / 1024 / 1024) {
invalid_param = true;
break;
}
params.longhaul_cache_bytes = gib * 1024 * 1024 * 1024;
} else if (arg == "-mg" || arg == "--main-gpu") {
if (++i >= argc) {
invalid_param = true;
@@ -1135,6 +1152,11 @@ static cmd_params parse_cmd_params(int argc, char ** argv) {
if (params.load_mode.empty()) {
params.load_mode = cmd_params_defaults.load_mode;
}
if (std::find(params.load_mode.begin(), params.load_mode.end(), LLAMA_LOAD_MODE_LONGHAUL) != params.load_mode.end() &&
params.longhaul_cache_bytes == 0) {
fprintf(stderr, "error: longhaul load mode requires --longhaul-cache N\n");
exit(1);
}
if (params.main_gpu.empty()) {
params.main_gpu = cmd_params_defaults.main_gpu;
}
@@ -1201,6 +1223,7 @@ struct cmd_params_instance {
int n_cpu_moe;
llama_split_mode split_mode;
llama_load_mode load_mode;
uint64_t longhaul_cache_bytes;
int main_gpu;
bool no_kv_offload;
llama_flash_attn_type flash_attn;
@@ -1222,6 +1245,7 @@ struct cmd_params_instance {
}
mparams.split_mode = split_mode;
mparams.load_mode = load_mode;
mparams.longhaul_cache_bytes = longhaul_cache_bytes;
mparams.main_gpu = main_gpu;
mparams.tensor_split = tensor_split.data();
mparams.no_host = no_host;
@@ -1269,7 +1293,8 @@ struct cmd_params_instance {
return model == other.model && n_gpu_layers == other.n_gpu_layers && n_cpu_moe == other.n_cpu_moe &&
split_mode == other.split_mode &&
main_gpu == other.main_gpu && tensor_split == other.tensor_split &&
load_mode == other.load_mode && devices == other.devices && no_host == other.no_host &&
load_mode == other.load_mode && longhaul_cache_bytes == other.longhaul_cache_bytes &&
devices == other.devices && no_host == other.no_host &&
vec_tensor_buft_override_equal(tensor_buft_overrides, other.tensor_buft_overrides);
}
@@ -1342,6 +1367,7 @@ static std::vector<cmd_params_instance> get_cmd_params_instances(const cmd_param
/* .n_cpu_moe = */ ncmoe,
/* .split_mode = */ sm,
/* .load_mode = */ lm,
/* .longhaul_cache_bytes = */ params.longhaul_cache_bytes,
/* .main_gpu = */ mg,
/* .no_kv_offload = */ nkvo,
/* .flash_attn = */ fa,
@@ -1378,6 +1404,7 @@ static std::vector<cmd_params_instance> get_cmd_params_instances(const cmd_param
/* .n_cpu_moe = */ ncmoe,
/* .split_mode = */ sm,
/* .load_mode = */ lm,
/* .longhaul_cache_bytes = */ params.longhaul_cache_bytes,
/* .main_gpu = */ mg,
/* .no_kv_offload = */ nkvo,
/* .flash_attn = */ fa,
@@ -1414,6 +1441,7 @@ static std::vector<cmd_params_instance> get_cmd_params_instances(const cmd_param
/* .n_cpu_moe = */ ncmoe,
/* .split_mode = */ sm,
/* .load_mode = */ lm,
/* .longhaul_cache_bytes = */ params.longhaul_cache_bytes,
/* .main_gpu = */ mg,
/* .no_kv_offload = */ nkvo,
/* .flash_attn = */ fa,
@@ -1455,6 +1483,7 @@ struct test {
int n_cpu_moe;
llama_split_mode split_mode;
llama_load_mode load_mode;
uint64_t longhaul_cache_bytes;
int main_gpu;
bool no_kv_offload;
llama_flash_attn_type flash_attn;
@@ -1494,6 +1523,7 @@ struct test {
n_cpu_moe = inst.n_cpu_moe;
split_mode = inst.split_mode;
load_mode = inst.load_mode;
longhaul_cache_bytes = inst.longhaul_cache_bytes;
main_gpu = inst.main_gpu;
no_kv_offload = inst.no_kv_offload;
flash_attn = inst.flash_attn;
@@ -1561,7 +1591,7 @@ struct test {
"n_ubatch", "n_threads", "cpu_mask", "cpu_strict", "poll",
"type_k", "type_v", "n_gpu_layers", "n_cpu_moe", "split_mode",
"main_gpu", "no_kv_offload", "flash_attn", "devices", "tensor_split",
"tensor_buft_overrides", "load_mode", "embeddings",
"tensor_buft_overrides", "load_mode", "longhaul_cache_bytes", "embeddings",
"no_op_offload", "no_host", "fit_target", "fit_min_ctx",
"n_prompt", "n_gen", "n_depth",
"test_time", "avg_ns", "stddev_ns", "avg_ts", "stddev_ts"
@@ -1576,7 +1606,8 @@ struct test {
field == "poll" || field == "model_size" || field == "model_n_params" || field == "n_gpu_layers" ||
field == "main_gpu" || field == "n_prompt" || field == "n_gen" || field == "n_depth" || field == "avg_ns" ||
field == "stddev_ns" || field == "no_op_offload" || field == "n_cpu_moe" ||
field == "fit_target" || field == "fit_min_ctx" || field == "flash_attn") {
field == "fit_target" || field == "fit_min_ctx" || field == "flash_attn" ||
field == "longhaul_cache_bytes") {
return INT;
}
if (field == "f16_kv" || field == "no_kv_offload" || field == "cpu_strict" ||
@@ -1656,6 +1687,7 @@ struct test {
tensor_split_str,
tensor_buft_overrides_str,
llama_load_mode_name(load_mode),
std::to_string(longhaul_cache_bytes),
std::to_string(embeddings),
std::to_string(no_op_offload),
std::to_string(no_host),
@@ -1970,6 +2002,9 @@ struct markdown_printer : public printer {
if (params.load_mode.size() > 1 || params.load_mode != cmd_params_defaults.load_mode) {
fields.emplace_back("load_mode");
}
if (params.longhaul_cache_bytes != cmd_params_defaults.longhaul_cache_bytes) {
fields.emplace_back("longhaul_cache_bytes");
}
if (params.embeddings.size() > 1 || params.embeddings != cmd_params_defaults.embeddings) {
fields.emplace_back("embeddings");
}