glm5-next : QSA without kq_mask (+ DSA) (#31009)
* glm5-next : QSA without kq_mask * DSA: dsa top-k mask change tested with layers 0,3 of DeepSeek-V3.2
This commit is contained in:
1 parent
fae9ac561f
commit
d78f40d9d3
2 files changed
+21
-7
No files matched your search
+14
-3
@@ -634,7 +634,10 @@ bool llm_graph_input_attn_kv_msa::can_reuse(const llm_graph_params & params) {
|
||||
void llm_graph_input_attn_k_dsa::set_input(const llama_ubatch * ubatch) {
|
||||
mctx->get_mla()->set_input_k_idxs(self_k_idxs_mla, ubatch);
|
||||
|
||||
mctx->get_mla()->set_input_kq_mask(self_kq_mask_mla, ubatch, cparams.causal_attn);
|
||||
// leave unallocated when the top-k layers use the indexer mask
|
||||
if (self_kq_mask_mla->buffer) {
|
||||
mctx->get_mla()->set_input_kq_mask(self_kq_mask_mla, ubatch, cparams.causal_attn);
|
||||
}
|
||||
|
||||
mctx->get_lid()->set_input_k_idxs(self_k_idxs_lid, ubatch);
|
||||
|
||||
@@ -1218,7 +1221,10 @@ bool llm_graph_input_mem_hybrid::can_reuse(const llm_graph_params & params) {
|
||||
void llm_graph_input_mem_hybrid_k::set_input(const llama_ubatch * ubatch) {
|
||||
mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch);
|
||||
|
||||
mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn);
|
||||
// leave the mask unallocated if all attention layers have their own (e.g. glm5-next DSA)
|
||||
if (inp_attn->self_kq_mask->buffer) {
|
||||
mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn);
|
||||
}
|
||||
|
||||
const int64_t n_rs = mctx->get_recr()->get_n_rs();
|
||||
|
||||
@@ -3441,7 +3447,12 @@ ggml_tensor * llm_graph_context::build_attn(
|
||||
ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il));
|
||||
}
|
||||
|
||||
const auto & kq_mask = inp->get_kq_mask_mla();
|
||||
// top_k indexes the cells, and the indexer mask of the same type is built the same way
|
||||
// let's use it, so when no other layer reads the MLA mask is not allocated
|
||||
ggml_tensor * kq_mask = inp->get_kq_mask_mla();
|
||||
if (inp->get_kq_mask_lid()->type == kq_mask->type) {
|
||||
kq_mask = inp->get_kq_mask_lid();
|
||||
}
|
||||
|
||||
ggml_tensor * kq_mask_top_k = nullptr;
|
||||
if (cparams.fused_top_k_mask) {
|
||||
|
||||
Reference in new issue
Block a user