DeepseekV4 MTP + DSpark (#25784)
This commit is contained in:
+62
-1
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user