Skip to content

Commit ef40457

Browse files
AnbeeldAnbeeld
committed
fix(bench): accept KVarN cache types
Teach llama-bench to parse kvarn2/kvarn3/kvarn4 pseudo cache types, normalize one-sided K/V selections the same way the common CLI does, and pass the selected KVarN preset through llama_context_params.kvarn. Keep benchmark output reporting KVarN names instead of the internal f16 storage fallback, and cover the path in the DFlash plumbing regression test.
1 parent 329dffb commit ef40457

3 files changed

Lines changed: 147 additions & 28 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
- Fixed the Gemma4 and Qwen3.6 DFlash + KVarN crash class seen after prompt-cache reuse. The server no longer asks KVarN-backed target memory to delete compressed historical groups during checkpoint restore, so repeated chat requests fall back to safe prefill instead of aborting in `seq_rm`.
1111
- Reduced KVarN side effects outside KVarN runs. KVarN is disabled for draft and auxiliary contexts, iSWA now keeps the SWA side on the compact upstream cache instead of allocating a full-size KVarN SWA cache, and DFlash backup streams are gated by target architecture so Gemma recurrent paths do not allocate unused backup state.
1212
- Simplified the public KVarN surface while it is still experimental. The extra tuning flags were removed; the only user-facing selection is the cache type itself. Internal defaults such as sink tokens and Sinkhorn iterations now come from the selected `kvarnN` type.
13-
- Improved KVarN CUDA behavior and validation. Cross-stream state copies are explicitly ordered, CUDA materialization group counts are precomputed, and regression coverage now checks stream synchronization, sequence-removal capability, prompt-cache rollback safety, iSWA sizing policy, unsupported runtime policy, and argument parsing.
13+
- Improved KVarN CUDA behavior and validation. Cross-stream state copies are explicitly ordered, CUDA materialization group counts are precomputed, `llama-bench` accepts KVarN cache-type names for benchmarking, and regression coverage now checks stream synchronization, sequence-removal capability, prompt-cache rollback safety, iSWA sizing policy, unsupported runtime policy, and argument parsing.
1414
- Updated preview/release maintenance for v0.3.2, including preview workflow routing, asset download handling, and avoiding cached release toolchains.
1515

1616
## v0.3.1

‎tests/test-dflash-plumbing.cpp‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,7 @@ int main(int argc, char ** argv) {
115115
const std::string arg_cpp = read_file(root + "/common/arg.cpp");
116116
const std::string common_h = read_file(root + "/common/common.h");
117117
const std::string common_cpp = read_file(root + "/common/common.cpp");
118+
const std::string llama_bench = read_file(root + "/tools/llama-bench/llama-bench.cpp");
118119
const std::string dflash_draft = read_file(root + "/src/models/dflash_draft.cpp");
119120
const std::string delta_net_base = read_file(root + "/src/models/delta-net-base.cpp");
120121
const std::string arch_cpp = read_file(root + "/src/llama-arch.cpp");
@@ -406,7 +407,7 @@ int main(int argc, char ** argv) {
406407
ggml_cpu_c.find("[GGML_TYPE_Q6_0]") != std::string::npos,
407408
"CPU backend traits must support q6_0 quantize and dot-product paths");
408409
ok &= expect(arg_cpp.find("GGML_TYPE_Q6_0") != std::string::npos &&
409-
read_file(root + "/tools/llama-bench/llama-bench.cpp").find("GGML_TYPE_Q6_0") != std::string::npos,
410+
llama_bench.find("GGML_TYPE_Q6_0") != std::string::npos,
410411
"CLI and benchmark cache-type parsers must expose q6_0");
411412
ok &= expect(cuda_common.find("ggml_cuda_type_traits<GGML_TYPE_Q6_0>") != std::string::npos &&
412413
cuda_dequantize.find("dequantize_q6_0") != std::string::npos &&
@@ -2238,6 +2239,13 @@ int main(int argc, char ** argv) {
22382239
ok &= expect(context_cpp.find("params.kvarn.type != LLAMA_KVARN_K4V2_G128") != std::string::npos &&
22392240
context_cpp.find("is experimental; only kvarn_k4v2_g128 is reference-aligned") != std::string::npos,
22402241
"non-k4v2 KVarN presets must be labeled experimental at runtime");
2242+
ok &= expect(llama_bench.find("bench_cache_type_from_name") != std::string::npos &&
2243+
llama_bench.find("kvarn_bits_from_cache_type") != std::string::npos &&
2244+
llama_bench.find("normalize_kvarn_cache_pair") != std::string::npos &&
2245+
llama_bench.find("cache_type_list_has_kvarn(params.type_k)") != std::string::npos &&
2246+
llama_bench.find("cparams.kvarn") != std::string::npos &&
2247+
llama_bench.find("bench_cache_type_name(type_k)") != std::string::npos,
2248+
"llama-bench must accept cache-type KVarN pseudo names and pass KVarN params to llama_context_params");
22412249
ok &= expect(ggml_cuda_kvarn.find("kvarn_live_groups_kernel") != std::string::npos &&
22422250
ggml_cuda_kvarn.find("ggml_cuda_pool_alloc<int> live_groups") != std::string::npos &&
22432251
ggml_cuda_kvarn.find("live_groups[out_stream]") != std::string::npos,

‎tools/llama-bench/llama-bench.cpp‎

Lines changed: 137 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -276,6 +276,37 @@ static std::string pair_str(const std::pair<int, int> & p) {
276276
return buf;
277277
}
278278

279+
struct bench_cache_type {
280+
ggml_type ggml;
281+
int32_t kvarn_bits;
282+
};
283+
284+
static bool operator==(const bench_cache_type & lhs, const bench_cache_type & rhs) {
285+
return lhs.ggml == rhs.ggml && lhs.kvarn_bits == rhs.kvarn_bits;
286+
}
287+
288+
static bool operator!=(const bench_cache_type & lhs, const bench_cache_type & rhs) {
289+
return !(lhs == rhs);
290+
}
291+
292+
static std::string bench_cache_type_name(const bench_cache_type & type) {
293+
if (type.kvarn_bits != 0) {
294+
return "kvarn" + std::to_string(type.kvarn_bits);
295+
}
296+
297+
return ggml_type_name(type.ggml);
298+
}
299+
300+
static bool cache_type_list_has_kvarn(const std::vector<bench_cache_type> & types) {
301+
for (const auto & type : types) {
302+
if (type.kvarn_bits != 0) {
303+
return true;
304+
}
305+
}
306+
307+
return false;
308+
}
309+
279310
static std::vector<int> parse_int_range(const std::string & s, bool allow_negative = false) {
280311
// first[-last[(+|*)step]]
281312
std::regex range_regex(allow_negative
@@ -329,8 +360,8 @@ struct cmd_params {
329360
std::vector<int> n_depth;
330361
std::vector<int> n_batch;
331362
std::vector<int> n_ubatch;
332-
std::vector<ggml_type> type_k;
333-
std::vector<ggml_type> type_v;
363+
std::vector<bench_cache_type> type_k;
364+
std::vector<bench_cache_type> type_v;
334365
std::vector<int> n_threads;
335366
std::vector<std::string> cpu_mask;
336367
std::vector<bool> cpu_strict;
@@ -373,8 +404,8 @@ static const cmd_params cmd_params_defaults = {
373404
/* n_depth */ { 0 },
374405
/* n_batch */ { 2048 },
375406
/* n_ubatch */ { 512 },
376-
/* type_k */ { GGML_TYPE_F16 },
377-
/* type_v */ { GGML_TYPE_F16 },
407+
/* type_k */ { { GGML_TYPE_F16, 0 } },
408+
/* type_v */ { { GGML_TYPE_F16, 0 } },
378409
/* n_threads */ { common_cpu_get_num_math() },
379410
/* cpu_mask */ { "0x0" },
380411
/* cpu_strict */ { false },
@@ -443,8 +474,8 @@ static void print_usage(int /* argc */, char ** argv) {
443474
printf(" -d, --n-depth <n> (default: %s)\n", join(cmd_params_defaults.n_depth, ",").c_str());
444475
printf(" -b, --batch-size <n> (default: %s)\n", join(cmd_params_defaults.n_batch, ",").c_str());
445476
printf(" -ub, --ubatch-size <n> (default: %s)\n", join(cmd_params_defaults.n_ubatch, ",").c_str());
446-
printf(" -ctk, --cache-type-k <t> (default: %s)\n", join(transform_to_str(cmd_params_defaults.type_k, ggml_type_name), ",").c_str());
447-
printf(" -ctv, --cache-type-v <t> (default: %s)\n", join(transform_to_str(cmd_params_defaults.type_v, ggml_type_name), ",").c_str());
477+
printf(" -ctk, --cache-type-k <t> (default: %s)\n", join(transform_to_str(cmd_params_defaults.type_k, bench_cache_type_name), ",").c_str());
478+
printf(" -ctv, --cache-type-v <t> (default: %s)\n", join(transform_to_str(cmd_params_defaults.type_v, bench_cache_type_name), ",").c_str());
448479
printf(" -t, --threads <n> (default: %s)\n", join(cmd_params_defaults.n_threads, ",").c_str());
449480
printf(" -C, --cpu-mask <hex,hex> (default: %s)\n", join(cmd_params_defaults.cpu_mask, ",").c_str());
450481
printf(" --cpu-strict <0|1> (default: %s)\n", join(cmd_params_defaults.cpu_strict, ",").c_str());
@@ -518,6 +549,81 @@ static ggml_type ggml_type_from_name(const std::string & s) {
518549
return GGML_TYPE_COUNT;
519550
}
520551

552+
static int32_t kvarn_bits_from_cache_type(const std::string & value) {
553+
if (value == "kvarn2") {
554+
return 2;
555+
}
556+
if (value == "kvarn3") {
557+
return 3;
558+
}
559+
if (value == "kvarn4") {
560+
return 4;
561+
}
562+
return 0;
563+
}
564+
565+
static llama_kvarn_type kvarn_type_from_bits(int32_t key_bits, int32_t value_bits) {
566+
switch (key_bits) {
567+
case 2:
568+
switch (value_bits) {
569+
case 2: return LLAMA_KVARN_K2V2_G128;
570+
case 3: return LLAMA_KVARN_K2V3_G128;
571+
case 4: return LLAMA_KVARN_K2V4_G128;
572+
}
573+
break;
574+
case 3:
575+
switch (value_bits) {
576+
case 2: return LLAMA_KVARN_K3V2_G128;
577+
case 3: return LLAMA_KVARN_K3V3_G128;
578+
case 4: return LLAMA_KVARN_K3V4_G128;
579+
}
580+
break;
581+
case 4:
582+
switch (value_bits) {
583+
case 2: return LLAMA_KVARN_K4V2_G128;
584+
case 3: return LLAMA_KVARN_K4V3_G128;
585+
case 4: return LLAMA_KVARN_K4V4_G128;
586+
}
587+
break;
588+
}
589+
590+
return LLAMA_KVARN_TYPE_INVALID;
591+
}
592+
593+
static bench_cache_type bench_cache_type_from_name(const std::string & name) {
594+
if (const int32_t kvarn_bits = kvarn_bits_from_cache_type(name)) {
595+
return { GGML_TYPE_F16, kvarn_bits };
596+
}
597+
598+
return { ggml_type_from_name(name), 0 };
599+
}
600+
601+
static void normalize_kvarn_cache_pair(bench_cache_type & key, bench_cache_type & value) {
602+
if (key.kvarn_bits == 0 && value.kvarn_bits == 0) {
603+
return;
604+
}
605+
606+
if (key.kvarn_bits == 0) {
607+
key = value;
608+
} else if (value.kvarn_bits == 0) {
609+
value = key;
610+
}
611+
}
612+
613+
static llama_kvarn_params kvarn_params_from_cache_pair(const bench_cache_type & key, const bench_cache_type & value) {
614+
if (key.kvarn_bits == 0 && value.kvarn_bits == 0) {
615+
return llama_kvarn_default_params();
616+
}
617+
618+
const llama_kvarn_type type = kvarn_type_from_bits(key.kvarn_bits, value.kvarn_bits);
619+
if (type == LLAMA_KVARN_TYPE_INVALID) {
620+
throw std::invalid_argument(string_format("invalid KVarN cache type combination: kvarn%d/kvarn%d",
621+
key.kvarn_bits, value.kvarn_bits));
622+
}
623+
624+
return llama_kvarn_params_for_type(type);
625+
}
626+
521627
static cmd_params parse_cmd_params(int argc, char ** argv) {
522628
cmd_params params;
523629
std::string arg;
@@ -629,14 +735,14 @@ static cmd_params parse_cmd_params(int argc, char ** argv) {
629735
}
630736
auto p = string_split<std::string>(argv[i], split_delim);
631737

632-
std::vector<ggml_type> types;
738+
std::vector<bench_cache_type> types;
633739
for (const auto & t : p) {
634-
ggml_type gt = ggml_type_from_name(t);
635-
if (gt == GGML_TYPE_COUNT) {
740+
bench_cache_type type = bench_cache_type_from_name(t);
741+
if (type.ggml == GGML_TYPE_COUNT) {
636742
invalid_param = true;
637743
break;
638744
}
639-
types.push_back(gt);
745+
types.push_back(type);
640746
}
641747
if (invalid_param) {
642748
break;
@@ -649,14 +755,14 @@ static cmd_params parse_cmd_params(int argc, char ** argv) {
649755
}
650756
auto p = string_split<std::string>(argv[i], split_delim);
651757

652-
std::vector<ggml_type> types;
758+
std::vector<bench_cache_type> types;
653759
for (const auto & t : p) {
654-
ggml_type gt = ggml_type_from_name(t);
655-
if (gt == GGML_TYPE_COUNT) {
760+
bench_cache_type type = bench_cache_type_from_name(t);
761+
if (type.ggml == GGML_TYPE_COUNT) {
656762
invalid_param = true;
657763
break;
658764
}
659-
types.push_back(gt);
765+
types.push_back(type);
660766
}
661767
if (invalid_param) {
662768
break;
@@ -1167,8 +1273,8 @@ struct cmd_params_instance {
11671273
int n_depth;
11681274
int n_batch;
11691275
int n_ubatch;
1170-
ggml_type type_k;
1171-
ggml_type type_v;
1276+
bench_cache_type type_k;
1277+
bench_cache_type type_v;
11721278
int n_threads;
11731279
std::string cpu_mask;
11741280
bool cpu_strict;
@@ -1259,8 +1365,9 @@ struct cmd_params_instance {
12591365
cparams.n_ctx = n_prompt + n_gen + n_depth;
12601366
cparams.n_batch = n_batch;
12611367
cparams.n_ubatch = n_ubatch;
1262-
cparams.type_k = type_k;
1263-
cparams.type_v = type_v;
1368+
cparams.type_k = type_k.ggml;
1369+
cparams.type_v = type_v.ggml;
1370+
cparams.kvarn = kvarn_params_from_cache_pair(type_k, type_v);
12641371
cparams.offload_kqv = !no_kv_offload;
12651372
cparams.flash_attn_type = flash_attn;
12661373
cparams.embeddings = embeddings;
@@ -1293,15 +1400,19 @@ static std::vector<cmd_params_instance> get_cmd_params_instances(const cmd_param
12931400
for (const auto & nopo : params.no_op_offload)
12941401
for (const auto & nb : params.n_batch)
12951402
for (const auto & nub : params.n_ubatch)
1296-
for (const auto & tk : params.type_k)
1297-
for (const auto & tv : params.type_v)
1403+
for (const auto & tk_arg : params.type_k)
1404+
for (const auto & tv_arg : params.type_v)
12981405
for (const auto & nkvo : params.no_kv_offload)
12991406
for (const auto & fa : params.flash_attn)
13001407
for (const auto & nt : params.n_threads)
13011408
for (const auto & cm : params.cpu_mask)
13021409
for (const auto & cs : params.cpu_strict)
13031410
for (const auto & nd : params.n_depth)
13041411
for (const auto & pl : params.poll) {
1412+
bench_cache_type tk = tk_arg;
1413+
bench_cache_type tv = tv_arg;
1414+
normalize_kvarn_cache_pair(tk, tv);
1415+
13051416
for (const auto & n_prompt : params.n_prompt) {
13061417
if (n_prompt == 0) {
13071418
continue;
@@ -1433,8 +1544,8 @@ struct test {
14331544
std::string cpu_mask;
14341545
bool cpu_strict;
14351546
int poll;
1436-
ggml_type type_k;
1437-
ggml_type type_v;
1547+
bench_cache_type type_k;
1548+
bench_cache_type type_v;
14381549
int n_gpu_layers;
14391550
int n_cpu_moe;
14401551
llama_split_mode split_mode;
@@ -1627,8 +1738,8 @@ struct test {
16271738
cpu_mask,
16281739
std::to_string(cpu_strict),
16291740
std::to_string(poll),
1630-
ggml_type_name(type_k),
1631-
ggml_type_name(type_v),
1741+
bench_cache_type_name(type_k),
1742+
bench_cache_type_name(type_v),
16321743
std::to_string(n_gpu_layers),
16331744
std::to_string(n_cpu_moe),
16341745
split_mode_str(split_mode),
@@ -1930,10 +2041,10 @@ struct markdown_printer : public printer {
19302041
if (params.n_ubatch.size() > 1 || params.n_ubatch != cmd_params_defaults.n_ubatch) {
19312042
fields.emplace_back("n_ubatch");
19322043
}
1933-
if (params.type_k.size() > 1 || params.type_k != cmd_params_defaults.type_k) {
2044+
if (params.type_k.size() > 1 || params.type_k != cmd_params_defaults.type_k || cache_type_list_has_kvarn(params.type_v)) {
19342045
fields.emplace_back("type_k");
19352046
}
1936-
if (params.type_v.size() > 1 || params.type_v != cmd_params_defaults.type_v) {
2047+
if (params.type_v.size() > 1 || params.type_v != cmd_params_defaults.type_v || cache_type_list_has_kvarn(params.type_k)) {
19372048
fields.emplace_back("type_v");
19382049
}
19392050
if (params.main_gpu.size() > 1 || params.main_gpu != cmd_params_defaults.main_gpu) {

0 commit comments

Comments
 (0)