Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
6 changes: 4 additions & 2 deletions common/speculative.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1420,7 +1420,9 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
llama_set_embeddings_nextn(ctx_tgt, true, /*masked*/ false);
llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true);

is_mem_shared = llama_get_ctx_other(ctx_dft) == ctx_tgt;
char arch[64] = {0};

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

see https://github.com/ggml-org/llama.cpp/pull/28243/changes#r4082502758. If we follow how other models support MTP, this will not be needed. e.g. #26725

llama_model_meta_val_str(llama_get_model(ctx_dft), "general.architecture", arch, sizeof(arch));
is_mem_shared = llama_get_ctx_other(ctx_dft) == ctx_tgt && std::strcmp(arch, "gemma4-assistant") == 0;
chain_heads = n_mtp_layers > 1 && !is_mem_shared;

if (chain_heads) {
Expand Down Expand Up @@ -2559,7 +2561,7 @@ common_speculative_init_result::common_speculative_init_result(
model_path = params.speculative.draft.mparams.path;
LOG_INF("%s: loading draft model '%s'\n", __func__, model_path.c_str());

llama_model * model_dft = llama_model_load_from_file(params.model.path.c_str(), mparams);
llama_model * model_dft = llama_model_load_from_file(model_path.c_str(), mparams);

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This isn't a fix, both expressions are the same string. I'd suggest just dropping this hunk from the PR.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ruixiang63 Are you sure about that? See 3fb9b98

@ruixiang63 ruixiang63 Sep 23, 2026 •

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Have you encountered any issues with MTP on the current master branch? If so, please open a separate issue.
This change does not seem relevant to the current Qwen3.8-Flash-Next MTP implementation.

if (model_dft == NULL) {
LOG_ERR("%s: failed to load draft model, '%s'\n", __func__, model_path.c_str());
return;
Expand Down
4 changes: 2 additions & 2 deletions conversion/bailingmoe3.py
Original file line number Diff line number Diff line change
Expand Up @@ -121,9 +121,9 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca

if is_mtp and cls.no_mtp:
return None
if cls.mtp_only and not is_mtp and name not in (
if cls.mtp_only and not is_mtp and (cls.mtp_shared_embd or name not in (

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Based on the suggestion below, I don’t think we need mtp_shared_embd anymore.

"model.word_embeddings.weight", "model.norm.weight", "lm_head.weight",
):
)):
return None

return super().filter_tensors((name, gen))
Expand Down
1 change: 1 addition & 0 deletions conversion/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,7 @@
supports_mtp_export: bool = False
mtp_only: bool = False
no_mtp: bool = False
mtp_shared_embd: bool = False

def __init__(self, dir_model: Path, ftype: gguf.LlamaFileType, fname_out: Path, *, is_big_endian: bool = False,
use_temp_file: bool = False, eager: bool = False,
Expand Down Expand Up @@ -251,7 +252,7 @@
if weight_map is None or not isinstance(weight_map, dict):
raise ValueError(f"Can't load 'weight_map' from {index_name!r}")
tensor_names_from_index.update(weight_map.keys())
part_dict: dict[str, None] = dict.fromkeys(weight_map.values(), None) # ty: ignore[invalid-assignment]

Check warning on line 255 in conversion/base.py

View workflow job for this annotation

GitHub Actions / python type-check

ty (unused-ignore-comment)

conversion/base.py:255:91: unused-ignore-comment: Unused `ty: ignore` directive help: Remove the unused suppression comment
part_names = sorted(part_dict.keys())
else:
weight_map = {}
Expand Down Expand Up @@ -1512,15 +1513,15 @@

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(self.dir_model)
vocab_size = self.hparams.get("vocab_size", len(tokenizer.vocab)) # ty: ignore[unresolved-attribute]

Check warning on line 1516 in conversion/base.py

View workflow job for this annotation

GitHub Actions / python type-check

ty (unused-ignore-comment)

conversion/base.py:1516:76: unused-ignore-comment: Unused `ty: ignore` directive help: Remove the unused suppression comment
assert max(tokenizer.vocab.values()) < vocab_size # ty: ignore[unresolved-attribute]

Check warning on line 1517 in conversion/base.py

View workflow job for this annotation

GitHub Actions / python type-check

ty (unused-ignore-comment)

conversion/base.py:1517:60: unused-ignore-comment: Unused `ty: ignore` directive help: Remove the unused suppression comment

tokpre = self.get_vocab_base_pre(tokenizer)

reverse_vocab = {id_: encoded_tok for encoded_tok, id_ in tokenizer.vocab.items()} # ty: ignore[unresolved-attribute]

Check warning on line 1521 in conversion/base.py

View workflow job for this annotation

GitHub Actions / python type-check

ty (unused-ignore-comment)

conversion/base.py:1521:93: unused-ignore-comment: Unused `ty: ignore` directive help: Remove the unused suppression comment
added_vocab = tokenizer.get_added_vocab() # ty: ignore[unresolved-attribute]

Check warning on line 1522 in conversion/base.py

View workflow job for this annotation

GitHub Actions / python type-check

ty (unused-ignore-comment)

conversion/base.py:1522:52: unused-ignore-comment: Unused `ty: ignore` directive help: Remove the unused suppression comment

added_tokens_decoder = tokenizer.added_tokens_decoder # ty: ignore[unresolved-attribute]

Check warning on line 1524 in conversion/base.py

View workflow job for this annotation

GitHub Actions / python type-check

ty (unused-ignore-comment)

conversion/base.py:1524:64: unused-ignore-comment: Unused `ty: ignore` directive help: Remove the unused suppression comment

for i in range(vocab_size):
if i not in reverse_vocab:
Expand All @@ -1533,7 +1534,7 @@
# To avoid unexpected issues - we make sure to normalize non-normalized tokens
if not added_tokens_decoder[i].normalized:
previous_token = token
token = tokenizer.decode(tokenizer.encode(token, add_special_tokens=False)) # ty: ignore[unresolved-attribute, invalid-assignment]

Check warning on line 1537 in conversion/base.py

View workflow job for this annotation

GitHub Actions / python type-check

ty (unused-ignore-comment)

conversion/base.py:1537:102: unused-ignore-comment: Unused `ty: ignore` directive help: Remove the unused suppression comment
if previous_token != token:
logger.info(f"{repr(previous_token)} is encoded and decoded back to {repr(token)} using AutoTokenizer")

Expand Down Expand Up @@ -1912,10 +1913,10 @@
def _set_vocab_hybriddna(self):
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True)
vocab_size = self.hparams.get("vocab_size", len(tokenizer.vocab)) # ty: ignore[unresolved-attribute]

Check warning on line 1916 in conversion/base.py

View workflow job for this annotation

GitHub Actions / python type-check

ty (unused-ignore-comment)

conversion/base.py:1916:76: unused-ignore-comment: Unused `ty: ignore` directive help: Remove the unused suppression comment
assert max(tokenizer.vocab.values()) < vocab_size # ty: ignore[unresolved-attribute]

Check warning on line 1917 in conversion/base.py

View workflow job for this annotation

GitHub Actions / python type-check

ty (unused-ignore-comment)

conversion/base.py:1917:60: unused-ignore-comment: Unused `ty: ignore` directive help: Remove the unused suppression comment

reverse_vocab = {id_: encoded_tok for encoded_tok, id_ in tokenizer.vocab.items()} # ty: ignore[unresolved-attribute]

Check warning on line 1919 in conversion/base.py

View workflow job for this annotation

GitHub Actions / python type-check

ty (unused-ignore-comment)

conversion/base.py:1919:93: unused-ignore-comment: Unused `ty: ignore` directive help: Remove the unused suppression comment
# k-mers can share text with a base-vocab BPE token (e.g. CCCCCC) and get
# dropped by get_vocab(); a reserved marker suffix (U+E000) keeps each
# k-mer's own id (llama.cpp strips it on detokenization)
Expand Down
4 changes: 2 additions & 2 deletions conversion/command_r.py
Original file line number Diff line number Diff line change
Expand Up @@ -131,9 +131,9 @@ def filter_tensors(cls, item):
is_mtp = (m := re.match(r"model\.layers\.(\d+)\.", name)) is not None and int(m.group(1)) >= cls._n_main_layers
if is_mtp and cls.no_mtp:
return None
if cls.mtp_only and not is_mtp and name not in (
if cls.mtp_only and not is_mtp and (cls.mtp_shared_embd or name not in (

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this will not be needed as well.

"model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",
):
)):
return None

return name, gen
Expand Down
4 changes: 2 additions & 2 deletions conversion/dots3.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,9 +99,9 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca
# --no-mtp: drop the NextN/MTP block; --mtp: keep only that block plus the shared embeddings/norm/lm_head
if is_mtp and cls.no_mtp:
return None
if cls.mtp_only and not is_mtp and name not in (
if cls.mtp_only and not is_mtp and (cls.mtp_shared_embd or name not in (
"model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",
):
)):
return None

return name, gen
Expand Down
12 changes: 6 additions & 6 deletions conversion/glm.py
Original file line number Diff line number Diff line change
Expand Up @@ -138,9 +138,9 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca

if is_mtp and cls.no_mtp:
return None
if cls.mtp_only and not is_mtp and name not in (
if cls.mtp_only and not is_mtp and (cls.mtp_shared_embd or name not in (
"model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",
):
)):
return None

return name, gen
Expand Down Expand Up @@ -292,9 +292,9 @@ def filter_tensors(cls, item):
is_mtp = match is not None and int(match.group(1)) >= cls._n_main_layers
if is_mtp and cls.no_mtp:
return None
if cls.mtp_only and not is_mtp and name not in (
if cls.mtp_only and not is_mtp and (cls.mtp_shared_embd or name not in (
"model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",
):
)):
return None

return name, gen
Expand Down Expand Up @@ -352,9 +352,9 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca
return None
# --mtp: keep ONLY NextN-block tensors plus the shared embeddings/
# norm/lm_head (so the resulting GGUF carries just the draft head).
if cls.mtp_only and not is_mtp and name not in (
if cls.mtp_only and not is_mtp and (cls.mtp_shared_embd or name not in (
"model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",
):
)):
return None

return name, gen
Expand Down
3 changes: 2 additions & 1 deletion conversion/qwen.py
Original file line number Diff line number Diff line change
Expand Up @@ -290,6 +290,7 @@ class _QwenMtpMixin:
tensor_map: gguf.TensorNameMap
no_mtp: bool
mtp_only: bool
mtp_shared_embd: bool
_original_block_count: int | None = None
opt_num_mtp_layers: int = 0

Expand Down Expand Up @@ -338,7 +339,7 @@ def filter_tensors(cls, item):
elif len(parts) == 3 and parts[1] in remapper:
name = f"model.layers.{cls._original_block_count}.{remapper[parts[1]]}.{parts[2]}"
elif cls.mtp_only:
keep = name in (
keep = not cls.mtp_shared_embd and name in (
"model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",
"embed_tokens.weight", "norm.weight",
)
Expand Down
53 changes: 42 additions & 11 deletions conversion/qwen4exp.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from __future__ import annotations

from typing import Iterable, cast
from typing import Callable, Iterable, cast

import torch
from torch import Tensor
Expand All @@ -21,20 +21,51 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
Shares the Qwen3.5 gated delta net and interleaved mrope, and adds three things:
hyper-connections in place of every layer norm, QSA sparse attention on the full
attention layers, and PLE n-gram hash embeddings on a single layer.
The checkpoint also carries a NextN/MTP draft head under `mtp.*`, exported as a
trailing block; pass --no-nextn to leave it out.
"""

model_arch = gguf.MODEL_ARCH.QWEN4EXP

# the MTP block is a separate draft head; vLLM drops it too
supports_mtp_export = False
no_mtp = True

def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# only the shard names, so the table itself is never held
self._ple_shards: dict[int, str] = {}
self._ple_row_dim: int | None = None

_MTP_MIXER_PREFIX = "mtp.hyper_connection_mixer."

@classmethod
def filter_tensors(cls, item):
name, gen = item
_, _, mixer_suffix = name.partition(cls._MTP_MIXER_PREFIX)
Comment on lines +37 to +42

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
_MTP_MIXER_PREFIX = "mtp.hyper_connection_mixer."
@classmethod
def filter_tensors(cls, item):
name, gen = item
_, _, mixer_suffix = name.partition(cls._MTP_MIXER_PREFIX)
@classmethod
def filter_tensors(cls, item):
name, gen = item
_, _, mixer_suffix = name.partition("mtp.hyper_connection_mixer.")

if mixer_suffix:
if cls.no_mtp:
return None
assert cls._original_block_count is not None
return f"model.layers.{cls._original_block_count}.hyper_connection_mixer.{mixer_suffix}", gen
return super().filter_tensors((name, gen))

def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
tensors = super().index_tensors(remote_hf_model_id=remote_hf_model_id)

emb = tensors.pop("mtp.fc_embedding.weight", None)
hid = tensors.pop("mtp.fc_hidden.weight", None)
if emb is None and hid is None:
return tensors
if emb is None or hid is None:
raise ValueError(
"the qwen4exp MTP combiner needs both mtp.fc_embedding.weight and "
"mtp.fc_hidden.weight; pass --no-nextn to convert without the draft head"
)

assert self._original_block_count is not None
# W_e@e + W_h@h == [W_e|W_h] @ concat(e, h); embedding first, matching the graph's concat.
name = f"model.layers.{self._original_block_count}.eh_proj.weight"
tensors[name] = lambda: torch.cat([emb(), hid()], dim=1)
return tensors

def _read_hash_constants(self, suffix: str) -> list[int]:
"""Read an int64 PLE constant straight from the checkpoint.
Expand Down Expand Up @@ -63,14 +94,14 @@ def set_gguf_parameters(self):
self.gguf_writer.add_indexer_top_k(hp["indexer_budget"])
ratio = hp["indexer_compress_ratio"]
layer_types = hp["layer_types"]
self.gguf_writer.add_attention_compress_ratios(
[ratio if layer_types[i] == "full_attention" else 0 for i in range(n_layer)]
)
ratios = [ratio if layer_types[i] == "full_attention" else 0 for i in range(n_layer)]
# 0 selects dense, which is how the MTP block attends.
ratios += [0] * (self.block_count - n_layer)
self.gguf_writer.add_attention_compress_ratios(ratios)

# ple_layer_ids is 1-based in the HF config; empty means no n-gram table,
# so emit no PLE keys rather than optional ones
# ple_layer_ids is 1-based in the HF config; empty means no n-gram table
ple_layers = [i - 1 for i in hp["ple_layer_ids"]]
if not ple_layers:
if not ple_layers or self.mtp_only:
return
self.gguf_writer.add_ple_layers(ple_layers)
self.gguf_writer.add_ple_ngram_size(hp["ngram_size"])
Expand Down
10 changes: 10 additions & 0 deletions convert_hf_to_gguf.py
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,10 @@ def parse_args() -> argparse.Namespace:
"--no-nextn", "--no-mtp", dest="no_mtp", action="store_true",
help="Exclude NextN speculative draft tensors from the converted GGUF. Pair with --mtp or --dspark on a second run to publish target and draft as two files.",
)
parser.add_argument(
"--mtp-shared-embd", action="store_true",
help="With --mtp, leave the token embeddings, output norm and LM head out of the draft and take them from the target model at load time. Much smaller draft, but it needs a llama.cpp new enough to read it.",
)
Comment on lines +128 to +131

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How much space do we save? My estimate is about 1GB of Q8_0?

If this is correct, my recommendation is to remove this sharing functionality in order to reduce the complexity. It is not worth it and additionally, am efficient MTP setup actually benefits from not sharing the target tensors and instead using fast low-bit quantizations such as Q4_0 for the token embeddings and the lm head.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I also don’t think we need to set this explicitly for MTP. I feel it should be less than 1 GB.

During MTP conversion, we can export the MTP head either together with the target model or separately. In the former case, the MTP head does not need its own token_embd or lm_head, since it can access them later through ctx_other. In the latter case, the exported one will be a self-contained draft model by design.

parser.add_argument(
"--dspark", action="store_true",
help="Export only the DeepSeek-V4 DSpark draft tensors as a separate GGUF.",
Expand Down Expand Up @@ -282,6 +286,12 @@ def main() -> None:
if args.mtp:
model_class.mtp_only = True

if args.mtp_shared_embd:

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same here.

if not args.mtp:
logger.error("--mtp-shared-embd only applies together with --mtp")
sys.exit(1)
model_class.mtp_shared_embd = True

model_instance = model_class(dir_model, output_type, fname_out,
is_big_endian=args.bigendian, use_temp_file=args.use_temp_file,
eager=args.no_lazy,
Expand Down
14 changes: 14 additions & 0 deletions gguf-py/gguf/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -1193,6 +1193,9 @@ class MODEL_TENSOR(IntEnum):
NEXTN_HNORM = auto()
NEXTN_SHARED_HEAD_HEAD = auto()
NEXTN_SHARED_HEAD_NORM = auto()
NEXTN_HC_HEAD_NORM = auto()
NEXTN_HC_HEAD_DOWN = auto()
NEXTN_HC_HEAD_UP = auto()
# eagle3
FC = auto() # feature fusion layer
D2T = auto() # draft to target vocabulary mapping
Expand Down Expand Up @@ -1979,6 +1982,9 @@ class MODEL_TENSOR(IntEnum):
MODEL_TENSOR.NEXTN_HNORM: "blk.{bid}.nextn.hnorm",
MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD: "blk.{bid}.nextn.shared_head_head",
MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM: "blk.{bid}.nextn.shared_head_norm",
MODEL_TENSOR.NEXTN_HC_HEAD_NORM: "blk.{bid}.nextn.hc_head_norm",
MODEL_TENSOR.NEXTN_HC_HEAD_DOWN: "blk.{bid}.nextn.hc_head_down",
MODEL_TENSOR.NEXTN_HC_HEAD_UP: "blk.{bid}.nextn.hc_head_up",
MODEL_TENSOR.FC: "fc",
MODEL_TENSOR.DSPARK_MARKOV_W1: "markov_w1",
MODEL_TENSOR.DSPARK_MARKOV_W2: "markov_w2",
Expand Down Expand Up @@ -2976,6 +2982,14 @@ class MODEL_TENSOR(IntEnum):
MODEL_TENSOR.PLE_NORM_QUERY,
MODEL_TENSOR.PLE_NORM_CONV,
MODEL_TENSOR.PLE_CONV1D,
MODEL_TENSOR.NEXTN_EH_PROJ,
MODEL_TENSOR.NEXTN_EMBED_TOKENS,
MODEL_TENSOR.NEXTN_ENORM,
MODEL_TENSOR.NEXTN_HNORM,
MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
MODEL_TENSOR.NEXTN_HC_HEAD_NORM,
MODEL_TENSOR.NEXTN_HC_HEAD_DOWN,
MODEL_TENSOR.NEXTN_HC_HEAD_UP,
],
MODEL_ARCH.PLAMO: [
MODEL_TENSOR.TOKEN_EMBD,
Expand Down
9 changes: 9 additions & 0 deletions gguf-py/gguf/tensor_mapping.py
Original file line number Diff line number Diff line change
Expand Up @@ -2803,6 +2803,15 @@ class TensorNameMap:
MODEL_TENSOR.HC_HEAD_UP: (
"model.hyper_connection_mixer.input_mix_weight_up",
),
MODEL_TENSOR.NEXTN_HC_HEAD_NORM: (
"model.layers.{bid}.hyper_connection_mixer.hc_norm",
),
MODEL_TENSOR.NEXTN_HC_HEAD_DOWN: (
"model.layers.{bid}.hyper_connection_mixer.input_mix_weight_down",
),
MODEL_TENSOR.NEXTN_HC_HEAD_UP: (
"model.layers.{bid}.hyper_connection_mixer.input_mix_weight_up",
),
MODEL_TENSOR.INDEXER_Q_NORM: (
"model.layers.{bid}.self_attn.indexer.q_layernorm",
),
Expand Down
6 changes: 6 additions & 0 deletions src/llama-arch.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -585,6 +585,9 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
{ LLM_TENSOR_NEXTN_HNORM, "blk.%d.nextn.hnorm" },
{ LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "blk.%d.nextn.shared_head_head" },
{ LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "blk.%d.nextn.shared_head_norm" },
{ LLM_TENSOR_NEXTN_HC_HEAD_NORM, "blk.%d.nextn.hc_head_norm" },
{ LLM_TENSOR_NEXTN_HC_HEAD_DOWN, "blk.%d.nextn.hc_head_down" },
{ LLM_TENSOR_NEXTN_HC_HEAD_UP, "blk.%d.nextn.hc_head_up" },
{ LLM_TENSOR_ATTN_SUB_NORM, "blk.%d.attn_sub_norm" },
{ LLM_TENSOR_FFN_SUB_NORM, "blk.%d.ffn_sub_norm" },
{ LLM_TENSOR_DEC_OUTPUT_NORM, "dec.output_norm" },
Expand Down Expand Up @@ -975,6 +978,9 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
{LLM_TENSOR_NEXTN_HNORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_NEXTN_HC_HEAD_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_NEXTN_HC_HEAD_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_NEXTN_HC_HEAD_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
// Nemotron 3 Super
// latent projections feed ggml_mul_mat, the buft probe must use MUL_MAT to keep them on GPU
{LLM_TENSOR_FFN_LATENT_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
Expand Down
3 changes: 3 additions & 0 deletions src/llama-arch.h
Original file line number Diff line number Diff line change
Expand Up @@ -696,6 +696,9 @@ enum llm_tensor {
LLM_TENSOR_NEXTN_HNORM,
LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD,
LLM_TENSOR_NEXTN_SHARED_HEAD_NORM,
LLM_TENSOR_NEXTN_HC_HEAD_NORM,
LLM_TENSOR_NEXTN_HC_HEAD_DOWN,
LLM_TENSOR_NEXTN_HC_HEAD_UP,
LLM_TENSOR_MASKED_EMBD_CENTROIDS,
LLM_TENSOR_MASKED_EMBD_ORDERING,
LLM_TENSOR_HRM_Z_L_INIT,
Expand Down
2 changes: 1 addition & 1 deletion src/llama-context.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -152,7 +152,7 @@ llama_context::llama_context(
cparams.ctx_other = params.ctx_other;
}

if (model.arch == LLM_ARCH_EAGLE3 || model.arch == LLM_ARCH_DFLASH) {
if (model.arch == LLM_ARCH_EAGLE3 || model.arch == LLM_ARCH_DFLASH || model.arch == LLM_ARCH_QWEN4EXP) {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here is not necessary as well. Let's follow the current MTP support design.

if (model.tok_embd == nullptr || model.output == nullptr) {
if (params.ctx_other == nullptr) {
throw std::runtime_error(model.arch_name() + " requires ctx_other to be set (this warning is normal during memory fitting)");
Expand Down
2 changes: 1 addition & 1 deletion src/llama-model.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -2539,7 +2539,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
const bool mtp_on_hybrid_qwen =
params.ctx_type == LLAMA_CONTEXT_TYPE_MTP &&
(arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE ||
arch == LLM_ARCH_BAILINGMOE3);
arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_QWEN4EXP);

const bool mtp_on_hybrid_nemotron =
params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && arch == LLM_ARCH_NEMOTRON_H_MOE;
Expand Down
4 changes: 4 additions & 0 deletions src/llama-model.h
Original file line number Diff line number Diff line change
Expand Up @@ -231,6 +231,10 @@ struct llama_layer_nextn {
struct ggml_tensor * shared_head_head_s = nullptr;
struct ggml_tensor * shared_head_head_in_s = nullptr;
struct ggml_tensor * shared_head_norm = nullptr;

struct ggml_tensor * hc_head_norm = nullptr;
struct ggml_tensor * hc_head_down = nullptr;
struct ggml_tensor * hc_head_up = nullptr;
};

struct llama_layer_switch_lora {
Expand Down
10 changes: 9 additions & 1 deletion src/models/models.h
Original file line number Diff line number Diff line change
Expand Up @@ -2389,7 +2389,11 @@ struct llama_model_qwen4exp : public llama_model_base {

struct graph : public llm_build_delta_net_base {
graph(const llama_model & model, const llm_graph_params & params);
private:
protected:
struct no_build_t {};
graph(const llama_model & model, const llm_graph_params & params, no_build_t) :
llm_build_delta_net_base(params), model(model) {}

// HC replaces every layer norm: residual is [n_embd, hc, n_tokens]
ggml_tensor * build_hc_mix(
ggml_tensor * x,
Expand Down Expand Up @@ -2481,6 +2485,10 @@ struct llama_model_qwen4exp : public llama_model_base {
const llama_model & model;
};

struct graph_mtp : public graph {
graph_mtp(const llama_model & model, const llm_graph_params & params);
};

std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};

Expand Down
Loading
Loading