DeepseekV4 MTP + DSpark (#25784)

This commit is contained in:
Aman Gupta
2026-08-02 20:55:34 +08:00
committed by GitHub
parent f5919bf458
commit 596a5795bd
14 changed files with 1534 additions and 225 deletions
+62 -1
View File
@@ -471,6 +471,45 @@ public:
const llama_kv_cache_iswa_context * mctx;
};
class llm_graph_input_attn_k_iswa : public llm_graph_input_i {
public:
llm_graph_input_attn_k_iswa(
const llama_hparams & hparams,
const llama_cparams & cparams,
const llama_kv_cache_iswa_context * mctx) :
hparams(hparams),
cparams(cparams),
mctx(mctx) {
}
~llm_graph_input_attn_k_iswa() = default;
void set_input(const llama_ubatch * ubatch) override;
bool can_reuse(const llm_graph_params & params) override;
ggml_tensor * get_k_idxs() const { return self_k_idxs; }
ggml_tensor * get_k_idxs_swa() const { return self_k_idxs_swa; }
ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }
ggml_tensor * get_kq_mask_swa() const { return self_kq_mask_swa_cnv; }
ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch]
ggml_tensor * self_k_idxs_swa = nullptr; // I64 [n_batch]
ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
ggml_tensor * self_kq_mask_swa = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
ggml_tensor * self_kq_mask_swa_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
ggml_tensor * self_k_rot = nullptr;
ggml_tensor * self_k_rot_swa = nullptr;
const llama_hparams hparams;
const llama_cparams cparams;
const llama_kv_cache_iswa_context * mctx;
};
// DSV4 raw graph inputs are SWA-only, but their mask may be stream-shaped
// so raw K can be concatenated with DSV4 compressed K in one attention op.
class llm_graph_input_dsv4_raw {
@@ -505,6 +544,10 @@ public:
ggml_tensor * state_pos = nullptr; // I32 [n_state]
ggml_tensor * state_persist_src_idxs = nullptr; // I32 [n_state_persist]
ggml_tensor * state_persist_dst_idxs = nullptr; // I32 [n_state_persist]
ggml_tensor * state_restore_src_idxs = nullptr; // I32 [n_state_restore]
ggml_tensor * state_restore_dst_idxs = nullptr; // I32 [n_state_restore]
ggml_tensor * state_snapshot_src_idxs = nullptr; // I32 [n_state_snapshot]
ggml_tensor * state_snapshot_dst_idxs = nullptr; // I32 [n_state_snapshot]
ggml_tensor * state_read_idxs = nullptr; // I32 [ratio*n_state_write]
ggml_tensor * state_write_idxs = nullptr; // I64 [n_state_write]
ggml_tensor * state_write_pos = nullptr; // I32 [n_state_write]
@@ -1068,7 +1111,7 @@ struct llm_graph_context {
ggml_tensor * build_attn_mha(
ggml_tensor * q, // [n_embd_head_q, n_head_q, n_tokens]
ggml_tensor * k, // [n_embd_head_k, n_head_k, n_tokens]
ggml_tensor * v, // [n_embd_head_v, n_head_v, n_tokens] (v_trans == false)
ggml_tensor * v, // [n_embd_head_v, n_head_v, n_tokens] (v_trans = false)
ggml_tensor * kq_b,
ggml_tensor * kq_mask,
ggml_tensor * sinks, // [n_head_q]
@@ -1160,6 +1203,24 @@ struct llm_graph_context {
float kq_scale,
int il) const;
llm_graph_input_attn_k_iswa * build_attn_inp_k_iswa() const;
// note: if k_cur is not provided, it will not be stored in the memory
// note: the K cache is used as V (MLA-style attention)
ggml_tensor * build_attn(
llm_graph_input_attn_k_iswa * inp,
ggml_tensor * wo,
ggml_tensor * wo_b,
ggml_tensor * wo_s,
ggml_tensor * q_cur, // [n_embd_head_q, n_head_q, n_tokens]
ggml_tensor * k_cur, // [n_embd_head_k, n_head_k, n_tokens] optional
ggml_tensor * v_cur, // [n_embd_head_v, n_head_v, n_tokens] optional
ggml_tensor * kq_b,
ggml_tensor * sinks, // [n_head_q]
ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v]
float kq_scale,
int il) const;
llm_graph_input_attn_cross * build_attn_inp_cross() const;
ggml_tensor * build_attn(