model : fix MTP context kv cache allocation for deepseek2, glm4moe, c…

[LoganChu]

Overview

Fixes #28626. MTP draft contexts on DEEPSEEK2, GLM4_MOE and COHERE2MOE allocate KV for every layer instead of only the appended NextN block. This commit adds the three architectures to the existing NextN filter condition and includes a corresponding regression test. This allows for correct kv cache allocation in both the main and draft contexts.

Additional information

test-llama-archs gains a '--mtp-kv' flag. It builds all architectures that are MTP-capable as a synthetic model with an appended NextN block. It tests with a 2-layer and an 8-layer trunk, and checks the MTP context's KV cache is the same size at both depths. The MTP graph only runs one NextN block regardless of depth, so filtered draft context kv cache allocation is depth-invariant.

                     @2-layer     @8-layer   growth
before  deepseek2      884736 B    2654208 B    3.00x   unfiltered
        glm4moe        786432 B    2359296 B    3.00x   unfiltered
        cohere2moe     786432 B    2359296 B    3.00x   unfiltered
after   deepseek2      294912 B     294912 B    1.00x
        glm4moe        262144 B     262144 B    1.00x
        cohere2moe     262144 B     262144 B    1.00x

Adding a NextN block raises block_count and different architectures disagree on the resulting per-layer array length (cohere2moe/mimo2 read the SWA pattern at n_layer(), step35 at n_layer_all). The test matches how their respective loader and converters handle it.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: Yes, for codebase investigation and reviewing and testing the code.

[CISC]

Why is this arch gated again?

[LoganChu]

It is arch-gated because otherwise the granite_switch and gemma4-assistant models would be incorrectly filtered. Granite repurposes n_layer_nextn (src/models/granite-switch.cpp:56-67) as a LoRA-adapter router layer instead of MTP. With its default context which is the only context it can have src/llama-context.cpp:3742-3745, it would hit the else branch and the router layer would get filtered out even though the decode graph still attends over it (src/models/granite-switch.cpp:265-277).

Gemma4-assistant declares its nextn count as its block count which the loader reads into hparams.n_layer_all (conversion/gemma.py:848-851, conversion/base.py:1363) so n_layer() reads as 0 (src/llama-hparams.cpp:347-349). With its default context type, it would hit the else branch and all layers would get filtered out.

The GLM_DSA/DEEPSEEK32 default entries are no-ops because they have their own case. I will remove them from the default gating in a follow-up commit in this PR.

[CISC]

Ah, forgot about those, though since they are the odd case out the gating should be inversed.