Skip to content

Commit 33a75f4

Browse files
authored
DeepseekV4: reduce graph splits (ggml-org#25702)
1 parent d3fba0c commit 33a75f4

1 file changed

Lines changed: 14 additions & 12 deletions

File tree

src/models/deepseek4.cpp

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -435,27 +435,29 @@ ggml_tensor * llama_model_deepseek4::graph::build_overlap_compressed_kv_from_sta
435435
kv_state = dsv4_append_zero_row(ctx0, kv_state, false);
436436
score_state = dsv4_append_zero_row(ctx0, score_state, true);
437437

438-
ggml_tensor * prev_idxs = dsv4_view_1d(ctx0, state_read_idxs, ratio*n_blocks, 0);
439-
ggml_tensor * cur_idxs = dsv4_view_1d(ctx0, state_read_idxs, ratio*n_blocks, ratio*n_blocks);
438+
const int64_t n_read = ratio*n_blocks;
440439

441-
ggml_tensor * kv_prev = ggml_get_rows(ctx0, kv_state, prev_idxs);
442-
kv_prev = ggml_cont(ctx0, ggml_view_2d(ctx0, kv_prev, n_embd_head, ratio*n_blocks, kv_prev->nb[1], 0));
440+
ggml_tensor * kv_rows = ggml_get_rows(ctx0, kv_state, state_read_idxs);
441+
ggml_tensor * score_rows = ggml_get_rows(ctx0, score_state, state_read_idxs);
442+
443+
ggml_tensor * kv_prev = ggml_cont(ctx0,
444+
ggml_view_2d(ctx0, kv_rows, n_embd_head, n_read, kv_rows->nb[1], 0));
443445
kv_prev = ggml_reshape_3d(ctx0, kv_prev, n_embd_head, ratio, n_blocks);
444446
cb(kv_prev, name, il);
445447

446-
ggml_tensor * score_prev = ggml_get_rows(ctx0, score_state, prev_idxs);
447-
score_prev = ggml_cont(ctx0, ggml_view_2d(ctx0, score_prev, n_embd_head, ratio*n_blocks, score_prev->nb[1], 0));
448+
ggml_tensor * score_prev = ggml_cont(ctx0,
449+
ggml_view_2d(ctx0, score_rows, n_embd_head, n_read, score_rows->nb[1], 0));
448450
score_prev = ggml_reshape_3d(ctx0, score_prev, n_embd_head, ratio, n_blocks);
449451
cb(score_prev, name, il);
450452

451-
ggml_tensor * kv_cur = ggml_get_rows(ctx0, kv_state, cur_idxs);
452-
kv_cur = ggml_cont(ctx0, ggml_view_2d(ctx0, kv_cur, n_embd_head, ratio*n_blocks, kv_cur->nb[1],
453-
ggml_row_size(kv_cur->type, n_embd_head)));
453+
ggml_tensor * kv_cur = ggml_cont(ctx0,
454+
ggml_view_2d(ctx0, kv_rows, n_embd_head, n_read, kv_rows->nb[1],
455+
n_read*kv_rows->nb[1] + ggml_row_size(kv_rows->type, n_embd_head)));
454456
kv_cur = ggml_reshape_3d(ctx0, kv_cur, n_embd_head, ratio, n_blocks);
455457

456-
ggml_tensor * score_cur = ggml_get_rows(ctx0, score_state, cur_idxs);
457-
score_cur = ggml_cont(ctx0, ggml_view_2d(ctx0, score_cur, n_embd_head, ratio*n_blocks, score_cur->nb[1],
458-
ggml_row_size(score_cur->type, n_embd_head)));
458+
ggml_tensor * score_cur = ggml_cont(ctx0,
459+
ggml_view_2d(ctx0, score_rows, n_embd_head, n_read, score_rows->nb[1],
460+
n_read*score_rows->nb[1] + ggml_row_size(score_rows->type, n_embd_head)));
459461
score_cur = ggml_reshape_3d(ctx0, score_cur, n_embd_head, ratio, n_blocks);
460462

461463
ggml_tensor * values = ggml_concat(ctx0, kv_prev, kv_cur, 1);

0 commit comments

Comments
 (0)