Skip to content

train: one training-context derivation; K/V format resolved and recorded (6edd495c slice 1) #60

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
joelteply wants to merge 1 commit into from
Closed

train: one training-context derivation; K/V format resolved and recorded (6edd495c slice 1) #60

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 8 additions & 36 deletions common/arg.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -301,34 +301,6 @@ struct handle_model_result {
std::string preset_path;
};

const std::vector<ggml_type> kv_cache_types = {
GGML_TYPE_F32,
GGML_TYPE_F16,
GGML_TYPE_BF16,
GGML_TYPE_Q8_0,
GGML_TYPE_Q4_0,
GGML_TYPE_Q4_1,
GGML_TYPE_IQ4_NL,
GGML_TYPE_Q5_0,
GGML_TYPE_Q5_1,
};

static ggml_type kv_cache_type_from_str(const std::string & s) {
for (const auto & type : kv_cache_types) {
if (ggml_type_name(type) == s) {
return type;
}
}
throw std::runtime_error("Unsupported cache type: " + s);
}

static std::string get_all_kv_cache_types() {
std::ostringstream msg;
for (const auto & type : kv_cache_types) {
msg << ggml_type_name(type) << (&type == &kv_cache_types.back() ? "" : ", ");
}
return msg.str();
}

static bool parse_bool_value(const std::string & value) {
if (is_truthy(value)) {
Expand Down Expand Up @@ -2428,11 +2400,11 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
"KV cache data type for K\n"
"allowed values: %s\n"
"(default: %s)",
get_all_kv_cache_types().c_str(),
common_kv_cache_types_str().c_str(),
ggml_type_name(params.cache_type_k)
),
[](common_params & params, const std::string & value) {
params.cache_type_k = kv_cache_type_from_str(value);
params.cache_type_k = common_kv_cache_type_from_str(value);
}
).set_env("LLAMA_ARG_CACHE_TYPE_K"));
add_opt(common_arg(
Expand All @@ -2441,11 +2413,11 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
"KV cache data type for V\n"
"allowed values: %s\n"
"(default: %s)",
get_all_kv_cache_types().c_str(),
common_kv_cache_types_str().c_str(),
ggml_type_name(params.cache_type_v)
),
[](common_params & params, const std::string & value) {
params.cache_type_v = kv_cache_type_from_str(value);
params.cache_type_v = common_kv_cache_type_from_str(value);
}
).set_env("LLAMA_ARG_CACHE_TYPE_V"));
add_opt(common_arg(
Expand Down Expand Up @@ -4052,11 +4024,11 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
"KV cache data type for K for the draft model\n"
"allowed values: %s\n"
"(default: %s)",
get_all_kv_cache_types().c_str(),
common_kv_cache_types_str().c_str(),
ggml_type_name(params.speculative.draft.cache_type_k)
),
[](common_params & params, const std::string & value) {
params.speculative.draft.cache_type_k = kv_cache_type_from_str(value);
params.speculative.draft.cache_type_k = common_kv_cache_type_from_str(value);
}
).set_env("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K"));
add_opt(common_arg(
Expand All @@ -4065,11 +4037,11 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
"KV cache data type for V for the draft model\n"
"allowed values: %s\n"
"(default: %s)",
get_all_kv_cache_types().c_str(),
common_kv_cache_types_str().c_str(),
ggml_type_name(params.speculative.draft.cache_type_v)
),
[](common_params & params, const std::string & value) {
params.speculative.draft.cache_type_v = kv_cache_type_from_str(value);
params.speculative.draft.cache_type_v = common_kv_cache_type_from_str(value);
}
).set_env("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V"));
add_opt(common_arg(
Expand Down
132 changes: 132 additions & 0 deletions common/common.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1721,6 +1721,138 @@ struct llama_model_params common_model_params_to_llama(common_params & params) {
return mparams;
}

const std::vector<ggml_type> & common_kv_cache_types() {
static const std::vector<ggml_type> types = {
GGML_TYPE_F32,
GGML_TYPE_F16,
GGML_TYPE_BF16,
GGML_TYPE_Q8_0,
GGML_TYPE_Q4_0,
GGML_TYPE_Q4_1,
GGML_TYPE_IQ4_NL,
GGML_TYPE_Q5_0,
GGML_TYPE_Q5_1,
};
return types;
}

std::string common_kv_cache_types_str() {
std::ostringstream msg;
const auto & types = common_kv_cache_types();
for (const auto & type : types) {
msg << ggml_type_name(type) << (&type == &types.back() ? "" : ", ");
}
return msg.str();
}

ggml_type common_kv_cache_type_from_str(const std::string & s) {
for (const auto & type : common_kv_cache_types()) {
if (ggml_type_name(type) == s) {
return type;
}
}
throw std::runtime_error("Unsupported cache type: " + s);
}

// The formats a training context has been RUN with, each with a measured learning curve: F16
// (c424b00c3) and Q8_0 (a8a235df5, 1.5B chunk 256: forward loss 0.74694 against the previous
// engine's 0.74658; eval 0.621 -> 0.454 over 3 epochs). The cache holds only constants (the context before a chunk, cast to F32
// where it joins the chunk), but a coarser format's effect on learning is a measurement, not an
// assumption: until its curve is posted beside these, it is not a training format.
static bool common_training_kv_measured(ggml_type t) {
return t == GGML_TYPE_F16 || t == GGML_TYPE_Q8_0;
}

bool common_training_kv_resolve(
const llama_model * model,
const common_params & serving,
const std::string & requested_k,
const std::string & requested_v,
common_training_kv & out,
std::string & why) {
// the training context always runs its context decodes with flash attention
const auto usable = [&](ggml_type k, ggml_type v, std::string & reason) {
if (!common_training_kv_measured(k) || !common_training_kv_measured(v)) {
reason = std::string("training has no measured run with a ") + ggml_type_name(k) + "/" + ggml_type_name(v)
+ " cache (measured: f16, q8_0)";
return false;
}
char buf[256];
if (!llama_kv_cache_types_check(model, k, v, /*flash_attn=*/ true, buf, sizeof(buf))) {
reason = buf;
return false;
}
return true;
};
if (requested_k.empty() != requested_v.empty()) {
why = "a training cache format names both cache_type_k and cache_type_v, or neither";
return false;
}
if (!requested_k.empty()) {
ggml_type k;
ggml_type v;
try {
k = common_kv_cache_type_from_str(requested_k);
v = common_kv_cache_type_from_str(requested_v);
} catch (const std::exception & e) {
why = std::string(e.what()) + " (allowed: " + common_kv_cache_types_str() + ")";
return false;
}
std::string reason;
if (!usable(k, v, reason)) {
why = "the requested training cache " + requested_k + "/" + requested_v + " can't be used: " + reason;
return false;
}
out = { k, v, "requested" };
return true;
}
std::string reason;
if (usable(serving.cache_type_k, serving.cache_type_v, reason)) {
out = { serving.cache_type_k, serving.cache_type_v, "serving" };
return true;
}
if (!usable(GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, reason)) {
why = "no training cache format fits this model: " + reason;
return false;
}
out = { GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, "default" };
return true;
}

struct llama_context_params common_context_params_for_training(
const common_params & serving,
uint32_t window,
uint32_t chunk,
const common_training_kv & kv) {
llama_context_params cparams = common_context_params_to_llama(serving);
// NO ROLLBACK SLOTS. Serving keeps recurrent-state snapshots for speculative decoding's
// rollback (n_rs_seq, from the MTP draft); a training context never rolls back, and the exact
// walk carries one recurrent state per sequence. Inherited, they crashed the 5090's serving
// process on Kimi's first exact run (2026-10-07 04:09Z: the delta-net invariant asserted).
cparams.n_rs_seq = 0;
// the context holds the whole window; one training chunk is one batch is one ubatch
cparams.n_ctx = window;
cparams.n_batch = chunk;
cparams.n_ubatch = chunk;
// two sequences: 0 is the walk; 1 is the scratch a recurrent state is snapshotted into
// around each training chunk (llama_context::opt_epoch_iter). Unified, so the attention
// cache stays one stream of the window, not one per sequence.
cparams.n_seq_max = 2;
cparams.kv_unified = true;
// Flash attention ON for the context: the walk's context decodes use it, and the cache then
// stores V un-transposed, which a quantized V cache requires. The training graphs build
// explicit attention regardless (build_attn_mha under cparams.training).
cparams.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
cparams.type_k = kv.type_k;
cparams.type_v = kv.type_v;
cparams.embeddings = false;
// training reads logits for every token of the window; the serving params cap outputs per
// ubatch to what sampling needs (a server-computed limit), which a training batch overruns
cparams.n_outputs_max = 0; // = n_batch
cparams.n_outputs_max_per_seq = 0; // = n_outputs_max
return cparams;
}

struct llama_context_params common_context_params_to_llama(const common_params & params) {
auto cparams = llama_context_default_params();

Expand Down
34 changes: 34 additions & 0 deletions common/common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -941,6 +941,40 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode
struct llama_model_params common_model_params_to_llama ( common_params & params);
struct llama_context_params common_context_params_to_llama(const common_params & params);

// K/V cache formats a context may be given, by name: the ONE parser (--cache-type-k/v and /train)
const std::vector<ggml_type> & common_kv_cache_types();
std::string common_kv_cache_types_str();
ggml_type common_kv_cache_type_from_str(const std::string & s); // throws on an unknown name

// The training context's K/V format and where it came from: "requested" (the /train caller),
// "serving" (serving's own resolved format) or "default" (neither was usable; Q8_0). Recorded in
// the run's receipt beside serving's format (continuum card 6edd495c).
struct common_training_kv {
ggml_type type_k = GGML_TYPE_Q8_0;
ggml_type type_v = GGML_TYPE_Q8_0;
std::string source;
};

// Resolve the training context's K/V format: the requested one (both names, or neither) if
// training has been measured with it and the model can hold it; else serving's if so; else the
// default. Only a REQUESTED format is refused (false, with `why`); serving's never is.
bool common_training_kv_resolve(
const llama_model * model,
const common_params & serving,
const std::string & requested_k,
const std::string & requested_v,
common_training_kv & out,
std::string & why);

// The training context: serving's params with only the training mode's differences
// (no rollback slots, one chunk per batch/ubatch, a scratch sequence in a unified cache,
// logits for every token, flash attention for the context decodes, no embeddings).
struct llama_context_params common_context_params_for_training(
const common_params & serving,
uint32_t window,
uint32_t chunk,
const common_training_kv & kv);

// clear LoRA adapters from context, then apply new list of adapters
void common_set_adapter_lora(struct llama_context * ctx, std::vector<common_adapter_lora_info> & lora);

Expand Down
10 changes: 10 additions & 0 deletions include/llama.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -532,6 +532,16 @@ extern "C" {

LLAMA_API void llama_model_free(struct llama_model * model);

// Whether a K/V cache format can hold this model's attention (flash_attn = the context's
// resolved setting). The same check context creation refuses by. On false, `why` names the rule.
LLAMA_API bool llama_kv_cache_types_check(
const struct llama_model * model,
enum ggml_type type_k,
enum ggml_type type_v,
bool flash_attn,
char * why,
size_t why_len);

LLAMA_API struct llama_context * llama_init_from_model(
struct llama_model * model,
struct llama_context_params params);
Expand Down
72 changes: 40 additions & 32 deletions src/llama-context.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3358,6 +3358,38 @@ llama_context_params llama_context_default_params() {
return result;
}

// THE ONE OWNER of whether a K/V cache format can hold this model's attention: context creation
// refuses by it, and a caller choosing a format (the training context, continuum card 6edd495c)
// asks it first instead of reading a null context, which also means out of memory.
bool llama_kv_cache_types_check(
const llama_model * model,
ggml_type type_k,
ggml_type type_v,
bool flash_attn,
char * why,
size_t why_len) {
std::string err;
if ((model->hparams.is_mla() || model->arch == LLM_ARCH_DEEPSEEK4) && type_k != type_v) {
err = std::string("model does not support different K (") + ggml_type_name(type_k) + ") and V (" + ggml_type_name(type_v) + ") cache types";
} else if (ggml_is_quantized(type_v) && !flash_attn) {
err = "quantized V cache requires flash_attn to be enabled";
} else if (flash_attn) {
for (uint32_t il = 0; il < model->hparams.n_layer() && err.empty(); ++il) {
if (ggml_is_quantized(type_k) && model->hparams.n_embd_head_k(il) % ggml_blck_size(type_k) != 0) {
err = std::string("K cache type ") + ggml_type_name(type_k) + " with block size " + std::to_string(ggml_blck_size(type_k))
+ " does not divide n_embd_head_k=" + std::to_string(model->hparams.n_embd_head_k(il));
} else if (ggml_is_quantized(type_v) && model->hparams.n_embd_head_v(il) % ggml_blck_size(type_v) != 0) {
err = std::string("V cache type ") + ggml_type_name(type_v) + " with block size " + std::to_string(ggml_blck_size(type_v))
+ " does not divide n_embd_head_v=" + std::to_string(model->hparams.n_embd_head_v(il));
}
}
}
if (why != nullptr && why_len > 0) {
snprintf(why, why_len, "%s", err.c_str());
}
return err.empty();
}

llama_context * llama_init_from_model(
llama_model * model,
llama_context_params params) {
Expand Down Expand Up @@ -3392,44 +3424,20 @@ llama_context * llama_init_from_model(
}
}

if ((model->hparams.is_mla() || model->arch == LLM_ARCH_DEEPSEEK4) && params.type_k != params.type_v) {
LLAMA_LOG_ERROR("%s: model does not support different K (%s) and V (%s) cache types\n", __func__, ggml_type_name(params.type_k), ggml_type_name(params.type_v));
return nullptr;
if (ggml_is_quantized(params.type_v) && params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_AUTO) {
LLAMA_LOG_INFO("%s: enabling flash_attn since it is required for quantized V cache\n", __func__);
params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
}

if (ggml_is_quantized(params.type_v) && params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_ENABLED) {
if (params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_AUTO) {
LLAMA_LOG_INFO("%s: enabling flash_attn since it is required for quantized V cache\n", __func__);
params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
}
if (params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_DISABLED) {
LLAMA_LOG_ERROR("%s: quantized V cache requires flash_attn to be enabled\n", __func__);
{
char why[256];
if (!llama_kv_cache_types_check(model, params.type_k, params.type_v,
params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED, why, sizeof(why))) {
LLAMA_LOG_ERROR("%s: %s\n", __func__, why);
return nullptr;
}
}

if (params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED && ggml_is_quantized(params.type_k)) {
const uint32_t blck_size = ggml_blck_size(params.type_k);
for (uint32_t il = 0; il < model->hparams.n_layer(); ++il) {
if (model->hparams.n_embd_head_k(il) % blck_size != 0) {
LLAMA_LOG_ERROR("%s: K cache type %s with block size %u does not divide n_embd_head_k=%u\n",
__func__, ggml_type_name(params.type_k), blck_size, model->hparams.n_embd_head_k(il));
return nullptr;
}
}
}

if (params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED && ggml_is_quantized(params.type_v)) {
const uint32_t blck_size = ggml_blck_size(params.type_v);
for (uint32_t il = 0; il < model->hparams.n_layer(); ++il) {
if (model->hparams.n_embd_head_v(il) % blck_size != 0) {
LLAMA_LOG_ERROR("%s: V cache type %s with block size %u does not divide n_embd_head_v=%u\n",
__func__, ggml_type_name(params.type_v), blck_size, model->hparams.n_embd_head_v(il));
return nullptr;
}
}
}

if (params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED &&
params.pooling_type != model->hparams.pooling_type) {
//user-specified pooling-type is different from the model default
Expand Down
Loading
Loading