mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-08-03 03:28:19 +02:00
Compare commits
6 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 221f0f6356 | |||
| 9d21b57f2e | |||
| 0ab9d6fed7 | |||
| fffbcbdb9d | |||
| bb4e0e1b3f | |||
| 3581ba0cf5 |
+35
-1
@@ -374,6 +374,10 @@ common_models_handler common_models_handler_init(const common_params & params, l
|
||||
params.speculative.types.end(),
|
||||
COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3) != params.speculative.types.end();
|
||||
|
||||
const bool spec_type_draft_dspark = std::find(params.speculative.types.begin(),
|
||||
params.speculative.types.end(),
|
||||
COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) != params.speculative.types.end();
|
||||
|
||||
// only download mmproj if the current example is using it
|
||||
bool use_mmproj = false;
|
||||
for (const auto & ex : mmproj_examples) {
|
||||
@@ -388,6 +392,7 @@ common_models_handler common_models_handler_init(const common_params & params, l
|
||||
opts.download_mtp = spec_type_draft_mtp;
|
||||
opts.download_eagle3 = spec_type_draft_eagle3;
|
||||
opts.download_dflash = spec_type_draft_dflash;
|
||||
opts.download_dspark = spec_type_draft_dspark;
|
||||
opts.download_mmproj = use_mmproj && !params.no_mmproj
|
||||
&& params.mmproj.path.empty() && params.mmproj.url.empty();
|
||||
|
||||
@@ -402,6 +407,7 @@ common_models_handler common_models_handler_init(const common_params & params, l
|
||||
opts_spec.download_mtp = true;
|
||||
opts_spec.download_dflash = true;
|
||||
opts_spec.download_eagle3 = true;
|
||||
opts_spec.download_dspark = true;
|
||||
}
|
||||
plan_spec = common_download_get_hf_plan(params.speculative.draft.mparams, opts_spec);
|
||||
}
|
||||
@@ -544,12 +550,19 @@ void common_models_handler_apply(common_models_handler & handler, common_params
|
||||
plan_spec.mtp = {};
|
||||
plan_spec.dflash = {};
|
||||
plan_spec.eagle3 = {};
|
||||
plan_spec.dspark = {};
|
||||
}
|
||||
|
||||
// infer the speculative type from the sidecar shipped by the draft repo when none is requested
|
||||
if (spec_types_is_default(params)) {
|
||||
if (!plan_spec.mtp.local_path.empty()) {
|
||||
params.speculative.types = { COMMON_SPECULATIVE_TYPE_DRAFT_MTP };
|
||||
plan_spec.dspark = {};
|
||||
plan_spec.dflash = {};
|
||||
plan_spec.eagle3 = {};
|
||||
} else if (!plan_spec.dspark.local_path.empty()) {
|
||||
// dspark outranks dflash, its sidecar carries the extra Markov head
|
||||
params.speculative.types = { COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK };
|
||||
plan_spec.dflash = {};
|
||||
plan_spec.eagle3 = {};
|
||||
} else if (!plan_spec.dflash.local_path.empty()) {
|
||||
@@ -563,7 +576,8 @@ void common_models_handler_apply(common_models_handler & handler, common_params
|
||||
// when a sidecar type is requested, the draft repo resolves to its sidecar instead of a full model
|
||||
const bool spec_sidecar_found = !plan_spec.mtp.local_path.empty() ||
|
||||
!plan_spec.dflash.local_path.empty() ||
|
||||
!plan_spec.eagle3.local_path.empty();
|
||||
!plan_spec.eagle3.local_path.empty() ||
|
||||
!plan_spec.dspark.local_path.empty();
|
||||
if (!plan_spec.mtp.local_path.empty() && !had_spec_url) {
|
||||
tasks.emplace_back(plan_spec.mtp, opts, [&]() {
|
||||
// only use the discovered MTP head when no draft path is set yet
|
||||
@@ -594,6 +608,16 @@ void common_models_handler_apply(common_models_handler & handler, common_params
|
||||
}
|
||||
});
|
||||
}
|
||||
if (!plan_spec.dspark.local_path.empty() && !had_spec_url) {
|
||||
tasks.emplace_back(plan_spec.dspark, opts, [&]() {
|
||||
// only use the discovered DSpark sidecar when no draft path is set yet
|
||||
if (params.speculative.draft.mparams.path.empty()) {
|
||||
params.speculative.draft.mparams.path = hf_cache::finalize_file(plan_spec.dspark);
|
||||
} else {
|
||||
hf_cache::finalize_file(plan_spec.dspark);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// a wired draft sidecar counts as an explicit draft for the main plan fallback below
|
||||
if (spec_sidecar_found) {
|
||||
@@ -649,6 +673,16 @@ void common_models_handler_apply(common_models_handler & handler, common_params
|
||||
}
|
||||
});
|
||||
}
|
||||
if (!plan.dspark.local_path.empty() && !had_spec_url) {
|
||||
tasks.emplace_back(plan.dspark, opts, [&]() {
|
||||
// only fall back to the discovered DSpark sidecar when no draft was explicitly provided
|
||||
if (params.speculative.draft.mparams.empty()) {
|
||||
params.speculative.draft.mparams.path = hf_cache::finalize_file(plan.dspark);
|
||||
} else {
|
||||
hf_cache::finalize_file(plan.dspark);
|
||||
}
|
||||
});
|
||||
}
|
||||
if (!plan.preset.local_path.empty()) {
|
||||
tasks.emplace_back(plan.preset, opts, [&]() {
|
||||
// if HF repo is a preset repo, we simply run server in router mode with the preset.ini file
|
||||
|
||||
+15
-4
@@ -656,6 +656,12 @@ static hf_cache::hf_file find_best_dflash(const hf_cache::hf_files & files,
|
||||
return find_best_sibling(files, model, "dflash-", tag);
|
||||
}
|
||||
|
||||
static hf_cache::hf_file find_best_dspark(const hf_cache::hf_files & files,
|
||||
const std::string & model,
|
||||
const std::string & tag = "") {
|
||||
return find_best_sibling(files, model, "dspark-", tag);
|
||||
}
|
||||
|
||||
static bool gguf_filename_is_model(const std::string & filepath) {
|
||||
if (!string_ends_with(filepath, ".gguf")) {
|
||||
return false;
|
||||
@@ -670,7 +676,8 @@ static bool gguf_filename_is_model(const std::string & filepath) {
|
||||
filename.find("imatrix") == std::string::npos &&
|
||||
filename.find("mtp-") == std::string::npos &&
|
||||
filename.find("eagle3-") == std::string::npos &&
|
||||
filename.find("dflash-") == std::string::npos;
|
||||
filename.find("dflash-") == std::string::npos &&
|
||||
filename.find("dspark-") == std::string::npos;
|
||||
}
|
||||
|
||||
static hf_cache::hf_file find_best_model(const hf_cache::hf_files & files,
|
||||
@@ -763,7 +770,7 @@ common_download_hf_plan common_download_get_hf_plan(const common_params_model &
|
||||
} else {
|
||||
primary = find_best_model(all, tag);
|
||||
// a requested sidecar can resolve on its own, without a full model of the same tag
|
||||
if (primary.path.empty() && !opts.download_mtp && !opts.download_dflash && !opts.download_eagle3) {
|
||||
if (primary.path.empty() && !opts.download_mtp && !opts.download_dflash && !opts.download_eagle3 && !opts.download_dspark) {
|
||||
LOG_ERR("%s: no GGUF files found in repository %s\n", __func__, repo.c_str());
|
||||
list_available_gguf_files(all);
|
||||
return plan;
|
||||
@@ -787,9 +794,12 @@ common_download_hf_plan common_download_get_hf_plan(const common_params_model &
|
||||
if (opts.download_eagle3) {
|
||||
plan.eagle3 = find_best_eagle3(all, primary.path, tag);
|
||||
}
|
||||
if (opts.download_dspark) {
|
||||
plan.dspark = find_best_dspark(all, primary.path, tag);
|
||||
}
|
||||
|
||||
if (primary.path.empty() &&
|
||||
plan.mtp.local_path.empty() && plan.dflash.local_path.empty() && plan.eagle3.local_path.empty()) {
|
||||
plan.mtp.local_path.empty() && plan.dflash.local_path.empty() && plan.eagle3.local_path.empty() && plan.dspark.local_path.empty()) {
|
||||
LOG_ERR("%s: no GGUF files found in repository %s\n", __func__, repo.c_str());
|
||||
list_available_gguf_files(all);
|
||||
}
|
||||
@@ -967,7 +977,8 @@ std::vector<common_cached_model_info> common_list_cached_models() {
|
||||
split.prefix.find("mmproj") != std::string::npos ||
|
||||
split.prefix.find("mtp-") != std::string::npos ||
|
||||
split.prefix.find("eagle3-") != std::string::npos ||
|
||||
split.prefix.find("dflash-") != std::string::npos) {
|
||||
split.prefix.find("dflash-") != std::string::npos ||
|
||||
split.prefix.find("dspark-") != std::string::npos) {
|
||||
continue;
|
||||
}
|
||||
if (seen.insert(f.repo_id + ":" + split.tag).second) {
|
||||
|
||||
@@ -59,6 +59,7 @@ struct common_download_opts {
|
||||
bool download_mtp = false;
|
||||
bool download_eagle3 = false;
|
||||
bool download_dflash = false;
|
||||
bool download_dspark = false;
|
||||
common_download_callback * callback = nullptr;
|
||||
};
|
||||
|
||||
@@ -110,6 +111,7 @@ struct common_download_hf_plan {
|
||||
hf_cache::hf_file mtp;
|
||||
hf_cache::hf_file eagle3;
|
||||
hf_cache::hf_file dflash;
|
||||
hf_cache::hf_file dspark;
|
||||
hf_cache::hf_file preset; // if set, only this file is downloaded
|
||||
};
|
||||
common_download_hf_plan common_download_get_hf_plan(const common_params_model & model, const common_download_opts & opts);
|
||||
|
||||
@@ -55,6 +55,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
|
||||
"DFlashDraftModel": "qwen",
|
||||
"Qwen3DSparkModel": "qwen",
|
||||
"DeepseekV4ForCausalLM": "deepseek",
|
||||
"DeepseekV4DSparkModel": "deepseek",
|
||||
"DistilBertForMaskedLM": "bert",
|
||||
"DistilBertForSequenceClassification": "bert",
|
||||
"DistilBertModel": "bert",
|
||||
|
||||
+104
-1
@@ -620,7 +620,8 @@ class DeepseekV4Model(TextModel):
|
||||
self.gguf_writer.add_hyper_connection_sinkhorn_iterations(hparams["hc_sinkhorn_iters"])
|
||||
self.gguf_writer.add_hyper_connection_epsilon(hparams["hc_eps"])
|
||||
self.gguf_writer.add_hash_layer_count(hparams["num_hash_layers"])
|
||||
self.gguf_writer.add_embedding_length_out(hparams["hidden_size"] * hparams["hc_mult"])
|
||||
if self.model_arch == gguf.MODEL_ARCH.DEEPSEEK4:
|
||||
self.gguf_writer.add_embedding_length_out(hparams["hidden_size"] * hparams["hc_mult"])
|
||||
if self.mtp_only and (num_nextn_predict_layers := hparams.get("num_nextn_predict_layers", 0)) > 0:
|
||||
self.gguf_writer.add_nextn_predict_layers(num_nextn_predict_layers)
|
||||
|
||||
@@ -878,3 +879,105 @@ class DeepseekV4Model(TextModel):
|
||||
super().prepare_tensors()
|
||||
self._is_mxfp4 = True
|
||||
self.ftype = gguf.LlamaFileType.MOSTLY_MXFP4_MOE
|
||||
|
||||
|
||||
@ModelBase.register("DeepseekV4DSparkModel")
|
||||
class DeepseekV4DSparkModel(DeepseekV4Model):
|
||||
model_arch = gguf.MODEL_ARCH.DFLASH
|
||||
|
||||
_DSPARK_ROOT_MAP: dict[str, tuple[gguf.MODEL_TENSOR, str]] = {
|
||||
"main_proj.weight": (gguf.MODEL_TENSOR.FC, ".weight"),
|
||||
"main_norm.weight": (gguf.MODEL_TENSOR.ENC_OUTPUT_NORM, ".weight"),
|
||||
"markov_head.markov_w1.weight": (gguf.MODEL_TENSOR.DSPARK_MARKOV_W1, ".weight"),
|
||||
"markov_head.markov_w2.weight": (gguf.MODEL_TENSOR.DSPARK_MARKOV_W2, ".weight"),
|
||||
"confidence_head.proj.weight": (gguf.MODEL_TENSOR.DSPARK_CONF_PROJ, ".weight"),
|
||||
}
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
self.block_count = 1 + max(
|
||||
int(match.group(1)) for name in self.model_tensors
|
||||
if (match := re.match(r"layers\.(\d+)\.", name))
|
||||
)
|
||||
self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
|
||||
|
||||
self.hparams["compress_ratios"] = [0] * self.block_count
|
||||
self.hparams["num_hash_layers"] = 0
|
||||
|
||||
def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
|
||||
if remote_hf_model_id is None:
|
||||
return super().index_tensors()
|
||||
|
||||
with open(self.dir_model / "model.safetensors.index.json", "r", encoding="utf-8") as f:
|
||||
weight_map = json.load(f)["weight_map"]
|
||||
|
||||
part_names = sorted({
|
||||
part_name for name, part_name in weight_map.items()
|
||||
if name.startswith("mtp.")
|
||||
})
|
||||
tensors: dict[str, Callable[[], Tensor]] = {}
|
||||
|
||||
for part_name in part_names:
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
logger.info("gguf: caching remote DSpark part '%s'", part_name)
|
||||
part_path = Path(hf_hub_download(repo_id=remote_hf_model_id, filename=part_name))
|
||||
with gguf.utility.SafetensorsLocal(part_path) as model_part:
|
||||
for name in model_part:
|
||||
data = model_part[name]
|
||||
data_gen = lambda data=data: LazyTorchTensor.from_local_tensor(data) # noqa: E731
|
||||
if titem := self.filter_tensors((name, data_gen)):
|
||||
tensor_name, tensor_gen = titem
|
||||
tensors[tensor_name] = tensor_gen
|
||||
|
||||
return tensors
|
||||
|
||||
@classmethod
|
||||
def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
|
||||
name, gen = item
|
||||
if not name.startswith("mtp."):
|
||||
return None
|
||||
return super().filter_tensors((cls._rekey_mtp_tensor_name(name), gen))
|
||||
|
||||
@staticmethod
|
||||
def _rekey_mtp_tensor_name(name: str) -> str:
|
||||
match = re.match(r"mtp\.(\d+)\.(.+)$", name)
|
||||
if match is None:
|
||||
raise ValueError(f"Unexpected DSpark tensor {name!r}")
|
||||
|
||||
stage, rest = match.group(1), match.group(2)
|
||||
root_names = (
|
||||
"main_proj.scale",
|
||||
"norm.weight",
|
||||
"hc_head_fn",
|
||||
"hc_head_base",
|
||||
"hc_head_scale",
|
||||
)
|
||||
if rest in DeepseekV4DSparkModel._DSPARK_ROOT_MAP or rest in root_names:
|
||||
return rest
|
||||
return f"layers.{stage}.{rest}"
|
||||
|
||||
def _map_dsv4_tensor_name(self, name: str, bid: int | None) -> tuple[gguf.MODEL_TENSOR, str]:
|
||||
if name in self._DSPARK_ROOT_MAP:
|
||||
return self._DSPARK_ROOT_MAP[name]
|
||||
return super()._map_dsv4_tensor_name(name, bid)
|
||||
|
||||
def set_vocab(self):
|
||||
if self.target_model_dir is None:
|
||||
raise ValueError("DeepSeek-V4 DSpark requires --target-model-dir with the target tokenizer")
|
||||
|
||||
original_dir = self.dir_model
|
||||
try:
|
||||
self.dir_model = self.target_model_dir
|
||||
super().set_vocab()
|
||||
finally:
|
||||
self.dir_model = original_dir
|
||||
|
||||
self.gguf_writer.add_mask_token_id(self.hparams["dspark_noise_token_id"])
|
||||
|
||||
def set_gguf_parameters(self):
|
||||
super().set_gguf_parameters()
|
||||
|
||||
self.gguf_writer.add_block_size(self.hparams["dspark_block_size"])
|
||||
self.gguf_writer.add_target_layers([layer + 1 for layer in self.hparams["dspark_target_layer_ids"]])
|
||||
|
||||
+16
-5
@@ -122,8 +122,12 @@ def parse_args() -> argparse.Namespace:
|
||||
help="Export only the multi-token prediction (MTP) head as a separate GGUF, suitable for use as a speculative draft. An 'mtp-' prefix will be added to the output file name.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-mtp", action="store_true",
|
||||
help="Exclude the multi-token prediction (MTP) head from the converted GGUF. Pair with --mtp on a second run to publish trunk and MTP as two files. Note: the split form duplicates embeddings, but even though the bundled default is more space-efficient overall, this allows differing quantization which may be more performant.",
|
||||
"--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(
|
||||
"--dspark", action="store_true",
|
||||
help="Export only the DeepSeek-V4 DSpark draft tensors as a separate GGUF.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mistral-format", action="store_true",
|
||||
@@ -254,13 +258,20 @@ def main() -> None:
|
||||
from conversion.mistral import MistralModel
|
||||
model_class = MistralModel
|
||||
|
||||
if args.mtp and args.no_mtp:
|
||||
logger.error("--mtp and --no-mtp are mutually exclusive")
|
||||
if sum((args.mtp, args.no_mtp, args.dspark)) > 1:
|
||||
logger.error("--mtp, --no-nextn, and --dspark are mutually exclusive")
|
||||
sys.exit(1)
|
||||
|
||||
if args.dspark:
|
||||
if is_mistral_format or model_architecture != "DeepseekV4ForCausalLM":
|
||||
logger.error("--dspark is only supported for DeepseekV4ForCausalLM")
|
||||
sys.exit(1)
|
||||
from conversion.deepseek import DeepseekV4DSparkModel
|
||||
model_class = DeepseekV4DSparkModel
|
||||
|
||||
if args.mtp or args.no_mtp:
|
||||
if not model_class.supports_mtp_export:
|
||||
logger.error("--mtp / --no-mtp are not supported for %s", model_architecture)
|
||||
logger.error("--mtp / --no-nextn are not supported for %s", model_architecture)
|
||||
sys.exit(1)
|
||||
if args.no_mtp:
|
||||
model_class.no_mtp = True
|
||||
|
||||
@@ -477,6 +477,24 @@ 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_dsv4_hc(ggml_metal_library_t lib, ggml_op op) {
|
||||
const char * name = nullptr;
|
||||
|
||||
switch (op) {
|
||||
case GGML_OP_DSV4_HC_COMB: name = "kernel_dsv4_hc_comb_f32"; break;
|
||||
case GGML_OP_DSV4_HC_PRE: name = "kernel_dsv4_hc_pre_f32"; break;
|
||||
case GGML_OP_DSV4_HC_POST: name = "kernel_dsv4_hc_post_f32"; break;
|
||||
default: GGML_ABORT("fatal error");
|
||||
}
|
||||
|
||||
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_ssm_conv(ggml_metal_library_t lib, const ggml_tensor * op) {
|
||||
GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32);
|
||||
@@ -2117,6 +2135,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_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);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
|
||||
@@ -1137,6 +1137,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:
|
||||
@@ -1181,6 +1189,7 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
|
||||
case GGML_OP_MUL:
|
||||
case GGML_OP_DIV:
|
||||
case GGML_OP_ADD_ID:
|
||||
return ggml_is_contiguous_rows(op->src[0]) && ggml_is_contiguous_rows(op->src[1]) && (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16) && (op->src[0]->type == op->src[1]->type);
|
||||
case GGML_OP_ACC:
|
||||
return ggml_is_contiguous_rows(op->src[0]) && ggml_is_contiguous_rows(op->src[1]) && op->src[0]->type == GGML_TYPE_F32;
|
||||
case GGML_OP_REPEAT:
|
||||
@@ -1299,6 +1308,42 @@ 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_DSV4_HC_COMB:
|
||||
return has_simdgroup_reduction &&
|
||||
op->src[0]->type == GGML_TYPE_F32 &&
|
||||
op->src[1]->type == GGML_TYPE_F32 &&
|
||||
op->src[2]->type == GGML_TYPE_F32 &&
|
||||
op->type == GGML_TYPE_F32 &&
|
||||
op->src[0]->ne[0] == 24 &&
|
||||
op->src[1]->ne[0] >= 3 &&
|
||||
op->src[2]->ne[0] == 24 &&
|
||||
ggml_is_contiguous_rows(op->src[0]) &&
|
||||
ggml_is_contiguous_rows(op->src[1]) &&
|
||||
ggml_is_contiguous_rows(op->src[2]);
|
||||
case GGML_OP_DSV4_HC_PRE:
|
||||
return has_simdgroup_reduction &&
|
||||
op->src[0]->type == GGML_TYPE_F32 &&
|
||||
op->src[1]->type == GGML_TYPE_F32 &&
|
||||
op->type == GGML_TYPE_F32 &&
|
||||
op->src[0]->ne[1] == 4 &&
|
||||
op->src[1]->ne[0] == 4 &&
|
||||
ggml_is_contiguous_rows(op->src[0]) &&
|
||||
ggml_is_contiguous_rows(op->src[1]);
|
||||
case GGML_OP_DSV4_HC_POST:
|
||||
return has_simdgroup_reduction &&
|
||||
op->src[0]->type == GGML_TYPE_F32 &&
|
||||
op->src[1]->type == GGML_TYPE_F32 &&
|
||||
op->src[2]->type == GGML_TYPE_F32 &&
|
||||
op->src[3]->type == GGML_TYPE_F32 &&
|
||||
op->type == GGML_TYPE_F32 &&
|
||||
op->src[1]->ne[1] == 4 &&
|
||||
op->src[2]->ne[0] == 4 &&
|
||||
op->src[3]->ne[0] == 4 &&
|
||||
op->src[3]->ne[1] == 4 &&
|
||||
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]);
|
||||
case GGML_OP_SSM_CONV:
|
||||
case GGML_OP_SSM_SCAN:
|
||||
return has_simdgroup_reduction;
|
||||
|
||||
@@ -1171,6 +1171,49 @@ typedef struct {
|
||||
int64_t val;
|
||||
} ggml_metal_kargs_memset;
|
||||
|
||||
typedef struct {
|
||||
int32_t n_tokens;
|
||||
int32_t n_iter;
|
||||
uint64_t nb_m0;
|
||||
uint64_t nb_m1;
|
||||
uint64_t nb_s0;
|
||||
uint64_t nb_b0;
|
||||
uint64_t nb_d0;
|
||||
uint64_t nb_d1;
|
||||
uint64_t nb_d2;
|
||||
float eps;
|
||||
} ggml_metal_kargs_dsv4_hc_comb;
|
||||
|
||||
typedef struct {
|
||||
int32_t n_embd;
|
||||
int32_t n_tokens;
|
||||
uint64_t nb_x0;
|
||||
uint64_t nb_x1;
|
||||
uint64_t nb_x2;
|
||||
uint64_t nb_w0;
|
||||
uint64_t nb_w1;
|
||||
uint64_t nb_d0;
|
||||
uint64_t nb_d1;
|
||||
} ggml_metal_kargs_dsv4_hc_pre;
|
||||
|
||||
typedef struct {
|
||||
int32_t n_embd;
|
||||
int32_t n_tokens;
|
||||
uint64_t nb_x0;
|
||||
uint64_t nb_x1;
|
||||
uint64_t nb_r0;
|
||||
uint64_t nb_r1;
|
||||
uint64_t nb_r2;
|
||||
uint64_t nb_p0;
|
||||
uint64_t nb_p1;
|
||||
uint64_t nb_c0;
|
||||
uint64_t nb_c1;
|
||||
uint64_t nb_c2;
|
||||
uint64_t nb_d0;
|
||||
uint64_t nb_d1;
|
||||
uint64_t nb_d2;
|
||||
} ggml_metal_kargs_dsv4_hc_post;
|
||||
|
||||
typedef struct {
|
||||
int32_t ne00;
|
||||
int32_t ne01;
|
||||
@@ -1222,4 +1265,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,12 @@ 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_DSV4_HC_COMB:
|
||||
case GGML_OP_DSV4_HC_PRE:
|
||||
case GGML_OP_DSV4_HC_POST:
|
||||
{
|
||||
n_fuse = ggml_metal_op_dsv4_hc(ctx, idx);
|
||||
} break;
|
||||
case GGML_OP_SOFT_MAX:
|
||||
{
|
||||
n_fuse = ggml_metal_op_soft_max(ctx, idx);
|
||||
@@ -1297,6 +1307,137 @@ int ggml_metal_op_diag(ggml_metal_op_t ctx, int idx) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
int ggml_metal_op_dsv4_hc(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_tensor * op = ctx->node(idx);
|
||||
|
||||
ggml_metal_encoder_t enc = ctx->enc;
|
||||
auto pipeline = ggml_metal_library_get_pipeline_dsv4_hc(ctx->lib, op->op);
|
||||
|
||||
ggml_metal_encoder_set_pipeline(enc, pipeline);
|
||||
|
||||
switch (op->op) {
|
||||
case GGML_OP_DSV4_HC_COMB:
|
||||
{
|
||||
const ggml_tensor * mixes = op->src[0];
|
||||
const ggml_tensor * scale = op->src[1];
|
||||
const ggml_tensor * base = op->src[2];
|
||||
|
||||
GGML_ASSERT(mixes->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(scale->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(base->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(op->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(mixes->ne[0] == 24);
|
||||
GGML_ASSERT(op->ne[0] == 4 && op->ne[1] == 4);
|
||||
|
||||
ggml_metal_kargs_dsv4_hc_comb args = {
|
||||
/*.n_tokens =*/ (int32_t) mixes->ne[1],
|
||||
/*.n_iter =*/ ggml_get_op_params_i32(op, 1),
|
||||
/*.nb_m0 =*/ mixes->nb[0],
|
||||
/*.nb_m1 =*/ mixes->nb[1],
|
||||
/*.nb_s0 =*/ scale->nb[0],
|
||||
/*.nb_b0 =*/ base->nb[0],
|
||||
/*.nb_d0 =*/ op->nb[0],
|
||||
/*.nb_d1 =*/ op->nb[1],
|
||||
/*.nb_d2 =*/ op->nb[2],
|
||||
/*.eps =*/ ggml_get_op_params_f32(op, 0),
|
||||
};
|
||||
|
||||
ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(mixes), 1);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(scale), 2);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(base), 3);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 4);
|
||||
|
||||
// One SIMDgroup owns one 4x4 Sinkhorn matrix. Packing up to four
|
||||
// independent tokens per threadgroup keeps both decode and prompt
|
||||
// dispatches compact without any threadgroup-memory synchronization.
|
||||
const int nsg = std::min(4, args.n_tokens);
|
||||
ggml_metal_encoder_dispatch_threadgroups(
|
||||
enc, (args.n_tokens + nsg - 1)/nsg, 1, 1, 32, nsg, 1);
|
||||
} break;
|
||||
case GGML_OP_DSV4_HC_PRE:
|
||||
{
|
||||
const ggml_tensor * x = op->src[0];
|
||||
const ggml_tensor * weights = op->src[1];
|
||||
|
||||
GGML_ASSERT(x->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(weights->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(op->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(x->ne[1] == 4);
|
||||
|
||||
ggml_metal_kargs_dsv4_hc_pre args = {
|
||||
/*.n_embd =*/ (int32_t) x->ne[0],
|
||||
/*.n_tokens =*/ (int32_t) x->ne[2],
|
||||
/*.nb_x0 =*/ x->nb[0],
|
||||
/*.nb_x1 =*/ x->nb[1],
|
||||
/*.nb_x2 =*/ x->nb[2],
|
||||
/*.nb_w0 =*/ weights->nb[0],
|
||||
/*.nb_w1 =*/ weights->nb[1],
|
||||
/*.nb_d0 =*/ op->nb[0],
|
||||
/*.nb_d1 =*/ op->nb[1],
|
||||
};
|
||||
|
||||
ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(x), 1);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(weights), 2);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 3);
|
||||
|
||||
const int n_tiles = (args.n_embd + 31)/32;
|
||||
const int nsg = std::min(4, n_tiles);
|
||||
ggml_metal_encoder_dispatch_threadgroups(
|
||||
enc, (n_tiles + nsg - 1)/nsg, args.n_tokens, 1, 32, nsg, 1);
|
||||
} break;
|
||||
case GGML_OP_DSV4_HC_POST:
|
||||
{
|
||||
const ggml_tensor * x = op->src[0];
|
||||
const ggml_tensor * residual = op->src[1];
|
||||
const ggml_tensor * post = op->src[2];
|
||||
const ggml_tensor * comb = op->src[3];
|
||||
|
||||
GGML_ASSERT(x->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(residual->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(post->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(comb->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(op->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(residual->ne[1] == 4);
|
||||
|
||||
ggml_metal_kargs_dsv4_hc_post args = {
|
||||
/*.n_embd =*/ (int32_t) x->ne[0],
|
||||
/*.n_tokens =*/ (int32_t) x->ne[1],
|
||||
/*.nb_x0 =*/ x->nb[0],
|
||||
/*.nb_x1 =*/ x->nb[1],
|
||||
/*.nb_r0 =*/ residual->nb[0],
|
||||
/*.nb_r1 =*/ residual->nb[1],
|
||||
/*.nb_r2 =*/ residual->nb[2],
|
||||
/*.nb_p0 =*/ post->nb[0],
|
||||
/*.nb_p1 =*/ post->nb[1],
|
||||
/*.nb_c0 =*/ comb->nb[0],
|
||||
/*.nb_c1 =*/ comb->nb[1],
|
||||
/*.nb_c2 =*/ comb->nb[2],
|
||||
/*.nb_d0 =*/ op->nb[0],
|
||||
/*.nb_d1 =*/ op->nb[1],
|
||||
/*.nb_d2 =*/ op->nb[2],
|
||||
};
|
||||
|
||||
ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(x), 1);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(residual), 2);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(post), 3);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(comb), 4);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 5);
|
||||
|
||||
const int n_tiles = (args.n_embd + 31)/32;
|
||||
const int nsg = std::min(4, n_tiles);
|
||||
ggml_metal_encoder_dispatch_threadgroups(
|
||||
enc, (n_tiles + nsg - 1)/nsg, args.n_tokens, 1, 32, nsg, 1);
|
||||
} break;
|
||||
default:
|
||||
GGML_ABORT("fatal error");
|
||||
}
|
||||
|
||||
return 1;
|
||||
}
|
||||
|
||||
int ggml_metal_op_soft_max(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_tensor * op = ctx->node(idx);
|
||||
|
||||
@@ -3197,9 +3338,6 @@ int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) {
|
||||
GGML_TENSOR_LOCALS( int32_t, ne, op, ne);
|
||||
GGML_TENSOR_LOCALS(uint64_t, nb, op, nb);
|
||||
|
||||
GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32);
|
||||
|
||||
GGML_ASSERT(ggml_is_contiguous_rows(op->src[0]));
|
||||
GGML_ASSERT(ggml_is_contiguous_rows(op->src[1]));
|
||||
|
||||
@@ -3339,6 +3477,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_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);
|
||||
int ggml_metal_op_ssm_scan (ggml_metal_op_t ctx, int idx);
|
||||
@@ -70,6 +71,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)]];
|
||||
@@ -1418,6 +1432,8 @@ typedef decltype(kernel_bin_fuse_impl<float, float, float>) kernel_bin_fuse_t;
|
||||
|
||||
template [[host_name("kernel_bin_fuse_f32_f32_f32")]] kernel kernel_bin_fuse_t kernel_bin_fuse_impl<float, float, float>;
|
||||
template [[host_name("kernel_bin_fuse_f32_f32_f32_4")]] kernel kernel_bin_fuse_t kernel_bin_fuse_impl<float4, float4, float4>;
|
||||
template [[host_name("kernel_bin_fuse_f16_f16_f16")]] kernel kernel_bin_fuse_t kernel_bin_fuse_impl<half, half, half>;
|
||||
template [[host_name("kernel_bin_fuse_f16_f16_f16_4")]] kernel kernel_bin_fuse_t kernel_bin_fuse_impl<half4, half4, half4>;
|
||||
|
||||
kernel void kernel_add_id(
|
||||
constant ggml_metal_kargs_add_id & args,
|
||||
@@ -11278,3 +11294,162 @@ kernel void kernel_count_equal(
|
||||
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>;
|
||||
|
||||
kernel void kernel_dsv4_hc_comb_f32(
|
||||
constant ggml_metal_kargs_dsv4_hc_comb & args,
|
||||
device const char * mixes,
|
||||
device const char * scale,
|
||||
device const char * base,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
constexpr ushort hc = 4;
|
||||
constexpr ushort comb_offset = 2*hc;
|
||||
|
||||
const int it = tgpig.x*ntg.y + sgitg;
|
||||
if (it >= args.n_tokens) {
|
||||
return;
|
||||
}
|
||||
|
||||
float scale_lane = 0.0f;
|
||||
if (tiisg == 0) {
|
||||
scale_lane = *(device const float *) (scale + 2*args.nb_s0);
|
||||
}
|
||||
const float scale_comb = simd_shuffle(scale_lane, 0);
|
||||
|
||||
float v = 0.0f;
|
||||
if (tiisg < hc*hc) {
|
||||
v = *(device const float *) (mixes + (comb_offset + tiisg)*args.nb_m0 + it*args.nb_m1)*scale_comb
|
||||
+ *(device const float *) (base + (comb_offset + tiisg)*args.nb_b0);
|
||||
}
|
||||
|
||||
// Softmax across destinations (the four contiguous lanes for each source).
|
||||
float vmax = max(v, simd_shuffle_xor(v, 1));
|
||||
vmax = max(vmax, simd_shuffle_xor(vmax, 2));
|
||||
v = exp(v - vmax);
|
||||
|
||||
float sum = v + simd_shuffle_xor(v, 1);
|
||||
sum += simd_shuffle_xor(sum, 2);
|
||||
v = v/sum + args.eps;
|
||||
|
||||
// Normalize columns: equal destination indices are four lanes apart.
|
||||
sum = v + simd_shuffle_xor(v, 4);
|
||||
sum += simd_shuffle_xor(sum, 8);
|
||||
v /= sum + args.eps;
|
||||
|
||||
for (int i = 1; i < args.n_iter; ++i) {
|
||||
sum = v + simd_shuffle_xor(v, 1);
|
||||
sum += simd_shuffle_xor(sum, 2);
|
||||
v /= sum + args.eps;
|
||||
|
||||
sum = v + simd_shuffle_xor(v, 4);
|
||||
sum += simd_shuffle_xor(sum, 8);
|
||||
v /= sum + args.eps;
|
||||
}
|
||||
|
||||
if (tiisg < hc*hc) {
|
||||
const ushort idst = tiisg & 3;
|
||||
const ushort isrc = tiisg >> 2;
|
||||
*(device float *) (dst + idst*args.nb_d0 + isrc*args.nb_d1 + it*args.nb_d2) = v;
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_dsv4_hc_pre_f32(
|
||||
constant ggml_metal_kargs_dsv4_hc_pre & args,
|
||||
device const char * x,
|
||||
device const char * weights,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
constexpr ushort hc = 4;
|
||||
|
||||
const int it = tgpig.y;
|
||||
const int i0 = ((int) tgpig.x*ntg.y + sgitg)*32 + tiisg;
|
||||
|
||||
float weight_lane = 0.0f;
|
||||
if (tiisg < hc) {
|
||||
weight_lane = *(device const float *) (weights + tiisg*args.nb_w0 + it*args.nb_w1);
|
||||
}
|
||||
|
||||
float w[hc];
|
||||
FOR_UNROLL (ushort ih = 0; ih < hc; ++ih) {
|
||||
w[ih] = simd_shuffle(weight_lane, ih);
|
||||
}
|
||||
|
||||
if (i0 >= args.n_embd) {
|
||||
return;
|
||||
}
|
||||
|
||||
device const char * xb = x + i0*args.nb_x0 + it*args.nb_x2;
|
||||
float result = 0.0f;
|
||||
FOR_UNROLL (ushort ih = 0; ih < hc; ++ih) {
|
||||
result = fma(*(device const float *) (xb + ih*args.nb_x1), w[ih], result);
|
||||
}
|
||||
|
||||
*(device float *) (dst + i0*args.nb_d0 + it*args.nb_d1) = result;
|
||||
}
|
||||
|
||||
kernel void kernel_dsv4_hc_post_f32(
|
||||
constant ggml_metal_kargs_dsv4_hc_post & args,
|
||||
device const char * x,
|
||||
device const char * residual,
|
||||
device const char * post,
|
||||
device const char * comb,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
constexpr ushort hc = 4;
|
||||
|
||||
const int it = tgpig.y;
|
||||
const int i0 = ((int) tgpig.x*ntg.y + sgitg)*32 + tiisg;
|
||||
|
||||
float coeff_lane = 0.0f;
|
||||
if (tiisg < hc) {
|
||||
coeff_lane = *(device const float *) (post + tiisg*args.nb_p0 + it*args.nb_p1);
|
||||
} else if (tiisg < hc + hc*hc) {
|
||||
const ushort idx = tiisg - hc;
|
||||
const ushort idst = idx & 3;
|
||||
const ushort isrc = idx >> 2;
|
||||
coeff_lane = *(device const float *) (comb + idst*args.nb_c0 + isrc*args.nb_c1 + it*args.nb_c2);
|
||||
}
|
||||
|
||||
float post_reg[hc];
|
||||
float comb_reg[hc][hc];
|
||||
FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
|
||||
post_reg[idst] = simd_shuffle(coeff_lane, idst);
|
||||
}
|
||||
FOR_UNROLL (ushort isrc = 0; isrc < hc; ++isrc) {
|
||||
FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
|
||||
comb_reg[isrc][idst] = simd_shuffle(coeff_lane, hc + idst + hc*isrc);
|
||||
}
|
||||
}
|
||||
|
||||
if (i0 >= args.n_embd) {
|
||||
return;
|
||||
}
|
||||
|
||||
const float xv = *(device const float *) (x + i0*args.nb_x0 + it*args.nb_x1);
|
||||
float result[hc];
|
||||
FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
|
||||
result[idst] = xv*post_reg[idst];
|
||||
}
|
||||
|
||||
device const char * rb = residual + i0*args.nb_r0 + it*args.nb_r2;
|
||||
FOR_UNROLL (ushort isrc = 0; isrc < hc; ++isrc) {
|
||||
const float rv = *(device const float *) (rb + isrc*args.nb_r1);
|
||||
FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
|
||||
result[idst] = fma(rv, comb_reg[isrc][idst], result[idst]);
|
||||
}
|
||||
}
|
||||
|
||||
FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
|
||||
*(device float *) (dst + i0*args.nb_d0 + idst*args.nb_d1 + it*args.nb_d2) = result[idst];
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -24260,7 +24260,7 @@ static void ggml_cl_glu(ggml_backend_t backend, const ggml_tensor * src0, const
|
||||
}
|
||||
|
||||
const size_t nrows = ggml_nrows(src0);
|
||||
size_t nth = 512;
|
||||
size_t nth = backend_ctx->max_workgroup_size < 512 ? backend_ctx->max_workgroup_size : 512;
|
||||
size_t global_work_size[] = {nrows*nth, 1, 1};
|
||||
size_t local_work_size[] = {nth, 1, 1};
|
||||
|
||||
|
||||
@@ -4383,10 +4383,35 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.ATTN_OUT,
|
||||
MODEL_TENSOR.ATTN_Q_NORM,
|
||||
MODEL_TENSOR.ATTN_K_NORM,
|
||||
MODEL_TENSOR.ATTN_SINKS,
|
||||
MODEL_TENSOR.ATTN_Q_A,
|
||||
MODEL_TENSOR.ATTN_Q_B,
|
||||
MODEL_TENSOR.ATTN_Q_A_NORM,
|
||||
MODEL_TENSOR.ATTN_KV,
|
||||
MODEL_TENSOR.ATTN_KV_NORM,
|
||||
MODEL_TENSOR.ATTN_OUT_A,
|
||||
MODEL_TENSOR.ATTN_OUT_B,
|
||||
MODEL_TENSOR.HC_ATTN_FN,
|
||||
MODEL_TENSOR.HC_ATTN_BASE,
|
||||
MODEL_TENSOR.HC_ATTN_SCALE,
|
||||
MODEL_TENSOR.HC_FFN_FN,
|
||||
MODEL_TENSOR.HC_FFN_BASE,
|
||||
MODEL_TENSOR.HC_FFN_SCALE,
|
||||
MODEL_TENSOR.HC_HEAD_FN,
|
||||
MODEL_TENSOR.HC_HEAD_BASE,
|
||||
MODEL_TENSOR.HC_HEAD_SCALE,
|
||||
MODEL_TENSOR.FFN_NORM,
|
||||
MODEL_TENSOR.FFN_GATE,
|
||||
MODEL_TENSOR.FFN_DOWN,
|
||||
MODEL_TENSOR.FFN_UP,
|
||||
MODEL_TENSOR.FFN_GATE_INP,
|
||||
MODEL_TENSOR.FFN_EXP_PROBS_B,
|
||||
MODEL_TENSOR.FFN_GATE_EXP,
|
||||
MODEL_TENSOR.FFN_DOWN_EXP,
|
||||
MODEL_TENSOR.FFN_UP_EXP,
|
||||
MODEL_TENSOR.FFN_GATE_SHEXP,
|
||||
MODEL_TENSOR.FFN_DOWN_SHEXP,
|
||||
MODEL_TENSOR.FFN_UP_SHEXP,
|
||||
MODEL_TENSOR.FC,
|
||||
MODEL_TENSOR.ENC_OUTPUT_NORM,
|
||||
# optional DSpark heads
|
||||
|
||||
@@ -8069,6 +8069,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
|
||||
test_cases.emplace_back(new test_dsv4_hc_comb(1, 1));
|
||||
test_cases.emplace_back(new test_dsv4_hc_comb(17, 4));
|
||||
test_cases.emplace_back(new test_dsv4_hc_comb(257, 8));
|
||||
test_cases.emplace_back(new test_dsv4_hc_comb(17, 20));
|
||||
|
||||
test_cases.emplace_back(new test_dsv4_hc_pre(1, 1));
|
||||
test_cases.emplace_back(new test_dsv4_hc_pre(31, 17));
|
||||
@@ -8078,6 +8079,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
|
||||
test_cases.emplace_back(new test_dsv4_hc_post(1, 1));
|
||||
test_cases.emplace_back(new test_dsv4_hc_post(31, 17));
|
||||
test_cases.emplace_back(new test_dsv4_hc_post(128, 257));
|
||||
test_cases.emplace_back(new test_dsv4_hc_post(4096, 21));
|
||||
|
||||
// glu ops
|
||||
for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) {
|
||||
|
||||
Reference in New Issue
Block a user