Skip to content

Commit 1c220db

Browse files
yrougyclaude
yrougy
and
claude
committed
models : skip MTP/nextn layers in qwen35, qwen35moe, qwen3next
Read nextn_predict_layers from GGUF metadata and mark the trailing MTP layers as TENSOR_SKIP so they are counted but not allocated. Set n_layer_kv_from_start to exclude them from KV cache, clear their recurrent flag, and cap the build loop to n_transformer_layers so they are never executed during inference. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 2cff104 commit 1c220db

3 files changed

Lines changed: 139 additions & 80 deletions

File tree

‎src/models/qwen35.cpp‎

Lines changed: 45 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,17 @@ void llama_model_qwen35::load_arch_hparams(llama_model_loader & ml) {
2121
}
2222
}
2323

24-
switch (hparams.n_layer) {
24+
// MTP (nextn predict) layers are standard attention layers appended after the main transformer
25+
ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false);
26+
if (hparams.nextn_predict_layers > 0) {
27+
GGML_ASSERT(hparams.nextn_predict_layers < hparams.n_layer);
28+
hparams.n_layer_kv_from_start = hparams.n_layer - hparams.nextn_predict_layers;
29+
for (uint32_t i = hparams.n_layer - hparams.nextn_predict_layers; i < hparams.n_layer; ++i) {
30+
hparams.recurrent_layer_arr[i] = false;
31+
}
32+
}
33+
34+
switch (hparams.n_layer - hparams.nextn_predict_layers) {
2535
case 24: type = hparams.n_embd == 1024 ? LLM_TYPE_0_8B : LLM_TYPE_2B; break;
2636
case 32: type = hparams.n_embd == 2560 ? LLM_TYPE_4B : LLM_TYPE_9B; break;
2737
case 64: type = LLM_TYPE_27B; break;
@@ -55,34 +65,44 @@ void llama_model_qwen35::load_arch_tensors(llama_model_loader &) {
5565
for (int i = 0; i < n_layer; ++i) {
5666
auto & layer = layers[i];
5767

58-
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, 0);
59-
layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), { n_embd }, 0);
68+
// MTP (nextn predict) layers are skipped entirely — they are not used in inference
69+
const int lflags = (hparams.nextn_predict_layers > 0 &&
70+
static_cast<uint32_t>(i) >= n_layer - hparams.nextn_predict_layers)
71+
? TENSOR_SKIP : 0;
72+
73+
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, lflags);
74+
layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), { n_embd }, lflags);
6075

6176
if (!hparams.is_recurrent(i)) {
6277
// Attention layers
63-
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, 0);
64-
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, 0);
65-
66-
// Q/K normalization for attention layers
67-
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, 0);
68-
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), { n_embd_head_k }, 0);
78+
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, lflags);
79+
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, lflags);
80+
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, lflags);
81+
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), { n_embd_head_k }, lflags);
6982
} else {
7083
// Linear attention (gated delta net) specific tensors
71-
// Create tensors with calculated dimensions
72-
layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), { n_embd, key_dim * 2 + value_dim }, TENSOR_NOT_REQUIRED);
73-
layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), { n_embd, value_dim }, TENSOR_NOT_REQUIRED);
74-
layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), { hparams.ssm_d_conv, conv_dim }, 0);
75-
layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), { hparams.ssm_dt_rank }, 0);
76-
layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, i), { hparams.ssm_dt_rank }, 0);
77-
layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), { n_embd, n_v_heads }, 0);
78-
layer.ssm_alpha = create_tensor(tn(LLM_TENSOR_SSM_ALPHA, "weight", i), { n_embd, n_v_heads }, 0);
79-
layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), { head_v_dim }, 0);
80-
layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), { value_dim, n_embd }, 0);
84+
layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), { n_embd, key_dim * 2 + value_dim }, TENSOR_NOT_REQUIRED | lflags);
85+
layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), { n_embd, value_dim }, TENSOR_NOT_REQUIRED | lflags);
86+
layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), { hparams.ssm_d_conv, conv_dim }, lflags);
87+
layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), { hparams.ssm_dt_rank }, lflags);
88+
layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, i), { hparams.ssm_dt_rank }, lflags);
89+
layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), { n_embd, n_v_heads }, lflags);
90+
layer.ssm_alpha = create_tensor(tn(LLM_TENSOR_SSM_ALPHA, "weight", i), { n_embd, n_v_heads }, lflags);
91+
layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), { head_v_dim }, lflags);
92+
layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), { value_dim, n_embd }, lflags);
8193
}
8294

83-
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
84-
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0);
85-
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
95+
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, lflags);
96+
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, lflags);
97+
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, lflags);
98+
99+
// NextN/MTP tensors are skipped (accounted for but not used in inference)
100+
if (lflags & TENSOR_SKIP) {
101+
create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2 * n_embd, n_embd}, TENSOR_SKIP);
102+
create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, TENSOR_SKIP);
103+
create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, TENSOR_SKIP);
104+
create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, TENSOR_SKIP);
105+
}
86106
}
87107
}
88108

@@ -111,7 +131,8 @@ llama_model_qwen35::graph::graph(const llama_model & model, const llm_graph_para
111131
ggml_tensor * inp_pos = build_inp_pos();
112132
ggml_tensor * inp_out_ids = build_inp_out_ids();
113133

114-
for (int il = 0; il < n_layer; ++il) {
134+
const int n_transformer_layers = n_layer - (int)hparams.nextn_predict_layers;
135+
for (int il = 0; il < n_transformer_layers; ++il) {
115136
ggml_tensor * inpSA = inpL;
116137

117138
cur = build_norm(inpL, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il);
@@ -128,7 +149,7 @@ llama_model_qwen35::graph::graph(const llama_model & model, const llm_graph_para
128149
cur = build_layer_attn(inp->get_attn(), cur, inp_pos, sections, il);
129150
}
130151

131-
if (il == n_layer - 1 && inp_out_ids) {
152+
if (il == n_transformer_layers - 1 && inp_out_ids) {
132153
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
133154
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
134155
}

‎src/models/qwen35moe.cpp‎

Lines changed: 47 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,16 @@ void llama_model_qwen35moe::load_arch_hparams(llama_model_loader & ml) {
2424
}
2525
}
2626

27-
switch (hparams.n_layer) {
27+
ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false);
28+
if (hparams.nextn_predict_layers > 0) {
29+
GGML_ASSERT(hparams.nextn_predict_layers < hparams.n_layer);
30+
hparams.n_layer_kv_from_start = hparams.n_layer - hparams.nextn_predict_layers;
31+
for (uint32_t i = hparams.n_layer - hparams.nextn_predict_layers; i < hparams.n_layer; ++i) {
32+
hparams.recurrent_layer_arr[i] = false;
33+
}
34+
}
35+
36+
switch (hparams.n_layer - hparams.nextn_predict_layers) {
2837
case 40: type = LLM_TYPE_35B_A3B; break;
2938
case 48: type = LLM_TYPE_122B_A10B; break;
3039
case 60: type = LLM_TYPE_397B_A17B; break;
@@ -60,42 +69,51 @@ void llama_model_qwen35moe::load_arch_tensors(llama_model_loader &) {
6069
for (int i = 0; i < n_layer; ++i) {
6170
auto & layer = layers[i];
6271

63-
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, 0);
64-
layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), { n_embd }, 0);
72+
const int lflags = (hparams.nextn_predict_layers > 0 &&
73+
static_cast<uint32_t>(i) >= n_layer - hparams.nextn_predict_layers)
74+
? TENSOR_SKIP : 0;
75+
76+
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, lflags);
77+
layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), { n_embd }, lflags);
6578

6679
if (!hparams.is_recurrent(i)) {
6780
// Attention layers
68-
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, 0);
69-
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, 0);
70-
71-
// Q/K normalization for attention layers
72-
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, 0);
73-
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), { n_embd_head_k }, 0);
81+
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, lflags);
82+
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, lflags);
83+
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, lflags);
84+
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), { n_embd_head_k }, lflags);
7485
} else {
7586
// Linear attention (gated delta net) specific tensors
76-
// Create tensors with calculated dimensions
77-
layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), { n_embd, key_dim * 2 + value_dim }, TENSOR_NOT_REQUIRED);
78-
layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), { n_embd, value_dim }, TENSOR_NOT_REQUIRED);
79-
layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), { hparams.ssm_d_conv, conv_dim }, 0);
80-
layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), { hparams.ssm_dt_rank }, 0);
81-
layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, i), { hparams.ssm_dt_rank }, 0);
82-
layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), { n_embd, n_v_heads }, 0);
83-
layer.ssm_alpha = create_tensor(tn(LLM_TENSOR_SSM_ALPHA, "weight", i), { n_embd, n_v_heads }, 0);
84-
layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), { head_v_dim }, 0);
85-
layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), { value_dim, n_embd }, 0);
87+
layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), { n_embd, key_dim * 2 + value_dim }, TENSOR_NOT_REQUIRED | lflags);
88+
layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), { n_embd, value_dim }, TENSOR_NOT_REQUIRED | lflags);
89+
layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), { hparams.ssm_d_conv, conv_dim }, lflags);
90+
layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), { hparams.ssm_dt_rank }, lflags);
91+
layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, i), { hparams.ssm_dt_rank }, lflags);
92+
layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), { n_embd, n_v_heads }, lflags);
93+
layer.ssm_alpha = create_tensor(tn(LLM_TENSOR_SSM_ALPHA, "weight", i), { n_embd, n_v_heads }, lflags);
94+
layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), { head_v_dim }, lflags);
95+
layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), { value_dim, n_embd }, lflags);
8696
}
8797

88-
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert }, 0);
89-
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), { n_ff_exp, n_embd, n_expert }, 0);
90-
create_tensor_gate_up_exps(layer, i, n_embd, n_ff_exp, n_expert, 0);
98+
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert }, lflags);
99+
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), { n_ff_exp, n_embd, n_expert }, lflags);
100+
create_tensor_gate_up_exps(layer, i, n_embd, n_ff_exp, n_expert, lflags);
91101

92102
// Shared experts
93103
const int64_t n_ff_shexp = hparams.n_ff_shexp ? hparams.n_ff_shexp : n_ff;
94104

95-
layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", i), { n_embd }, 0);
96-
layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), { n_embd, n_ff_shexp }, 0);
97-
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), { n_embd, n_ff_shexp }, 0);
98-
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_shexp, n_embd }, 0);
105+
layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", i), { n_embd }, lflags);
106+
layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), { n_embd, n_ff_shexp }, lflags);
107+
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), { n_embd, n_ff_shexp }, lflags);
108+
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_shexp, n_embd }, lflags);
109+
110+
// NextN/MTP tensors are skipped (accounted for but not used in inference)
111+
if (lflags & TENSOR_SKIP) {
112+
create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2 * n_embd, n_embd}, TENSOR_SKIP);
113+
create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, TENSOR_SKIP);
114+
create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, TENSOR_SKIP);
115+
create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, TENSOR_SKIP);
116+
}
99117
}
100118
}
101119

@@ -124,7 +142,8 @@ llama_model_qwen35moe::graph::graph(const llama_model & model, const llm_graph_p
124142
ggml_tensor * inp_pos = build_inp_pos();
125143
ggml_tensor * inp_out_ids = build_inp_out_ids();
126144

127-
for (int il = 0; il < n_layer; ++il) {
145+
const int n_transformer_layers = n_layer - (int)hparams.nextn_predict_layers;
146+
for (int il = 0; il < n_transformer_layers; ++il) {
128147
ggml_tensor * inpSA = inpL;
129148

130149
cur = build_norm(inpL, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il);
@@ -141,7 +160,7 @@ llama_model_qwen35moe::graph::graph(const llama_model & model, const llm_graph_p
141160
cur = build_layer_attn(inp->get_attn(), cur, inp_pos, sections, il);
142161
}
143162

144-
if (il == n_layer - 1 && inp_out_ids) {
163+
if (il == n_transformer_layers - 1 && inp_out_ids) {
145164
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
146165
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
147166
}

0 commit comments

Comments
 (0)