spec: properly handle mtmd input for mtp (#30257)

* spec: properly handle mtmd input for mtp

* nits
This commit is contained in:
Xuan-Son Nguyen authored and GitHub committed 2026-10-10 11:22:38 +02:00
1 parent 0fd868cbca
commit 781dbc5ac9
27 files changed
+248 -99

No files matched your search

+13 -7
View File
@@ -1132,7 +1132,7 @@ static void test_compat(testing & t) {
}
static void test_mtp_embd_width(testing & t) {
t.test("mtp_uses_n_embd_out", [&](testing & t) {
t.test("mtp_keeps_n_embd_inp_and_takes_state_at_n_embd_out", [&](testing & t) {
llama_hparams hparams = {};
hparams.n_embd = 64;
hparams.n_deepstack_layers = 2; // makes n_embd_inp() = 64 + 64*2 = 192
@@ -1141,16 +1141,22 @@ static void test_mtp_embd_width(testing & t) {
t.assert_equal("default context uses n_embd_inp (deepstack-aware)",
(size_t) 192, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_DEFAULT, LLM_ARCH_LLAMA, hparams));
t.assert_equal("MTP context uses n_embd_out instead (target-model hidden state width)",
(size_t) 96, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_MTP, LLM_ARCH_LLAMA, hparams));
t.assert_equal("MTP context keeps n_embd_inp for the token embeddings",
(size_t) 192, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_MTP, LLM_ARCH_LLAMA, hparams));
t.assert_equal("MTP context takes the target hidden state at n_embd_out",
(size_t) 96, llama_batch_ext_select_n_embd_state(LLAMA_CONTEXT_TYPE_MTP, hparams));
t.assert_equal("default context takes no state",
(size_t) 0, llama_batch_ext_select_n_embd_state(LLAMA_CONTEXT_TYPE_DEFAULT, hparams));
});
t.test("mtp_falls_back_to_n_embd_when_no_override", [&](testing & t) {
t.test("mtp_state_falls_back_to_n_embd_when_no_override", [&](testing & t) {
llama_hparams hparams = {};
hparams.n_embd = 64; // no deepstack, no n_embd_out_impl override
t.assert_equal((size_t) 64, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_DEFAULT, LLM_ARCH_LLAMA, hparams));
t.assert_equal((size_t) 64, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_MTP, LLM_ARCH_LLAMA, hparams));
t.assert_equal((size_t) 64, llama_batch_ext_select_n_embd_state(LLAMA_CONTEXT_TYPE_MTP, hparams));
});
t.test("dflash_uses_n_embd_inp_enc", [&](testing & t) {
@@ -1165,8 +1171,8 @@ static void test_mtp_embd_width(testing & t) {
t.assert_equal("other archs ignore n_embd_inp_enc",
(size_t) 64, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_DEFAULT, LLM_ARCH_LLAMA, hparams));
t.assert_equal("MTP takes precedence over DFlash",
(size_t) 96, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_MTP, LLM_ARCH_DFLASH, hparams));
t.assert_equal("MTP context does not change the DFlash input width",
(size_t) 128, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_MTP, LLM_ARCH_DFLASH, hparams));
});
}