Skip to content

Commit 6af1783

Browse files
fairydreamingsszymczyCISCggerganov
authored
model : support for DeepseekV32ForCausalLM with generic DeepSeek Sparse Attention (DSA) implementation (ggml-org#23346)
* llama : support DeepSeek V3.2 model family (with DSA lightning indexer) * convert : handle DeepseekV32ForCausalLM architecture * ggml : support for f16 GGML_OP_FILL * memory : separate hparams argument in llama_kv_cache constructor * memory : add llama_kv_cache_dsa memory (KV cache + lightning indexer cache) * llama : support for LLM_ARCH_DEEPSEEK32 * model : llama_model_deepseek32 implementation * model : merge two scale operations into one in DSA lightning indexer implementation * chore : remove unused code * model : support NVFP4 in DeepSeek V3.2 Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@scala.com> * memory : refactoring TODO Co-authored-by: ggerganov <ggerganov@users.noreply.github.com> --------- Co-authored-by: Stanisław Szymczyk <sszymczy@gmail.com> Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@scala.com> Co-authored-by: ggerganov <ggerganov@users.noreply.github.com>
1 parent 4c14d45 commit 6af1783

22 files changed

Lines changed: 1261 additions & 7 deletions

‎conversion/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@
4747
"DeepseekForCausalLM": "deepseek",
4848
"DeepseekV2ForCausalLM": "deepseek",
4949
"DeepseekV3ForCausalLM": "deepseek",
50+
"DeepseekV32ForCausalLM": "deepseek",
5051
"DistilBertForMaskedLM": "bert",
5152
"DistilBertForSequenceClassification": "bert",
5253
"DistilBertModel": "bert",

‎conversion/base.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -915,6 +915,8 @@ def load():
915915
gguf.MODEL_TENSOR.SSM_CONV1D_Q,
916916
gguf.MODEL_TENSOR.SSM_CONV1D_K,
917917
gguf.MODEL_TENSOR.SSM_CONV1D_V,
918+
# DSA indexer weights should be F32
919+
gguf.MODEL_TENSOR.INDEXER_PROJ,
918920
)
919921
)
920922
or new_name[-7:] not in (".weight", ".lora_a", ".lora_b")

‎conversion/deepseek.py‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -386,3 +386,32 @@ def prepare_tensors(self):
386386
experts = [k for d in self._experts for k in d.keys()]
387387
if len(experts) > 0:
388388
raise ValueError(f"Unprocessed experts: {experts}")
389+
390+
391+
@ModelBase.register("DeepseekV32ForCausalLM")
392+
class DeepseekV32Model(DeepseekV2Model):
393+
model_arch = gguf.MODEL_ARCH.DEEPSEEK32
394+
skip_mtp = False
395+
396+
def __init__(self, *args, **kwargs):
397+
super().__init__(*args, **kwargs)
398+
self.block_count = self.hparams["num_hidden_layers"] + self.hparams.get("num_nextn_predict_layers", 0)
399+
self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
400+
401+
def set_vocab(self):
402+
from transformers import AutoTokenizer
403+
tokenizer = AutoTokenizer.from_pretrained(self.dir_model)
404+
assert getattr(tokenizer, "add_bos_token", False), "Change value of add_bos_token to true in tokenizer_config.json file."
405+
self._set_vocab_gpt2()
406+
407+
def set_gguf_parameters(self):
408+
super().set_gguf_parameters()
409+
410+
# NextN/MTP prediction layers
411+
if (num_nextn_predict_layers := self.hparams.get("num_nextn_predict_layers")) is not None:
412+
self.gguf_writer.add_nextn_predict_layers(num_nextn_predict_layers)
413+
414+
# DSA indexer parameters
415+
self.gguf_writer.add_indexer_head_count(self.hparams["index_n_heads"])
416+
self.gguf_writer.add_indexer_key_length(self.hparams["index_head_dim"])
417+
self.gguf_writer.add_indexer_top_k(self.hparams["index_topk"])

‎ggml/src/ggml-cpu/ops.cpp‎

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2235,8 +2235,42 @@ static void ggml_compute_forward_fill_f32(const ggml_compute_params * params, gg
22352235
}
22362236
}
22372237

2238+
static void ggml_compute_forward_fill_f16(const ggml_compute_params * params, ggml_tensor * dst) {
2239+
const ggml_fp16_t c = GGML_CPU_FP32_TO_FP16(ggml_get_op_params_f32(dst, 0));
2240+
2241+
GGML_TENSOR_LOCALS(int64_t, ne, dst, ne);
2242+
GGML_TENSOR_LOCALS(size_t, nb, dst, nb);
2243+
2244+
const auto [ir0, ir1] = get_thread_range(params, dst);
2245+
2246+
for (int64_t ir = ir0; ir < ir1; ++ir) {
2247+
const int64_t i03 = ir/(ne2*ne1);
2248+
const int64_t i02 = (ir - i03*ne2*ne1)/ne1;
2249+
const int64_t i01 = (ir - i03*ne2*ne1 - i02*ne1);
2250+
2251+
ggml_fp16_t * dst_ptr = (ggml_fp16_t *) ((char *) dst->data + i03*nb3 + i02*nb2 + i01*nb1);
2252+
2253+
ggml_vec_set_f16(ne0, dst_ptr, c);
2254+
}
2255+
}
2256+
22382257
void ggml_compute_forward_fill(const ggml_compute_params * params, ggml_tensor * dst) {
2239-
ggml_compute_forward_fill_f32(params, dst);
2258+
const ggml_tensor * src0 = dst->src[0];
2259+
2260+
switch (src0->type) {
2261+
case GGML_TYPE_F32:
2262+
{
2263+
ggml_compute_forward_fill_f32(params, dst);
2264+
} break;
2265+
case GGML_TYPE_F16:
2266+
{
2267+
ggml_compute_forward_fill_f16(params, dst);
2268+
} break;
2269+
default:
2270+
{
2271+
GGML_ABORT("unsupported type for ggml_compute_forward_fill: %s", ggml_type_name(src0->type));
2272+
}
2273+
}
22402274
}
22412275

22422276
// ggml_compute_tri

‎ggml/src/ggml.c‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5223,7 +5223,7 @@ static struct ggml_tensor * ggml_fill_impl(
52235223
struct ggml_tensor * a,
52245224
float c,
52255225
bool inplace) {
5226-
GGML_ASSERT(a->type == GGML_TYPE_F32);
5226+
GGML_ASSERT(a->type == GGML_TYPE_F32 || a->type == GGML_TYPE_F16);
52275227
GGML_ASSERT(ggml_is_contiguous(a));
52285228

52295229
struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a);

‎gguf-py/gguf/constants.py‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -451,6 +451,7 @@ class MODEL_ARCH(IntEnum):
451451
DEEPSEEK = auto()
452452
DEEPSEEK2 = auto()
453453
DEEPSEEK2OCR = auto()
454+
DEEPSEEK32 = auto()
454455
CHATGLM = auto()
455456
GLM4 = auto()
456457
GLM4_MOE = auto()
@@ -967,6 +968,7 @@ class MODEL_TENSOR(IntEnum):
967968
MODEL_ARCH.DEEPSEEK: "deepseek",
968969
MODEL_ARCH.DEEPSEEK2: "deepseek2",
969970
MODEL_ARCH.DEEPSEEK2OCR: "deepseek2-ocr",
971+
MODEL_ARCH.DEEPSEEK32: "deepseek32",
970972
MODEL_ARCH.CHATGLM: "chatglm",
971973
MODEL_ARCH.GLM4: "glm4",
972974
MODEL_ARCH.GLM4_MOE: "glm4moe",
@@ -2930,6 +2932,46 @@ class MODEL_TENSOR(IntEnum):
29302932
MODEL_TENSOR.FFN_UP_SHEXP,
29312933
MODEL_TENSOR.FFN_EXP_PROBS_B,
29322934
],
2935+
MODEL_ARCH.DEEPSEEK32: [
2936+
MODEL_TENSOR.TOKEN_EMBD,
2937+
MODEL_TENSOR.OUTPUT_NORM,
2938+
MODEL_TENSOR.OUTPUT,
2939+
MODEL_TENSOR.ROPE_FREQS,
2940+
MODEL_TENSOR.ATTN_NORM,
2941+
MODEL_TENSOR.ATTN_Q,
2942+
MODEL_TENSOR.ATTN_Q_A,
2943+
MODEL_TENSOR.ATTN_Q_B,
2944+
MODEL_TENSOR.ATTN_KV_A_MQA,
2945+
MODEL_TENSOR.ATTN_K_B,
2946+
MODEL_TENSOR.ATTN_V_B,
2947+
MODEL_TENSOR.ATTN_Q_A_NORM,
2948+
MODEL_TENSOR.ATTN_KV_A_NORM,
2949+
MODEL_TENSOR.ATTN_OUT,
2950+
MODEL_TENSOR.ATTN_ROT_EMBD,
2951+
MODEL_TENSOR.FFN_GATE_INP,
2952+
MODEL_TENSOR.FFN_NORM,
2953+
MODEL_TENSOR.FFN_GATE,
2954+
MODEL_TENSOR.FFN_DOWN,
2955+
MODEL_TENSOR.FFN_UP,
2956+
MODEL_TENSOR.FFN_GATE_EXP,
2957+
MODEL_TENSOR.FFN_DOWN_EXP,
2958+
MODEL_TENSOR.FFN_UP_EXP,
2959+
MODEL_TENSOR.FFN_GATE_SHEXP,
2960+
MODEL_TENSOR.FFN_DOWN_SHEXP,
2961+
MODEL_TENSOR.FFN_UP_SHEXP,
2962+
MODEL_TENSOR.FFN_EXP_PROBS_B,
2963+
MODEL_TENSOR.INDEXER_K_NORM,
2964+
MODEL_TENSOR.INDEXER_PROJ,
2965+
MODEL_TENSOR.INDEXER_ATTN_K,
2966+
MODEL_TENSOR.INDEXER_ATTN_Q_B,
2967+
# NextN/MTP tensors - preserved but unused
2968+
MODEL_TENSOR.NEXTN_EH_PROJ,
2969+
MODEL_TENSOR.NEXTN_EMBED_TOKENS,
2970+
MODEL_TENSOR.NEXTN_ENORM,
2971+
MODEL_TENSOR.NEXTN_HNORM,
2972+
MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
2973+
MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
2974+
],
29332975
MODEL_ARCH.ERNIE4_5_MOE: [
29342976
MODEL_TENSOR.TOKEN_EMBD,
29352977
MODEL_TENSOR.OUTPUT_NORM,
@@ -4077,6 +4119,10 @@ class MODEL_TENSOR(IntEnum):
40774119
MODEL_TENSOR.ROPE_FREQS,
40784120
MODEL_TENSOR.ATTN_ROT_EMBD,
40794121
],
4122+
MODEL_ARCH.DEEPSEEK32: [
4123+
MODEL_TENSOR.ROPE_FREQS,
4124+
MODEL_TENSOR.ATTN_ROT_EMBD,
4125+
],
40804126
MODEL_ARCH.CHATGLM: [
40814127
MODEL_TENSOR.ROPE_FREQS,
40824128
],

‎src/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ add_library(llama
2424
llama-io.cpp
2525
llama-kv-cache.cpp
2626
llama-kv-cache-iswa.cpp
27+
llama-kv-cache-dsa.cpp
2728
llama-memory.cpp
2829
llama-memory-hybrid.cpp
2930
llama-memory-hybrid-iswa.cpp

‎src/llama-arch.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
7575
{ LLM_ARCH_DEEPSEEK, "deepseek" },
7676
{ LLM_ARCH_DEEPSEEK2, "deepseek2" },
7777
{ LLM_ARCH_DEEPSEEK2OCR, "deepseek2-ocr" },
78+
{ LLM_ARCH_DEEPSEEK32, "deepseek32" },
7879
{ LLM_ARCH_CHATGLM, "chatglm" },
7980
{ LLM_ARCH_GLM4, "glm4" },
8081
{ LLM_ARCH_GLM4_MOE, "glm4moe" },
@@ -904,6 +905,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
904905
case LLM_ARCH_OLMO2:
905906
case LLM_ARCH_OLMOE:
906907
case LLM_ARCH_DEEPSEEK2:
908+
case LLM_ARCH_DEEPSEEK32:
907909
case LLM_ARCH_GLM_DSA:
908910
case LLM_ARCH_BITNET:
909911
case LLM_ARCH_T5:

‎src/llama-arch.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@ enum llm_arch {
7979
LLM_ARCH_DEEPSEEK,
8080
LLM_ARCH_DEEPSEEK2,
8181
LLM_ARCH_DEEPSEEK2OCR,
82+
LLM_ARCH_DEEPSEEK32,
8283
LLM_ARCH_CHATGLM,
8384
LLM_ARCH_GLM4,
8485
LLM_ARCH_GLM4_MOE,

‎src/llama-graph.cpp‎

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
#include "llama-kv-cache.h"
99
#include "llama-kv-cache-iswa.h"
10+
#include "llama-kv-cache-dsa.h"
1011
#include "llama-memory-hybrid.h"
1112
#include "llama-memory-hybrid-iswa.h"
1213
#include "llama-memory-recurrent.h"
@@ -531,6 +532,34 @@ bool llm_graph_input_attn_k::can_reuse(const llm_graph_params & params) {
531532
return res;
532533
}
533534

535+
void llm_graph_input_attn_k_dsa::set_input(const llama_ubatch * ubatch) {
536+
mctx->get_mla()->set_input_k_idxs(self_k_idxs_mla, ubatch);
537+
538+
mctx->get_mla()->set_input_kq_mask(self_kq_mask_mla, ubatch, cparams.causal_attn);
539+
540+
mctx->get_lid()->set_input_k_idxs(self_k_idxs_lid, ubatch);
541+
542+
mctx->get_lid()->set_input_kq_mask(self_kq_mask_lid, ubatch, cparams.causal_attn);
543+
544+
mctx->get_lid()->set_input_k_rot(self_k_rot_lid);
545+
}
546+
547+
bool llm_graph_input_attn_k_dsa::can_reuse(const llm_graph_params & params) {
548+
const auto * mctx = static_cast<const llama_kv_cache_dsa_context *>(params.mctx);
549+
550+
this->mctx = mctx;
551+
552+
bool res = true;
553+
554+
res &= self_k_idxs_mla->ne[0] == params.ubatch.n_tokens;
555+
res &= self_k_idxs_lid->ne[0] == params.ubatch.n_tokens;
556+
557+
res &= can_reuse_kq_mask(self_kq_mask_mla, mctx->get_mla(), params.ubatch, params.cparams);
558+
res &= can_reuse_kq_mask(self_kq_mask_lid, mctx->get_lid(), params.ubatch, params.cparams);
559+
560+
return res;
561+
}
562+
534563
void llm_graph_input_attn_kv_iswa::set_input(const llama_ubatch * ubatch) {
535564
// base tensors may not be allocated if there are no non-SWA attention layers
536565
if (self_k_idxs && self_k_idxs->buffer) {
@@ -2396,6 +2425,82 @@ ggml_tensor * llm_graph_context::build_attn(
23962425
return cur;
23972426
}
23982427

2428+
ggml_tensor * llm_graph_context::build_attn(
2429+
llm_graph_input_attn_k_dsa * inp,
2430+
ggml_tensor * wo,
2431+
ggml_tensor * wo_b,
2432+
ggml_tensor * wo_s,
2433+
ggml_tensor * q_cur,
2434+
ggml_tensor * k_cur,
2435+
ggml_tensor * v_cur,
2436+
ggml_tensor * kq_b,
2437+
ggml_tensor * sinks,
2438+
ggml_tensor * v_mla,
2439+
ggml_tensor * top_k,
2440+
float kq_scale,
2441+
int il) const {
2442+
// these nodes are added to the graph together so that they are not reordered
2443+
// by doing so, the number of splits in the graph is reduced
2444+
// expand k later to enable rope fusion which directly writes into k-v cache
2445+
ggml_build_forward_expand(gf, q_cur);
2446+
ggml_build_forward_expand(gf, v_cur);
2447+
ggml_build_forward_expand(gf, k_cur);
2448+
2449+
const auto * mctx_cur = inp->mctx->get_mla();
2450+
2451+
// store to KV cache
2452+
{
2453+
const auto & k_idxs = inp->get_k_idxs_mla();
2454+
2455+
ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il));
2456+
}
2457+
2458+
const auto & kq_mask = inp->get_kq_mask_mla();
2459+
2460+
// prepare new kq mask - starts filled with -INFINITY
2461+
ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);
2462+
2463+
// reshape KQ mask into tensor with rows of size 1:
2464+
// [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream]
2465+
kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0);
2466+
2467+
// reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1]
2468+
ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0);
2469+
2470+
// prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream]
2471+
// this will be our source of zero values for unmasking top k mask elements
2472+
ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]);
2473+
zeros = ggml_fill(ctx0, zeros, 0.0f);
2474+
2475+
// modify KQ mask by unmasking elements that are in top_k indices
2476+
// ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1])
2477+
ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d);
2478+
2479+
// reshape to restore the original shape of KQ mask:
2480+
// [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream]
2481+
kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0);
2482+
2483+
// combine with the original kq mask
2484+
kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask);
2485+
2486+
ggml_tensor * q = q_cur;
2487+
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
2488+
ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0);
2489+
2490+
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, sinks, v_mla, kq_scale, il);
2491+
cb(cur, "kqv_out", il);
2492+
2493+
if (wo) {
2494+
cur = build_lora_mm(wo, cur, wo_s);
2495+
}
2496+
2497+
if (wo_b) {
2498+
cur = ggml_add(ctx0, cur, wo_b);
2499+
}
2500+
2501+
return cur;
2502+
}
2503+
23992504
ggml_tensor * llm_graph_context::build_attn(
24002505
llm_graph_input_attn_kv_iswa * inp,
24012506
ggml_tensor * wo,
@@ -2542,6 +2647,30 @@ ggml_tensor * llm_graph_context::build_attn(
25422647
return cur;
25432648
}
25442649

2650+
llm_graph_input_attn_k_dsa * llm_graph_context::build_attn_inp_k_dsa() const {
2651+
const auto * mctx_cur = static_cast<const llama_kv_cache_dsa_context *>(mctx);
2652+
2653+
auto inp = std::make_unique<llm_graph_input_attn_k_dsa>(hparams, cparams, mctx_cur);
2654+
2655+
{
2656+
inp->self_k_idxs_mla = mctx_cur->get_mla()->build_input_k_idxs(ctx0, ubatch);
2657+
2658+
inp->self_kq_mask_mla = build_attn_inp_kq_mask(ctx0, mctx_cur->get_mla(), ubatch, cparams);
2659+
inp->self_kq_mask_mla_cnv = cparams.flash_attn ? ggml_cast(ctx0, inp->self_kq_mask_mla, GGML_TYPE_F16) : inp->self_kq_mask_mla;
2660+
}
2661+
2662+
{
2663+
inp->self_k_idxs_lid = mctx_cur->get_lid()->build_input_k_idxs(ctx0, ubatch);
2664+
2665+
inp->self_kq_mask_lid = build_attn_inp_kq_mask(ctx0, mctx_cur->get_lid(), ubatch, cparams);
2666+
inp->self_kq_mask_lid_cnv = cparams.flash_attn ? ggml_cast(ctx0, inp->self_kq_mask_lid, GGML_TYPE_F16) : inp->self_kq_mask_lid;
2667+
2668+
inp->self_k_rot_lid = mctx_cur->get_lid()->build_input_k_rot(ctx0);
2669+
}
2670+
2671+
return (llm_graph_input_attn_k_dsa *) res->add_input(std::move(inp));
2672+
}
2673+
25452674
// TODO: maybe separate the inner implementation into a separate function
25462675
// like with the non-sliding window equivalent
25472676
// once sliding-window hybrid caches are a thing.

0 commit comments

Comments
 (0)