mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-08-03 11:40:14 +02:00
Compare commits
5 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| f2b52a87e8 | |||
| 4ed2b13f75 | |||
| 2b63e0610b | |||
| 1464c62d88 | |||
| 221f0f6356 |
+33
-2
@@ -447,12 +447,43 @@ class DeepseekV2Model(TextModel):
|
||||
class DeepseekV32Model(DeepseekV2Model):
|
||||
model_arch = gguf.MODEL_ARCH.DEEPSEEK32
|
||||
skip_mtp = False
|
||||
supports_mtp_export = True
|
||||
_n_main_layers: int | None = None
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.block_count = self.hparams["num_hidden_layers"] + self.hparams.get("num_nextn_predict_layers", 0)
|
||||
self.block_count = self.hparams["num_hidden_layers"]
|
||||
if not self.no_mtp:
|
||||
self.block_count += self.hparams.get("num_nextn_predict_layers", 0)
|
||||
self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
|
||||
|
||||
def index_tensors(self, remote_hf_model_id: str | None = None):
|
||||
type(self)._n_main_layers = self.hparams["num_hidden_layers"]
|
||||
return super().index_tensors(remote_hf_model_id=remote_hf_model_id)
|
||||
|
||||
@classmethod
|
||||
def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
|
||||
if (titem := super().filter_tensors(item)) is None:
|
||||
return None
|
||||
name, gen = titem
|
||||
|
||||
# DeepSeek V3.2 appends the NextN/MTP block past num_hidden_layers
|
||||
# (model.layers.61 -> blk.61 in the 62-block file).
|
||||
assert cls._n_main_layers is not None
|
||||
is_mtp = (m := re.match(r"model\.layers\.(\d+)\.", name)) is not None and int(m.group(1)) >= cls._n_main_layers
|
||||
|
||||
# --no-mtp: drop the appended NextN block entirely.
|
||||
if is_mtp and cls.no_mtp:
|
||||
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 (
|
||||
"model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",
|
||||
):
|
||||
return None
|
||||
|
||||
return name, gen
|
||||
|
||||
def set_vocab(self):
|
||||
from transformers import AutoTokenizer
|
||||
tokenizer = AutoTokenizer.from_pretrained(self.dir_model)
|
||||
@@ -463,7 +494,7 @@ class DeepseekV32Model(DeepseekV2Model):
|
||||
super().set_gguf_parameters()
|
||||
|
||||
# NextN/MTP prediction layers
|
||||
if (num_nextn_predict_layers := self.hparams.get("num_nextn_predict_layers")) is not None:
|
||||
if not self.no_mtp and (num_nextn_predict_layers := self.hparams.get("num_nextn_predict_layers")) is not None:
|
||||
self.gguf_writer.add_nextn_predict_layers(num_nextn_predict_layers)
|
||||
|
||||
# DSA indexer parameters
|
||||
|
||||
+96
-97
@@ -268,8 +268,101 @@ class Qwen3MoeModel(Qwen2MoeModel):
|
||||
super().set_vocab()
|
||||
|
||||
|
||||
class _QwenMtpMixin:
|
||||
"""Shared MTP wiring for Qwen3-Next and Qwen3.5/3.6 text variants. The HF
|
||||
config carries the MTP block under `mtp_num_hidden_layers` (computed from
|
||||
the checkpoint when absent, e.g. Qwen3-Next) and the tensors under
|
||||
`mtp.*`; we extend block_count, emit the nextn metadata key, and remap
|
||||
`mtp.*` to the standard layer-indexed nextn naming so the existing
|
||||
tensor_map handles them."""
|
||||
|
||||
supports_mtp_export = True
|
||||
hparams: dict[str, Any]
|
||||
model_arch: gguf.MODEL_ARCH
|
||||
gguf_writer: gguf.GGUFWriter
|
||||
block_count: int
|
||||
tensor_map: gguf.TensorNameMap
|
||||
no_mtp: bool
|
||||
mtp_only: bool
|
||||
_original_block_count: int | None = None
|
||||
opt_num_mtp_layers: int = 0
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.block_count = self.hparams["num_hidden_layers"]
|
||||
if not self.no_mtp:
|
||||
n_mtp = self.hparams.get("mtp_num_hidden_layers", 0)
|
||||
# Qwen-3-Next doesn't include `mtp_num_hidden_layers` in config.
|
||||
if n_mtp == 0:
|
||||
assert self.opt_num_mtp_layers != 0
|
||||
n_mtp = self.opt_num_mtp_layers
|
||||
self.block_count += n_mtp
|
||||
self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
|
||||
|
||||
def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
|
||||
hparams = {**self.hparams, **self.hparams.get("text_config", {})}
|
||||
key = next((k for k in ["n_layers", "num_hidden_layers", "n_layer", "num_layers"] if k in hparams), None)
|
||||
type(self)._original_block_count = hparams.get(key)
|
||||
type(self).opt_num_mtp_layers = 0
|
||||
return super().index_tensors(remote_hf_model_id=remote_hf_model_id) # ty: ignore[unresolved-attribute]
|
||||
|
||||
@classmethod
|
||||
def filter_tensors(cls, item):
|
||||
assert cls._original_block_count is not None
|
||||
# TODO: change TextModel to super()
|
||||
if (titem := TextModel.filter_tensors(item)) is None:
|
||||
return None
|
||||
name, gen = titem
|
||||
if name.startswith("model.mtp."):
|
||||
name = name.replace("model.", "", 1)
|
||||
if name.startswith("mtp."):
|
||||
if cls.no_mtp:
|
||||
return None
|
||||
remapper = {
|
||||
"fc": "eh_proj",
|
||||
"pre_fc_norm_embedding": "enorm",
|
||||
"pre_fc_norm_hidden": "hnorm",
|
||||
"norm": "shared_head.norm",
|
||||
}
|
||||
parts = name.split(".", 3)
|
||||
if len(parts) == 4 and parts[1] == "layers" and parts[2].isdecimal():
|
||||
mtp_idx = int(parts[2])
|
||||
name = f"model.layers.{cls._original_block_count + mtp_idx}.{parts[3]}"
|
||||
cls.opt_num_mtp_layers = max(cls.opt_num_mtp_layers, mtp_idx + 1)
|
||||
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 (
|
||||
"model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",
|
||||
"embed_tokens.weight", "norm.weight",
|
||||
)
|
||||
if not keep:
|
||||
return None
|
||||
return name, gen
|
||||
|
||||
def set_gguf_parameters(self):
|
||||
super().set_gguf_parameters() # ty: ignore[unresolved-attribute]
|
||||
if self.no_mtp:
|
||||
return
|
||||
if (n := self.block_count - self.hparams["num_hidden_layers"]) > 0:
|
||||
self.gguf_writer.add_nextn_predict_layers(n)
|
||||
|
||||
def prepare_metadata(self, vocab_only: bool):
|
||||
from_dir = self.fname_out.is_dir()
|
||||
super().prepare_metadata(vocab_only=vocab_only) # ty: ignore[unresolved-attribute]
|
||||
|
||||
if not self.mtp_only or not from_dir:
|
||||
return
|
||||
|
||||
output_type: str = self.ftype.name.partition("_")[2] # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
|
||||
fname_default: str = gguf.naming_convention(
|
||||
self.metadata.name, self.metadata.basename, self.metadata.finetune, # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
|
||||
self.metadata.version, size_label=None, output_type=output_type, model_type=None) # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
|
||||
self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"
|
||||
|
||||
|
||||
@ModelBase.register("Qwen3NextForCausalLM")
|
||||
class Qwen3NextModel(Qwen2MoeModel):
|
||||
class Qwen3NextModel(_QwenMtpMixin, Qwen2MoeModel):
|
||||
model_arch = gguf.MODEL_ARCH.QWEN3NEXT
|
||||
|
||||
def set_gguf_parameters(self):
|
||||
@@ -284,16 +377,6 @@ class Qwen3NextModel(Qwen2MoeModel):
|
||||
rope_dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"]
|
||||
self.gguf_writer.add_rope_dimension_count(int(rope_dim * self.rope_parameters.get("partial_rotary_factor", 0.25)))
|
||||
|
||||
@classmethod
|
||||
def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
|
||||
name, gen = item
|
||||
|
||||
if name.startswith("mtp"):
|
||||
# ignore MTP layers for now
|
||||
return None
|
||||
|
||||
return super().filter_tensors(item)
|
||||
|
||||
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
|
||||
if name.endswith(".A_log"):
|
||||
data_torch = -torch.exp(data_torch)
|
||||
@@ -536,97 +619,13 @@ class _Qwen35MRopeMixin:
|
||||
self.gguf_writer.add_rope_dimension_sections(self._QWEN35_DEFAULT_MROPE_SECTION)
|
||||
|
||||
|
||||
class _Qwen35MtpMixin:
|
||||
"""Shared MTP wiring for Qwen3.5/3.6 text variants. The HF config carries
|
||||
the MTP block under `mtp_num_hidden_layers` and the tensors under
|
||||
`mtp.*`; we extend block_count, emit the nextn metadata key, and remap
|
||||
`mtp.*` to the standard layer-indexed nextn naming so the existing
|
||||
tensor_map handles them."""
|
||||
|
||||
supports_mtp_export = True
|
||||
hparams: dict[str, Any]
|
||||
model_arch: gguf.MODEL_ARCH
|
||||
gguf_writer: gguf.GGUFWriter
|
||||
block_count: int
|
||||
tensor_map: gguf.TensorNameMap
|
||||
no_mtp: bool
|
||||
mtp_only: bool
|
||||
_original_block_count: int | None = None
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.block_count = self.hparams["num_hidden_layers"]
|
||||
if not self.no_mtp:
|
||||
self.block_count += self.hparams.get("mtp_num_hidden_layers", 0)
|
||||
self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
|
||||
|
||||
def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
|
||||
hparams = {**self.hparams, **self.hparams.get("text_config", {})}
|
||||
key = next((k for k in ["n_layers", "num_hidden_layers", "n_layer", "num_layers"] if k in hparams), None)
|
||||
type(self)._original_block_count = hparams.get(key)
|
||||
return super().index_tensors(remote_hf_model_id=remote_hf_model_id) # ty: ignore[unresolved-attribute]
|
||||
|
||||
@classmethod
|
||||
def filter_tensors(cls, item):
|
||||
assert cls._original_block_count is not None
|
||||
# TODO: change TextModel to super()
|
||||
if (titem := TextModel.filter_tensors(item)) is None:
|
||||
return None
|
||||
name, gen = titem
|
||||
if name.startswith("model.mtp."):
|
||||
name = name.replace("model.", "", 1)
|
||||
if name.startswith("mtp."):
|
||||
if cls.no_mtp:
|
||||
return None
|
||||
remapper = {
|
||||
"fc": "eh_proj",
|
||||
"pre_fc_norm_embedding": "enorm",
|
||||
"pre_fc_norm_hidden": "hnorm",
|
||||
"norm": "shared_head.norm",
|
||||
}
|
||||
parts = name.split(".", 3)
|
||||
if len(parts) == 4 and parts[1] == "layers" and parts[2].isdecimal():
|
||||
mtp_idx = int(parts[2])
|
||||
name = f"model.layers.{cls._original_block_count + mtp_idx}.{parts[3]}"
|
||||
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 (
|
||||
"model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",
|
||||
"embed_tokens.weight", "norm.weight",
|
||||
)
|
||||
if not keep:
|
||||
return None
|
||||
return name, gen
|
||||
|
||||
def set_gguf_parameters(self):
|
||||
super().set_gguf_parameters() # ty: ignore[unresolved-attribute]
|
||||
if self.no_mtp:
|
||||
return
|
||||
if (n := self.hparams.get("mtp_num_hidden_layers", 0)) > 0:
|
||||
self.gguf_writer.add_nextn_predict_layers(n)
|
||||
|
||||
def prepare_metadata(self, vocab_only: bool):
|
||||
from_dir = self.fname_out.is_dir()
|
||||
super().prepare_metadata(vocab_only=vocab_only) # ty: ignore[unresolved-attribute]
|
||||
|
||||
if not self.mtp_only or not from_dir:
|
||||
return
|
||||
|
||||
output_type: str = self.ftype.name.partition("_")[2] # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
|
||||
fname_default: str = gguf.naming_convention(
|
||||
self.metadata.name, self.metadata.basename, self.metadata.finetune, # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
|
||||
self.metadata.version, size_label=None, output_type=output_type, model_type=None) # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
|
||||
self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"
|
||||
|
||||
|
||||
@ModelBase.register("Qwen3_5ForConditionalGeneration", "Qwen3_5ForCausalLM")
|
||||
class Qwen3_5TextModel(_Qwen35MtpMixin, _Qwen35MRopeMixin, _LinearAttentionVReorderBase):
|
||||
class Qwen3_5TextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
|
||||
model_arch = gguf.MODEL_ARCH.QWEN35
|
||||
|
||||
|
||||
@ModelBase.register("Qwen3_5MoeForConditionalGeneration", "Qwen3_5MoeForCausalLM")
|
||||
class Qwen3_5MoeTextModel(_Qwen35MtpMixin, _Qwen35MRopeMixin, _LinearAttentionVReorderBase):
|
||||
class Qwen3_5MoeTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
|
||||
model_arch = gguf.MODEL_ARCH.QWEN35MOE
|
||||
|
||||
|
||||
|
||||
@@ -477,6 +477,23 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_soft_max(ggml_me
|
||||
return res;
|
||||
}
|
||||
|
||||
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_lightning_indexer(
|
||||
ggml_metal_library_t lib,
|
||||
const ggml_tensor * op) {
|
||||
GGML_ASSERT(op->op == GGML_OP_LIGHTNING_INDEXER);
|
||||
|
||||
char name[256];
|
||||
|
||||
snprintf(name, 256, "kernel_lightning_indexer_%s", ggml_type_name(op->src[1]->type));
|
||||
|
||||
ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
|
||||
if (!res.pipeline) {
|
||||
res = ggml_metal_library_compile_pipeline(lib, name, name, nullptr);
|
||||
}
|
||||
|
||||
return res;
|
||||
}
|
||||
|
||||
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc(ggml_metal_library_t lib, ggml_op op) {
|
||||
const char * name = nullptr;
|
||||
|
||||
@@ -2135,6 +2152,23 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_opt_step_sgd(ggm
|
||||
return res;
|
||||
}
|
||||
|
||||
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_silu_back(ggml_metal_library_t lib, const ggml_tensor * op) {
|
||||
assert(op->op == GGML_OP_SILU_BACK);
|
||||
|
||||
char base[256];
|
||||
char name[256];
|
||||
|
||||
snprintf(base, 256, "kernel_silu_back_%s", ggml_type_name(op->src[0]->type));
|
||||
snprintf(name, 256, "%s", base);
|
||||
|
||||
ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
|
||||
if (!res.pipeline) {
|
||||
res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
|
||||
}
|
||||
|
||||
return res;
|
||||
}
|
||||
|
||||
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_memset(ggml_metal_library_t lib, const ggml_tensor * op) {
|
||||
GGML_ASSERT(op->type == GGML_TYPE_I64);
|
||||
|
||||
|
||||
@@ -117,6 +117,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_diag
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_repeat (ggml_metal_library_t lib, enum ggml_type tsrc);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_concat (ggml_metal_library_t lib, enum ggml_type tsrc);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_unary (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_silu_back (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glu (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum_rows (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
@@ -124,6 +125,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_bl
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_add (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_tri (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_soft_max (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_lightning_indexer (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc (ggml_metal_library_t lib, enum ggml_op op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_batched (ggml_metal_library_t lib, const struct ggml_tensor * op, int ssm_conv_bs);
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
#import "ggml-impl.h"
|
||||
#import "ggml-backend-impl.h"
|
||||
#import "ggml-metal-impl.h"
|
||||
|
||||
#include <Foundation/Foundation.h>
|
||||
|
||||
@@ -1137,6 +1138,14 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
|
||||
default:
|
||||
return false;
|
||||
}
|
||||
case GGML_OP_SILU_BACK:
|
||||
return (op->src[0]->type == GGML_TYPE_F32) &&
|
||||
(op->src[1]->type == GGML_TYPE_F32) &&
|
||||
(op->type == GGML_TYPE_F32) &&
|
||||
ggml_is_contiguous(op->src[0]) &&
|
||||
ggml_is_contiguous(op->src[1]) &&
|
||||
ggml_is_contiguous(op) &&
|
||||
ggml_are_same_shape(op->src[0], op->src[1]);
|
||||
case GGML_OP_GLU:
|
||||
switch (ggml_get_glu_op(op)) {
|
||||
case GGML_GLU_OP_REGLU:
|
||||
@@ -1300,6 +1309,36 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
|
||||
return false;
|
||||
}
|
||||
return has_simdgroup_mm; // TODO: over-restricted for vec-kernels
|
||||
case GGML_OP_LIGHTNING_INDEXER:
|
||||
if (op->src[0]->ne[0] != OP_LIGHTNING_INDEXER_DK ||
|
||||
op->src[0]->ne[1] != OP_LIGHTNING_INDEXER_NH) {
|
||||
return false;
|
||||
}
|
||||
if (!has_simdgroup_mm ||
|
||||
op->src[0]->type != GGML_TYPE_F32 ||
|
||||
op->src[2]->type != GGML_TYPE_F32 ||
|
||||
op->src[3]->type != GGML_TYPE_F16 ||
|
||||
op->type != GGML_TYPE_F32 ||
|
||||
!ggml_is_contiguous_rows(op->src[0]) ||
|
||||
!ggml_is_contiguous_rows(op->src[1]) ||
|
||||
!ggml_is_contiguous_rows(op->src[2]) ||
|
||||
!ggml_is_contiguous_rows(op->src[3])) {
|
||||
return false;
|
||||
}
|
||||
switch (op->src[1]->type) {
|
||||
case GGML_TYPE_F32:
|
||||
case GGML_TYPE_F16:
|
||||
case GGML_TYPE_Q4_0:
|
||||
case GGML_TYPE_Q4_1:
|
||||
case GGML_TYPE_Q5_0:
|
||||
case GGML_TYPE_Q5_1:
|
||||
case GGML_TYPE_Q8_0:
|
||||
return true;
|
||||
case GGML_TYPE_BF16:
|
||||
return has_bfloat;
|
||||
default:
|
||||
return false;
|
||||
}
|
||||
case GGML_OP_DSV4_HC_COMB:
|
||||
return has_simdgroup_reduction &&
|
||||
op->src[0]->type == GGML_TYPE_F32 &&
|
||||
|
||||
@@ -112,6 +112,13 @@
|
||||
#define OP_FLASH_ATTN_EXT_VEC_NQPSG 1
|
||||
#define OP_FLASH_ATTN_EXT_VEC_NCPSG 32
|
||||
|
||||
#define OP_LIGHTNING_INDEXER_DK 128
|
||||
#define OP_LIGHTNING_INDEXER_NH 64
|
||||
#define OP_LIGHTNING_INDEXER_NHPTG 8
|
||||
#define OP_LIGHTNING_INDEXER_NKPSG 8
|
||||
#define OP_LIGHTNING_INDEXER_NSG 8
|
||||
#define OP_LIGHTNING_INDEXER_NBPTG 8
|
||||
|
||||
#define OP_UNARY_NUM_SCALE 10
|
||||
#define OP_UNARY_NUM_FILL 11
|
||||
#define OP_UNARY_NUM_CLAMP 12
|
||||
@@ -1171,6 +1178,23 @@ typedef struct {
|
||||
int64_t val;
|
||||
} ggml_metal_kargs_memset;
|
||||
|
||||
typedef struct {
|
||||
int32_t n_kv;
|
||||
int32_t n_batch;
|
||||
int32_t mask_ne3;
|
||||
uint64_t nb1;
|
||||
uint64_t nb3;
|
||||
uint64_t nbq1;
|
||||
uint64_t nbq2;
|
||||
uint64_t nbq3;
|
||||
uint64_t nbk2;
|
||||
uint64_t nbk3;
|
||||
uint64_t nbw1;
|
||||
uint64_t nbw3;
|
||||
uint64_t nbm1;
|
||||
uint64_t nbm3;
|
||||
} ggml_metal_kargs_lightning_indexer;
|
||||
|
||||
typedef struct {
|
||||
int32_t n_tokens;
|
||||
int32_t n_iter;
|
||||
@@ -1265,4 +1289,8 @@ typedef struct {
|
||||
int64_t np;
|
||||
} ggml_metal_kargs_opt_step_sgd;
|
||||
|
||||
typedef struct {
|
||||
int64_t ne;
|
||||
} ggml_metal_kargs_silu_back;
|
||||
|
||||
#endif // GGML_METAL_IMPL
|
||||
|
||||
@@ -299,6 +299,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) {
|
||||
{
|
||||
n_fuse = ggml_metal_op_unary(ctx, idx);
|
||||
} break;
|
||||
case GGML_OP_SILU_BACK:
|
||||
{
|
||||
n_fuse = ggml_metal_op_silu_back(ctx, idx);
|
||||
} break;
|
||||
case GGML_OP_GLU:
|
||||
{
|
||||
n_fuse = ggml_metal_op_glu(ctx, idx);
|
||||
@@ -316,6 +320,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) {
|
||||
{
|
||||
n_fuse = ggml_metal_op_cumsum(ctx, idx);
|
||||
} break;
|
||||
case GGML_OP_LIGHTNING_INDEXER:
|
||||
{
|
||||
n_fuse = ggml_metal_op_lightning_indexer(ctx, idx);
|
||||
} break;
|
||||
case GGML_OP_DSV4_HC_COMB:
|
||||
case GGML_OP_DSV4_HC_PRE:
|
||||
case GGML_OP_DSV4_HC_POST:
|
||||
@@ -1303,6 +1311,72 @@ int ggml_metal_op_diag(ggml_metal_op_t ctx, int idx) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
int ggml_metal_op_lightning_indexer(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_tensor * op = ctx->node(idx);
|
||||
|
||||
ggml_metal_encoder_t enc = ctx->enc;
|
||||
|
||||
GGML_ASSERT(op->op == GGML_OP_LIGHTNING_INDEXER);
|
||||
|
||||
const ggml_tensor * q = op->src[0];
|
||||
const ggml_tensor * k = op->src[1];
|
||||
const ggml_tensor * w = op->src[2];
|
||||
const ggml_tensor * m = op->src[3];
|
||||
|
||||
GGML_ASSERT(q->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(k->type == GGML_TYPE_F32 ||
|
||||
k->type == GGML_TYPE_F16 ||
|
||||
k->type == GGML_TYPE_BF16 ||
|
||||
k->type == GGML_TYPE_Q4_0 ||
|
||||
k->type == GGML_TYPE_Q4_1 ||
|
||||
k->type == GGML_TYPE_Q5_0 ||
|
||||
k->type == GGML_TYPE_Q5_1 ||
|
||||
k->type == GGML_TYPE_Q8_0);
|
||||
GGML_ASSERT(w->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(m->type == GGML_TYPE_F16);
|
||||
GGML_ASSERT(op->type == GGML_TYPE_F32);
|
||||
|
||||
GGML_ASSERT(q->ne[0] == OP_LIGHTNING_INDEXER_DK);
|
||||
GGML_ASSERT(q->ne[1] == OP_LIGHTNING_INDEXER_NH);
|
||||
|
||||
ggml_metal_kargs_lightning_indexer args = {
|
||||
/*.n_kv =*/ (int32_t) k->ne[2],
|
||||
/*.n_batch =*/ (int32_t) q->ne[2],
|
||||
/*.mask_ne3 =*/ (int32_t) m->ne[3],
|
||||
/*.nb1 =*/ op->nb[1],
|
||||
/*.nb3 =*/ op->nb[3],
|
||||
/*.nbq1 =*/ q->nb[1],
|
||||
/*.nbq2 =*/ q->nb[2],
|
||||
/*.nbq3 =*/ q->nb[3],
|
||||
/*.nbk2 =*/ k->nb[2],
|
||||
/*.nbk3 =*/ k->nb[3],
|
||||
/*.nbw1 =*/ w->nb[1],
|
||||
/*.nbw3 =*/ w->nb[3],
|
||||
/*.nbm1 =*/ m->nb[1],
|
||||
/*.nbm3 =*/ m->nb[3],
|
||||
};
|
||||
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(q), 1);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(k), 2);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(w), 3);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(m), 4);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 5);
|
||||
|
||||
const int nsg = OP_LIGHTNING_INDEXER_NSG;
|
||||
const int nkptg = OP_LIGHTNING_INDEXER_NKPSG*nsg;
|
||||
const int nbptg = OP_LIGHTNING_INDEXER_NBPTG;
|
||||
|
||||
auto pipeline = ggml_metal_library_get_pipeline_lightning_indexer(ctx->lib, op);
|
||||
ggml_metal_encoder_set_pipeline(enc, pipeline);
|
||||
ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0);
|
||||
ggml_metal_encoder_dispatch_threadgroups(enc,
|
||||
(k->ne[2] + nkptg - 1)/nkptg,
|
||||
(q->ne[2] + nbptg - 1)/nbptg,
|
||||
q->ne[3], 32, nsg, 1);
|
||||
|
||||
return 1;
|
||||
}
|
||||
|
||||
int ggml_metal_op_dsv4_hc(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_tensor * op = ctx->node(idx);
|
||||
|
||||
@@ -3473,6 +3547,36 @@ int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) {
|
||||
return n_fuse;
|
||||
}
|
||||
|
||||
int ggml_metal_op_silu_back(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_tensor * op = ctx->node(idx);
|
||||
|
||||
ggml_metal_library_t lib = ctx->lib;
|
||||
ggml_metal_encoder_t enc = ctx->enc;
|
||||
|
||||
auto pipeline = ggml_metal_library_get_pipeline_silu_back(lib, op);
|
||||
|
||||
const int64_t ne = ggml_nelements(op);
|
||||
|
||||
ggml_metal_kargs_silu_back args = {
|
||||
/*.ne =*/ ne,
|
||||
};
|
||||
|
||||
int arg_idx{0};
|
||||
|
||||
ggml_metal_encoder_set_pipeline(enc, pipeline);
|
||||
ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), arg_idx++);
|
||||
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), arg_idx++);
|
||||
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), arg_idx++);
|
||||
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), arg_idx++);
|
||||
|
||||
const int nth = std::min<int64_t>(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), ne);
|
||||
const int64_t n = (ne + nth - 1) / nth;
|
||||
|
||||
ggml_metal_encoder_dispatch_threadgroups(enc, n, 1, 1, nth, 1, 1);
|
||||
|
||||
return 1;
|
||||
}
|
||||
|
||||
int ggml_metal_op_l2_norm(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_tensor * op = ctx->node(idx);
|
||||
|
||||
|
||||
@@ -54,6 +54,7 @@ int ggml_metal_op_cumsum (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_get_rows (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_set_rows (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_diag (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_lightning_indexer (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_dsv4_hc (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_soft_max (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_ssm_conv (ggml_metal_op_t ctx, int idx);
|
||||
@@ -71,6 +72,7 @@ int ggml_metal_op_mul_mat_id (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_add_id (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_flash_attn_ext (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_bin (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_silu_back (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_l2_norm (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_group_norm (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_norm (ggml_metal_op_t ctx, int idx);
|
||||
|
||||
@@ -1255,6 +1255,20 @@ template [[host_name("kernel_unary_f32_f32_4")]] kernel kernel_unary_t kernel_un
|
||||
template [[host_name("kernel_unary_f16_f16")]] kernel kernel_unary_t kernel_unary_impl<half, half, float>;
|
||||
template [[host_name("kernel_unary_f16_f16_4")]] kernel kernel_unary_t kernel_unary_impl<half4, half4, float4>;
|
||||
|
||||
kernel void kernel_silu_back_f32(
|
||||
constant ggml_metal_kargs_silu_back & args,
|
||||
device const float * dy,
|
||||
device const float * x,
|
||||
device float * dx,
|
||||
uint gid [[thread_position_in_grid]]) {
|
||||
if (gid >= args.ne) {
|
||||
return;
|
||||
}
|
||||
|
||||
const float s = 1.0f / (1.0f + exp(-x[gid]));
|
||||
dx[gid] = dy[gid] * s * (1.0f + x[gid] * (1.0f - s));
|
||||
}
|
||||
|
||||
// OP: 0 - add, 1 - sub, 2 - mul, 3 - div
|
||||
constant short FC_bin_op [[function_constant(FC_BIN + 0)]];
|
||||
constant short FC_bin_f [[function_constant(FC_BIN + 1)]];
|
||||
@@ -11281,6 +11295,155 @@ typedef decltype(kernel_count_equal<int32_t>) kernel_count_equal_t;
|
||||
|
||||
template [[host_name("kernel_count_equal_i32")]] kernel kernel_count_equal_t kernel_count_equal<int32_t>;
|
||||
|
||||
template<
|
||||
typename kd4x4_t,
|
||||
short nl_k,
|
||||
void (*deq_k)(device const kd4x4_t *, short, thread half4x4 &)>
|
||||
kernel void kernel_lightning_indexer(
|
||||
constant ggml_metal_kargs_lightning_indexer & args,
|
||||
device const char * q,
|
||||
device const char * k,
|
||||
device const char * w,
|
||||
device const char * m,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiitg[[thread_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
||||
constexpr short DK = OP_LIGHTNING_INDEXER_DK;
|
||||
constexpr short NH = OP_LIGHTNING_INDEXER_NH;
|
||||
constexpr short NHPTG = OP_LIGHTNING_INDEXER_NHPTG;
|
||||
constexpr short NKPSG = OP_LIGHTNING_INDEXER_NKPSG;
|
||||
constexpr short NSG = OP_LIGHTNING_INDEXER_NSG;
|
||||
constexpr short NBPTG = OP_LIGHTNING_INDEXER_NBPTG;
|
||||
|
||||
constexpr short DK4 = DK/4;
|
||||
constexpr short DK8 = DK/8;
|
||||
constexpr short DK16 = DK/16;
|
||||
|
||||
constexpr short NK = NKPSG*NSG; // keys per threadgroup
|
||||
constexpr short NTG = 32*NSG; // threads per threadgroup
|
||||
|
||||
const int i_stream = tgpig.z;
|
||||
const int i_kv_0 = tgpig.x*NK; // first key of this threadgroup
|
||||
const int i_kv = i_kv_0 + sgitg*NKPSG; // first key of this simdgroup
|
||||
|
||||
threadgroup half4x4 sk4x4[NK*DK16];
|
||||
threadgroup half * sk = (threadgroup half *) sk4x4;
|
||||
|
||||
for (short i = tiitg; i < NK*DK16; i += NTG) {
|
||||
const short ik = i/DK16;
|
||||
const short i16 = i%DK16;
|
||||
|
||||
half4x4 tmp;
|
||||
|
||||
if (i_kv_0 + ik < args.n_kv) {
|
||||
device const kd4x4_t * kr = (device const kd4x4_t *) (k + (i_kv_0 + ik)*args.nbk2 + i_stream*args.nbk3);
|
||||
|
||||
deq_k(kr + i16/nl_k, i16%nl_k, tmp);
|
||||
} else {
|
||||
FOR_UNROLL (short j = 0; j < 4; ++j) {
|
||||
tmp[j] = half4(0.0h);
|
||||
}
|
||||
}
|
||||
|
||||
sk4x4[i] = tmp;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// K tile of this simdgroup, transposed to [DK, NKPSG]
|
||||
simdgroup_half8x8 mk[DK8];
|
||||
|
||||
FOR_UNROLL (short i = 0; i < DK8; ++i) {
|
||||
simdgroup_load(mk[i], sk + sgitg*NKPSG*DK + 8*i, DK, 0, true);
|
||||
}
|
||||
|
||||
threadgroup half4 sq4[NHPTG*DK4];
|
||||
threadgroup half * sq = (threadgroup half *) sq4;
|
||||
|
||||
threadgroup float sw [NHPTG];
|
||||
threadgroup float sqk[NSG*NHPTG*NKPSG];
|
||||
|
||||
const int i_batch_0 = tgpig.y*NBPTG;
|
||||
const int n_batch = min((int) NBPTG, args.n_batch - i_batch_0);
|
||||
|
||||
for (short ib = 0; ib < n_batch; ++ib) {
|
||||
const int i_batch = i_batch_0 + ib;
|
||||
|
||||
device const char * pq = q + i_batch*args.nbq2 + i_stream*args.nbq3;
|
||||
device const char * pw = w + i_batch*args.nbw1 + i_stream*args.nbw3;
|
||||
|
||||
float score = 0.0f;
|
||||
|
||||
FOR_UNROLL (short i_head = 0; i_head < NH; i_head += NHPTG) {
|
||||
// stage the Q tile [DK, NHPTG] and the (prescaled) head weights
|
||||
for (short i = tiitg; i < NHPTG*DK4; i += NTG) {
|
||||
const short ih = i/DK4;
|
||||
const short i4 = i%DK4;
|
||||
|
||||
device const float4 * q4 = (device const float4 *) (pq + (i_head + ih)*args.nbq1);
|
||||
|
||||
sq4[ih*DK4 + i4] = half4(q4[i4]);
|
||||
}
|
||||
|
||||
if (tiitg < NHPTG) {
|
||||
sw[tiitg] = ((device const float *) pw)[i_head + tiitg];
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
simdgroup_float8x8 mqk = make_filled_simdgroup_matrix<float, 8>(0.0f);
|
||||
|
||||
FOR_UNROLL (short i = 0; i < DK8; ++i) {
|
||||
simdgroup_half8x8 mq;
|
||||
|
||||
simdgroup_load(mq, sq + 8*i, DK, 0, false);
|
||||
simdgroup_multiply_accumulate(mqk, mq, mk[i], mqk);
|
||||
}
|
||||
|
||||
threadgroup float * pqk = sqk + sgitg*NHPTG*NKPSG;
|
||||
|
||||
simdgroup_store(mqk, pqk, NKPSG, 0, false);
|
||||
simdgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// one lane per key: ReLU, apply the head weight and accumulate over the head tile
|
||||
if (tiisg < NKPSG) {
|
||||
FOR_UNROLL (short ih = 0; ih < NHPTG; ++ih) {
|
||||
score += max(pqk[ih*NKPSG + tiisg], 0.0f)*sw[ih];
|
||||
}
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
if (tiisg < NKPSG) {
|
||||
const int ik = i_kv + tiisg;
|
||||
if (ik < args.n_kv) {
|
||||
device const half * pm = (device const half *) (m + i_batch*args.nbm1 + (i_stream % args.mask_ne3)*args.nbm3);
|
||||
device float * pd = (device float *) (dst + i_batch*args.nb1 + i_stream*args.nb3);
|
||||
|
||||
pd[ik] = score + (float) pm[ik];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_lightning_indexer<half4x4, 1, dequantize_f16>) kernel_lightning_indexer_t;
|
||||
|
||||
template [[host_name("kernel_lightning_indexer_f32")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<float4x4, 1, dequantize_f32>;
|
||||
template [[host_name("kernel_lightning_indexer_f16")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<half4x4, 1, dequantize_f16>;
|
||||
|
||||
#if defined(GGML_METAL_HAS_BF16)
|
||||
template [[host_name("kernel_lightning_indexer_bf16")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<bfloat4x4, 1, dequantize_bf16>;
|
||||
#endif
|
||||
|
||||
template [[host_name("kernel_lightning_indexer_q4_0")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<block_q4_0, 2, dequantize_q4_0>;
|
||||
template [[host_name("kernel_lightning_indexer_q4_1")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<block_q4_1, 2, dequantize_q4_1>;
|
||||
template [[host_name("kernel_lightning_indexer_q5_0")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<block_q5_0, 2, dequantize_q5_0>;
|
||||
template [[host_name("kernel_lightning_indexer_q5_1")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<block_q5_1, 2, dequantize_q5_1>;
|
||||
template [[host_name("kernel_lightning_indexer_q8_0")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<block_q8_0, 2, dequantize_q8_0>;
|
||||
|
||||
kernel void kernel_dsv4_hc_comb_f32(
|
||||
constant ggml_metal_kargs_dsv4_hc_comb & args,
|
||||
device const char * mixes,
|
||||
@@ -11438,4 +11601,3 @@ kernel void kernel_dsv4_hc_post_f32(
|
||||
*(device float *) (dst + i0*args.nb_d0 + idst*args.nb_d1 + it*args.nb_d2) = result[idst];
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -2329,7 +2329,13 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.SSM_NORM,
|
||||
MODEL_TENSOR.SSM_IN,
|
||||
MODEL_TENSOR.SSM_BETA_ALPHA,
|
||||
MODEL_TENSOR.SSM_OUT
|
||||
MODEL_TENSOR.SSM_OUT,
|
||||
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_SHARED_HEAD_NORM,
|
||||
],
|
||||
MODEL_ARCH.QWEN3VL: [
|
||||
MODEL_TENSOR.TOKEN_EMBD,
|
||||
|
||||
+8
-24
@@ -2071,24 +2071,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
|
||||
{
|
||||
res = nullptr;
|
||||
} break;
|
||||
case LLM_ARCH_DEEPSEEK32:
|
||||
{
|
||||
res = new llama_kv_cache_dsa(
|
||||
*this,
|
||||
params.type_k,
|
||||
params.type_v,
|
||||
!cparams.flash_attn,
|
||||
cparams.offload_kqv,
|
||||
cparams.kv_unified,
|
||||
cparams.n_ctx_seq,
|
||||
cparams.n_seq_max,
|
||||
1,
|
||||
hparams.n_swa,
|
||||
hparams.swa_type,
|
||||
nullptr,
|
||||
nullptr);
|
||||
} break;
|
||||
case LLM_ARCH_GLM_DSA:
|
||||
case LLM_ARCH_DEEPSEEK32:
|
||||
{
|
||||
if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && hparams.n_layer_nextn > 0) {
|
||||
// The NextN/MTP draft head runs dense MLA (no DSA indexer), so the
|
||||
@@ -2211,11 +2195,11 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
|
||||
// checks
|
||||
default:
|
||||
{
|
||||
// The MTP head is dense-attention only on hybrid Qwen3.5/3.6, so use a plain
|
||||
// The MTP head is dense-attention only on hybrid Qwen3-Next/3.5/3.6, so use a plain
|
||||
// attention KV cache for the MTP context instead of the hybrid wrapper.
|
||||
const bool mtp_on_hybrid_qwen35 =
|
||||
const bool mtp_on_hybrid_qwen =
|
||||
params.ctx_type == LLAMA_CONTEXT_TYPE_MTP &&
|
||||
(arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE);
|
||||
(arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE);
|
||||
|
||||
if (llm_arch_is_recurrent(arch)) {
|
||||
res = new llama_memory_recurrent(
|
||||
@@ -2227,7 +2211,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
|
||||
cparams.n_seq_max,
|
||||
cparams.n_rs_seq,
|
||||
nullptr);
|
||||
} else if (llm_arch_is_hybrid(arch) && !mtp_on_hybrid_qwen35) {
|
||||
} else if (llm_arch_is_hybrid(arch) && !mtp_on_hybrid_qwen) {
|
||||
// The main difference between hybrid architectures is the
|
||||
// layer filters, so pick the right one here
|
||||
llama_memory_hybrid::layer_filter_cb filter_attn = nullptr;
|
||||
@@ -2242,7 +2226,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
|
||||
filter_recr = [&](uint32_t il) {
|
||||
return hparams.is_recr(il) && hparams.n_ff(il) == 0;
|
||||
};
|
||||
} else if (arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE) {
|
||||
} else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE) {
|
||||
filter_attn = [&](uint32_t il) {
|
||||
return il < hparams.n_layer() && !hparams.is_recr(il);
|
||||
};
|
||||
@@ -2308,12 +2292,12 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
|
||||
};
|
||||
}
|
||||
|
||||
if (mtp_on_hybrid_qwen35) {
|
||||
if (mtp_on_hybrid_qwen) {
|
||||
filter = [&](uint32_t il) { return il >= hparams.n_layer(); };
|
||||
}
|
||||
|
||||
if ((arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_HY_V3 || arch == LLM_ARCH_GLM_DSA ||
|
||||
arch == LLM_ARCH_MIMO2) &&
|
||||
arch == LLM_ARCH_MIMO2 || arch == LLM_ARCH_DEEPSEEK32) &&
|
||||
hparams.n_layer_nextn > 0) {
|
||||
if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP) {
|
||||
filter = [&](uint32_t il) { return il >= hparams.n_layer(); };
|
||||
|
||||
+269
-10
@@ -44,13 +44,24 @@ void llama_model_deepseek32::load_arch_hparams(llama_model_loader & ml) {
|
||||
GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer");
|
||||
|
||||
switch (hparams.n_layer()) {
|
||||
case 62: type = LLM_TYPE_685B_A37B; break;
|
||||
case 61: type = LLM_TYPE_685B_A37B; break;
|
||||
default: type = LLM_TYPE_UNKNOWN;
|
||||
}
|
||||
}
|
||||
|
||||
void llama_model_deepseek32::load_arch_tensors(llama_model_loader &) {
|
||||
void llama_model_deepseek32::load_arch_tensors(llama_model_loader & ml) {
|
||||
LLAMA_LOAD_LOCALS;
|
||||
|
||||
const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
|
||||
const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
|
||||
const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
|
||||
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
|
||||
int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
|
||||
|
||||
if (!ml.load_mtp) {
|
||||
mtp_flags |= TENSOR_SKIP;
|
||||
}
|
||||
|
||||
const bool is_mla = hparams.is_mla();
|
||||
if (!is_mla) {
|
||||
throw std::runtime_error("DEEPSEEK32 architecture requires MLA");
|
||||
@@ -80,12 +91,7 @@ void llama_model_deepseek32::load_arch_tensors(llama_model_loader &) {
|
||||
}
|
||||
|
||||
for (int i = 0; i < n_layer_all; ++i) {
|
||||
int flags = 0;
|
||||
if (i >= n_layer) {
|
||||
// skip all tensors in the NextN layers
|
||||
// TODO @ngxson : TENSOR_NOT_REQUIRED was a hack, need to remove it later
|
||||
flags |= TENSOR_SKIP | TENSOR_NOT_REQUIRED;
|
||||
}
|
||||
const int flags = (i >= n_layer) ? mtp_flags : trunk_flags;
|
||||
|
||||
auto & layer = layers[i];
|
||||
|
||||
@@ -138,7 +144,7 @@ void llama_model_deepseek32::load_arch_tensors(llama_model_loader &) {
|
||||
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags);
|
||||
}
|
||||
|
||||
// NextN/MTP tensors (preserved but unused) - conditionally load for last nextn_predict_layers
|
||||
// NextN/MTP tensors - conditionally load for last nextn_predict_layers
|
||||
if (i >= n_layer) {
|
||||
layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), { 2 * n_embd, n_embd }, flags);
|
||||
layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), { n_embd }, flags);
|
||||
@@ -153,6 +159,9 @@ void llama_model_deepseek32::load_arch_tensors(llama_model_loader &) {
|
||||
}
|
||||
|
||||
std::unique_ptr<llm_graph_context> llama_model_deepseek32::build_arch_graph(const llm_graph_params & params) const {
|
||||
if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
|
||||
return std::make_unique<graph_mtp>(*this, params);
|
||||
}
|
||||
return std::make_unique<graph>(*this, params);
|
||||
}
|
||||
|
||||
@@ -430,7 +439,9 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_
|
||||
Qcur, Kcur, Vcur, nullptr, nullptr, model.layers[il].wv_b, top_k, kq_scale, il);
|
||||
}
|
||||
}
|
||||
if (il == n_layer - 1 && inp_out_ids) {
|
||||
// when unmasked nextn embeddings are requested, t_h_nextn must keep all rows,
|
||||
// so the early output masking has to be skipped (it is applied after the final norm instead)
|
||||
if (il == n_layer - 1 && inp_out_ids && (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked)) {
|
||||
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
|
||||
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
|
||||
}
|
||||
@@ -493,6 +504,14 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_
|
||||
|
||||
cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
|
||||
|
||||
// post-norm hidden state feeds the NextN/MTP draft head
|
||||
cb(cur, "h_nextn", -1);
|
||||
res->t_h_nextn = cur;
|
||||
|
||||
if (cparams.embeddings_nextn && !cparams.embeddings_nextn_masked && inp_out_ids) {
|
||||
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
|
||||
}
|
||||
|
||||
cb(cur, "result_norm", -1);
|
||||
res->t_embd = cur;
|
||||
|
||||
@@ -504,3 +523,243 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_
|
||||
|
||||
ggml_build_forward_expand(gf, cur);
|
||||
}
|
||||
|
||||
// LLM_GRAPH_TYPE_DECODER_MTP draft head for DeepSeek V3.2 (DEEPSEEK32).
|
||||
// Semantics mirror the deepseek-family NextN/MTP layer:
|
||||
// enorm(embed) + hnorm(prev_hidden) -> concat(e, h) -> eh_proj ->
|
||||
// full deepseek32 decoder block (dense MLA attention + sigmoid-gated MoE FFN
|
||||
// with shared expert, exactly as the trunk deepseek2 graph builds it) ->
|
||||
// shared_head_norm (fallback output_norm) -> shared LM head.
|
||||
// The DSA indexer is not used at runtime.
|
||||
llama_model_deepseek32::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)
|
||||
: llm_graph_context(params) {
|
||||
GGML_ASSERT(hparams.n_layer_nextn > 0 && "DEEPSEEK32 MTP requires n_layer_nextn > 0");
|
||||
GGML_ASSERT(hparams.n_layer_nextn == 1 && "DEEPSEEK32 MTP currently only supports a single MTP block");
|
||||
GGML_ASSERT(hparams.is_mla() && "DEEPSEEK32 MTP requires MLA");
|
||||
|
||||
const int il = hparams.n_layer() + cparams.nextn_layer_offset;
|
||||
GGML_ASSERT(cparams.nextn_layer_offset >= 0 &&
|
||||
cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&
|
||||
"nextn_layer_offset out of range [0, n_layer_nextn)");
|
||||
const auto & layer = model.layers[il];
|
||||
|
||||
GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj");
|
||||
GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm");
|
||||
GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm");
|
||||
GGML_ASSERT(layer.ffn_gate_inp && "MTP block missing ffn_gate_inp");
|
||||
|
||||
// note: these are the actual head sizes you get when treating as MHA or after "decompression" using wv_b for MLA
|
||||
const int64_t n_embd_head_k = hparams.n_embd_head_k_mla();
|
||||
|
||||
const int64_t n_embd_head_qk_rope = hparams.n_rot();
|
||||
const int64_t n_embd_head_qk_nope = n_embd_head_k - n_embd_head_qk_rope;
|
||||
|
||||
const uint32_t kv_lora_rank = hparams.n_lora_kv;
|
||||
|
||||
// We have to pre-scale kq_scale and attn_factor to make the YaRN RoPE work correctly.
|
||||
// See the deepseek2 trunk graph for the detailed explanation - this must match it EXACTLY.
|
||||
GGML_ASSERT(ext_factor >= 0.0f);
|
||||
const float attn_factor_org = attn_factor * (1.0f + 0.1f * logf(1.0f / freq_scale));
|
||||
|
||||
const float mscale = attn_factor_org * (1.0f + 0.1f * hparams.rope_yarn_log_mul * logf(1.0f / freq_scale));
|
||||
const float kq_scale = 1.0f * mscale * mscale / sqrtf(float(n_embd_head_k));
|
||||
|
||||
// TODO: extract in a common llm_graph_context::build_inp_embd_h()
|
||||
auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd);
|
||||
|
||||
inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
|
||||
ggml_set_input(inp->tokens);
|
||||
|
||||
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);
|
||||
ggml_set_input(inp->embd);
|
||||
|
||||
ggml_tensor * tok_embd;
|
||||
if (ubatch.token) {
|
||||
ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
|
||||
|
||||
tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);
|
||||
} else {
|
||||
tok_embd = inp->embd;
|
||||
}
|
||||
cb(tok_embd, "mtp_tok_embd", il);
|
||||
|
||||
inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
|
||||
ggml_set_input(inp->h);
|
||||
ggml_set_name(inp->h, "mtp_h_input");
|
||||
|
||||
ggml_tensor * h_embd = inp->h;
|
||||
|
||||
res->add_input(std::move(inp));
|
||||
|
||||
ggml_tensor * inp_pos = build_inp_pos();
|
||||
ggml_tensor * inp_out_ids = build_inp_out_ids();
|
||||
|
||||
// MLA with the absorption optimization uses a K-only cache (V is a view of K)
|
||||
auto * inp_attn = build_attn_inp_k();
|
||||
|
||||
ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
|
||||
cb(h_norm, "mtp_hnorm", il);
|
||||
|
||||
ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
|
||||
cb(e_norm, "mtp_enorm", il);
|
||||
|
||||
ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0);
|
||||
cb(concat, "mtp_concat", il);
|
||||
|
||||
ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);
|
||||
cb(cur, "mtp_eh_proj", il);
|
||||
|
||||
ggml_tensor * inpSA = cur;
|
||||
|
||||
cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
|
||||
cb(cur, "mtp_attn_norm", il);
|
||||
|
||||
// self-attention: dense MLA, same construction as the deepseek2 trunk graph
|
||||
{
|
||||
ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_a, cur);
|
||||
cb(q, "mtp_q", il);
|
||||
|
||||
q = build_norm(q, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
|
||||
cb(q, "mtp_q", il);
|
||||
|
||||
q = ggml_mul_mat(ctx0, layer.wq_b, q);
|
||||
cb(q, "mtp_q", il);
|
||||
|
||||
// split into {n_embd_head_qk_nope, n_head, n_tokens}
|
||||
ggml_tensor * q_nope =
|
||||
ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k),
|
||||
ggml_row_size(q->type, n_embd_head_k) * n_head, 0);
|
||||
cb(q_nope, "mtp_q_nope", il);
|
||||
|
||||
// and {n_embd_head_qk_rope, n_head, n_tokens}
|
||||
ggml_tensor * q_pe = ggml_view_3d(
|
||||
ctx0, q, n_embd_head_qk_rope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k),
|
||||
ggml_row_size(q->type, n_embd_head_k) * n_head, ggml_row_size(q->type, n_embd_head_qk_nope));
|
||||
cb(q_pe, "mtp_q_pe", il);
|
||||
|
||||
ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur);
|
||||
cb(kv_cmpr_pe, "mtp_kv_cmpr_pe", il);
|
||||
|
||||
// split into {kv_lora_rank, n_tokens}
|
||||
ggml_tensor * kv_cmpr =
|
||||
ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens,
|
||||
ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0);
|
||||
cb(kv_cmpr, "mtp_kv_cmpr", il);
|
||||
|
||||
// and {n_embd_head_qk_rope, 1, n_tokens}
|
||||
ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens,
|
||||
ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),
|
||||
ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),
|
||||
ggml_row_size(kv_cmpr_pe->type, kv_lora_rank));
|
||||
cb(k_pe, "mtp_k_pe", il);
|
||||
|
||||
q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
|
||||
ext_factor, attn_factor, beta_fast, beta_slow);
|
||||
cb(q_pe, "mtp_q_pe", il);
|
||||
|
||||
k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
|
||||
ext_factor, attn_factor, beta_fast, beta_slow);
|
||||
cb(k_pe, "mtp_k_pe", il);
|
||||
|
||||
kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);
|
||||
cb(kv_cmpr, "mtp_kv_cmpr", il);
|
||||
|
||||
// {n_embd_head_qk_nope, n_tokens, n_head}
|
||||
q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
|
||||
cb(q_nope, "mtp_q_nope_perm", il);
|
||||
|
||||
// {n_embd_head_qk_nope, kv_lora_rank, n_head} x {n_embd_head_qk_nope, n_tokens, n_head}
|
||||
ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope);
|
||||
cb(q_nope_absorbed, "mtp_q_nope_absorbed", il);
|
||||
|
||||
// {kv_lora_rank, n_head, n_tokens}
|
||||
q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3);
|
||||
cb(q_nope_absorbed, "mtp_q_nope_absorbed_perm", il);
|
||||
|
||||
// {n_embd_head_qk_rope + kv_lora_rank, n_head, n_tokens}
|
||||
// note: rope must go first for in-place context shifting in build_rope_shift()
|
||||
ggml_tensor * Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0);
|
||||
cb(Qcur, "mtp_Qcur", il);
|
||||
|
||||
kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);
|
||||
cb(kv_cmpr, "mtp_kv_cmpr_reshape", il);
|
||||
|
||||
// {n_embd_head_qk_rope + kv_lora_rank, 1, n_tokens}
|
||||
ggml_tensor * Kcur = ggml_concat(ctx0, kv_cmpr, k_pe, 0);
|
||||
cb(Kcur, "mtp_Kcur", il);
|
||||
|
||||
// {kv_lora_rank, 1, n_tokens}
|
||||
ggml_tensor * Vcur = kv_cmpr;
|
||||
cb(Vcur, "mtp_Vcur", il);
|
||||
|
||||
// note: MLA with the absorption optimization converts into MQA (ie: GQA with 1 group)
|
||||
cur = build_attn(inp_attn,
|
||||
layer.wo, NULL, layer.wo_s,
|
||||
Qcur, Kcur, Vcur, nullptr, nullptr, layer.wv_b, kq_scale, il);
|
||||
cb(cur, "mtp_attn_out", il);
|
||||
}
|
||||
|
||||
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
|
||||
cb(ffn_inp, "mtp_ffn_inp", il);
|
||||
|
||||
cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il);
|
||||
cb(cur, "mtp_ffn_norm", il);
|
||||
|
||||
// MoE FFN with shared expert - same construction as the deepseek2 trunk graph
|
||||
ggml_tensor * moe_out = build_moe_ffn(cur,
|
||||
layer.ffn_gate_inp,
|
||||
layer.ffn_up_exps,
|
||||
layer.ffn_gate_exps,
|
||||
layer.ffn_down_exps,
|
||||
layer.ffn_exp_probs_b,
|
||||
n_expert, n_expert_used,
|
||||
LLM_FFN_SILU, hparams.expert_weights_norm,
|
||||
hparams.expert_weights_scale,
|
||||
(llama_expert_gating_func_type) hparams.expert_gating_func,
|
||||
il,
|
||||
nullptr,
|
||||
layer.ffn_gate_up_exps,
|
||||
layer.ffn_up_exps_s,
|
||||
layer.ffn_gate_exps_s,
|
||||
layer.ffn_down_exps_s);
|
||||
cb(moe_out, "mtp_ffn_moe_out", il);
|
||||
|
||||
// FFN shared expert
|
||||
ggml_tensor * ffn_shexp =
|
||||
build_ffn(cur,
|
||||
layer.ffn_up_shexp, NULL, layer.ffn_up_shexp_s,
|
||||
layer.ffn_gate_shexp, NULL, layer.ffn_gate_shexp_s,
|
||||
layer.ffn_down_shexp, NULL, layer.ffn_down_shexp_s,
|
||||
NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
|
||||
cb(ffn_shexp, "mtp_ffn_shexp", il);
|
||||
|
||||
cur = ggml_add(ctx0, moe_out, ffn_shexp);
|
||||
cb(cur, "mtp_ffn_out", il);
|
||||
|
||||
cur = ggml_add(ctx0, cur, ffn_inp);
|
||||
cb(cur, "mtp_post_ffn", il);
|
||||
|
||||
// shared_head_norm applied after the decoder block, before the shared LM head.
|
||||
// The post-norm hidden state seeds the next MTP step.
|
||||
ggml_tensor * head_norm_w = layer.nextn.shared_head_norm
|
||||
? layer.nextn.shared_head_norm
|
||||
: model.output_norm;
|
||||
GGML_ASSERT(head_norm_w && "DEEPSEEK32 MTP: missing both nextn.shared_head_norm and output_norm");
|
||||
cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1);
|
||||
|
||||
cb(cur, "h_nextn", -1);
|
||||
res->t_h_nextn = cur;
|
||||
|
||||
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
|
||||
cb(cur, "mtp_shared_head_norm", -1);
|
||||
|
||||
ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;
|
||||
ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;
|
||||
GGML_ASSERT(head_w && "DEEPSEEK32 MTP: missing LM head (nextn.shared_head_head or model.output)");
|
||||
cur = build_lora_mm(head_w, cur, head_s);
|
||||
cb(cur, "result_output", -1);
|
||||
|
||||
res->t_logits = cur;
|
||||
ggml_build_forward_expand(gf, cur);
|
||||
}
|
||||
|
||||
|
||||
@@ -1097,6 +1097,10 @@ struct llama_model_deepseek32 : public llama_model_base {
|
||||
graph(const llama_model & model, const llm_graph_params & params);
|
||||
};
|
||||
|
||||
struct graph_mtp : public llm_graph_context {
|
||||
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;
|
||||
};
|
||||
|
||||
@@ -2033,6 +2037,10 @@ struct llama_model_qwen3next : public llama_model_base {
|
||||
const llama_model & model;
|
||||
};
|
||||
|
||||
struct graph_mtp : public llm_graph_context {
|
||||
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;
|
||||
};
|
||||
|
||||
|
||||
+276
-48
@@ -13,7 +13,11 @@ void llama_model_qwen3next::load_arch_hparams(llama_model_loader & ml) {
|
||||
ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank);
|
||||
ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group);
|
||||
|
||||
// Mark recurrent layers (linear attention layers)
|
||||
// NextN/MTP: extra decoder block appended beyond the main stack
|
||||
ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false);
|
||||
GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer_all");
|
||||
|
||||
// Mark recurrent layers (linear attention layers).
|
||||
if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) {
|
||||
uint32_t full_attn_interval = 4;
|
||||
ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false);
|
||||
@@ -28,13 +32,17 @@ void llama_model_qwen3next::load_arch_hparams(llama_model_loader & ml) {
|
||||
}
|
||||
}
|
||||
|
||||
void llama_model_qwen3next::load_arch_tensors(llama_model_loader &) {
|
||||
void llama_model_qwen3next::load_arch_tensors(llama_model_loader & ml) {
|
||||
LLAMA_LOAD_LOCALS;
|
||||
|
||||
if (n_expert == 0) {
|
||||
throw std::runtime_error(arch_name() + " model cannot have zero experts");
|
||||
}
|
||||
|
||||
const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
|
||||
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
|
||||
int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
|
||||
|
||||
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
|
||||
|
||||
// output
|
||||
@@ -61,49 +69,73 @@ void llama_model_qwen3next::load_arch_tensors(llama_model_loader &) {
|
||||
const int64_t qkvz_dim = key_dim * 2 + value_dim * 2;
|
||||
const int64_t ba_dim = n_v_heads * 2;
|
||||
|
||||
for (int i = 0; i < n_layer; ++i) {
|
||||
auto & layer = layers[i];
|
||||
const uint32_t n_ff_shexp = hparams.n_ff_shexp > 0 ? hparams.n_ff_shexp : hparams.n_ff(i);
|
||||
auto load_block_trunk = [&](int il, int flags) {
|
||||
auto & layer = layers[il];
|
||||
const uint32_t n_ff_shexp = hparams.n_ff_shexp > 0 ? hparams.n_ff_shexp : hparams.n_ff(il);
|
||||
|
||||
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, 0);
|
||||
layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), { n_embd }, 0);
|
||||
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, flags);
|
||||
layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, flags);
|
||||
|
||||
if (!hparams.is_recr(i)) {
|
||||
if (!hparams.is_recr(il)) {
|
||||
// Attention layers
|
||||
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, 0);
|
||||
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, 0);
|
||||
|
||||
create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, flags);
|
||||
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, flags);
|
||||
// Q/K normalization for attention layers
|
||||
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, 0);
|
||||
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), { n_embd_head_k }, 0);
|
||||
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, flags);
|
||||
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, flags);
|
||||
} else {
|
||||
// Linear attention (gated delta net) specific tensors
|
||||
// Create tensors with calculated dimensions
|
||||
// note: ssm_in is used by legacy GGUF
|
||||
layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", i), { n_embd, qkvz_dim }, TENSOR_NOT_REQUIRED);
|
||||
layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), { n_embd, key_dim * 2 + value_dim }, TENSOR_NOT_REQUIRED);
|
||||
layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), { n_embd, value_dim }, TENSOR_NOT_REQUIRED);
|
||||
layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), { hparams.ssm_d_conv, conv_dim }, 0);
|
||||
layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), { hparams.ssm_dt_rank }, 0);
|
||||
layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, i), { hparams.ssm_dt_rank }, 0);
|
||||
layer.ssm_beta_alpha = create_tensor(tn(LLM_TENSOR_SSM_BETA_ALPHA, "weight", i), { n_embd, ba_dim }, 0);
|
||||
layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), { head_v_dim }, 0);
|
||||
layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), { value_dim, n_embd }, 0);
|
||||
layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", il), { n_embd, qkvz_dim }, TENSOR_NOT_REQUIRED | flags);
|
||||
layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", il), { n_embd, key_dim * 2 + value_dim }, TENSOR_NOT_REQUIRED | flags);
|
||||
layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, value_dim }, TENSOR_NOT_REQUIRED | flags);
|
||||
layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", il), { hparams.ssm_d_conv, conv_dim }, flags);
|
||||
layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { hparams.ssm_dt_rank }, flags);
|
||||
layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, il), { hparams.ssm_dt_rank }, flags);
|
||||
layer.ssm_beta_alpha = create_tensor(tn(LLM_TENSOR_SSM_BETA_ALPHA, "weight", il), { n_embd, ba_dim }, flags);
|
||||
layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), { head_v_dim }, flags);
|
||||
layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", il), { value_dim, n_embd }, flags);
|
||||
}
|
||||
|
||||
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert }, 0);
|
||||
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), { n_ff_exp, n_embd, n_expert }, 0);
|
||||
create_tensor_gate_up_exps(layer, i, n_embd, n_ff_exp, n_expert, 0);
|
||||
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, flags);
|
||||
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, flags);
|
||||
create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, flags);
|
||||
|
||||
// Shared experts
|
||||
layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", i), { n_embd }, 0);
|
||||
layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), { n_embd, n_ff_shexp }, 0);
|
||||
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), { n_embd, n_ff_shexp }, 0);
|
||||
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_shexp, n_embd }, 0);
|
||||
layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, flags);
|
||||
layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, flags);
|
||||
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_shexp }, flags);
|
||||
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_shexp, n_embd }, flags);
|
||||
};
|
||||
|
||||
auto load_block_mtp = [&](int il) {
|
||||
// MTP head is identical to the trunk block (full attention + FFN)
|
||||
load_block_trunk(il, mtp_flags);
|
||||
|
||||
auto & layer = layers[il];
|
||||
|
||||
// NextN-specific tensors that define the MTP block.
|
||||
layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, mtp_flags);
|
||||
layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, mtp_flags);
|
||||
layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { n_embd }, mtp_flags);
|
||||
layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), { n_embd, n_vocab }, mtp_flags | TENSOR_NOT_REQUIRED);
|
||||
layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_vocab }, mtp_flags | TENSOR_NOT_REQUIRED);
|
||||
layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd }, mtp_flags | TENSOR_NOT_REQUIRED);
|
||||
};
|
||||
|
||||
for (int i = 0; i < n_layer; i++) {
|
||||
load_block_trunk(i, trunk_flags);
|
||||
}
|
||||
for (int i = n_layer; i < n_layer_all; i++) {
|
||||
load_block_mtp(i);
|
||||
}
|
||||
}
|
||||
|
||||
std::unique_ptr<llm_graph_context> llama_model_qwen3next::build_arch_graph(const llm_graph_params & params) const {
|
||||
if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
|
||||
return std::make_unique<graph_mtp>(*this, params);
|
||||
}
|
||||
return std::make_unique<graph>(*this, params);
|
||||
}
|
||||
|
||||
@@ -120,6 +152,7 @@ llama_model_qwen3next::graph::graph(const llama_model & model, const llm_graph_p
|
||||
ggml_tensor * inp_pos = build_inp_pos();
|
||||
ggml_tensor * inp_out_ids = build_inp_out_ids();
|
||||
|
||||
// MTP/NextN layers are loaded as extra decoder blocks but not executed in the main pass.
|
||||
for (int il = 0; il < n_layer; ++il) {
|
||||
res->t_layer_inp[il] = inpL;
|
||||
|
||||
@@ -139,7 +172,7 @@ llama_model_qwen3next::graph::graph(const llama_model & model, const llm_graph_p
|
||||
cur = build_layer_attn(inp->get_attn(), cur, inp_pos, il);
|
||||
}
|
||||
|
||||
if (il == n_layer - 1 && inp_out_ids) {
|
||||
if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
|
||||
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
|
||||
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
|
||||
}
|
||||
@@ -171,9 +204,16 @@ llama_model_qwen3next::graph::graph(const llama_model & model, const llm_graph_p
|
||||
}
|
||||
cur = inpL;
|
||||
|
||||
// Final norm
|
||||
// post-norm hidden state is input to both the LM head and the MTP head
|
||||
cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
|
||||
|
||||
cb(cur, "h_nextn", -1);
|
||||
res->t_h_nextn = cur;
|
||||
|
||||
if (!cparams.embeddings_nextn_masked && inp_out_ids) {
|
||||
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
|
||||
}
|
||||
|
||||
cb(cur, "result_norm", -1);
|
||||
res->t_embd = cur;
|
||||
|
||||
@@ -186,14 +226,6 @@ llama_model_qwen3next::graph::graph(const llama_model & model, const llm_graph_p
|
||||
ggml_build_forward_expand(gf, cur);
|
||||
}
|
||||
|
||||
// utility to get one slice from the third dimension
|
||||
// input dim: [x, y, c, b]
|
||||
// output dim: [x, y, 1, b]
|
||||
static ggml_tensor * get_slice_2d(ggml_context * ctx0, ggml_tensor * t, int64_t c) {
|
||||
return ggml_view_4d(ctx0, t, t->ne[0], t->ne[1], 1, t->ne[3],
|
||||
t->nb[1], t->nb[2], t->nb[3], t->nb[2] * c);
|
||||
}
|
||||
|
||||
ggml_tensor * llama_model_qwen3next::graph::build_norm_gated(
|
||||
ggml_tensor * input,
|
||||
ggml_tensor * weights,
|
||||
@@ -216,7 +248,7 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_attn(
|
||||
// Order: joint QG projection, QG split, Q norm, KV projection, K norm, RoPE, attention
|
||||
|
||||
// Qwen3Next uses a single Q projection that outputs query + gate
|
||||
ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur);
|
||||
ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s);
|
||||
cb(Qcur_full, "Qcur_full", il);
|
||||
|
||||
Qcur_full = ggml_reshape_4d(ctx0, Qcur_full, n_embd_head * 2, n_head, n_tokens, 1);
|
||||
@@ -232,10 +264,10 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_attn(
|
||||
Qcur_full->nb[1], Qcur_full->nb[2], Qcur_full->nb[3], n_embd_head * ggml_element_size(Qcur_full));
|
||||
cb(gate, "gate", il);
|
||||
|
||||
ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur);
|
||||
ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur, model.layers[il].wk_s);
|
||||
cb(Kcur, "Kcur", il);
|
||||
|
||||
ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur);
|
||||
ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur, model.layers[il].wv_s);
|
||||
cb(Vcur, "Vcur", il);
|
||||
|
||||
Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
|
||||
@@ -274,8 +306,6 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_attn(
|
||||
gate = ggml_sigmoid(ctx0, gate);
|
||||
cb(gate, "gate_sigmoid", il);
|
||||
|
||||
gate = ggml_reshape_2d(ctx0, gate, n_embd_head * n_head, n_tokens);
|
||||
|
||||
cur = ggml_mul(ctx0, cur, gate);
|
||||
cb(cur, "attn_gated", il);
|
||||
|
||||
@@ -550,16 +580,19 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_ffn(ggml_tensor * cur, c
|
||||
LLM_FFN_SILU, true,
|
||||
hparams.expert_weights_scale,
|
||||
LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, il,
|
||||
nullptr, model.layers[il].ffn_gate_up_exps);
|
||||
nullptr, model.layers[il].ffn_gate_up_exps,
|
||||
model.layers[il].ffn_up_exps_s,
|
||||
model.layers[il].ffn_gate_exps_s,
|
||||
model.layers[il].ffn_down_exps_s);
|
||||
cb(moe_out, "ffn_moe_out", il);
|
||||
|
||||
// Add shared experts if present - following Qwen3Next reference implementation
|
||||
if (model.layers[il].ffn_up_shexp != nullptr) {
|
||||
ggml_tensor * ffn_shexp =
|
||||
build_ffn(cur,
|
||||
model.layers[il].ffn_up_shexp, NULL, NULL,
|
||||
model.layers[il].ffn_gate_shexp, NULL, NULL,
|
||||
model.layers[il].ffn_down_shexp, NULL, NULL,
|
||||
model.layers[il].ffn_up_shexp, NULL, model.layers[il].ffn_up_shexp_s,
|
||||
model.layers[il].ffn_gate_shexp, NULL, model.layers[il].ffn_gate_shexp_s,
|
||||
model.layers[il].ffn_down_shexp, NULL, model.layers[il].ffn_down_shexp_s,
|
||||
NULL,
|
||||
LLM_FFN_SILU, LLM_FFN_PAR, il);
|
||||
cb(ffn_shexp, "ffn_shexp", il);
|
||||
@@ -593,3 +626,198 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_ffn(ggml_tensor * cur, c
|
||||
}
|
||||
return cur;
|
||||
}
|
||||
|
||||
// LLM_GRAPH_TYPE_DECODER_MTP draft head for Qwen3-Next
|
||||
llama_model_qwen3next::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)
|
||||
: llm_graph_context(params) {
|
||||
GGML_ASSERT(hparams.n_layer_nextn > 0 && "QWEN3NEXT MTP requires n_layer_nextn > 0");
|
||||
GGML_ASSERT(hparams.n_layer_nextn == 1 && "QWEN3NEXT MTP currently only supports a single MTP block");
|
||||
|
||||
const int64_t n_embd_head = hparams.n_embd_head_v();
|
||||
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
|
||||
|
||||
const int il = hparams.n_layer();
|
||||
const auto & layer = model.layers[il];
|
||||
|
||||
GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj");
|
||||
GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm");
|
||||
GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm");
|
||||
GGML_ASSERT(layer.ffn_gate_inp && "MTP block missing ffn_gate_inp");
|
||||
|
||||
// TODO: extract in a common llm_graph_context::build_inp_embd_h()
|
||||
auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd);
|
||||
|
||||
inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
|
||||
ggml_set_input(inp->tokens);
|
||||
|
||||
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);
|
||||
ggml_set_input(inp->embd);
|
||||
|
||||
// TODO: make static using `ggml_build_forward_select()`
|
||||
// see llm_graph_context::build_inp_embd() for reference
|
||||
ggml_tensor * tok_embd;
|
||||
if (ubatch.token) {
|
||||
ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
|
||||
|
||||
tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);
|
||||
} else {
|
||||
tok_embd = inp->embd;
|
||||
}
|
||||
cb(tok_embd, "mtp_tok_embd", il);
|
||||
|
||||
inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
|
||||
ggml_set_input(inp->h);
|
||||
ggml_set_name(inp->h, "mtp_h_input");
|
||||
|
||||
ggml_tensor * h_embd = inp->h;
|
||||
|
||||
res->add_input(std::move(inp));
|
||||
|
||||
ggml_tensor * inp_pos = build_inp_pos();
|
||||
ggml_tensor * inp_out_ids = build_inp_out_ids();
|
||||
|
||||
auto * inp_attn = build_attn_inp_kv();
|
||||
|
||||
ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
|
||||
cb(h_norm, "mtp_hnorm", il);
|
||||
|
||||
ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
|
||||
cb(e_norm, "mtp_enorm", il);
|
||||
|
||||
ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0);
|
||||
cb(concat, "mtp_concat", il);
|
||||
|
||||
ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);
|
||||
cb(cur, "mtp_eh_proj", il);
|
||||
|
||||
ggml_tensor * inpSA = cur;
|
||||
|
||||
cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
|
||||
cb(cur, "mtp_attn_norm", il);
|
||||
|
||||
ggml_tensor * Qcur_full = build_lora_mm(layer.wq, cur, layer.wq_s);
|
||||
cb(Qcur_full, "mtp_Qcur_full", il);
|
||||
|
||||
ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full,
|
||||
n_embd_head, n_head, n_tokens,
|
||||
ggml_element_size(Qcur_full) * n_embd_head * 2,
|
||||
ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head,
|
||||
0);
|
||||
Qcur = build_norm(Qcur, layer.attn_q_norm, nullptr, LLM_NORM_RMS, il);
|
||||
cb(Qcur, "mtp_Qcur_normed", il);
|
||||
|
||||
ggml_tensor * Kcur = build_lora_mm(layer.wk, cur, layer.wk_s);
|
||||
Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
|
||||
Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il);
|
||||
cb(Kcur, "mtp_Kcur_normed", il);
|
||||
|
||||
ggml_tensor * Vcur = build_lora_mm(layer.wv, cur, layer.wv_s);
|
||||
Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);
|
||||
|
||||
Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr,
|
||||
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
|
||||
ext_factor, attn_factor, beta_fast, beta_slow);
|
||||
Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr,
|
||||
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
|
||||
ext_factor, attn_factor, beta_fast, beta_slow);
|
||||
|
||||
cb(Qcur, "mtp_Qcur", il);
|
||||
cb(Kcur, "mtp_Kcur", il);
|
||||
cb(Vcur, "mtp_Vcur", il);
|
||||
|
||||
const float kq_scale = hparams.f_attention_scale == 0.0f
|
||||
? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
|
||||
|
||||
cur = build_attn(inp_attn,
|
||||
nullptr, nullptr, nullptr,
|
||||
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
|
||||
cb(cur, "mtp_attn_pregate", il);
|
||||
|
||||
ggml_tensor * gate = ggml_view_3d(ctx0, Qcur_full,
|
||||
n_embd_head, n_head, n_tokens,
|
||||
ggml_element_size(Qcur_full) * n_embd_head * 2,
|
||||
ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head,
|
||||
ggml_element_size(Qcur_full) * n_embd_head);
|
||||
|
||||
// TODO: CUDA is missing non-contiguous unary ops. when implemented: remove this cont
|
||||
gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens);
|
||||
cb(gate, "mtp_gate", il);
|
||||
|
||||
cur = ggml_mul(ctx0, cur, ggml_sigmoid(ctx0, gate));
|
||||
cur = build_lora_mm(layer.wo, cur, layer.wo_s);
|
||||
cb(cur, "mtp_attn_out", il);
|
||||
|
||||
if (inp_out_ids) {
|
||||
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
|
||||
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
|
||||
}
|
||||
|
||||
cur = ggml_add(ctx0, cur, inpSA);
|
||||
cb(cur, "mtp_attn_residual", il);
|
||||
|
||||
ggml_tensor * ffn_residual = cur;
|
||||
cur = build_norm(cur, layer.attn_post_norm, nullptr, LLM_NORM_RMS, il);
|
||||
cb(cur, "mtp_attn_post_norm", il);
|
||||
|
||||
// MoE FFN — routed experts plus gated shared expert (mirrors the trunk).
|
||||
ggml_tensor * moe_out =
|
||||
build_moe_ffn(cur,
|
||||
layer.ffn_gate_inp,
|
||||
layer.ffn_up_exps,
|
||||
layer.ffn_gate_exps,
|
||||
layer.ffn_down_exps,
|
||||
nullptr,
|
||||
n_expert, n_expert_used,
|
||||
LLM_FFN_SILU, true,
|
||||
hparams.expert_weights_scale,
|
||||
LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, il,
|
||||
nullptr, layer.ffn_gate_up_exps,
|
||||
layer.ffn_up_exps_s,
|
||||
layer.ffn_gate_exps_s,
|
||||
layer.ffn_down_exps_s);
|
||||
cb(moe_out, "mtp_ffn_moe_out", il);
|
||||
|
||||
if (layer.ffn_up_shexp != nullptr) {
|
||||
ggml_tensor * ffn_shexp =
|
||||
build_ffn(cur,
|
||||
layer.ffn_up_shexp, nullptr, layer.ffn_up_shexp_s,
|
||||
layer.ffn_gate_shexp, nullptr, layer.ffn_gate_shexp_s,
|
||||
layer.ffn_down_shexp, nullptr, layer.ffn_down_shexp_s,
|
||||
nullptr,
|
||||
LLM_FFN_SILU, LLM_FFN_PAR, il);
|
||||
cb(ffn_shexp, "mtp_ffn_shexp", il);
|
||||
|
||||
ggml_tensor * shared_gate = build_lora_mm(layer.ffn_gate_inp_shexp, cur);
|
||||
shared_gate = ggml_sigmoid(ctx0, shared_gate);
|
||||
cb(shared_gate, "mtp_shared_expert_gate_sigmoid", il);
|
||||
|
||||
ffn_shexp = ggml_mul(ctx0, ffn_shexp, shared_gate);
|
||||
cb(ffn_shexp, "mtp_ffn_shexp_gated", il);
|
||||
|
||||
cur = ggml_add(ctx0, moe_out, ffn_shexp);
|
||||
} else {
|
||||
cur = moe_out;
|
||||
}
|
||||
cb(cur, "mtp_ffn_out", il);
|
||||
|
||||
cur = ggml_add(ctx0, cur, ffn_residual);
|
||||
cb(cur, "mtp_post_ffn", il);
|
||||
|
||||
ggml_tensor * head_norm_w = layer.nextn.shared_head_norm
|
||||
? layer.nextn.shared_head_norm
|
||||
: model.output_norm;
|
||||
GGML_ASSERT(head_norm_w && "QWEN3NEXT MTP: missing both nextn.shared_head_norm and output_norm");
|
||||
cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1);
|
||||
|
||||
cb(cur, "h_nextn", -1);
|
||||
res->t_h_nextn = cur;
|
||||
|
||||
ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;
|
||||
ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;
|
||||
GGML_ASSERT(head_w && "QWEN3NEXT MTP: missing LM head (nextn.shared_head_head or model.output)");
|
||||
cur = build_lora_mm(head_w, cur, head_s);
|
||||
cb(cur, "result_output", -1);
|
||||
|
||||
res->t_logits = cur;
|
||||
ggml_build_forward_expand(gf, cur);
|
||||
}
|
||||
|
||||
@@ -9731,6 +9731,12 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
|
||||
}
|
||||
}
|
||||
|
||||
for (int kv : { 1, 7, 8, 63, 64, 65 }) {
|
||||
for (ggml_type type_K : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0}) {
|
||||
test_cases.emplace_back(new test_lightning_indexer(128, 64, kv, 32, 4, 1, type_K));
|
||||
}
|
||||
}
|
||||
|
||||
return test_cases;
|
||||
}
|
||||
#ifdef _MSC_VER
|
||||
|
||||
@@ -199,6 +199,9 @@ Invoke a tool call, request body is a JSON object with:
|
||||
- `tool` (string): the name of the tool
|
||||
- `params` (object): a mapping from argument name (string) to argument value
|
||||
|
||||
Headers:
|
||||
- `x-tool-cwd`: optional; if set, use as the CWD for tool; this is not part of tool's params because it's meant to be set by the runtime, not the LLM itself
|
||||
|
||||
Returns JSON object. There are two response formats (MCP tools use the same two formats: their result content is concatenated into `plain_text_response`, and RPC or tool errors are surfaced as the `error` string):
|
||||
|
||||
Format 1: Plain text. The text will be placed into a field called `plain_text_response`, example:
|
||||
|
||||
@@ -64,24 +64,27 @@ public:
|
||||
|
||||
class tools_io_basic : public tools_io {
|
||||
public:
|
||||
// cwd, if non-empty, is used to resolve relative paths and as the working directory for run()
|
||||
explicit tools_io_basic(std::string cwd = "") : cwd(std::move(cwd)) {}
|
||||
|
||||
bool is_directory(const std::string & path) const override {
|
||||
std::error_code ec;
|
||||
return fs::is_directory(path, ec) && !ec;
|
||||
return fs::is_directory(resolve(path), ec) && !ec;
|
||||
}
|
||||
|
||||
bool is_regular_file(const std::string & path) const override {
|
||||
std::error_code ec;
|
||||
return fs::is_regular_file(path, ec) && !ec;
|
||||
return fs::is_regular_file(resolve(path), ec) && !ec;
|
||||
}
|
||||
|
||||
bool file_size(const std::string & path, uintmax_t & out_size) const override {
|
||||
std::error_code ec;
|
||||
out_size = fs::file_size(path, ec);
|
||||
out_size = fs::file_size(resolve(path), ec);
|
||||
return !ec;
|
||||
}
|
||||
|
||||
bool read_file(const std::string & path, std::string & out) const override {
|
||||
std::ifstream f(path, std::ios::binary);
|
||||
std::ifstream f(resolve(path), std::ios::binary);
|
||||
if (!f) return false;
|
||||
std::ostringstream ss;
|
||||
ss << f.rdbuf();
|
||||
@@ -91,12 +94,12 @@ public:
|
||||
|
||||
bool write_file(const std::string & path, const std::string & content) const override {
|
||||
std::error_code ec;
|
||||
fs::path fpath(path);
|
||||
fs::path fpath(resolve(path));
|
||||
if (fpath.has_parent_path()) {
|
||||
fs::create_directories(fpath.parent_path(), ec);
|
||||
if (ec) return false;
|
||||
}
|
||||
std::ofstream f(path, std::ios::binary);
|
||||
std::ofstream f(fpath, std::ios::binary);
|
||||
if (!f) return false;
|
||||
f << content;
|
||||
return (bool) f;
|
||||
@@ -104,13 +107,14 @@ public:
|
||||
|
||||
std::vector<std::string> list_files(const std::string & base, std::string & err) const override {
|
||||
err.clear();
|
||||
std::string abs_base = resolve(base);
|
||||
if (!is_directory(base)) {
|
||||
err = "path does not exist or is not a directory: " + base;
|
||||
return {};
|
||||
}
|
||||
|
||||
auto res = run(
|
||||
{"git", "-C", base, "ls-files", "--cached", "--others", "--exclude-standard"},
|
||||
{"git", "-C", abs_base, "ls-files", "--cached", "--others", "--exclude-standard"},
|
||||
SERVER_TOOL_GIT_LS_FILES_MAX_OUTPUT, SERVER_TOOL_GIT_LS_FILES_TIMEOUT);
|
||||
|
||||
if (res.exit_code == 0 && !res.timed_out) {
|
||||
@@ -128,7 +132,7 @@ public:
|
||||
return result;
|
||||
}
|
||||
|
||||
return list_files_fallback(base);
|
||||
return list_files_fallback(abs_base);
|
||||
}
|
||||
|
||||
exec_result run(
|
||||
@@ -145,7 +149,7 @@ public:
|
||||
| subprocess_option_inherit_environment
|
||||
| subprocess_option_search_user_path;
|
||||
|
||||
if (!proc.create(args, options)) {
|
||||
if (!proc.create(args, options, {}, cwd.empty() ? nullptr : cwd.c_str())) {
|
||||
res.output = "failed to spawn process";
|
||||
return res;
|
||||
}
|
||||
@@ -205,6 +209,16 @@ public:
|
||||
}
|
||||
|
||||
private:
|
||||
std::string cwd;
|
||||
|
||||
// resolves `path` against `cwd` if `path` is relative and `cwd` is set; otherwise returns `path` unchanged
|
||||
std::string resolve(const std::string & path) const {
|
||||
if (cwd.empty() || fs::path(path).is_absolute()) {
|
||||
return path;
|
||||
}
|
||||
return (fs::path(cwd) / path).string();
|
||||
}
|
||||
|
||||
static const std::unordered_set<std::string> & junk_dir_names() {
|
||||
static const std::unordered_set<std::string> names = {
|
||||
".git", ".svn", ".hg", "node_modules", "__pycache__",
|
||||
@@ -244,8 +258,8 @@ private:
|
||||
};
|
||||
|
||||
static std::unique_ptr<tools_io> make_tools_io(const json & params) {
|
||||
GGML_UNUSED(params); // TODO in follow-up PR
|
||||
return std::make_unique<tools_io_basic>();
|
||||
std::string cwd = json_value(params, "cwd", std::string());
|
||||
return std::make_unique<tools_io_basic>(cwd);
|
||||
}
|
||||
|
||||
// no '/' in pattern -> match basename at any depth; else match full relative path
|
||||
@@ -1188,6 +1202,22 @@ static std::vector<std::unique_ptr<server_tool>> build_tools() {
|
||||
return tools;
|
||||
}
|
||||
|
||||
static std::string str_to_lower(const std::string & value) {
|
||||
std::string lowered(value.size(), '\0');
|
||||
std::transform(value.begin(), value.end(), lowered.begin(), [](unsigned char c) { return std::tolower(c); });
|
||||
return lowered;
|
||||
}
|
||||
|
||||
static std::string get_header(const std::map<std::string, std::string> & headers, const std::string & key, std::string default_value = "") {
|
||||
const auto lowered_key = str_to_lower(key);
|
||||
for (const auto & h : headers) {
|
||||
if (str_to_lower(h.first) == lowered_key) {
|
||||
return h.second;
|
||||
}
|
||||
}
|
||||
return default_value;
|
||||
}
|
||||
|
||||
void server_tools::setup(const std::vector<std::string> & enabled_tools,
|
||||
server_mcp & mcp_mgr) {
|
||||
if (!enabled_tools.empty()) {
|
||||
@@ -1271,6 +1301,12 @@ void server_tools::setup(const std::vector<std::string> & enabled_tools,
|
||||
json params = body.value("params", json::object());
|
||||
bool stream = body.value("stream", false);
|
||||
|
||||
// accept x-tool-cwd header to override of the process
|
||||
auto cwd = get_header(req.headers, "x-tool-cwd");
|
||||
if (!cwd.empty()) {
|
||||
params["cwd"] = cwd;
|
||||
}
|
||||
|
||||
server_tool & tool = find_tool(tools, tool_name, stream);
|
||||
|
||||
if (stream) {
|
||||
|
||||
@@ -19,8 +19,8 @@ def create_server():
|
||||
server.server_tools = "all"
|
||||
|
||||
|
||||
def call_tool(name: str, params: dict) -> dict:
|
||||
res = server.make_request("POST", "/tools", data={"tool": name, "params": params})
|
||||
def call_tool(name: str, params: dict, headers: dict | None = None) -> dict:
|
||||
res = server.make_request("POST", "/tools", data={"tool": name, "params": params}, headers=headers)
|
||||
assert res.status_code == 200, res.body
|
||||
assert "error" not in res.body, res.body
|
||||
return res.body
|
||||
@@ -123,6 +123,29 @@ def test_tools_builtin_exec_shell_command_stream():
|
||||
assert "[exit code: 0]" in chunks
|
||||
|
||||
|
||||
def test_tools_builtin_cwd_header():
|
||||
global server
|
||||
server.start()
|
||||
|
||||
cwd_dir = os.path.join(PROJECT_ROOT, "tools", "server", "tests", "unit")
|
||||
headers = {"x-tool-cwd": cwd_dir}
|
||||
|
||||
res = call_tool("read_file", {"path": "test_tools_builtin.py"}, headers=headers)
|
||||
assert GREP_MARKER in res["plain_text_response"]
|
||||
|
||||
# exec_shell_command should also run with that directory as its working directory:
|
||||
# writing to a relative filename must land inside cwd_dir
|
||||
marker_name = "llama_cpp_test_tools_builtin_cwd_marker.txt"
|
||||
marker_path = os.path.join(cwd_dir, marker_name)
|
||||
try:
|
||||
command = f"echo hello > {marker_name}"
|
||||
call_tool("exec_shell_command", {"command": command}, headers=headers)
|
||||
assert os.path.exists(marker_path)
|
||||
finally:
|
||||
if os.path.exists(marker_path):
|
||||
os.remove(marker_path)
|
||||
|
||||
|
||||
def test_tools_builtin_edit_file_rejects_overlapping_edits():
|
||||
global server
|
||||
server.start()
|
||||
|
||||
Reference in New Issue
Block a user