Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
47 commits
Select commit Hold shift + click to select a range
87cc29b
Add preliminary MiniMax-M3 support
danielhanchen Jun 12, 2026
53c81dd
MiniMax-M3 vision tower (mmproj + clip graph)
timkhronos Jun 17, 2026
f07f1d4
Delete m3_vision_ref.py
timkhronos Jun 17, 2026
4a3206f
Update clip.cpp
timkhronos Jun 17, 2026
bbf1a80
MSA
timkhronos Jun 19, 2026
09657dd
Update constants.py
timkhronos Jun 19, 2026
8fe2e01
Update minimax.py
timkhronos Jun 19, 2026
8c953a9
Cache creation. Working withotu flash attention
timkhronos Jun 19, 2026
ea6fbd6
Added flash attention for sparse layers
timkhronos Jun 19, 2026
2e82759
Decomposed slow cpu OP into GPU + CPU ops. Massive speedup over long ctx
timkhronos Jun 20, 2026
0152226
Rewrote indexer op to be cuda native. Modified flash attention to mat…
timkhronos Jun 21, 2026
d1a04f7
Implement sparse attention calc out of stock ops.
timkhronos Jun 21, 2026
b1b174e
Fix a cache allocation and cont issue
timkhronos Jun 22, 2026
cea714a
Fixed -fa auto crash, flagged debug spots
timkhronos Jun 22, 2026
69c958d
Delete vocab.json
timkhronos Jun 22, 2026
afc09f1
Delete model.safetensors.index.json
timkhronos Jun 22, 2026
f40d0f5
Delete generation_config.json
timkhronos Jun 22, 2026
714bbe9
Delete Minimax directory
timkhronos Jun 22, 2026
8136a9c
Handled multi stream case to fall back on Dense Attention
timkhronos Jun 23, 2026
3ed9b18
Development scaffolding cleanup. No functional change to the decode or
timkhronos Jun 23, 2026
79a6eec
Remove redundant comment from minimax-m3.cpp
timkhronos Jun 23, 2026
35990be
Changed 3 Gelu Ops for vision into Gelu_erf ops
timkhronos Jun 25, 2026
fa15850
Assert that n_kv is multiple of 128
timkhronos Jun 25, 2026
eae55e2
Rename MSA index tensors to indexer convention
timkhronos Jun 26, 2026
2bb7eeb
Fix incorrect Assert
timkhronos Jun 26, 2026
d6f9426
Review driven changes (#3)
timkhronos Jun 28, 2026
0a7b2dc
Remove comment from conversion minimax.py
timkhronos Jun 28, 2026
636143a
Remove whitespaces from constants.py
timkhronos Jun 28, 2026
7b7ff65
Tighten comment in minimax.py
timkhronos Jun 28, 2026
1cd03ef
inherit MiniMax-M3 from MiniMax-M2
timkhronos Jun 28, 2026
c70c8a9
drop dead text_config fallbacks
timkhronos Jun 28, 2026
25199fa
Add indexer writer methods
timkhronos Jun 28, 2026
d54b4eb
Reuse LLM_FFN_SWIGLU_OAI_MOE
timkhronos Jun 28, 2026
f57efce
Remove duplicate indexer setters, add only block_size/local_blocks, …
timkhronos Jun 28, 2026
618e145
Fix conversion error /gguf_writer.py
timkhronos Jun 28, 2026
5a24782
Update gguf-py/gguf/gguf_writer.py
timkhronos Jun 28, 2026
5dfe838
Update gguf-py/gguf/tensor_mapping.py
timkhronos Jun 28, 2026
b99a8d5
Update conversion/minimax.py
timkhronos Jun 28, 2026
56ba541
Update conversion/minimax.py
timkhronos Jun 28, 2026
044391a
Remove whitespace in src/llama-kv-cache.cpp
timkhronos Jun 28, 2026
c314faf
Remove Whitespace in Update src/llama-model.h
timkhronos Jun 28, 2026
5015a6b
Remove whitespace in src/llama-hparams.h
timkhronos Jun 28, 2026
c97a2ac
remove multimodal code upon maintainer request. Will be made as a sep…
timkhronos Jun 28, 2026
fcb9ed2
Whitespace clean in tensor_mapping.py
timkhronos Jun 29, 2026
51e7bf8
Merge branch 'master' into MSA
timkhronos Jun 30, 2026
6b14fea
Merge branch 'ggml-org:master' into MSA
timkhronos Jul 1, 2026
d2fe995
Log cache size on launch, block ctx shift, support prompt caching
timkhronos Jul 3, 2026
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions conversion/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -155,6 +155,8 @@
"MiniCPMForCausalLM": "minicpm",
"MiniCPMV4_6ForConditionalGeneration": "minicpm",
"MiniMaxM2ForCausalLM": "minimax",
"MiniMaxM3SparseForCausalLM": "minimax",
"MiniMaxM3SparseForConditionalGeneration": "minimax",
"Ministral3ForCausalLM": "mistral3",
"Mistral3ForConditionalGeneration": "mistral3",
"MistralForCausalLM": "llama",
Expand Down
4 changes: 2 additions & 2 deletions conversion/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -1154,7 +1154,7 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca
or "projector." in name or "pre_mm_projector_norm" in name \
or "image_newline" in name or "view_seperator" in name \
or "patch_embed" in name or "patch_embedding" in name \
or "patch_merger." in name or "model.connector." in name:
or "patch_merger." in name or "patch_merge_mlp." in name or "model.connector." in name:
return None

return super().filter_tensors(item)
Expand Down Expand Up @@ -1201,7 +1201,7 @@ def set_gguf_parameters(self):
self.gguf_writer.add_embedding_length(n_embd)
logger.info(f"gguf: embedding length = {n_embd}")

if (n_ff := self.find_hparam(["prefix_dense_intermediate_size", "intermediate_size", "n_inner", "hidden_dim"], optional=True)) is not None:
if (n_ff := self.find_hparam(["prefix_dense_intermediate_size", "dense_intermediate_size", "intermediate_size", "n_inner", "hidden_dim"], optional=True)) is not None:
self.gguf_writer.add_feed_forward_length(n_ff)
logger.info(f"gguf: feed forward length = {n_ff}")

Expand Down
36 changes: 35 additions & 1 deletion conversion/minimax.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ def set_gguf_parameters(self):

def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None):
# merge expert weights
if 'experts' in name:
if "block_sparse_moe.experts." in name:
n_experts = self.find_hparam(["num_local_experts", "num_experts"])
assert bid is not None

Expand Down Expand Up @@ -52,3 +52,37 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None):
return

yield from super().modify_tensors(data_torch, name, bid)


@ModelBase.register("MiniMaxM3SparseForCausalLM", "MiniMaxM3SparseForConditionalGeneration")
class MiniMaxM3Model(MiniMaxM2Model):
model_arch = gguf.MODEL_ARCH.MINIMAXM3
def set_gguf_parameters(self):
super().set_gguf_parameters()

self.gguf_writer.add_expert_shared_count(self.find_hparam(["n_shared_experts"]))
self.gguf_writer.add_expert_weights_scale(self.find_hparam(["routed_scaling_factor"]))
self.gguf_writer.add_expert_weights_norm(True)

sac = self.find_hparam(["sparse_attention_config"])
self.gguf_writer.add_indexer_head_count(sac["sparse_num_index_heads"])
self.gguf_writer.add_indexer_key_length(sac["sparse_index_dim"])
self.gguf_writer.add_indexer_top_k(sac["sparse_topk_blocks"])
self.gguf_writer.add_indexer_block_size(sac["sparse_block_size"])
self.gguf_writer.add_indexer_local_blocks(sac["sparse_local_block"])

moe_layer_freq = self.find_hparam(["moe_layer_freq"])
n_dense = 0
for v in moe_layer_freq:
if v == 0:
n_dense += 1
else:
break
self.gguf_writer.add_leading_dense_block_count(n_dense)

def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None):
# Gemma-style (1 + w) RMSNorm: bake the +1 in so llama.cpp can use plain RMSNorm
if name.endswith("norm.weight"):
data_torch = data_torch + 1.0

yield from super().modify_tensors(data_torch, name, bid)
38 changes: 38 additions & 0 deletions gguf-py/gguf/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -200,6 +200,8 @@ class Indexer:
HEAD_COUNT = "{arch}.attention.indexer.head_count"
KEY_LENGTH = "{arch}.attention.indexer.key_length"
TOP_K = "{arch}.attention.indexer.top_k"
BLOCK_SIZE = "{arch}.attention.indexer.block_size" #MSA
LOCAL_BLOCKS = "{arch}.attention.indexer.local_blocks" #MSA

class HyperConnection:
COUNT = "{arch}.hyper_connection.count"
Expand Down Expand Up @@ -525,6 +527,7 @@ class MODEL_ARCH(IntEnum):
APERTUS = auto()
COGVLM = auto()
MINIMAXM2 = auto()
MINIMAXM3 = auto()
RND1 = auto()
PANGU_EMBED = auto()
MISTRAL3 = auto()
Expand Down Expand Up @@ -771,6 +774,9 @@ class MODEL_TENSOR(IntEnum):
INDEXER_PROJ = auto()
INDEXER_ATTN_K = auto()
INDEXER_ATTN_Q_B = auto()
INDEXER_Q_PROJ = auto()
INDEXER_K_PROJ = auto()
INDEXER_Q_NORM = auto()
INDEXER_COMPRESSOR_WKV = auto()
INDEXER_COMPRESSOR_WGATE = auto()
INDEXER_COMPRESSOR_APE = auto()
Expand Down Expand Up @@ -1105,6 +1111,7 @@ class MODEL_TENSOR(IntEnum):
MODEL_ARCH.GROVEMOE: "grovemoe",
MODEL_ARCH.APERTUS: "apertus",
MODEL_ARCH.MINIMAXM2: "minimax-m2",
MODEL_ARCH.MINIMAXM3: "minimax-m3",
MODEL_ARCH.COGVLM: "cogvlm",
MODEL_ARCH.RND1: "rnd1",
MODEL_ARCH.PANGU_EMBED: "pangu-embedded",
Expand Down Expand Up @@ -1350,6 +1357,9 @@ class MODEL_TENSOR(IntEnum):
MODEL_TENSOR.INDEXER_PROJ: "blk.{bid}.indexer.proj",
MODEL_TENSOR.INDEXER_ATTN_K: "blk.{bid}.indexer.attn_k",
MODEL_TENSOR.INDEXER_ATTN_Q_B: "blk.{bid}.indexer.attn_q_b",
MODEL_TENSOR.INDEXER_Q_PROJ: "blk.{bid}.indexer.q_proj",
MODEL_TENSOR.INDEXER_K_PROJ: "blk.{bid}.indexer.k_proj",
MODEL_TENSOR.INDEXER_Q_NORM: "blk.{bid}.indexer.q_norm",
MODEL_TENSOR.INDEXER_COMPRESSOR_WKV: "blk.{bid}.indexer_compressor_kv",
MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE: "blk.{bid}.indexer_compressor_gate",
MODEL_TENSOR.INDEXER_COMPRESSOR_APE: "blk.{bid}.indexer_compressor_ape",
Expand Down Expand Up @@ -4102,6 +4112,34 @@ class MODEL_TENSOR(IntEnum):
MODEL_TENSOR.FFN_UP_EXP,
MODEL_TENSOR.FFN_EXP_PROBS_B,
],
MODEL_ARCH.MINIMAXM3: [
MODEL_TENSOR.TOKEN_EMBD,
MODEL_TENSOR.OUTPUT_NORM,
MODEL_TENSOR.OUTPUT,
MODEL_TENSOR.ATTN_NORM,
MODEL_TENSOR.ATTN_Q,
MODEL_TENSOR.ATTN_Q_NORM,
MODEL_TENSOR.ATTN_K,
MODEL_TENSOR.ATTN_K_NORM,
MODEL_TENSOR.ATTN_V,
MODEL_TENSOR.ATTN_OUT,
MODEL_TENSOR.FFN_NORM,
MODEL_TENSOR.FFN_GATE_INP,
MODEL_TENSOR.FFN_EXP_PROBS_B,
MODEL_TENSOR.FFN_GATE_EXP,
MODEL_TENSOR.FFN_DOWN_EXP,
MODEL_TENSOR.FFN_UP_EXP,
MODEL_TENSOR.FFN_GATE_SHEXP,
MODEL_TENSOR.FFN_DOWN_SHEXP,
MODEL_TENSOR.FFN_UP_SHEXP,
MODEL_TENSOR.FFN_GATE,
MODEL_TENSOR.FFN_DOWN,
MODEL_TENSOR.FFN_UP,
MODEL_TENSOR.INDEXER_Q_PROJ,
MODEL_TENSOR.INDEXER_K_PROJ,
MODEL_TENSOR.INDEXER_Q_NORM,
MODEL_TENSOR.INDEXER_K_NORM,
],
MODEL_ARCH.COGVLM: [
MODEL_TENSOR.TOKEN_EMBD,
MODEL_TENSOR.OUTPUT_NORM,
Expand Down
6 changes: 6 additions & 0 deletions gguf-py/gguf/gguf_writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -793,6 +793,12 @@ def add_indexer_key_length(self, length: int) -> None:
def add_indexer_top_k(self, top_k: int) -> None:
self.add_uint32(Keys.Attention.Indexer.TOP_K.format(arch=self.arch), top_k)

def add_indexer_block_size(self, block_size: int) -> None:
self.add_uint32(Keys.Attention.Indexer.BLOCK_SIZE.format(arch=self.arch), block_size)

def add_indexer_local_blocks(self, local_blocks: int) -> None:
self.add_uint32(Keys.Attention.Indexer.LOCAL_BLOCKS.format(arch=self.arch), local_blocks)

def add_max_alibi_bias(self, bias: float) -> None:
self.add_float32(Keys.Attention.MAX_ALIBI_BIAS.format(arch=self.arch), bias)

Expand Down
15 changes: 14 additions & 1 deletion gguf-py/gguf/tensor_mapping.py
Original file line number Diff line number Diff line change
Expand Up @@ -1263,7 +1263,8 @@ class TensorNameMap:
),

MODEL_TENSOR.INDEXER_K_NORM: (
"model.layers.{bid}.self_attn.indexer.k_norm", # DSA
"model.layers.{bid}.self_attn.indexer.k_norm", # DSA
"model.layers.{bid}.self_attn.index_k_norm", # MSA
),

MODEL_TENSOR.INDEXER_PROJ: (
Expand All @@ -1278,6 +1279,18 @@ class TensorNameMap:
"model.layers.{bid}.self_attn.indexer.wq_b", # DSA
),

MODEL_TENSOR.INDEXER_Q_PROJ: (
"model.layers.{bid}.self_attn.index_q_proj", # MSA
),

MODEL_TENSOR.INDEXER_K_PROJ: (
"model.layers.{bid}.self_attn.index_k_proj", # MSA
),

MODEL_TENSOR.INDEXER_Q_NORM: (
"model.layers.{bid}.self_attn.index_q_norm", # MSA
),

############################################################################
# TODO: these do not belong to block_mappings_cfg - move them to mappings_cfg
MODEL_TENSOR.ENC_OUTPUT_NORM: (
Expand Down
10 changes: 10 additions & 0 deletions src/llama-arch.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
{ LLM_ARCH_GROVEMOE, "grovemoe" },
{ LLM_ARCH_APERTUS, "apertus" },
{ LLM_ARCH_MINIMAX_M2, "minimax-m2" },
{ LLM_ARCH_MINIMAX_M3, "minimax-m3" },
{ LLM_ARCH_COGVLM, "cogvlm" },
{ LLM_ARCH_RND1, "rnd1" },
{ LLM_ARCH_PANGU_EMBED, "pangu-embedded" },
Expand Down Expand Up @@ -251,6 +252,8 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
{ LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, "%s.attention.indexer.head_count" },
{ LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, "%s.attention.indexer.key_length" },
{ LLM_KV_ATTENTION_INDEXER_TOP_K, "%s.attention.indexer.top_k" },
{ LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, "%s.attention.indexer.block_size" },
{ LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, "%s.attention.indexer.local_blocks" },
{ LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, "%s.attention.output_group_count" },
{ LLM_KV_ATTENTION_OUTPUT_LORA_RANK, "%s.attention.output_lora_rank" },
{ LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, "%s.attention.compress_rope_freq_base" },
Expand Down Expand Up @@ -594,6 +597,9 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
{ LLM_TENSOR_INDEXER_PROJ, "blk.%d.indexer.proj" },
{ LLM_TENSOR_INDEXER_ATTN_K, "blk.%d.indexer.attn_k" },
{ LLM_TENSOR_INDEXER_ATTN_Q_B, "blk.%d.indexer.attn_q_b" },
{ LLM_TENSOR_INDEXER_Q_PROJ, "blk.%d.indexer.q_proj" },
{ LLM_TENSOR_INDEXER_K_PROJ, "blk.%d.indexer.k_proj" },
{ LLM_TENSOR_INDEXER_Q_NORM, "blk.%d.indexer.q_norm" },
{ LLM_TENSOR_INDEXER_COMPRESSOR_WKV, "blk.%d.indexer_compressor_kv" },
{ LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "blk.%d.indexer_compressor_gate" },
{ LLM_TENSOR_INDEXER_COMPRESSOR_APE, "blk.%d.indexer_compressor_ape" },
Expand Down Expand Up @@ -829,6 +835,9 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
{LLM_TENSOR_INDEXER_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_ATTN_K, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_ATTN_Q_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_Q_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_K_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_Q_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_INDEXER_COMPRESSOR_WKV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_COMPRESSOR_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}},
Expand Down Expand Up @@ -998,6 +1007,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
case LLM_ARCH_LFM2:
case LLM_ARCH_LFM2MOE:
case LLM_ARCH_MINIMAX_M2:
case LLM_ARCH_MINIMAX_M3:
case LLM_ARCH_MISTRAL4:
case LLM_ARCH_KIMI_LINEAR:
return false;
Expand Down
6 changes: 6 additions & 0 deletions src/llama-arch.h
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,7 @@ enum llm_arch {
LLM_ARCH_TALKIE,
LLM_ARCH_MELLUM,
LLM_ARCH_EAGLE3,
LLM_ARCH_MINIMAX_M3,
LLM_ARCH_DFLASH,
LLM_ARCH_UNKNOWN,
};
Expand Down Expand Up @@ -256,6 +257,8 @@ enum llm_kv {
LLM_KV_ATTENTION_INDEXER_HEAD_COUNT,
LLM_KV_ATTENTION_INDEXER_KEY_LENGTH,
LLM_KV_ATTENTION_INDEXER_TOP_K,
LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE,
LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS,
LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT,
LLM_KV_ATTENTION_OUTPUT_LORA_RANK,
LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE,
Expand Down Expand Up @@ -594,6 +597,9 @@ enum llm_tensor {
LLM_TENSOR_INDEXER_PROJ,
LLM_TENSOR_INDEXER_ATTN_K,
LLM_TENSOR_INDEXER_ATTN_Q_B,
LLM_TENSOR_INDEXER_Q_NORM,
LLM_TENSOR_INDEXER_K_PROJ,
LLM_TENSOR_INDEXER_Q_PROJ,
LLM_TENSOR_INDEXER_COMPRESSOR_WKV,
LLM_TENSOR_INDEXER_COMPRESSOR_WGATE,
LLM_TENSOR_INDEXER_COMPRESSOR_APE,
Expand Down
3 changes: 2 additions & 1 deletion src/llama-context.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -2325,7 +2325,8 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
model.arch == LLM_ARCH_KIMI_LINEAR ||
model.arch == LLM_ARCH_QWEN35 ||
model.arch == LLM_ARCH_QWEN35MOE ||
model.arch == LLM_ARCH_DEEPSEEK4) {
model.arch == LLM_ARCH_DEEPSEEK4 ||
model.arch == LLM_ARCH_MINIMAX_M3) {
return std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
}
uint32_t res = std::max<uint32_t>(1024u, 8u*model.n_tensors());
Expand Down
16 changes: 14 additions & 2 deletions src/llama-graph.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1700,6 +1700,17 @@ ggml_tensor * llm_graph_context::build_ffn(
cur = ggml_swiglu(ctx0, cur);
cb(cur, "ffn_swiglu", il);
} break;
case LLM_FFN_SWIGLU_OAI_MOE:
if (gate && type_gate == LLM_FFN_PAR) {
//Same constants as LLM_FFN_SWIGLU_OAI_MOE
const float alpha = 1.702f;
const float limit = 7.0f;
cur = ggml_swiglu_oai(ctx0, cur, tmp, alpha, limit);
cb(cur, "ffn_swiglu_oai", il);
type_gate = LLM_FFN_SEQ;
} else {
GGML_ABORT("LLM_FFN_SWIGLU_OAI requires a parallel gate");
} break;
case LLM_FFN_GEGLU:
{
cur = ggml_geglu(ctx0, cur);
Expand Down Expand Up @@ -2629,7 +2640,8 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_tensor * sinks,
ggml_tensor * v_mla, // TODO: remove
float kq_scale,
int il) const {
int il,
ggml_tensor * kq_mask_override) const {
GGML_ASSERT(v_mla == nullptr);

if (inp->self_k_rot) {
Expand Down Expand Up @@ -2659,7 +2671,7 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il));
}

const auto & kq_mask = inp->get_kq_mask();
ggml_tensor * kq_mask = kq_mask_override ? kq_mask_override : inp->get_kq_mask();

ggml_tensor * q = q_cur;
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
Expand Down
3 changes: 2 additions & 1 deletion src/llama-graph.h
Original file line number Diff line number Diff line change
Expand Up @@ -1085,7 +1085,8 @@ struct llm_graph_context {
ggml_tensor * sinks, // [n_head_q]
ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] // TODO: remove
float kq_scale,
int il) const;
int il,
ggml_tensor * kq_mask_override = nullptr) const;

llm_graph_input_attn_k * build_attn_inp_k() const;

Expand Down
10 changes: 10 additions & 0 deletions src/llama-hparams.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,16 @@ uint32_t llama_hparams::n_embd_v_gqa_max() const {
return val;
}

uint32_t llama_hparams::n_embd_k_idx(uint32_t il) const {
if (indexer_head_size == 0) {
return 0; // arch without MSA
}
if (il < n_layer_dense_lead) {
return 0; // leading dense layers carry no indexer
}
return indexer_head_size; // 128
}

uint32_t llama_hparams::n_embd_r() const {
if (wkv_head_size != 0) {
// for RWKV models
Expand Down
6 changes: 6 additions & 0 deletions src/llama-hparams.h
Original file line number Diff line number Diff line change
Expand Up @@ -226,6 +226,9 @@ struct llama_hparams {
uint32_t indexer_n_head = 0;
uint32_t indexer_head_size = 0;
uint32_t indexer_top_k = 0;
// MSA
uint32_t indexer_block_size = 0;
uint32_t indexer_local_blocks = 0;

// DeepSeek-V4
uint32_t dsv4_o_group_count = 0;
Expand Down Expand Up @@ -344,6 +347,9 @@ struct llama_hparams {
uint32_t n_embd_k_gqa_max() const;
uint32_t n_embd_v_gqa_max() const;

// dimension of the single-head MSA indexer key stream
uint32_t n_embd_k_idx(uint32_t il = 0) const;

// dimension of the rolling state embeddings
// corresponds to Mamba's conv_states size or RWKV's token_shift states size
uint32_t n_embd_r() const;
Expand Down
Loading