diff --git a/conversion/__init__.py b/conversion/__init__.py index 2c38123df..2a87bd75b 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -74,6 +74,7 @@ TEXT_MODEL_MAP: dict[str, str] = { "Gemma3nForCausalLM": "gemma", "Gemma3nForConditionalGeneration": "gemma", "Gemma4ForConditionalGeneration": "gemma", + "Gemma4ForCausalLM": "gemma", "GemmaForCausalLM": "gemma", "Glm4ForCausalLM": "glm", "Glm4MoeForCausalLM": "glm", @@ -215,6 +216,7 @@ TEXT_MODEL_MAP: dict[str, str] = { "T5EncoderModel": "t5", "T5ForConditionalGeneration": "t5", "T5WithLMHeadModel": "t5", + "TalkieForCausalLM": "talkie", "UMT5ForConditionalGeneration": "t5", "UMT5Model": "t5", "UltravoxModel": "ultravox", diff --git a/conversion/base.py b/conversion/base.py index d8f050ed3..1d3554ea2 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -467,7 +467,14 @@ class ModelBase: elif quant_method == "compressed-tensors": quant_format = quant_config["format"] groups = quant_config["config_groups"] - if len(groups) > 1: + nvfp4_compressed_tensors = ( + quant_format == "nvfp4-pack-quantized" + or quant_format == "mixed-precision" + and bool(groups) + and all(g.get("format") == "nvfp4-pack-quantized" for g in groups.values() if isinstance(g, dict)) + ) + + if len(groups) > 1 and not nvfp4_compressed_tensors: raise NotImplementedError("Can't handle multiple config groups for compressed-tensors yet") weight_config = tuple(groups.values())[0]["weights"] @@ -505,6 +512,9 @@ class ModelBase: tensors_to_remove += [base_name + n for n in ("_packed", "_shape", "_scale")] if (base_name + "_zero_point") in self.model_tensors: tensors_to_remove.append(base_name + "_zero_point") + elif nvfp4_compressed_tensors: + # Don't error from compressed-tensors, we'll handle them in _generate_nvfp4_tensors + pass else: raise NotImplementedError(f"Quant format {quant_format!r} for method {quant_method!r} is not yet supported") elif quant_method == "modelopt": @@ -746,10 +756,13 @@ class ModelBase: del experts, merged def prepare_tensors(self): - # detect NVFP4 quantization (ModelOpt format) - quant_algo = (self.hparams.get("quantization_config") or {}).get("quant_algo") - quant_method = (self.hparams.get("quantization_config") or {}).get("quant_method") - quant_layers = (self.hparams.get("quantization_config") or {}).get("quantized_layers") or {} + # detect NVFP4 quantization (ModelOpt and Compressed-tensors formats) + quantization_config = self.hparams.get("quantization_config") or {} + quant_algo = quantization_config.get("quant_algo") + quant_method = quantization_config.get("quant_method") + quant_format = quantization_config.get("format") + quant_groups = quantization_config.get("config_groups") or {} + quant_layers = quantization_config.get("quantized_layers") or {} quant_config_file = self.dir_model / "hf_quant_config.json" if (not quant_algo or not quant_layers) and quant_config_file.is_file(): @@ -760,13 +773,25 @@ class ModelBase: producer_name = (producer.get("name") or "").lower() if quant_method is None: self.hparams.setdefault("quantization_config", {})["quant_method"] = producer_name + quant_method = producer_name quant_algo = quant_config.get("quant_algo", quant_algo) + quant_method = quant_config.get("quant_method", quant_method) + quant_format = quant_config.get("format", quant_format) + quant_groups = quant_config.get("config_groups", quant_groups) or {} quant_layers = quant_config.get("quantized_layers", quant_layers) or {} # Some models use per-tensor quant_algo (e.g. "MIXED_PRECISION" with # per-layer NVFP4/FP8) instead of a single global "NVFP4" value. + nvfp4_compressed_tensors = quant_method == "compressed-tensors" and ( + quant_format == "nvfp4-pack-quantized" + or quant_format == "mixed-precision" + and bool(quant_groups) + and all(g.get("format") == "nvfp4-pack-quantized" for g in quant_groups.values() if isinstance(g, dict)) + ) if quant_algo != "NVFP4": - if any(v.get("quant_algo") == "NVFP4" for v in quant_layers.values() if isinstance(v, dict)): + if nvfp4_compressed_tensors: + quant_algo = "NVFP4" + elif any(v.get("quant_algo") == "NVFP4" for v in quant_layers.values() if isinstance(v, dict)): quant_algo = "NVFP4" self._is_nvfp4 = quant_algo == "NVFP4" @@ -776,6 +801,28 @@ class ModelBase: # This must run before dequant_model so NVFP4 tensors are removed # from model_tensors, leaving only non-NVFP4 (e.g. FP8) for dequant. if self._is_nvfp4: + if nvfp4_compressed_tensors: + # Convert compressed-tensors 'global' scales into the reciprocal + def inverse_scale(gen): + def load(): + scale = LazyTorchTensor.to_eager(gen()).float() + return 1.0 / scale + return load + + # Change the compressed-tensors names to the ModelOpt names for handling consistently later + for name in list(self.model_tensors.keys()): + if name.endswith(".weight_packed"): + weight_name = name.removesuffix("_packed") + if weight_name not in self.model_tensors: + self.model_tensors[weight_name] = self.model_tensors.pop(name) + elif name.endswith(".weight_global_scale"): + scale2_name = name.replace(".weight_global_scale", ".weight_scale_2") + if scale2_name not in self.model_tensors: + self.model_tensors[scale2_name] = inverse_scale(self.model_tensors.pop(name)) + elif name.endswith(".input_global_scale"): + input_scale_name = name.replace(".input_global_scale", ".input_scale") + if input_scale_name not in self.model_tensors: + self.model_tensors[input_scale_name] = inverse_scale(self.model_tensors.pop(name)) self._generate_nvfp4_tensors() self.dequant_model() @@ -1575,6 +1622,9 @@ class TextModel(ModelBase): if chkhsh == "62f6fb0a6fd5098caeabb19b07a5c1099cafc8b9c40eab6ea89ece4ec02fbc57": # ref: https://huggingface.co/sarvamai/sarvam-30b res = "sarvam-moe" + if chkhsh == "f728162c1315c26e40249849799b4ba3fe584c32084b4795b03eb295e63cb5af": + # ref: https://huggingface.co/lewtun/talkie-1930-13b-it-hf + res = "talkie" if res is None: logger.warning("\n") diff --git a/conversion/gemma.py b/conversion/gemma.py index a6e14fbcb..1b427a30c 100644 --- a/conversion/gemma.py +++ b/conversion/gemma.py @@ -614,7 +614,7 @@ class Gemma3NModel(Gemma3Model): yield from super().modify_tensors(data_torch, name, bid) -@ModelBase.register("Gemma4ForConditionalGeneration") +@ModelBase.register("Gemma4ForConditionalGeneration", "Gemma4ForCausalLM") class Gemma4Model(Gemma3Model): model_arch = gguf.MODEL_ARCH.GEMMA4 diff --git a/conversion/qwen.py b/conversion/qwen.py index 45d1f98c2..7eb135c83 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -1,6 +1,5 @@ from __future__ import annotations -from pathlib import Path from typing import Any, Callable, Iterable, TYPE_CHECKING import torch @@ -549,6 +548,7 @@ class _Qwen35MtpMixin: 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) @@ -557,22 +557,44 @@ class _Qwen35MtpMixin: 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): - name, _ = 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 - return item - if cls.mtp_only: - canonical = name.replace("language_model.", "") - keep = canonical in ( + 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 super().filter_tensors(item) # ty: ignore[unresolved-attribute] + return name, gen def set_gguf_parameters(self): super().set_gguf_parameters() # ty: ignore[unresolved-attribute] @@ -594,29 +616,6 @@ class _Qwen35MtpMixin: 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" - def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: - if name.startswith("mtp."): - n_layer = self.hparams["num_hidden_layers"] - if name.find("layers.") != -1: - assert bid is not None - name = name.replace(f"mtp.layers.{bid}", f"model.layers.{bid + n_layer}") - bid = bid + n_layer - else: - remapper = { - "mtp.fc": "model.layers.{bid}.eh_proj", - "mtp.pre_fc_norm_embedding": "model.layers.{bid}.enorm", - "mtp.pre_fc_norm_hidden": "model.layers.{bid}.hnorm", - "mtp.norm": "model.layers.{bid}.shared_head.norm", - } - stem = Path(name).stem - suffix = Path(name).suffix - tmpl = remapper[stem] + suffix - for b in range(n_layer, self.block_count): - yield from super().modify_tensors(data_torch, tmpl.format(bid=b), b) # ty: ignore[unresolved-attribute] - return - - yield from super().modify_tensors(data_torch, name, bid) # ty: ignore[unresolved-attribute] - @ModelBase.register("Qwen3_5ForConditionalGeneration", "Qwen3_5ForCausalLM") class Qwen3_5TextModel(_Qwen35MtpMixin, _Qwen35MRopeMixin, _LinearAttentionVReorderBase): diff --git a/conversion/talkie.py b/conversion/talkie.py new file mode 100644 index 000000000..a970b32d3 --- /dev/null +++ b/conversion/talkie.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +from typing import Iterable, TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from torch import Tensor + +from .base import LazyTorchTensor, ModelBase, TextModel, gguf + + +@ModelBase.register("TalkieForCausalLM") +class TalkieModel(TextModel): + model_arch = gguf.MODEL_ARCH.TALKIE + + def set_gguf_parameters(self): + super().set_gguf_parameters() + # Talkie used F.rms_norm without an explicit eps + self.gguf_writer.add_layer_norm_rms_eps(torch.finfo(torch.float32).eps) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + prefix = f"model.blocks.{bid}." if bid is not None else "" + suffix = name.removeprefix(prefix) + + if suffix == "attn_gain.a_g": + yield self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_OUT, bid, ".scale"), data_torch + return + elif suffix == "mlp_gain.a_g": + yield self.format_tensor_name(gguf.MODEL_TENSOR.FFN_DOWN, bid, ".scale"), data_torch + return + elif suffix == "lm_head_gain.w_g": + self.gguf_writer.add_logit_scale(LazyTorchTensor.to_eager(data_torch).item()) + return + elif suffix in ("attn.attn_query.weight", "attn.attn_key.weight"): + # absorb inverse rope + head_dim = self.hparams["head_dim"] + shape = data_torch.shape + data_torch = torch.reshape(data_torch, (-1, head_dim, shape[-1])) + signs = torch.ones((1, head_dim, 1), dtype=data_torch.dtype) + signs[:, head_dim // 2 :, :] = -1 + if self.lazy: + signs = LazyTorchTensor.from_eager(signs) + # (n_head, head_dim, n_in) -> (n_out, n_in) + data_torch = torch.reshape(data_torch * signs, shape) + elif suffix == "attn.head_gain.head_g": + # allow head gain to broadcast + data_torch = data_torch.unsqueeze(-1) + + if not name.endswith(".weight"): + name += ".weight" + + yield from super().modify_tensors(data_torch, name, bid) diff --git a/convert_hf_to_gguf_update.py b/convert_hf_to_gguf_update.py index 8b2a9454f..8bfe04b3d 100755 --- a/convert_hf_to_gguf_update.py +++ b/convert_hf_to_gguf_update.py @@ -156,6 +156,7 @@ models = [ {"name": "kanana2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/kakaocorp/kanana-2-30b-a3b-instruct-2601", }, {"name": "f2llmv2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/codefuse-ai/F2LLM-v2-4B", }, {"name": "sarvam-moe", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/sarvamai/sarvam-30b", }, + {"name": "talkie", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/lewtun/talkie-1930-13b-it-hf", }, ] # some models are known to be broken upstream, so we will skip them as exceptions diff --git a/convert_lora_to_gguf.py b/convert_lora_to_gguf.py index 81658ba03..9a6437bea 100755 --- a/convert_lora_to_gguf.py +++ b/convert_lora_to_gguf.py @@ -208,6 +208,16 @@ class LoraTorchTensor: def to(self, *args, **kwargs): return LoraTorchTensor(self._lora_A.to(*args, **kwargs), self._lora_B.to(*args, **kwargs)) + def __mul__(self, other) -> LoraTorchTensor: + # Only output-side multiplication for now + # W = B @ A, so M_out * W == (M_out * B) @ A + if not isinstance(other, (int, float)) and other.shape and other.shape[-1] != 1: + raise NotImplementedError + return LoraTorchTensor(self._lora_A, self._lora_B * other) + + def __rmul__(self, other) -> LoraTorchTensor: + return self * other + @classmethod def __torch_function__(cls, func: Callable, types, args=(), kwargs=None): del types # unused diff --git a/embd_res/klite.embd b/embd_res/klite.embd index c585668f6..5d832a4bd 100644 --- a/embd_res/klite.embd +++ b/embd_res/klite.embd @@ -6022,7 +6022,7 @@ Current version indicated by LITEVER below. inputtxt = replaceAll(inputtxt,instructsysplaceholder_end.trim(),get_instruct_systag_end(false)); return inputtxt; } - function replace_noninstruct_placeholders(inputtxt,escape=false,isgametext=true,persistent_countmap=null) + function replace_chat_name_placeholders(inputtxt, escape=false) { let firstopponent = localsettings.chatopponent; if (firstopponent && firstopponent.includes("||$||")) { @@ -6040,7 +6040,11 @@ Current version indicated by LITEVER below. inputtxt = replaceAll(inputtxt,"{{user}}",(localsettings.chatname?localsettings.chatname:"User"),true); inputtxt = replaceAll(inputtxt,"{{char}}",(localsettings.chatopponent?firstopponent:defaultchatopponent),true); } - + return inputtxt; + } + function replace_noninstruct_placeholders(inputtxt,escape=false,isgametext=true,persistent_countmap=null) + { + inputtxt = replace_chat_name_placeholders(inputtxt,escape); inputtxt = replaceAll(inputtxt,"{{original}}","",true); //redundant in Lite; use sysprompt or assistant jailbreak instead for(let i=0;i { + if (userInput.match(/character\//i)) + { + userInput = userInput.split("#")[0].split("?")[0]; + userInput = userInput.endsWith('/') ? userInput.slice(0, -1) : userInput; + let tmp = userInput.split("character/")[1]; + return `https://botbooru.com/download/png/${tmp}`; + } + else + { + userInput = userInput.split("#")[0].split("?")[0]; + userInput = userInput.endsWith('/') ? userInput.slice(0, -1) : userInput; + return userInput; + } + }, + fetch: (userInput) => { + temp_scenario = { + "title":"", + "desc": "", + "opmode":3, + "chatopponent": "", + "gui_type":1, + "prompt":"", + "memory": "", + "authorsnote": "", + "worldinfo": [], + }; + + finalurl = apply_proxy_url(userInput,true); + //try to obtain the full portrait image + return fetch(finalurl, { + method: 'GET', + redirect: 'follow', + referrerPolicy: 'no-referrer', + }) + .then(rb => { + if(rb.ok) + { + return rb.blob(); + }else{ + throw new Error('Cannot fetch tavern image'); + } + }) + .then(blob => { + preview_temp_scenario(); + + readTavernPngFromBlob(blob,(obj)=>{ + load_temp_scenario_from_tavernobj(obj); + }); + + const objectURL = URL.createObjectURL(blob); + const compressedImg = compressImage(objectURL, (compressedImageURI, aspectratio)=>{ + temp_scenario.image = compressedImageURI; + temp_scenario.image_aspect = aspectratio; + preview_temp_scenario(); + }, true, AVATAR_PX); + }) + .catch(error => { + throw new Error("Selected scenario is invalid."); + console.error(error); + }); + } + }, + //raw png URL download + { + name: "Raw URL", + urlParam: "rawurl", + matchers: ["/"], + inputBox: { + text: "Enter the direct URL to the character card file", + placeholder: "https://example.com/file.png", + }, + extraction: (userInput) => { + return userInput; + }, + fetch: (userInput) => { + temp_scenario = { + "title":"", + "desc": "", + "opmode":3, + "chatopponent": "", + "gui_type":1, + "prompt":"", + "memory": "", + "authorsnote": "", + "worldinfo": [], + }; + + finalurl = apply_proxy_url(userInput,false); + //try to obtain the full portrait image + return fetch(finalurl, { + method: 'GET', + redirect: 'follow', + referrerPolicy: 'no-referrer', + }) + .then(rb => { + if(rb.ok) + { + return rb.blob(); + }else{ + throw new Error('Cannot fetch tavern image'); + } + }) + .then(blob => { + preview_temp_scenario(); + + readTavernPngFromBlob(blob,(obj)=>{ + load_temp_scenario_from_tavernobj(obj); + }); + + const objectURL = URL.createObjectURL(blob); + const compressedImg = compressImage(objectURL, (compressedImageURI, aspectratio)=>{ + temp_scenario.image = compressedImageURI; + temp_scenario.image_aspect = aspectratio; + preview_temp_scenario(); + }, true, AVATAR_PX); + }) + .catch(error => { + throw new Error("Selected scenario is invalid."); + console.error(error); + }); + } + }, // char-archive has shut down // { @@ -12308,7 +12450,7 @@ Current version indicated by LITEVER below. else inputBox( scenario_source.inputBox.text, "Import from " + scenario_source.name, - "", + getInputBoxValue(), scenario_source.inputBox.placeholder, (value) => { const input = getInputBoxValue().trim(); @@ -12317,6 +12459,57 @@ Current version indicated by LITEVER below. ); } + function import_scenario_unified_dropdown() + { + let dropdown = document.getElementById("importscenariounifieddropdown"); + if(dropdown) + { + let picked = dropdown.value; + if(picked!="-1") + { + import_scenario(scenario_sources[picked]); + } + } + } + function import_scenario_unified() //a single button entry point to import from any source + { + let combined_txt = "Please enter the full URL of the character card to import, or pick a specific website."; + let selbox = `Source: "; + combined_txt += "
" + selbox; + + inputBox(combined_txt,"Import character card from URL","","Enter full URL to character card", + (value) => { + const input = getInputBoxValue().trim(); + let found = -1; + for(let i=0;iLoad Character Card or JSON File`; scenarios += ``; scenarios += ``; - - for(let i=0;i${scenario_sources[i].name}` - } + scenarios += `` for(let i=0;in_free_blocks; i++) { + for (int i = idx; i < chunk->n_free_blocks - 1; i++) { chunk->free_blocks[i] = chunk->free_blocks[i+1]; } chunk->n_free_blocks--; diff --git a/ggml/src/ggml-cuda/fwht.cu b/ggml/src/ggml-cuda/fwht.cu new file mode 100644 index 000000000..184dc254c --- /dev/null +++ b/ggml/src/ggml-cuda/fwht.cu @@ -0,0 +1,101 @@ +#include "common.cuh" +#include "fwht.cuh" + +template +__launch_bounds__(4*ggml_cuda_get_physical_warp_size(), 1) +__global__ void fwht_cuda(const float * src, float * dst, const int64_t n_rows, const float scale) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + + const int64_t r = (int64_t) blockIdx.x * blockDim.y + threadIdx.y; + + if (r >= n_rows) { + return; + } + + src += r * N; + dst += r * N; + + static constexpr int el_w = N / warp_size; + float reg[el_w]; + const int lane = threadIdx.x; + + ggml_cuda_pdl_sync(); +#pragma unroll + for (int i = 0; i < el_w; ++i) { + reg[i] = src[i * warp_size + lane] * scale; + } + +#pragma unroll + for (int h = 1; h < warp_size; h *= 2) { +#pragma unroll + for (int j = 0; j < el_w; j++) { + const float val = reg[j]; + const float val2 = __shfl_xor_sync(0xFFFFFFFF, val, h, warp_size); + + reg[j] = (lane & h) == 0 ? val + val2 : val2 - val; + } + } + +#pragma unroll + for (int h = warp_size; h < N; h *= 2) { + const int step = h / warp_size; +#pragma unroll + for (int j = 0; j < el_w; j += 2 * step) { +#pragma unroll + for (int k = 0; k < step; k++) { + const float x = reg[j + k]; + const float y = reg[j + k + step]; + + reg[j + k] = x + y; + reg[j + k + step] = x - y; + } + } + } + +#pragma unroll + for (int i = 0; i < el_w; ++i) { + dst[i * warp_size + lane] = reg[i]; + } +} + +bool ggml_cuda_op_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst) { + GGML_ASSERT(ggml_are_same_shape(src, dst)); + if (!ggml_is_contiguous(src) || !ggml_is_contiguous(dst)) { + return false; + } + const int n = src->ne[0]; + const int64_t rows = ggml_nrows(src); + + const float * src_d = (const float *) src->data; + float * dst_d = (float *) dst->data; + + const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; + const int rows_per_block = 4; + + const int64_t num_blocks = (rows + rows_per_block - 1) / rows_per_block; + + cudaStream_t stream = ctx.stream(); + dim3 grid_dims(num_blocks, 1, 1); + dim3 block_dims(warp_size, rows_per_block, 1); + const ggml_cuda_kernel_launch_params launch_params = + ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, stream); + + const float scale = 1 / sqrtf(n); + + switch (n) { + case 64: + ggml_cuda_kernel_launch(fwht_cuda<64>, launch_params, src_d, dst_d, rows, scale); + return true; + case 128: + ggml_cuda_kernel_launch(fwht_cuda<128>, launch_params, src_d, dst_d, rows, scale); + return true; + case 256: + ggml_cuda_kernel_launch(fwht_cuda<256>, launch_params, src_d, dst_d, rows, scale); + return true; + case 512: + ggml_cuda_kernel_launch(fwht_cuda<512>, launch_params, src_d, dst_d, rows, scale); + return true; + default: + return false; + } +} diff --git a/ggml/src/ggml-cuda/fwht.cuh b/ggml/src/ggml-cuda/fwht.cuh new file mode 100644 index 000000000..cf3df94ca --- /dev/null +++ b/ggml/src/ggml-cuda/fwht.cuh @@ -0,0 +1,4 @@ +#include "common.cuh" + +// Returns whether the Fast Walsh-Hadamard transform could be used. +bool ggml_cuda_op_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 25bcbda57..a0e3cd30a 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -26,6 +26,7 @@ bool g_mul_mat_q = true; #include "ggml-cuda/diagmask.cuh" #include "ggml-cuda/diag.cuh" #include "ggml-cuda/fattn.cuh" +#include "ggml-cuda/fwht.cuh" #include "ggml-cuda/getrows.cuh" #include "ggml-cuda/im2col.cuh" #include "ggml-cuda/mmf.cuh" @@ -2594,6 +2595,11 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor bool use_batched_cublas_bf16 = src0->type == GGML_TYPE_BF16 && bf16_mma_hardware_available(cc); bool use_batched_cublas_f32 = src0->type == GGML_TYPE_F32; + const int32_t hint = ggml_get_op_params_i32(dst, 1); + if (hint == GGML_HINT_SRC0_IS_HADAMARD && !split && ggml_cuda_op_fwht(ctx, src1, dst)) { + return; + } + if(ggml_cuda_highest_compiled_arch(cc) <= GGML_CUDA_CC_TURING) { //kcpp: https://github.com/ggml-org/llama.cpp/pull/14361 broke oldpc mode without this. diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 1f212a92f..4a3ebb556 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -215,6 +215,30 @@ void ggml_metal_rsets_free(ggml_metal_rsets_t rsets); // device // +enum ggml_metal_device_id { + GGML_METAL_DEVICE_GENERIC = 0, + + GGML_METAL_DEVICE_M1, + GGML_METAL_DEVICE_M1_PRO, + GGML_METAL_DEVICE_M1_MAX, + GGML_METAL_DEVICE_M1_ULTRA, + GGML_METAL_DEVICE_M2, + GGML_METAL_DEVICE_M2_PRO, + GGML_METAL_DEVICE_M2_MAX, + GGML_METAL_DEVICE_M2_ULTRA, + GGML_METAL_DEVICE_M3, + GGML_METAL_DEVICE_M3_PRO, + GGML_METAL_DEVICE_M3_MAX, + GGML_METAL_DEVICE_M3_ULTRA, + GGML_METAL_DEVICE_M4, + GGML_METAL_DEVICE_M4_PRO, + GGML_METAL_DEVICE_M4_MAX, + GGML_METAL_DEVICE_M5, + GGML_METAL_DEVICE_M5_PRO, + GGML_METAL_DEVICE_M5_MAX, + GGML_METAL_DEVICE_M5_ULTRA, +}; + struct ggml_metal_device_props { int device; char name[128]; @@ -234,6 +258,8 @@ struct ggml_metal_device_props { bool supports_gpu_family_apple7; + enum ggml_metal_device_id device_id; + int op_offload_min_batch_size; }; diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index c2a4e6905..9138678b9 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -634,6 +634,50 @@ void ggml_metal_rsets_free(ggml_metal_rsets_t rsets) { free(rsets); } +static enum ggml_metal_device_id ggml_metal_device_id_parse(const char * name) { + if (!name) { + return GGML_METAL_DEVICE_GENERIC; + } + + static const char prefix[] = "Apple "; + if (strncmp(name, prefix, sizeof(prefix) - 1) != 0) { + return GGML_METAL_DEVICE_GENERIC; + } + const char * suffix = name + sizeof(prefix) - 1; + + static const struct { + const char * name; + enum ggml_metal_device_id id; + } table[] = { + {"M1", GGML_METAL_DEVICE_M1}, + {"M1 Pro", GGML_METAL_DEVICE_M1_PRO}, + {"M1 Max", GGML_METAL_DEVICE_M1_MAX}, + {"M1 Ultra", GGML_METAL_DEVICE_M1_ULTRA}, + {"M2", GGML_METAL_DEVICE_M2}, + {"M2 Pro", GGML_METAL_DEVICE_M2_PRO}, + {"M2 Max", GGML_METAL_DEVICE_M2_MAX}, + {"M2 Ultra", GGML_METAL_DEVICE_M2_ULTRA}, + {"M3", GGML_METAL_DEVICE_M3}, + {"M3 Pro", GGML_METAL_DEVICE_M3_PRO}, + {"M3 Max", GGML_METAL_DEVICE_M3_MAX}, + {"M3 Ultra", GGML_METAL_DEVICE_M3_ULTRA}, + {"M4", GGML_METAL_DEVICE_M4}, + {"M4 Pro", GGML_METAL_DEVICE_M4_PRO}, + {"M4 Max", GGML_METAL_DEVICE_M4_MAX}, + {"M5", GGML_METAL_DEVICE_M5}, + {"M5 Pro", GGML_METAL_DEVICE_M5_PRO}, + {"M5 Max", GGML_METAL_DEVICE_M5_MAX}, + {"M5 Ultra", GGML_METAL_DEVICE_M5_ULTRA}, + }; + + for (size_t i = 0; i < sizeof(table)/sizeof(table[0]); ++i) { + if (strcmp(suffix, table[i].name) == 0) { + return table[i].id; + } + } + return GGML_METAL_DEVICE_GENERIC; +} + ggml_metal_device_t ggml_metal_device_init(int device) { ggml_metal_device_t dev = calloc(1, sizeof(struct ggml_metal_device)); @@ -801,6 +845,8 @@ ggml_metal_device_t ggml_metal_device_init(int device) { dev->props.supports_gpu_family_apple7 = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7]; + dev->props.device_id = ggml_metal_device_id_parse([[dev->mtl_device name] UTF8String]); + dev->props.op_offload_min_batch_size = getenv("GGML_OP_OFFLOAD_MIN_BATCH") ? atoi(getenv("GGML_OP_OFFLOAD_MIN_BATCH")) : 32; dev->props.max_buffer_size = dev->mtl_device.maxBufferLength; diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 7fdcf03d7..0189f6f03 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -505,6 +505,7 @@ class MODEL_ARCH(IntEnum): LLAMA_EMBED = auto() MAINCODER = auto() KIMI_LINEAR = auto() + TALKIE = auto() class VISION_PROJECTOR_TYPE(IntEnum): @@ -1021,6 +1022,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { MODEL_ARCH.LLAMA_EMBED: "llama-embed", MODEL_ARCH.MAINCODER: "maincoder", MODEL_ARCH.KIMI_LINEAR: "kimi-linear", + MODEL_ARCH.TALKIE: "talkie", } VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = { @@ -4013,6 +4015,19 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.FFN_DOWN_SHEXP, MODEL_TENSOR.FFN_UP_SHEXP, ], + MODEL_ARCH.TALKIE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.LAYER_OUT_SCALE, + ], # TODO } diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index c2235cb3b..ecc3c05f9 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -34,6 +34,7 @@ class TensorNameMap: "encoder", # neobert "model.transformer.wte", # llada "embed_tokens", # qwen3-embedding + "model.embed", # talkie ), # Token type embeddings @@ -259,6 +260,7 @@ class TensorNameMap: "model.transformer.blocks.{bid}.q_proj", # llada "layers.{bid}.self_attn.q_proj", # qwen3-embedding "backbone.layers.{bid}.mixer.q_proj", # nemotron-h + "model.blocks.{bid}.attn.attn_query", # talkie ), # Attention key @@ -279,6 +281,7 @@ class TensorNameMap: "model.transformer.blocks.{bid}.k_proj", # llada "layers.{bid}.self_attn.k_proj", # qwen3-embedding "backbone.layers.{bid}.mixer.k_proj", # nemotron-h + "model.blocks.{bid}.attn.attn_key", # talkie ), # Attention value @@ -298,6 +301,7 @@ class TensorNameMap: "model.transformer.blocks.{bid}.v_proj", # llada "layers.{bid}.self_attn.v_proj", # qwen3-embedding "backbone.layers.{bid}.mixer.v_proj", # nemotron-h + "model.blocks.{bid}.attn.attn_value", # talkie ), # Attention output @@ -336,6 +340,7 @@ class TensorNameMap: "layers.{bid}.self_attn.o_proj", # qwen3-embedding "backbone.layers.{bid}.mixer.o_proj", # nemotron-h "model.layers.{bid}.self_attn.language_expert_dense", # cogvlm + "model.blocks.{bid}.attn.attn_resid", # talkie ), # Attention output norm @@ -508,6 +513,7 @@ class TensorNameMap: "layers.{bid}.mlp.up_proj", # qwen3-embedding "backbone.layers.{bid}.mixer.up_proj", # nemotron-h "model.layers.{bid}.mlp.language_mlp.up_proj", # cogvlm + "model.blocks.{bid}.mlp.mlp_linear", # talkie ), MODEL_TENSOR.FFN_UP_EXP: ( @@ -561,6 +567,7 @@ class TensorNameMap: "model.transformer.blocks.{bid}.ff_proj", # llada "layers.{bid}.mlp.gate_proj", # qwen3-embedding "model.layers.{bid}.mlp.language_mlp.gate_proj", # cogvlm + "model.blocks.{bid}.mlp.mlp_gate", # talkie ), MODEL_TENSOR.FFN_GATE_EXP: ( @@ -636,6 +643,7 @@ class TensorNameMap: "layers.{bid}.mlp.down_proj", # qwen3-embedding "backbone.layers.{bid}.mixer.down_proj", # nemotron-h "model.layers.{bid}.mlp.language_mlp.down_proj", # cogvlm + "model.blocks.{bid}.mlp.mlp_resid", # talkie ), MODEL_TENSOR.FFN_DOWN_EXP: ( @@ -682,6 +690,7 @@ class TensorNameMap: "model.layers.layers.{bid}.mixer.q_norm", # plamo3 "layers.{bid}.self_attn.q_norm", # qwen3-embedding "model.layers.{bid}.attention.query_layernorm", # apertus + "model.blocks.{bid}.attn.head_gain.head_g", # talkie ), MODEL_TENSOR.ATTN_K_NORM: ( @@ -716,6 +725,7 @@ class TensorNameMap: MODEL_TENSOR.LAYER_OUT_SCALE: ( "model.layers.{bid}.layer_scalar", # gemma4 + "model.blocks.{bid}.embed_skip.a_g", # talkie ), MODEL_TENSOR.PER_LAYER_TOKEN_EMBD: ( diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index c9eead18a..e95ba6daa 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -133,6 +133,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_LLAMA_EMBED, "llama-embed" }, { LLM_ARCH_MAINCODER, "maincoder" }, { LLM_ARCH_KIMI_LINEAR, "kimi-linear" }, + { LLM_ARCH_TALKIE, "talkie" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; @@ -767,8 +768,9 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, // Nemotron 3 Super - {LLM_TENSOR_FFN_LATENT_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, - {LLM_TENSOR_FFN_LATENT_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + // latent projections feed ggml_mul_mat, the buft probe must use MUL_MAT to keep them on GPU + {LLM_TENSOR_FFN_LATENT_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_FFN_LATENT_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, }; LLM_KV::LLM_KV(llm_arch arch, const char * suffix) : arch(arch), suffix(suffix) {} diff --git a/src/llama-arch.h b/src/llama-arch.h index 89cf16cc3..7c1dcc4d6 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -137,6 +137,7 @@ enum llm_arch { LLM_ARCH_LLAMA_EMBED, LLM_ARCH_MAINCODER, LLM_ARCH_KIMI_LINEAR, + LLM_ARCH_TALKIE, LLM_ARCH_UNKNOWN, }; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index be63c2701..1a10389c3 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -161,6 +161,7 @@ #include "models/step35.cpp" #include "models/t5.cpp" #include "models/t5encoder.cpp" +#include "models/talkie.cpp" #include "models/wavtokenizer-dec.cpp" #include "models/xverse.cpp" @@ -174,6 +175,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_llama_embed(params); case LLM_ARCH_MAINCODER: return new llama_model_maincoder(params); + case LLM_ARCH_TALKIE: + return new llama_model_talkie(params); case LLM_ARCH_DECI: return new llama_model_deci(params); case LLM_ARCH_BAICHUAN: @@ -2484,6 +2487,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_QWEN3NEXT: case LLM_ARCH_MIMO2: case LLM_ARCH_STEP35: + case LLM_ARCH_TALKIE: return LLAMA_ROPE_TYPE_NEOX; case LLM_ARCH_QWEN2VL: diff --git a/src/llama-model.h b/src/llama-model.h index 398a0aa72..b797b8966 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -488,7 +488,7 @@ struct llama_layer { struct ggml_tensor * indexer_attn_k = nullptr; struct ggml_tensor * indexer_attn_q_b = nullptr; // note: for lora a/b, not bias - // gemma4 layer output scale + // gemma4 layer output scale, reused for talkie embedding skip scale struct ggml_tensor * out_scale = nullptr; struct llama_layer_posnet posnet; diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index c0a7cc3ad..acdb42cd4 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -2432,7 +2432,8 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { } else if ( tokenizer_pre == "gpt-4o" || tokenizer_pre == "llama4" || - tokenizer_pre == "kanana2") { + tokenizer_pre == "kanana2" || + tokenizer_pre == "talkie") { pre_type = LLAMA_VOCAB_PRE_TYPE_GPT4O; clean_spaces = false; } else if ( diff --git a/src/models/mistral3.cpp b/src/models/mistral3.cpp index 4e6ebef82..1ac5a95cc 100644 --- a/src/models/mistral3.cpp +++ b/src/models/mistral3.cpp @@ -177,9 +177,9 @@ llama_model_mistral3::graph::graph(const llama_model & model, const llm_graph_pa cb(cur, "ffn_norm", il); cur = build_ffn(cur, - model.layers[il].ffn_up, model.layers[il].ffn_up_b, NULL, - model.layers[il].ffn_gate, model.layers[il].ffn_gate_b, NULL, - model.layers[il].ffn_down, model.layers[il].ffn_down_b, NULL, + model.layers[il].ffn_up, model.layers[il].ffn_up_b, model.layers[il].ffn_up_s, + model.layers[il].ffn_gate, model.layers[il].ffn_gate_b, model.layers[il].ffn_gate_s, + model.layers[il].ffn_down, model.layers[il].ffn_down_b, model.layers[il].ffn_down_s, NULL, LLM_FFN_SILU, LLM_FFN_PAR, il); cb(cur, "ffn_out", il); @@ -200,7 +200,11 @@ llama_model_mistral3::graph::graph(const llama_model & model, const llm_graph_pa LLM_FFN_SILU, true, hparams.expert_weights_scale, LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, - il); + il, + nullptr, nullptr, + model.layers[il].ffn_up_exps_s, + model.layers[il].ffn_gate_exps_s, + model.layers[il].ffn_down_exps_s); cb(cur, "ffn_moe_out", il); } cur = ggml_add(ctx0, cur, ffn_inp); diff --git a/src/models/models.h b/src/models/models.h index 7e551eb96..db228865d 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -186,6 +186,19 @@ struct llama_model_maincoder : public llama_model_base { }; +struct llama_model_talkie : public llama_model_base { + llama_model_talkie(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + struct graph : public llm_graph_context { + graph(const llama_model & model, const llm_graph_params & params); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_deci : public llama_model_base { llama_model_deci(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/src/models/talkie.cpp b/src/models/talkie.cpp new file mode 100644 index 000000000..1258eeb19 --- /dev/null +++ b/src/models/talkie.cpp @@ -0,0 +1,149 @@ +#include "models.h" + +void llama_model_talkie::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale); + + switch (hparams.n_layer) { + case 40: type = LLM_TYPE_13B; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_talkie::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0); + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_gqa, n_embd_gqa, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0); + + // no k gain + layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {1, n_head}, 0); + + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0); + + layer.out_scale = create_tensor(tn(LLM_TENSOR_LAYER_OUT_SCALE, "weight", i), {1}, 0); + } +} + +std::unique_ptr llama_model_talkie::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +llama_model_talkie::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { + const int64_t n_embd_head = hparams.n_embd_head_k(); + + GGML_ASSERT(n_embd_head == hparams.n_embd_head_v()); + GGML_ASSERT(n_embd_head == n_rot); + + ggml_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + inpL = build_norm(inpL, nullptr, nullptr, LLM_NORM_RMS, -1); + cb(inpL, "inp_norm", -1); + + ggml_tensor * embd_skip = inpL; + + // inp_pos - contains the positions + ggml_tensor * inp_pos = build_inp_pos(); + + auto * inp_attn = build_attn_inp_kv(); + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + const float kq_scale = 1.0f / sqrtf(float(n_embd_head)); + + for (int il = 0; il < n_layer; ++il) { + ggml_tensor * inpSA = inpL; + ggml_tensor * inp_skip = embd_skip; + + cur = build_norm(inpL, nullptr, nullptr, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + // self-attention + { + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head, n_head, n_head_kv, il); + + 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); + + // reference applies qknorm after rope + Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il); + cb(Qcur, "Qcur_norm", il); + + Kcur = build_norm(Kcur, nullptr, nullptr, LLM_NORM_RMS, il); + cb(Kcur, "Kcur_norm", il); + + cb(Vcur, "Vcur", il); + + cur = build_attn(inp_attn, + model.layers[il].wo, nullptr, model.layers[il].wo_s, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + cb(cur, "attn_out", il); + } + + if (il == n_layer - 1 && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); + inp_skip = ggml_get_rows(ctx0, inp_skip, inp_out_ids); + } + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cb(ffn_inp, "ffn_inp", il); + + cur = build_norm(ffn_inp, nullptr, nullptr, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + cur = build_ffn(cur, + model.layers[il].ffn_up, nullptr, nullptr, + model.layers[il].ffn_gate, nullptr, nullptr, + model.layers[il].ffn_down, nullptr, model.layers[il].ffn_down_s, + nullptr, + LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(cur, "ffn_out", il); + + cur = ggml_add(ctx0, cur, ffn_inp); + + ggml_tensor * skip = ggml_mul(ctx0, inp_skip, model.layers[il].out_scale); + cb(skip, "embd_skip", il); + + cur = ggml_add(ctx0, cur, skip); + + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + + // input for next layer + inpL = cur; + } + + cur = inpL; + + cur = build_norm(cur, nullptr, nullptr, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + + res->t_embd = cur; + + cur = build_lora_mm(model.output, cur); + cur = ggml_scale(ctx0, cur, hparams.f_logit_scale); + cb(cur, "result_output", -1); + + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} diff --git a/tools/ui/src/lib/enums/agentic.enums.ts b/tools/ui/src/lib/enums/agentic.enums.ts index b96d244cd..9ad7b4f1a 100644 --- a/tools/ui/src/lib/enums/agentic.enums.ts +++ b/tools/ui/src/lib/enums/agentic.enums.ts @@ -16,3 +16,12 @@ export enum AgenticSectionType { REASONING = 'reasoning', REASONING_PENDING = 'reasoning_pending' } + +/** + * How a Continue click on an assistant message resumes generation. + */ +export enum ContinueIntentKind { + APPEND_TEXT = 'append_text', + RERUN_TURN = 'rerun_turn', + NEXT_TURN = 'next_turn' +} diff --git a/tools/ui/src/lib/enums/index.ts b/tools/ui/src/lib/enums/index.ts index a17cca1d8..b80b5b61e 100644 --- a/tools/ui/src/lib/enums/index.ts +++ b/tools/ui/src/lib/enums/index.ts @@ -6,7 +6,7 @@ export { AttachmentItemVisibleWhen } from './attachment.enums'; -export { AgenticSectionType, ToolCallType } from './agentic.enums'; +export { AgenticSectionType, ContinueIntentKind, ToolCallType } from './agentic.enums'; export { ChatMessageStatsView, diff --git a/tools/ui/src/lib/stores/chat.svelte.ts b/tools/ui/src/lib/stores/chat.svelte.ts index 61ea4c892..5b2644826 100644 --- a/tools/ui/src/lib/stores/chat.svelte.ts +++ b/tools/ui/src/lib/stores/chat.svelte.ts @@ -33,6 +33,7 @@ import { isAbortError, generateConversationTitle } from '$lib/utils'; +import { classifyContinueIntent } from '$lib/utils/agentic'; import { MAX_INACTIVE_CONVERSATION_STATES, INACTIVE_CONVERSATION_STATE_MAX_AGE_MS, @@ -51,7 +52,7 @@ import type { DatabaseMessage, DatabaseMessageExtra } from '$lib/types'; -import { ErrorDialogType, MessageRole, MessageType } from '$lib/enums'; +import { ContinueIntentKind, ErrorDialogType, MessageRole, MessageType } from '$lib/enums'; interface ConversationStateEntry { lastAccessed: number; @@ -1259,6 +1260,57 @@ class ChatStore { } } + /** + * Open a fresh assistant turn anchored at the last tool result of a resolved + * agentic round and let streamChatCompletion route through runAgenticFlow. + * Used by continueAssistantMessage when classifyContinueIntent returns + * next_turn, meaning the target assistant already has its tool_calls paired + * with trailing tool results and the next thing to generate is a brand new + * turn rather than a token level continuation. + */ + private async continueAsNextAgenticTurn(anchorIndex: number): Promise { + const activeConv = conversationsStore.activeConversation; + if (!activeConv) return; + const anchor = conversationsStore.activeMessages[anchorIndex]; + if (!anchor) return; + this.cancelPreEncode(); + this.setChatLoading(activeConv.id, true); + this.clearChatStreaming(activeConv.id); + try { + const allMessages = await conversationsStore.getConversationMessages(activeConv.id); + const anchorMessage = findMessageById(allMessages, anchor.id); + if (!anchorMessage) { + this.setChatLoading(activeConv.id, false); + return; + } + const newAssistantMessage = await DatabaseService.createMessageBranch( + { + convId: activeConv.id, + type: MessageType.TEXT, + timestamp: Date.now(), + role: MessageRole.ASSISTANT, + content: '', + toolCalls: '', + children: [], + model: null + }, + anchorMessage.id + ); + await conversationsStore.updateCurrentNode(newAssistantMessage.id); + conversationsStore.updateConversationTimestamp(); + await conversationsStore.refreshActiveMessages(); + const conversationPath = filterByLeafNodeId( + allMessages, + anchorMessage.id, + false + ) as DatabaseMessage[]; + await this.streamChatCompletion(conversationPath, newAssistantMessage); + } catch (error) { + if (!isAbortError(error)) console.error('Failed to continue agentic turn:', error); + this.setChatLoading(activeConv.id, false); + } + } + async continueAssistantMessage(messageId: string): Promise { const activeConv = conversationsStore.activeConversation; if (!activeConv || this.isChatLoadingInternal(activeConv.id)) return; @@ -1268,6 +1320,18 @@ class ChatStore { const { message: msg, index: idx } = result; + // Decide which resume path applies. tool_calls without tool results can + // not be resumed mid sequence by continue_final_message, branch instead. + // tool_calls already paired with tool results need a fresh next turn, + // not a token level continuation of the target assistant. + const intent = classifyContinueIntent(conversationsStore.activeMessages, idx); + if (intent.kind === ContinueIntentKind.RERUN_TURN) { + return this.regenerateMessageWithBranching(messageId); + } + if (intent.kind === ContinueIntentKind.NEXT_TURN) { + return this.continueAsNextAgenticTurn(intent.truncateAfter); + } + try { this.showErrorDialog(null); this.setChatLoading(activeConv.id, true); @@ -1283,15 +1347,11 @@ class ChatStore { const originalContent = dbMessage.content; const originalReasoning = dbMessage.reasoningContent || ''; - const conversationContext = conversationsStore.activeMessages.slice(0, idx); - const contextWithContinue = [ - ...conversationContext, - { - role: MessageRole.ASSISTANT as const, - content: originalContent, - reasoning_content: originalReasoning || undefined - } - ]; + // Hand the persisted DatabaseMessage straight to sendMessage so its + // internal converter preserves tool_calls and extras when present. + // Reconstructing a bare {role, content} here would drop those fields + // and break continue_final_message for messages with tool calls. + const contextWithContinue = conversationsStore.activeMessages.slice(0, idx + 1); let appendedContent = ''; let appendedReasoning = ''; diff --git a/tools/ui/src/lib/utils/agentic.ts b/tools/ui/src/lib/utils/agentic.ts index 549a1c9a0..52ff35793 100644 --- a/tools/ui/src/lib/utils/agentic.ts +++ b/tools/ui/src/lib/utils/agentic.ts @@ -1,4 +1,4 @@ -import { AgenticSectionType, MessageRole } from '$lib/enums'; +import { AgenticSectionType, ContinueIntentKind, MessageRole } from '$lib/enums'; import { ATTACHMENT_SAVED_REGEX, NEWLINE_SEPARATOR } from '$lib/constants'; import type { ApiChatCompletionToolCall } from '$lib/types/api'; import type { @@ -225,3 +225,62 @@ export function hasAgenticContent( return toolMessages.length > 0; } + +/** + * Classification of how a Continue click on an assistant message should resume + * generation. The caller dispatches the resume path based on this value. + * + * append_text -> the target is a plain text turn, resume with + * continue_final_message and rehydrate the persisted + * tool_calls and attachments through the regular DB to API + * message converter. + * rerun_turn -> the target carries tool_calls that were never resolved by + * tool result messages. The agentic stream was cut mid turn, + * so we drop the target and rerun the loop from the previous + * history. truncateAfter is the last kept index, inclusive. + * next_turn -> the target's tool_calls were already resolved by trailing + * tool results. Hand the history up to and including the + * last consecutive tool result back to the agentic loop so it + * starts the next turn naturally. truncateAfter points at + * that last tool result. + */ +export type ContinueIntent = + | { kind: ContinueIntentKind.APPEND_TEXT } + | { kind: ContinueIntentKind.RERUN_TURN; truncateAfter: number } + | { kind: ContinueIntentKind.NEXT_TURN; truncateAfter: number }; + +/** + * Decide how a Continue click on messages[idx] should resume generation. + * Pure function over the persisted history snapshot. + */ +export function classifyContinueIntent(messages: DatabaseMessage[], idx: number): ContinueIntent { + const target = messages[idx]; + + // Defensive default: callers already filter by role, stay deterministic. + if (!target || target.role !== MessageRole.ASSISTANT) { + return { kind: ContinueIntentKind.APPEND_TEXT }; + } + + const hasToolCalls = parseToolCalls(target.toolCalls).length > 0; + if (!hasToolCalls) { + return { kind: ContinueIntentKind.APPEND_TEXT }; + } + + // Walk consecutive trailing tool results. The agentic loop only emits tool + // messages directly after the assistant turn that owns them, so the first + // non tool message marks the boundary. + let lastTrailingTool = idx; + for (let i = idx + 1; i < messages.length; i++) { + if (messages[i].role === MessageRole.TOOL) { + lastTrailingTool = i; + } else { + break; + } + } + + if (lastTrailingTool > idx) { + return { kind: ContinueIntentKind.NEXT_TURN, truncateAfter: lastTrailingTool }; + } + + return { kind: ContinueIntentKind.RERUN_TURN, truncateAfter: idx - 1 }; +} diff --git a/tools/ui/tests/unit/continue-intent.test.ts b/tools/ui/tests/unit/continue-intent.test.ts new file mode 100644 index 000000000..76539c76a --- /dev/null +++ b/tools/ui/tests/unit/continue-intent.test.ts @@ -0,0 +1,166 @@ +import { describe, it, expect } from 'vitest'; +import { classifyContinueIntent } from '$lib/utils/agentic'; +import { ContinueIntentKind, MessageRole, MessageType } from '$lib/enums'; +import type { DatabaseMessage } from '$lib/types/database'; + +/** + * Tests for the Continue button intent classifier. + * + * The classifier walks the persisted message history to decide which of three + * resume paths a Continue click should take: + * + * A. append_text -> plain text assistant turn, resume with + * continue_final_message. + * B. rerun_turn -> assistant turn with tool_calls but no tool results yet, + * the stream was cut mid turn and the tool_calls are + * unrecoverable as a token level continuation. Drop the + * target and rerun from the previous history. + * C. next_turn -> assistant turn with tool_calls that were already + * resolved by trailing tool results. Hand the history + * back to the agentic loop so it starts the next turn. + */ + +let nextId = 0; +function makeMsg(role: MessageRole, opts: Partial = {}): DatabaseMessage { + nextId++; + return { + id: `msg-${nextId}`, + convId: 'conv-1', + type: MessageType.TEXT, + timestamp: nextId, + role, + content: '', + parent: null, + children: [], + ...opts + }; +} + +function toolCall(id: string, name: string, args: string = '{}'): string { + return JSON.stringify([{ id, type: 'function', function: { name, arguments: args } }]); +} + +describe('classifyContinueIntent', () => { + it('returns append_text for a plain text assistant turn at the tail', () => { + const messages = [ + makeMsg(MessageRole.USER, { content: 'hello' }), + makeMsg(MessageRole.ASSISTANT, { content: 'hi there' }) + ]; + + const intent = classifyContinueIntent(messages, 1); + + expect(intent).toEqual({ kind: ContinueIntentKind.APPEND_TEXT }); + }); + + it('returns append_text for a plain text assistant turn in the middle', () => { + const messages = [ + makeMsg(MessageRole.USER, { content: 'q1' }), + makeMsg(MessageRole.ASSISTANT, { content: 'a1' }), + makeMsg(MessageRole.USER, { content: 'q2' }), + makeMsg(MessageRole.ASSISTANT, { content: 'a2' }) + ]; + + expect(classifyContinueIntent(messages, 1)).toEqual({ kind: ContinueIntentKind.APPEND_TEXT }); + }); + + it('returns rerun_turn when the assistant has tool_calls without results', () => { + const messages = [ + makeMsg(MessageRole.USER, { content: 'list files' }), + makeMsg(MessageRole.ASSISTANT, { + content: '', + toolCalls: toolCall('call_1', 'bash_tool', '{"command":"ls"}') + }) + ]; + + const intent = classifyContinueIntent(messages, 1); + + expect(intent).toEqual({ kind: ContinueIntentKind.RERUN_TURN, truncateAfter: 0 }); + }); + + it('returns next_turn when trailing tool results resolve the tool_calls', () => { + const messages = [ + makeMsg(MessageRole.USER, { content: 'list files' }), + makeMsg(MessageRole.ASSISTANT, { + content: '', + toolCalls: toolCall('call_1', 'bash_tool') + }), + makeMsg(MessageRole.TOOL, { content: 'file1\nfile2', toolCallId: 'call_1' }) + ]; + + const intent = classifyContinueIntent(messages, 1); + + expect(intent).toEqual({ kind: ContinueIntentKind.NEXT_TURN, truncateAfter: 2 }); + }); + + it('next_turn keeps all consecutive trailing tool results, not just one', () => { + const messages = [ + makeMsg(MessageRole.USER, { content: 'do many things' }), + makeMsg(MessageRole.ASSISTANT, { + content: '', + toolCalls: JSON.stringify([ + { id: 'call_1', type: 'function', function: { name: 'a', arguments: '{}' } }, + { id: 'call_2', type: 'function', function: { name: 'b', arguments: '{}' } } + ]) + }), + makeMsg(MessageRole.TOOL, { content: 'r1', toolCallId: 'call_1' }), + makeMsg(MessageRole.TOOL, { content: 'r2', toolCallId: 'call_2' }) + ]; + + const intent = classifyContinueIntent(messages, 1); + + expect(intent).toEqual({ kind: ContinueIntentKind.NEXT_TURN, truncateAfter: 3 }); + }); + + it('next_turn stops at the first non tool message after the target', () => { + const messages = [ + makeMsg(MessageRole.USER, { content: 'go' }), + makeMsg(MessageRole.ASSISTANT, { + content: '', + toolCalls: toolCall('call_1', 'a') + }), + makeMsg(MessageRole.TOOL, { content: 'r1', toolCallId: 'call_1' }), + makeMsg(MessageRole.USER, { content: 'wait' }), + makeMsg(MessageRole.TOOL, { content: 'late', toolCallId: 'call_1' }) + ]; + + const intent = classifyContinueIntent(messages, 1); + + // truncateAfter must point at the contiguous tool block, not jump over + // the user message to grab the dangling late tool. + expect(intent).toEqual({ kind: ContinueIntentKind.NEXT_TURN, truncateAfter: 2 }); + }); + + it('returns append_text when toolCalls is set but parses to empty array', () => { + const messages = [ + makeMsg(MessageRole.USER, { content: 'q' }), + makeMsg(MessageRole.ASSISTANT, { content: 'a', toolCalls: '[]' }) + ]; + + expect(classifyContinueIntent(messages, 1)).toEqual({ kind: ContinueIntentKind.APPEND_TEXT }); + }); + + it('returns append_text when toolCalls is malformed JSON', () => { + const messages = [ + makeMsg(MessageRole.USER, { content: 'q' }), + makeMsg(MessageRole.ASSISTANT, { content: 'a', toolCalls: '{not json' }) + ]; + + expect(classifyContinueIntent(messages, 1)).toEqual({ kind: ContinueIntentKind.APPEND_TEXT }); + }); + + it('returns append_text defensively when idx points at a non assistant message', () => { + const messages = [ + makeMsg(MessageRole.USER, { content: 'q' }), + makeMsg(MessageRole.ASSISTANT, { content: 'a' }) + ]; + + expect(classifyContinueIntent(messages, 0)).toEqual({ kind: ContinueIntentKind.APPEND_TEXT }); + }); + + it('returns append_text defensively when idx is out of bounds', () => { + const messages = [makeMsg(MessageRole.ASSISTANT, { content: 'a' })]; + + expect(classifyContinueIntent(messages, 5)).toEqual({ kind: ContinueIntentKind.APPEND_TEXT }); + expect(classifyContinueIntent([], 0)).toEqual({ kind: ContinueIntentKind.APPEND_TEXT }); + }); +});