Skip to content

Commit 3a20879

Browse files
committed
cont : naming + TODOs
1 parent e87eca1 commit 3a20879

15 files changed

Lines changed: 61 additions & 40 deletions

src/llama-hparams.cpp

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,19 @@ void llama_hparams::set_swa_pattern(uint32_t n_pattern, bool dense_first) {
1717
}
1818
}
1919

20+
// TODO: implement
21+
//void llama_hparams::set_recr_pattern(uint32_t n_pattern, bool dense_first) {
22+
// if (dense_first) {
23+
// for (uint32_t il = 0; il < n_layer; ++il) {
24+
// is_recr_impl[il] = n_pattern == 0 || (il % n_pattern != 0);
25+
// }
26+
// } else {
27+
// for (uint32_t il = 0; il < n_layer; ++il) {
28+
// is_recr_impl[il] = n_pattern == 0 || (il % n_pattern < (n_pattern - 1));
29+
// }
30+
// }
31+
//}
32+
2033
bool llama_hparams::is_swa_any() const {
2134
for (uint32_t il = 0; il < n_layer; ++il) {
2235
if (is_swa_impl[il]) {
@@ -193,7 +206,7 @@ uint32_t llama_hparams::n_embd_s() const {
193206
return ssm_d_state * ssm_d_inner;
194207
}
195208

196-
bool llama_hparams::is_recurrent(uint32_t il) const {
209+
bool llama_hparams::is_recr(uint32_t il) const {
197210
if (il < n_layer) {
198211
return is_recr_impl[il];
199212
}

src/llama-hparams.h

Lines changed: 16 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,9 @@ struct llama_hparams_convnext {
3434
};
3535

3636
struct llama_hparams {
37+
// note: use the `_impl` suffix to avoid name conflict between members and getters
38+
// for example: n_embd_out() vs n_embd_out_impl
39+
3740
bool vocab_only;
3841
bool no_alloc;
3942
bool rope_finetuned;
@@ -43,7 +46,7 @@ struct llama_hparams {
4346
uint32_t n_ctx_train; // context size the model was trained on
4447
uint32_t n_embd;
4548
uint32_t n_layer;
46-
int32_t n_layer_kv_from_start = -1; // if non-negative, the first n_layer_kv_from_start layers have KV cache
49+
int32_t n_layer_kv_from_start = -1; // if non-negative, the first n_layer_kv_from_start layers have KV cache
4750
uint32_t n_expert = 0;
4851
uint32_t n_expert_used = 0;
4952
uint32_t n_rel_attn_bkts = 0;
@@ -134,12 +137,16 @@ struct llama_hparams {
134137
llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE;
135138
// the size of the sliding window (0 - no SWA)
136139
uint32_t n_swa = 0;
140+
137141
// if is_swa_impl[il] == 1, then layer il is SWA
138142
// if is_swa_impl[il] == 0, then layer il is dense (i.e. non-SWA)
139143
// by default, all layers are dense
140144
// note: using uint32_t type for compatibility reason
141145
std::array<uint32_t, LLAMA_MAX_LAYERS> is_swa_impl;
142146

147+
// for hybrid state space models
148+
std::array<uint32_t, LLAMA_MAX_LAYERS> is_recr_impl;
149+
143150
// for State Space Models
144151
uint32_t ssm_d_conv = 0;
145152
uint32_t ssm_d_inner = 0;
@@ -150,9 +157,6 @@ struct llama_hparams {
150157
// for Kimi Linear KDA
151158
uint32_t n_embd_head_kda = 0;
152159

153-
// for hybrid state space models
154-
std::array<uint32_t, LLAMA_MAX_LAYERS> is_recr_impl;
155-
156160
bool ssm_dt_b_c_rms = false;
157161

158162
float f_clamp_kqv = 0.0f;
@@ -255,6 +259,14 @@ struct llama_hparams {
255259
// return true if one of the layers is SWA
256260
bool is_swa_any() const;
257261

262+
bool is_swa(uint32_t il) const;
263+
264+
// TODO: implement
265+
//void set_recr_pattern(uint32_t n_pattern, bool dense_first = false);
266+
267+
// whether or not the given layer is recurrent (for hybrid models)
268+
bool is_recr(uint32_t il) const;
269+
258270
uint32_t n_head(uint32_t il = 0) const;
259271

260272
uint32_t n_head_kv(uint32_t il = 0) const;
@@ -296,13 +308,8 @@ struct llama_hparams {
296308
// dimension of the recurrent state embeddings
297309
uint32_t n_embd_s() const;
298310

299-
// whether or not the given layer is recurrent (for hybrid models)
300-
bool is_recurrent(uint32_t il) const;
301-
302311
uint32_t n_pos_per_embd() const;
303312

304-
bool is_swa(uint32_t il) const;
305-
306313
// note: currently only support if either all or none of the layers are MLA
307314
bool is_mla() const;
308315

src/llama-memory-hybrid-iswa.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ llama_memory_hybrid_iswa::llama_memory_hybrid_iswa(
4444
n_ubatch,
4545
n_pad,
4646
filter_attn == nullptr ?
47-
[&](int32_t il) { return !hparams.is_recurrent(il); }
47+
[&](int32_t il) { return !hparams.is_recr(il); }
4848
: filter_attn,
4949
nullptr
5050
)),
@@ -57,7 +57,7 @@ llama_memory_hybrid_iswa::llama_memory_hybrid_iswa(
5757
n_seq_max,
5858
n_rs_seq,
5959
filter_recr == nullptr ?
60-
[&](int32_t il) { return hparams.is_recurrent(il); }
60+
[&](int32_t il) { return hparams.is_recr(il); }
6161
: filter_recr
6262
)) {}
6363

src/llama-memory-hybrid.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ llama_memory_hybrid::llama_memory_hybrid(
4545
n_swa,
4646
swa_type,
4747
filter_attn == nullptr ?
48-
[&](int32_t il) { return !hparams.is_recurrent(il); }
48+
[&](int32_t il) { return !hparams.is_recr(il); }
4949
: filter_attn,
5050
nullptr
5151
)),
@@ -58,7 +58,7 @@ llama_memory_hybrid::llama_memory_hybrid(
5858
n_seq_max,
5959
n_rs_seq,
6060
filter_recr == nullptr ?
61-
[&](int32_t il) { return hparams.is_recurrent(il); }
61+
[&](int32_t il) { return hparams.is_recr(il); }
6262
: filter_recr
6363
)) {}
6464

src/llama-model-saver.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -243,6 +243,7 @@ void llama_model_saver::add_kv_from_model() {
243243
add_kv(LLM_KV_EMBEDDING_SCALE, hparams.f_embedding_scale);
244244
add_kv(LLM_KV_TOKEN_SHIFT_COUNT, hparams.token_shift_count);
245245
add_kv(LLM_KV_INTERLEAVE_MOE_LAYER_STEP, hparams.n_moe_layer_step);
246+
// add_kv(LLM_KV_FULL_ATTENTION_INTERVAL, ???); // saved as LLM_KV_ATTENTION_RECURRENT_LAYERS instead
246247

247248
add_kv(LLM_KV_ATTENTION_HEAD_COUNT, hparams.n_head_arr, true);
248249
add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, hparams.n_head_kv_arr, true);

src/llama-model.cpp

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -371,10 +371,10 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
371371
// count only the same type of previous layers to avoid this
372372
auto get_il_eff = [&](const size_t il){
373373
size_t ret = 0;
374-
const bool il_is_recurrent = hparams.is_recurrent(il);
375-
const bool il_is_swa = hparams.is_swa(il);
374+
const bool il_is_recr = hparams.is_recr(il);
375+
const bool il_is_swa = hparams.is_swa(il);
376376
for (size_t il_prev = 0; il_prev < il; il_prev++) {
377-
ret += hparams.is_recurrent(il_prev) == il_is_recurrent && hparams.is_swa(il_prev) == il_is_swa;
377+
ret += hparams.is_recr(il_prev) == il_is_recr && hparams.is_swa(il_prev) == il_is_swa;
378378
}
379379
return ret;
380380
};
@@ -551,7 +551,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
551551
};
552552

553553
auto get_split_granularity = [&](int64_t blck_size, uint32_t il, const std::vector<std::pair<int64_t, uint32_t>> & segments) -> std::vector<int64_t> {
554-
if (hparams.is_recurrent(il)) {
554+
if (hparams.is_recr(il)) {
555555
// linear attention
556556
const int64_t head_dim = hparams.ssm_d_state;
557557
const int64_t granularity_qkv = std::lcm(blck_size, head_dim);
@@ -2009,18 +2009,18 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
20092009
filter_recr = [&](int32_t) { return true; };
20102010
} else if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) {
20112011
filter_attn = [&](int32_t il) {
2012-
return !hparams.is_recurrent(il) && hparams.n_ff(il) == 0;
2012+
return !hparams.is_recr(il) && hparams.n_ff(il) == 0;
20132013
};
20142014
filter_recr = [&](int32_t il) {
2015-
return hparams.is_recurrent(il) && hparams.n_ff(il) == 0;
2015+
return hparams.is_recr(il) && hparams.n_ff(il) == 0;
20162016
};
20172017
} else if (arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE) {
20182018
const uint32_t n_main = hparams.n_layer - hparams.nextn_predict_layers;
20192019
filter_attn = [&, n_main](int32_t il) {
2020-
return (uint32_t)il < n_main && !hparams.is_recurrent(il);
2020+
return (uint32_t)il < n_main && !hparams.is_recr(il);
20212021
};
20222022
filter_recr = [&, n_main](int32_t il) {
2023-
return (uint32_t)il < n_main && hparams.is_recurrent(il);
2023+
return (uint32_t)il < n_main && hparams.is_recr(il);
20242024
};
20252025
}
20262026

src/models/granite-hybrid.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ void llama_model_granite_hybrid::load_arch_tensors(llama_model_loader &) {
7171
// norm
7272
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
7373

74-
if (hparams.is_recurrent(i)) {
74+
if (hparams.is_recr(i)) {
7575
// ssm layers
7676
layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", i), {n_embd, d_in_proj}, 0);
7777

@@ -158,7 +158,7 @@ llama_model_granite_hybrid::graph::graph(const llama_model & model, const llm_gr
158158
cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
159159
cb(cur, "attn_norm", il);
160160

161-
if (hparams.is_recurrent(il)) {
161+
if (hparams.is_recr(il)) {
162162
// ssm layer //
163163
cur = build_mamba2_layer(inp->get_recr(), cur, model, ubatch, il);
164164
} else {

src/models/kimi-linear.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ void llama_model_kimi_linear::load_arch_tensors(llama_model_loader &) {
5353
const int64_t n_embd_head_v_kda = hparams.n_embd_head_kda;
5454
const int64_t ssm_d_conv = hparams.ssm_d_conv;
5555

56-
if (hparams.is_recurrent(i)) {
56+
if (hparams.is_recr(i)) {
5757
// Conv1d weights: try 4D first, then 3D (quantization may remove trailing 1)
5858
// 4D: [d_conv, 1, d_inner, 1], 3D: [d_conv, 1, d_inner]
5959
layer.ssm_q_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_Q, "weight", i), {ssm_d_conv, 1, n_embd_head_k_kda * n_head, 1}, TENSOR_NOT_REQUIRED);
@@ -285,7 +285,7 @@ llama_model_kimi_linear::graph::graph(const llama_model & model, const llm_graph
285285

286286
ggml_build_forward_expand(gf, cur);
287287

288-
if (hparams.is_recurrent(il)) {
288+
if (hparams.is_recr(il)) {
289289
// === KDA Layer (Kimi Delta Attention) with Recurrent State ===
290290
// Reference: vLLM kda.py
291291
const auto * mctx_cur = inp_rs->mctx;

src/models/lfm2.cpp

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ void llama_model_lfm2::load_arch_tensors(llama_model_loader &) {
5959
// for operator_norm
6060
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
6161

62-
if (!hparams.is_recurrent(i)) {
62+
if (!hparams.is_recr(i)) {
6363
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0);
6464
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0);
6565
GGML_ASSERT(n_embd_v_gqa == n_embd_k_gqa);
@@ -235,8 +235,8 @@ llama_model_lfm2::graph<iswa>::graph(const llama_model & model, const llm_graph_
235235
cur = build_norm(cur, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
236236
cb(cur, "model.layers.{}.operator_norm", il);
237237

238-
cur = hparams.is_recurrent(il) ? build_shortconv_block(cur, inp_hybrid->get_recr(), il) :
239-
build_attn_block(cur, inp_pos, inp_hybrid->get_attn(), il);
238+
cur = hparams.is_recr(il) ? build_shortconv_block(cur, inp_hybrid->get_recr(), il) :
239+
build_attn_block(cur, inp_pos, inp_hybrid->get_attn(), il);
240240

241241
if (il == n_layer - 1 && inp_out_ids) {
242242
cur = ggml_get_rows(ctx0, cur, inp_out_ids);

src/models/lfm2moe.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ void llama_model_lfm2moe::load_arch_tensors(llama_model_loader &) {
5555
// for operator_norm
5656
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
5757

58-
if (!hparams.is_recurrent(i)) {
58+
if (!hparams.is_recr(i)) {
5959
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0);
6060
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0);
6161
GGML_ASSERT(n_embd_v_gqa == n_embd_k_gqa);

0 commit comments

Comments
 (0)