From 9cb6ae64818b8c8a8097276155ab405d57455942 Mon Sep 17 00:00:00 2001 From: Nicholas Sparks Date: Sat, 2 May 2026 04:05:43 +0000 Subject: [PATCH] model : add DeepSeek V4 architecture support Adds runtime and conversion support for the DeepSeek V4 architecture (deepseek-ai/DeepSeek-V4-Flash and friends). The model uses a latent KV with a compression-window indexer + MLA-style nope/rope split + an Hadamard-chunk grouped output projection, and a 256-expert sparse MoE with per-token softplus+sqrt routing. The custom llama_memory_deepseek4 module owns the per-layer attention KV (latent, head_dim-wide), the compression-window state (F32 KV / score state) and the optional indexer state, since none of the existing KV cache layouts cover the V4 indexer's compression-window scoring. Conversion: DeepseekV4Model subclasses DeepseekV2Model, merges the per-expert FFN tensors into a single 3D tensor per layer, and emits the V4 hparams. FP8 source checkpoints are dequantized via the existing parent-class fp8 path, so the converter emits standard ftypes (F16, Q4_K, Q8_0, ...) that current upstream readers handle. The chat template covers the OpenAI-API subset of roles (system/user/assistant) with an enable_thinking switch and drop_thinking history rewriting, validated against the official encoding/encoding_dsv4.py test fixtures. --- convert_hf_to_gguf.py | 243 +++++++-- gguf-py/gguf/constants.py | 87 +++ gguf-py/gguf/tensor_mapping.py | 104 +++- src/CMakeLists.txt | 1 + src/llama-arch.cpp | 43 ++ src/llama-arch.h | 22 + src/llama-context.cpp | 8 + src/llama-memory-deepseek4.cpp | 689 ++++++++++++++++++++++++ src/llama-memory-deepseek4.h | 109 ++++ src/llama-model-saver.cpp | 7 +- src/llama-model.cpp | 195 ++++++- src/llama-model.h | 25 + src/models/deepseek4.cpp | 957 +++++++++++++++++++++++++++++++++ src/models/models.h | 4 + tests/test-llama-archs.cpp | 5 + 15 files changed, 2448 insertions(+), 51 deletions(-) create mode 100644 src/llama-memory-deepseek4.cpp create mode 100644 src/llama-memory-deepseek4.h create mode 100644 src/models/deepseek4.cpp diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index 5287c4df9410..a7ad744baf46 100755 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -167,7 +167,7 @@ def __init__(self, dir_model: Path, ftype: gguf.LlamaFileType, fname_out: Path, logger.info("heuristics unable to detect tensor dtype, defaulting to --outtype f16") # Configure GGUF Writer - self.gguf_writer = gguf.GGUFWriter(path=None, arch=gguf.MODEL_ARCH_NAMES[self.model_arch], endianess=self.endianess, use_temp_file=self.use_temp_file, + self.gguf_writer = gguf.GGUFWriter(path=fname_out, arch=gguf.MODEL_ARCH_NAMES[self.model_arch], endianess=self.endianess, use_temp_file=self.use_temp_file, split_max_tensors=split_max_tensors, split_max_size=split_max_size, dry_run=dry_run, small_first_shard=small_first_shard) # Mistral specific @@ -728,9 +728,6 @@ def _flush_nvfp4_experts(self, key, expert_blocks, expert_scales, expert_input_s del experts, merged - def _needs_nvfp4_processing(self) -> bool: - return True - def prepare_tensors(self): # detect NVFP4 quantization (ModelOpt format) quant_algo = (self.hparams.get("quantization_config") or {}).get("quant_algo") @@ -761,7 +758,7 @@ def prepare_tensors(self): # NVFP4 weights are repacked and written directly to gguf_writer. # 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 and self._needs_nvfp4_processing(): + if self._is_nvfp4: self._generate_nvfp4_tensors() self.dequant_model() @@ -780,7 +777,8 @@ def prepare_tensors(self): old_dtype = data_torch.dtype # convert any unsupported data types to float32 - if data_torch.dtype not in (torch.float16, torch.float32): + preserve_integer_tensor = name.endswith(".ffn.gate.tid2eid") + if data_torch.dtype not in (torch.float16, torch.float32) and not preserve_integer_tensor: data_torch = data_torch.to(torch.float32) # use the first number-like part of the tensor name as the block id @@ -791,6 +789,13 @@ def prepare_tensors(self): break for new_name, data_torch in (self.modify_tensors(data_torch, name, bid)): + if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.FFN_GATE_TID2EID, bid, suffix=""): + data = LazyTorchTensor.to_eager(data_torch).to(torch.int32).numpy() + shape_str = f"{{{', '.join(str(n) for n in reversed(data.shape))}}}" + logger.info(f"{f'%-{max_name_len}s' % f'{new_name},'} {old_dtype} --> I32, shape = {shape_str}") + self.gguf_writer.add_tensor(new_name, data) + continue + # TODO: why do we squeeze here? # data = data_torch.squeeze().numpy() data = data_torch.numpy() @@ -2193,10 +2198,6 @@ def __init__(self, *args, **kwargs): # merge configs self.preprocessor_config = {**self.preprocessor_config, **cfg} - def _needs_nvfp4_processing(self) -> bool: - # nvfp4 quantization applies to the text model only. - return False - def get_vision_config(self) -> dict[str, Any] | None: config_name = "vision_config" if not self.is_mistral_format else "vision_encoder" return self.global_config.get(config_name) @@ -4457,12 +4458,6 @@ def get_vision_config(self) -> dict[str, Any] | None: } return vision_config - def dequant_model(self): - if self._is_nvfp4: - # Skip nvfp4 quantization for vision/audio model. - return - super().dequant_model() - def set_gguf_parameters(self): if "image_mean" not in self.preprocessor_config: self.preprocessor_config["image_mean"] = [0.485, 0.456, 0.406] @@ -4486,10 +4481,6 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter if "input_conditioner" in name: return - # mtmd does not support video yet so skip tensors related to video. - if "radio_model.model.patch_generator.video_embedder" in name: - return - # RADIO's pos_embed doesn't have .weight suffix, but clip.cpp expects it if "patch_generator.pos_embed" in name: if not name.endswith(".weight"): @@ -6658,7 +6649,7 @@ def _xlmroberta_set_vocab(self) -> None: tokens: list[bytes] = [f"[PAD{i}]".encode("utf-8") for i in range(vocab_size)] scores: list[float] = [-10000.0] * vocab_size - toktypes: list[int] = [SentencePieceTokenTypes.UNUSED] * vocab_size + toktypes: list[int] = [SentencePieceTokenTypes.UNUSED] * vocab_size # ty: ignore[invalid-assignment] if isinstance(tokenizer, SentencePieceProcessor): for token_id in range(tokenizer.vocab_size()): @@ -9199,6 +9190,186 @@ def prepare_tensors(self): raise ValueError(f"Unprocessed experts: {experts}") +@ModelBase.register("DeepseekV4ForCausalLM") +class DeepseekV4Model(DeepseekV2Model): + model_arch = gguf.MODEL_ARCH.DEEPSEEK4 + skip_mtp = True + merge_expert = True + # Chat template: basic system / user / assistant turns with proper + # `<｜Assistant｜>` framing, an `enable_thinking` switch (chat vs thinking + # mode), and `drop_thinking` semantics (only the assistant turn after the + # last user turn keeps its `reasoning_content`; earlier assistant turns + # are emitted as `` only). This matches the `thinking_mode` and + # `drop_thinking=True` default of the official `encoding/encoding_dsv4.py` + # for the OpenAI-API subset of roles. Tool calling, the `developer` / + # `latest_reminder` roles, and quick-instruction tasks are not included. + # Validated byte-for-byte against `encoding/tests/test_input_2.json` -> + # `test_output_2.txt` (basic chat with interleaved thinking and + # drop_thinking applied). + chat_template = ( + "{%- if not add_generation_prompt is defined -%}" + "{%- set add_generation_prompt = false -%}" + "{%- endif -%}" + "{%- if enable_thinking is defined -%}" + "{%- set thinking = enable_thinking -%}" + "{%- elif thinking is not defined -%}" + "{%- set thinking = false -%}" + "{%- endif -%}" + "{%- set ns = namespace(last_user_idx=-1) -%}" + "{%- for message in messages -%}" + "{%- if message['role'] in ['user', 'tool'] -%}" + "{%- set ns.last_user_idx = loop.index0 -%}" + "{%- endif -%}" + "{%- endfor -%}" + "{{- '<｜begin▁of▁sentence｜>' -}}" + "{%- for message in messages -%}" + "{%- if message['role'] == 'system' -%}" + "{{- message['content'] -}}" + "{%- elif message['role'] == 'user' -%}" + "{{- '<｜User｜>' + message['content'] -}}" + "{%- elif message['role'] == 'assistant' -%}" + "{{- '<｜Assistant｜>' -}}" + "{%- if thinking and loop.index0 > ns.last_user_idx and message['reasoning_content'] is defined and message['reasoning_content'] -%}" + "{{- '' + message['reasoning_content'] + '' -}}" + "{%- else -%}" + "{{- '' -}}" + "{%- endif -%}" + "{{- message['content'] + '<｜end▁of▁sentence｜>' -}}" + "{%- endif -%}" + "{%- endfor -%}" + "{%- if add_generation_prompt -%}" + "{{- '<｜Assistant｜>' -}}" + "{%- if thinking -%}" + "{{- '' -}}" + "{%- else -%}" + "{{- '' -}}" + "{%- endif -%}" + "{%- endif -%}" + ) + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # State for merging per-expert weights into a single 3D tensor per + # layer; populated lazily in modify_tensors(). + self._expert_buffers: list[dict[str, Tensor]] | None = None + self._expert_seen: list[dict[str, set[int]]] | None = None + + def set_gguf_parameters(self): + self.hparams["num_key_value_heads"] = self.find_hparam(["num_key_value_heads"], optional=True) or 1 + self.hparams["rms_norm_eps"] = self.find_hparam(["rms_norm_eps", "norm_eps"], optional=True) or 1e-6 + + score_func_keys = {} + for key in ("scoring_func", "score_func"): + if key in self.hparams: + score_func_keys[key] = self.hparams.pop(key) + + try: + TextModel.set_gguf_parameters(self) + finally: + self.hparams.update(score_func_keys) + + self.gguf_writer.add_chat_template(self.chat_template) + + self.gguf_writer.add_vocab_size(self.find_hparam(["vocab_size"])) + + if (q_lora_rank := self.find_hparam(["q_lora_rank"], optional=True)) is not None: + self.gguf_writer.add_q_lora_rank(q_lora_rank) + + if (rope_dim := self.find_hparam(["qk_rope_head_dim"], optional=True)) is not None: + self.gguf_writer.add_rope_dimension_count(rope_dim) + + if (sliding_window := self.find_hparam(["sliding_window"], optional=True)) is not None: + self.gguf_writer.add_sliding_window(sliding_window) + + if (compress_rope_theta := self.find_hparam(["compress_rope_theta"], optional=True)) is not None: + self.gguf_writer.add_rope_freq_base_swa(compress_rope_theta) + + self.gguf_writer.add_leading_dense_block_count(0) + + self.gguf_writer.add_expert_feed_forward_length(self.find_hparam(["moe_intermediate_size"])) + + if (n_routed_experts := self.find_hparam(["n_routed_experts"], optional=True)) is not None: + self.gguf_writer.add_expert_count(n_routed_experts) + + if (n_shared_experts := self.find_hparam(["n_shared_experts"], optional=True)) is not None: + self.gguf_writer.add_expert_shared_count(n_shared_experts) + + if (routed_scaling_factor := self.find_hparam(["routed_scaling_factor"], optional=True)) is not None: + self.gguf_writer.add_expert_weights_scale(routed_scaling_factor) + + if self.find_hparam(["scoring_func"], optional=True) != "softmax": + self.gguf_writer.add_expert_weights_norm(True) + + if (swiglu_limit := self.find_hparam(["swiglu_limit"], optional=True)) is not None: + self.gguf_writer.add_swiglu_clamp_exp([float(swiglu_limit)] * self.block_count) + + if (index_n_heads := self.find_hparam(["index_n_heads"], optional=True)) is not None: + self.gguf_writer.add_indexer_head_count(index_n_heads) + + if (index_head_dim := self.find_hparam(["index_head_dim"], optional=True)) is not None: + self.gguf_writer.add_indexer_key_length(index_head_dim) + + if (index_topk := self.find_hparam(["index_topk"], optional=True)) is not None: + self.gguf_writer.add_indexer_top_k(index_topk) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name.startswith("mtp."): + return + + if self.hparams.get("tie_word_embeddings", False) and name == "head.weight": + logger.info("Skipping tied output layer 'head.weight' (will use token_embd.weight)") + return + + if self.merge_expert and ".ffn.experts." in name: + n_experts = self.hparams["n_routed_experts"] + assert bid is not None + + match = re.fullmatch(r"layers\.(\d+)\.ffn\.experts\.(\d+)\.(w[123])\.weight", name) + if match is None: + raise ValueError(f"Unexpected DeepSeek V4 expert tensor name: {name}") + + xid = int(match.group(2)) + w_name = match.group(3) + if xid >= n_experts: + raise ValueError(f"Unexpected DeepSeek V4 expert id {xid} for tensor {name}") + + if self._expert_buffers is None: + self._expert_buffers = [{} for _ in range(self.block_count)] + self._expert_seen = [{} for _ in range(self.block_count)] + assert self._expert_seen is not None + + layer_buffers = self._expert_buffers[bid] + layer_seen = self._expert_seen[bid] + + seen = layer_seen.setdefault(w_name, set()) + if xid in seen: + raise ValueError(f"Duplicate DeepSeek V4 expert tensor: {name}") + + if w_name not in layer_buffers: + layer_buffers[w_name] = torch.empty((n_experts, *data_torch.shape), dtype=data_torch.dtype) + elif layer_buffers[w_name].shape[1:] != data_torch.shape: + raise ValueError( + f"Unexpected DeepSeek V4 expert shape {tuple(data_torch.shape)} for tensor {name}; " + f"expected {tuple(layer_buffers[w_name].shape[1:])}" + ) + + layer_buffers[w_name][xid].copy_(data_torch) + seen.add(xid) + + if all(len(layer_seen.get(done_w_name, set())) >= n_experts for done_w_name in ("w2", "w1", "w3")): + for done_w_name in ["w2", "w1", "w3"]: + merged = layer_buffers.pop(done_w_name) + del layer_seen[done_w_name] + merged_name = f"layers.{bid}.ffn.experts.{done_w_name}.weight" + for new_name, data_torch in TextModel.modify_tensors(self, merged, merged_name, bid): + yield new_name, data_torch + return + else: + return + + yield from TextModel.modify_tensors(self, data_torch, name, bid) + + @ModelBase.register( "Mistral3ForConditionalGeneration", "Ministral3ForCausalLM", @@ -10837,11 +11008,7 @@ def __init__(self, *args, **kwargs): # uses self.model_arch to build the tensor name map, and all MoE-specific # mappings would be missed if it were called with the default non-MoE arch. hparams = ModelBase.load_hparams(args[0], self.is_mistral_format) - has_moe_params = ( - "num_experts_per_tok" in hparams - or (isinstance(hparams.get("llm_config"), dict) and "num_experts_per_tok" in hparams["llm_config"]) - ) - if has_moe_params: + if "num_experts_per_tok" in hparams: self.model_arch = gguf.MODEL_ARCH.NEMOTRON_H_MOE self.is_moe = True @@ -10988,11 +11155,6 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter if name.startswith(("vision_model.", "mlp1.")): return - if name.startswith(("sound_encoder.")): - return - if name.startswith(("sound_projection.")): - return - # Strip language_model. prefix for VLM models (e.g., Nemotron Nano 12B v2 VL) if name.startswith("language_model."): name = name[len("language_model."):] @@ -13334,6 +13496,11 @@ def __torch_function__(cls, func, types, args=(), kwargs=None): return cls._wrap_fn(func)(*args, **kwargs) +if (torch_float8_e8m0fnu := getattr(torch, "float8_e8m0fnu", None)) is not None: + LazyTorchTensor._dtype_byteswap_map[torch_float8_e8m0fnu] = np.uint8 + LazyTorchTensor._dtype_str_map["F8_E8M0"] = torch_float8_e8m0fnu + + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Convert a huggingface model to a GGML compatible file") @@ -13346,8 +13513,8 @@ def parse_args() -> argparse.Namespace: help="path to write to; default: based on input. {ftype} will be replaced by the outtype.", ) parser.add_argument( - "--outtype", type=str, choices=["f32", "f16", "bf16", "q8_0", "tq1_0", "tq2_0", "auto"], default="auto", - help="output format - use f32 for float32, f16 for float16, bf16 for bfloat16, q8_0 for Q8_0, tq1_0 or tq2_0 for ternary, and auto for the highest-fidelity 16-bit float type", + "--outtype", type=str, choices=["f32", "f16", "bf16", "q8_0", "tq1_0", "tq2_0", "native", "auto"], default="auto", + help="output format - use f32 for float32, f16 for float16, bf16 for bfloat16, q8_0 for Q8_0, tq1_0 or tq2_0 for ternary, native to preserve supported source quantization formats, and auto for the highest-fidelity 16-bit float type", ) parser.add_argument( "--bigendian", action="store_true", @@ -13374,6 +13541,10 @@ def parse_args() -> argparse.Namespace: "--verbose", action="store_true", help="increase output verbosity", ) + parser.add_argument( + "--torch-threads", type=int, default=None, + help="number of PyTorch CPU threads to use for tensor conversion operations", + ) parser.add_argument( "--split-max-tensors", type=int, default=0, help="max tensors in each split", @@ -13495,6 +13666,12 @@ def main() -> None: else: logging.basicConfig(level=logging.INFO) + if args.torch_threads is not None: + if args.torch_threads <= 0: + raise ValueError("--torch-threads must be a positive integer") + torch.set_num_threads(args.torch_threads) + logger.info(f"PyTorch tensor conversion threads: {torch.get_num_threads()}") + if args.remote: hf_repo_id = args.model from huggingface_hub import snapshot_download diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 83ae51ce9ce3..689e68c8dc7c 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -442,6 +442,7 @@ class MODEL_ARCH(IntEnum): DEEPSEEK = auto() DEEPSEEK2 = auto() DEEPSEEK2OCR = auto() + DEEPSEEK4 = auto() CHATGLM = auto() GLM4 = auto() GLM4_MOE = auto() @@ -708,6 +709,27 @@ class MODEL_TENSOR(IntEnum): INDEXER_PROJ = auto() INDEXER_ATTN_K = auto() INDEXER_ATTN_Q_B = auto() + ATTN_KV_LATENT = auto() + ATTN_OUT_A = auto() + ATTN_OUT_B = auto() + ATTN_COMPRESS_APE = auto() + ATTN_COMPRESS_NORM = auto() + ATTN_COMPRESS_KV = auto() + ATTN_COMPRESS_GATE = auto() + INDEXER_COMPRESS_APE = auto() + INDEXER_COMPRESS_NORM = auto() + INDEXER_COMPRESS_KV = auto() + INDEXER_COMPRESS_GATE = auto() + HC_HEAD_BASE = auto() + HC_HEAD_FN = auto() + HC_HEAD_SCALE = auto() + HC_ATTN_BASE = auto() + HC_ATTN_FN = auto() + HC_ATTN_SCALE = auto() + HC_FFN_BASE = auto() + HC_FFN_FN = auto() + HC_FFN_SCALE = auto() + FFN_GATE_TID2EID = auto() # vision V_MMPROJ = auto() V_MMPROJ_FC = auto() @@ -928,6 +950,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.DEEPSEEK: "deepseek", MODEL_ARCH.DEEPSEEK2: "deepseek2", MODEL_ARCH.DEEPSEEK2OCR: "deepseek2-ocr", + MODEL_ARCH.DEEPSEEK4: "deepseek4", MODEL_ARCH.CHATGLM: "chatglm", MODEL_ARCH.GLM4: "glm4", MODEL_ARCH.GLM4_MOE: "glm4moe", @@ -1193,6 +1216,27 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.INDEXER_PROJ: "blk.{bid}.indexer.proj", MODEL_TENSOR.INDEXER_ATTN_K: "blk.{bid}.indexer.attn_k", MODEL_TENSOR.INDEXER_ATTN_Q_B: "blk.{bid}.indexer.attn_q_b", + MODEL_TENSOR.ATTN_KV_LATENT: "blk.{bid}.attn_kv_latent", + MODEL_TENSOR.ATTN_OUT_A: "blk.{bid}.attn_output_a", + MODEL_TENSOR.ATTN_OUT_B: "blk.{bid}.attn_output_b", + MODEL_TENSOR.ATTN_COMPRESS_APE: "blk.{bid}.attn_compress_ape", + MODEL_TENSOR.ATTN_COMPRESS_NORM: "blk.{bid}.attn_compress_norm", + MODEL_TENSOR.ATTN_COMPRESS_KV: "blk.{bid}.attn_compress_kv", + MODEL_TENSOR.ATTN_COMPRESS_GATE: "blk.{bid}.attn_compress_gate", + MODEL_TENSOR.INDEXER_COMPRESS_APE: "blk.{bid}.indexer.compress_ape", + MODEL_TENSOR.INDEXER_COMPRESS_NORM: "blk.{bid}.indexer.compress_norm", + MODEL_TENSOR.INDEXER_COMPRESS_KV: "blk.{bid}.indexer.compress_kv", + MODEL_TENSOR.INDEXER_COMPRESS_GATE: "blk.{bid}.indexer.compress_gate", + MODEL_TENSOR.HC_HEAD_BASE: "hc_head_base", + MODEL_TENSOR.HC_HEAD_FN: "hc_head_fn", + MODEL_TENSOR.HC_HEAD_SCALE: "hc_head_scale", + MODEL_TENSOR.HC_ATTN_BASE: "blk.{bid}.hc_attn_base", + MODEL_TENSOR.HC_ATTN_FN: "blk.{bid}.hc_attn_fn", + MODEL_TENSOR.HC_ATTN_SCALE: "blk.{bid}.hc_attn_scale", + MODEL_TENSOR.HC_FFN_BASE: "blk.{bid}.hc_ffn_base", + MODEL_TENSOR.HC_FFN_FN: "blk.{bid}.hc_ffn_fn", + MODEL_TENSOR.HC_FFN_SCALE: "blk.{bid}.hc_ffn_scale", + MODEL_TENSOR.FFN_GATE_TID2EID: "blk.{bid}.ffn_gate_tid2eid", # vision MODEL_TENSOR.V_MMPROJ: "mm.{bid}", MODEL_TENSOR.V_MMPROJ_FC: "mm.model.fc", @@ -2816,6 +2860,49 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_UP_SHEXP, MODEL_TENSOR.FFN_EXP_PROBS_B, ], + MODEL_ARCH.DEEPSEEK4: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.HC_HEAD_BASE, + MODEL_TENSOR.HC_HEAD_FN, + MODEL_TENSOR.HC_HEAD_SCALE, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_KV_LATENT, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.ATTN_OUT_A, + MODEL_TENSOR.ATTN_OUT_B, + MODEL_TENSOR.ATTN_SINKS, + MODEL_TENSOR.ATTN_COMPRESS_APE, + MODEL_TENSOR.ATTN_COMPRESS_NORM, + MODEL_TENSOR.ATTN_COMPRESS_KV, + MODEL_TENSOR.ATTN_COMPRESS_GATE, + MODEL_TENSOR.INDEXER_PROJ, + MODEL_TENSOR.INDEXER_ATTN_Q_B, + MODEL_TENSOR.INDEXER_COMPRESS_APE, + MODEL_TENSOR.INDEXER_COMPRESS_NORM, + MODEL_TENSOR.INDEXER_COMPRESS_KV, + MODEL_TENSOR.INDEXER_COMPRESS_GATE, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_TID2EID, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.HC_ATTN_BASE, + MODEL_TENSOR.HC_ATTN_FN, + MODEL_TENSOR.HC_ATTN_SCALE, + MODEL_TENSOR.HC_FFN_BASE, + MODEL_TENSOR.HC_FFN_FN, + MODEL_TENSOR.HC_FFN_SCALE, + ], MODEL_ARCH.ERNIE4_5_MOE: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 01a9b236000b..2ec2e6be17d6 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -36,6 +36,7 @@ class TensorNameMap: "encoder", # neobert "model.transformer.wte", # llada "embed_tokens", # qwen3-embedding + "embed", # deepseek-v4 ), # Token type embeddings @@ -196,6 +197,7 @@ class TensorNameMap: "layers.{bid}.input_layernorm", # qwen3-embedding "model.layers.{bid}.attention_layernorm", # apertus "model.layers.{bid}.pre_attention_layernorm", # kormo + "layers.{bid}.attn_norm", # deepseek-v4 ), # Attention norm 2 @@ -357,6 +359,7 @@ class TensorNameMap: MODEL_TENSOR.ATTN_SINKS: ( "model.layers.{bid}.self_attn.sinks", # openai-moe "model.layers.{bid}.self_attn.attention_sink_bias", # mimov2 + "layers.{bid}.attn.attn_sink", # deepseek-v4 ), MODEL_TENSOR.ATTN_GATE: ( @@ -390,7 +393,8 @@ class TensorNameMap: "layers.{bid}.post_attention_layernorm", # qwen3-embedding "model.layers.{bid}.feedforward_layernorm", # apertus "model.layers.{bid}.pre_mlp_layernorm", # kormo - "layers.{bid}.mlp_norm" # modern-bert + "layers.{bid}.mlp_norm", # modern-bert + "layers.{bid}.ffn_norm", # deepseek-v4 ), # Pre feed-forward norm @@ -441,6 +445,7 @@ class TensorNameMap: "backbone.layers.{bid}.mixer.gate", # nemotron-h-moe "model.layers.{bid}.moe.gate", # step3.5 "model.layers.{bid}.router.proj", # gemma4 + "layers.{bid}.ffn.gate", # deepseek-v4 ), MODEL_TENSOR.FFN_GATE_INP_SHEXP: ( @@ -458,6 +463,7 @@ class TensorNameMap: "model.layers.{bid}.mlp.e_score_correction", # exaone-moe "model.layers.{bid}.block_sparse_moe.gate.e_score_correction", # kimi "model.layers.{bid}.moe.router_bias", # step3.5 expert selection bias + "layers.{bid}.ffn.gate.bias", # deepseek-v4 ), # Feed-forward up @@ -513,6 +519,7 @@ class TensorNameMap: "encoder.layers.{bid}.mlp.experts.mlp.w1", # nomic-bert-moe "model.layers.{bid}.block_sparse_moe.experts.up", # smallthinker "model.layers.{bid}.moe.up_proj", # step3.5 + "layers.{bid}.ffn.experts.w3", # deepseek-v4 (merged) ), MODEL_TENSOR.FFN_UP_SHEXP: ( @@ -525,6 +532,7 @@ class TensorNameMap: "backbone.layers.{bid}.mixer.shared_experts.up_proj", # nemotron-h-moe "model.layers.{bid}.block_sparse_moe.shared_experts.up_proj", # kimi "model.layers.{bid}.share_expert.up_proj", # step3.5 + "layers.{bid}.ffn.shared_experts.w3", # deepseek-v4 ), MODEL_TENSOR.FFN_UP_CHEXP: ( @@ -565,6 +573,7 @@ class TensorNameMap: "model.layers.{bid}.feed_forward.experts.gate_proj", # llama4 "model.layers.{bid}.block_sparse_moe.experts.gate", # smallthinker "model.layers.{bid}.moe.gate_proj", # step3.5 + "layers.{bid}.ffn.experts.w1", # deepseek-v4 (merged) ), MODEL_TENSOR.FFN_GATE_SHEXP: ( @@ -575,6 +584,7 @@ class TensorNameMap: "layers.{bid}.shared_experts.w1", # mistral-large "model.layers.{bid}.block_sparse_moe.shared_experts.gate_proj", # kimi "model.layers.{bid}.share_expert.gate_proj", # step3.5 + "layers.{bid}.ffn.shared_experts.w1", # deepseek-v4 ), MODEL_TENSOR.FFN_GATE_CHEXP: ( @@ -644,6 +654,7 @@ class TensorNameMap: "model.layers.{bid}.block_sparse_moe.experts.down", # smallthinker "model.layers.{bid}.moe.down_proj", # step3.5 "model.layers.{bid}.experts.down_proj", # gemma4 + "layers.{bid}.ffn.experts.w2", # deepseek-v4 (merged) ), MODEL_TENSOR.FFN_DOWN_SHEXP: ( @@ -656,6 +667,7 @@ class TensorNameMap: "backbone.layers.{bid}.mixer.shared_experts.down_proj", # nemotron-h-moe "model.layers.{bid}.block_sparse_moe.shared_experts.down_proj", # kimi "model.layers.{bid}.share_expert.down_proj", # step3.5 + "layers.{bid}.ffn.shared_experts.w2", # deepseek-v4 ), MODEL_TENSOR.FFN_DOWN_CHEXP: ( @@ -1064,11 +1076,13 @@ class TensorNameMap: MODEL_TENSOR.ATTN_Q_A: ( "model.layers.{bid}.self_attn.q_a_proj", # deepseek2 "layers.{bid}.attention.wq_a", # mistral-large + "layers.{bid}.attn.wq_a", # deepseek-v4 ), MODEL_TENSOR.ATTN_Q_B: ( "model.layers.{bid}.self_attn.q_b_proj", # deepseek2 "layers.{bid}.attention.wq_b", # mistral-large + "layers.{bid}.attn.wq_b", # deepseek-v4 ), MODEL_TENSOR.ATTN_KV_A_MQA: ( @@ -1093,11 +1107,97 @@ class TensorNameMap: MODEL_TENSOR.ATTN_Q_A_NORM: ( "model.layers.{bid}.self_attn.q_a_layernorm", # deepseek2 "layers.{bid}.attention.q_a_norm", # mistral-large + "layers.{bid}.attn.q_norm", # deepseek-v4 ), MODEL_TENSOR.ATTN_KV_A_NORM: ( "model.layers.{bid}.self_attn.kv_a_layernorm", # deepseek2 "layers.{bid}.attention.kv_a_norm", # mistral-large + "layers.{bid}.attn.kv_norm", # deepseek-v4 + ), + + MODEL_TENSOR.ATTN_KV_LATENT: ( + "layers.{bid}.attn.wkv", # deepseek-v4 + ), + + MODEL_TENSOR.ATTN_OUT_A: ( + "layers.{bid}.attn.wo_a", # deepseek-v4 + ), + + MODEL_TENSOR.ATTN_OUT_B: ( + "layers.{bid}.attn.wo_b", # deepseek-v4 + ), + + MODEL_TENSOR.ATTN_COMPRESS_APE: ( + "layers.{bid}.attn.compressor.ape", # deepseek-v4 + ), + + MODEL_TENSOR.ATTN_COMPRESS_NORM: ( + "layers.{bid}.attn.compressor.norm", # deepseek-v4 + ), + + MODEL_TENSOR.ATTN_COMPRESS_KV: ( + "layers.{bid}.attn.compressor.wkv", # deepseek-v4 + ), + + MODEL_TENSOR.ATTN_COMPRESS_GATE: ( + "layers.{bid}.attn.compressor.wgate", # deepseek-v4 + ), + + MODEL_TENSOR.INDEXER_COMPRESS_APE: ( + "layers.{bid}.attn.indexer.compressor.ape", # deepseek-v4 + ), + + MODEL_TENSOR.INDEXER_COMPRESS_NORM: ( + "layers.{bid}.attn.indexer.compressor.norm", # deepseek-v4 + ), + + MODEL_TENSOR.INDEXER_COMPRESS_KV: ( + "layers.{bid}.attn.indexer.compressor.wkv", # deepseek-v4 + ), + + MODEL_TENSOR.INDEXER_COMPRESS_GATE: ( + "layers.{bid}.attn.indexer.compressor.wgate", # deepseek-v4 + ), + + MODEL_TENSOR.HC_HEAD_BASE: ( + "hc_head_base", # deepseek-v4 + ), + + MODEL_TENSOR.HC_HEAD_FN: ( + "hc_head_fn", # deepseek-v4 + ), + + MODEL_TENSOR.HC_HEAD_SCALE: ( + "hc_head_scale", # deepseek-v4 + ), + + MODEL_TENSOR.HC_ATTN_BASE: ( + "layers.{bid}.hc_attn_base", # deepseek-v4 + ), + + MODEL_TENSOR.HC_ATTN_FN: ( + "layers.{bid}.hc_attn_fn", # deepseek-v4 + ), + + MODEL_TENSOR.HC_ATTN_SCALE: ( + "layers.{bid}.hc_attn_scale", # deepseek-v4 + ), + + MODEL_TENSOR.HC_FFN_BASE: ( + "layers.{bid}.hc_ffn_base", # deepseek-v4 + ), + + MODEL_TENSOR.HC_FFN_FN: ( + "layers.{bid}.hc_ffn_fn", # deepseek-v4 + ), + + MODEL_TENSOR.HC_FFN_SCALE: ( + "layers.{bid}.hc_ffn_scale", # deepseek-v4 + ), + + MODEL_TENSOR.FFN_GATE_TID2EID: ( + "layers.{bid}.ffn.gate.tid2eid", # deepseek-v4 ), MODEL_TENSOR.ATTN_SUB_NORM: ( @@ -1244,6 +1344,7 @@ class TensorNameMap: MODEL_TENSOR.INDEXER_PROJ: ( "model.layers.{bid}.self_attn.indexer.weights_proj", # DSA + "layers.{bid}.attn.indexer.weights_proj", # deepseek-v4 ), MODEL_TENSOR.INDEXER_ATTN_K: ( @@ -1252,6 +1353,7 @@ class TensorNameMap: MODEL_TENSOR.INDEXER_ATTN_Q_B: ( "model.layers.{bid}.self_attn.indexer.wq_b", # DSA + "layers.{bid}.attn.indexer.wq_b", # deepseek-v4 ), ############################################################################ diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 7b1fcfca0ada..6f8eae4d11aa 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -25,6 +25,7 @@ add_library(llama llama-kv-cache.cpp llama-kv-cache-iswa.cpp llama-memory.cpp + llama-memory-deepseek4.cpp llama-memory-hybrid.cpp llama-memory-hybrid-iswa.cpp llama-memory-recurrent.cpp diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 633a66fc6651..f2cf5e21f578 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -75,6 +75,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_DEEPSEEK, "deepseek" }, { LLM_ARCH_DEEPSEEK2, "deepseek2" }, { LLM_ARCH_DEEPSEEK2OCR, "deepseek2-ocr" }, + { LLM_ARCH_DEEPSEEK4, "deepseek4" }, { LLM_ARCH_CHATGLM, "chatglm" }, { LLM_ARCH_GLM4, "glm4" }, { LLM_ARCH_GLM4_MOE, "glm4moe" }, @@ -547,6 +548,27 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_INDEXER_PROJ, "blk.%d.indexer.proj" }, { LLM_TENSOR_INDEXER_ATTN_K, "blk.%d.indexer.attn_k" }, { LLM_TENSOR_INDEXER_ATTN_Q_B, "blk.%d.indexer.attn_q_b" }, + { LLM_TENSOR_ATTN_KV_LATENT, "blk.%d.attn_kv_latent" }, + { LLM_TENSOR_ATTN_OUT_A, "blk.%d.attn_output_a" }, + { LLM_TENSOR_ATTN_OUT_B, "blk.%d.attn_output_b" }, + { LLM_TENSOR_ATTN_COMPRESS_APE, "blk.%d.attn_compress_ape" }, + { LLM_TENSOR_ATTN_COMPRESS_NORM, "blk.%d.attn_compress_norm" }, + { LLM_TENSOR_ATTN_COMPRESS_KV, "blk.%d.attn_compress_kv" }, + { LLM_TENSOR_ATTN_COMPRESS_GATE, "blk.%d.attn_compress_gate" }, + { LLM_TENSOR_INDEXER_COMPRESS_APE, "blk.%d.indexer.compress_ape" }, + { LLM_TENSOR_INDEXER_COMPRESS_NORM, "blk.%d.indexer.compress_norm" }, + { LLM_TENSOR_INDEXER_COMPRESS_KV, "blk.%d.indexer.compress_kv" }, + { LLM_TENSOR_INDEXER_COMPRESS_GATE, "blk.%d.indexer.compress_gate" }, + { LLM_TENSOR_HC_HEAD_BASE, "hc_head_base" }, + { LLM_TENSOR_HC_HEAD_FN, "hc_head_fn" }, + { LLM_TENSOR_HC_HEAD_SCALE, "hc_head_scale" }, + { LLM_TENSOR_HC_ATTN_BASE, "blk.%d.hc_attn_base" }, + { LLM_TENSOR_HC_ATTN_FN, "blk.%d.hc_attn_fn" }, + { LLM_TENSOR_HC_ATTN_SCALE, "blk.%d.hc_attn_scale" }, + { LLM_TENSOR_HC_FFN_BASE, "blk.%d.hc_ffn_base" }, + { LLM_TENSOR_HC_FFN_FN, "blk.%d.hc_ffn_fn" }, + { LLM_TENSOR_HC_FFN_SCALE, "blk.%d.hc_ffn_scale" }, + { LLM_TENSOR_FFN_GATE_TID2EID, "blk.%d.ffn_gate_tid2eid" }, }; // declare information about the model weight tensors: @@ -756,6 +778,27 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_INDEXER_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_INDEXER_ATTN_K, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_INDEXER_ATTN_Q_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_ATTN_KV_LATENT, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_ATTN_OUT_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_ATTN_OUT_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_ATTN_COMPRESS_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}}, + {LLM_TENSOR_ATTN_COMPRESS_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_ATTN_COMPRESS_KV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_ATTN_COMPRESS_GATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_INDEXER_COMPRESS_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}}, + {LLM_TENSOR_INDEXER_COMPRESS_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_INDEXER_COMPRESS_KV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_INDEXER_COMPRESS_GATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_HEAD_BASE, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_ADD}}, + {LLM_TENSOR_HC_HEAD_FN, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_HEAD_SCALE, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_SCALE}}, + {LLM_TENSOR_HC_ATTN_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}}, + {LLM_TENSOR_HC_ATTN_FN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_ATTN_SCALE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SCALE}}, + {LLM_TENSOR_HC_FFN_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}}, + {LLM_TENSOR_HC_FFN_FN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_FFN_SCALE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SCALE}}, + {LLM_TENSOR_FFN_GATE_TID2EID, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}}, // NextN/MTP tensors are currently ignored (reserved for future MTP support) // These tensors only exist in the last layer(s) and are treated as output tensors {LLM_TENSOR_NEXTN_EH_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, diff --git a/src/llama-arch.h b/src/llama-arch.h index 8f335f5c7b3e..9438d9bb1d33 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -79,6 +79,7 @@ enum llm_arch { LLM_ARCH_DEEPSEEK, LLM_ARCH_DEEPSEEK2, LLM_ARCH_DEEPSEEK2OCR, + LLM_ARCH_DEEPSEEK4, LLM_ARCH_CHATGLM, LLM_ARCH_GLM4, LLM_ARCH_GLM4_MOE, @@ -548,6 +549,27 @@ enum llm_tensor { LLM_TENSOR_INDEXER_PROJ, LLM_TENSOR_INDEXER_ATTN_K, LLM_TENSOR_INDEXER_ATTN_Q_B, + LLM_TENSOR_ATTN_KV_LATENT, + LLM_TENSOR_ATTN_OUT_A, + LLM_TENSOR_ATTN_OUT_B, + LLM_TENSOR_ATTN_COMPRESS_APE, + LLM_TENSOR_ATTN_COMPRESS_NORM, + LLM_TENSOR_ATTN_COMPRESS_KV, + LLM_TENSOR_ATTN_COMPRESS_GATE, + LLM_TENSOR_INDEXER_COMPRESS_APE, + LLM_TENSOR_INDEXER_COMPRESS_NORM, + LLM_TENSOR_INDEXER_COMPRESS_KV, + LLM_TENSOR_INDEXER_COMPRESS_GATE, + LLM_TENSOR_HC_HEAD_BASE, + LLM_TENSOR_HC_HEAD_FN, + LLM_TENSOR_HC_HEAD_SCALE, + LLM_TENSOR_HC_ATTN_BASE, + LLM_TENSOR_HC_ATTN_FN, + LLM_TENSOR_HC_ATTN_SCALE, + LLM_TENSOR_HC_FFN_BASE, + LLM_TENSOR_HC_FFN_FN, + LLM_TENSOR_HC_FFN_SCALE, + LLM_TENSOR_FFN_GATE_TID2EID, LLM_TENSOR_NEXTN_EH_PROJ, LLM_TENSOR_NEXTN_EMBED_TOKENS, LLM_TENSOR_NEXTN_ENORM, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 8126249e1436..263c7a31c383 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -469,6 +469,11 @@ void llama_context::sched_reserve() { if (cparams.auto_fgdn) { LLAMA_LOG_INFO("%s: resolving fused Gated Delta Net support:\n", __func__); + if (model.arch == LLM_ARCH_DEEPSEEK4) { + cparams.fused_gdn_ar = false; + cparams.fused_gdn_ch = false; + } + if (cparams.fused_gdn_ar) { auto * gf = graph_reserve(1, n_seqs, n_outputs, mctx.get(), true); if (!gf) { @@ -2073,6 +2078,9 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_KIMI_LINEAR || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE) { return std::max(n_tokens * 40, 32u * model.n_tensors()); } + if (model.arch == LLM_ARCH_DEEPSEEK4) { + return std::max(n_tokens * 256, 128u * model.n_tensors()); + } uint32_t res = std::max(1024u, 8u*model.n_tensors()); for (const auto & lora : model.loras) { res += lora->get_n_nodes(); diff --git a/src/llama-memory-deepseek4.cpp b/src/llama-memory-deepseek4.cpp new file mode 100644 index 000000000000..662ecea234c5 --- /dev/null +++ b/src/llama-memory-deepseek4.cpp @@ -0,0 +1,689 @@ +#include "llama-memory-deepseek4.h" + +#include "llama-impl.h" +#include "llama-model.h" +#include "llama-context.h" +#include "llama-io.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +// v1: every cache tensor was serialized with its full ggml_nbytes(), regardless of how +// many slots were populated. With n_ctx in the millions this made each checkpoint +// several GiB even for short conversations; the server's per-turn checkpoint restore +// (triggered because DeepSeek4 only supports full-removal seq_rm) became dominant. +// v2: only the active row prefix of n_ctx-scaling tensors (attn_kv, indexer_kv) is +// written. On read the active prefix bytes are restored and the remaining tail is +// explicitly zeroed via ggml_backend_tensor_memset, preserving the +// "untouched-slot == zero" invariant the compute graph relies on. +static constexpr uint32_t DEEPSEEK4_STATE_VERSION = 2; + +static llama_ubatch make_dummy_ubatch() { + llama_ubatch ubatch = {}; + ubatch.data = std::make_shared(); + + ubatch.b_equal_seqs = 1; + ubatch.n_tokens = 1; + ubatch.n_seq_tokens = 1; + ubatch.n_seqs = 1; + ubatch.n_seqs_unq = 1; + ubatch.n_pos = 1; + + ubatch.data->token = { 0 }; + ubatch.data->pos = { 0 }; + ubatch.data->n_seq_id = { 1 }; + ubatch.data->seq_id_unq = { 0 }; + ubatch.data->seq_idx.assign(LLAMA_MAX_SEQ, -1); + ubatch.data->seq_idx[0] = 0; + ubatch.data->output = { 0 }; + ubatch.data->seq_id_data = { 0 }; + ubatch.data->seq_id = { ubatch.data->seq_id_data.data() }; + + ubatch.token = ubatch.data->token.data(); + ubatch.embd = nullptr; + ubatch.pos = ubatch.data->pos.data(); + ubatch.n_seq_id = ubatch.data->n_seq_id.data(); + ubatch.seq_id = ubatch.data->seq_id.data(); + ubatch.seq_id_unq = ubatch.data->seq_id_unq.data(); + ubatch.seq_idx = ubatch.data->seq_idx.data(); + ubatch.output = ubatch.data->output.data(); + + return ubatch; +} + +static uint32_t deepseek4_compress_ratio(const llama_layer & layer) { + return layer.attn_compress_ape ? static_cast(layer.attn_compress_ape->ne[1]) : 0; +} + +static uint32_t deepseek4_comp_slots(const ggml_tensor * ape, uint32_t head_dim) { + if (!ape || head_dim == 0) { + return 0; + } + + return static_cast(ape->ne[0] / head_dim); +} + +static void deepseek4_fill_f32_tensor(ggml_tensor * tensor, float value) { + if (!tensor) { + return; + } + + GGML_ASSERT(tensor->type == GGML_TYPE_F32); + std::vector data(ggml_nelements(tensor), value); + ggml_backend_tensor_set(tensor, data.data(), 0, ggml_nbytes(tensor)); +} + +static void deepseek4_write_tensor(llama_io_write_i & io, const ggml_tensor * tensor, uint64_t active_bytes_override = UINT64_MAX) { + const uint32_t present = tensor != nullptr; + io.write(&present, sizeof(present)); + + if (!present) { + return; + } + + const int32_t type = static_cast(tensor->type); + const uint32_t n_dims = ggml_n_dims(tensor); + int64_t ne[GGML_MAX_DIMS] = {}; + for (uint32_t i = 0; i < GGML_MAX_DIMS; ++i) { + ne[i] = tensor->ne[i]; + } + const uint64_t total_bytes = ggml_nbytes(tensor); + const uint64_t active_bytes = active_bytes_override == UINT64_MAX + ? total_bytes + : std::min(active_bytes_override, total_bytes); + + io.write(&type, sizeof(type)); + io.write(&n_dims, sizeof(n_dims)); + io.write(ne, sizeof(ne)); + io.write(&active_bytes, sizeof(active_bytes)); + io.write(&total_bytes, sizeof(total_bytes)); + if (active_bytes > 0) { + io.write_tensor(tensor, 0, active_bytes); + } +} + +static void deepseek4_read_tensor(llama_io_read_i & io, ggml_tensor * tensor) { + uint32_t present; + io.read_to(&present, sizeof(present)); + + if (!present) { + if (tensor != nullptr) { + throw std::runtime_error("DeepSeek4 state is missing a runtime tensor"); + } + return; + } + + if (tensor == nullptr) { + throw std::runtime_error("DeepSeek4 state contains an unexpected runtime tensor"); + } + + int32_t type_ref; + uint32_t n_dims_ref; + int64_t ne_ref[GGML_MAX_DIMS]; + uint64_t active_bytes_ref; + uint64_t total_bytes_ref; + + io.read_to(&type_ref, sizeof(type_ref)); + io.read_to(&n_dims_ref, sizeof(n_dims_ref)); + io.read_to(ne_ref, sizeof(ne_ref)); + io.read_to(&active_bytes_ref, sizeof(active_bytes_ref)); + io.read_to(&total_bytes_ref, sizeof(total_bytes_ref)); + + if (type_ref != static_cast(tensor->type)) { + throw std::runtime_error("DeepSeek4 state tensor type mismatch"); + } + if (n_dims_ref != static_cast(ggml_n_dims(tensor))) { + throw std::runtime_error("DeepSeek4 state tensor rank mismatch"); + } + for (uint32_t i = 0; i < GGML_MAX_DIMS; ++i) { + if (ne_ref[i] != tensor->ne[i]) { + throw std::runtime_error("DeepSeek4 state tensor shape mismatch"); + } + } + + const uint64_t total_bytes = ggml_nbytes(tensor); + if (total_bytes_ref != total_bytes) { + throw std::runtime_error("DeepSeek4 state tensor size mismatch"); + } + if (active_bytes_ref > total_bytes) { + throw std::runtime_error("DeepSeek4 state tensor active range exceeds tensor size"); + } + + if (active_bytes_ref > 0) { + ggml_backend_tensor_set(tensor, io.read(active_bytes_ref), 0, active_bytes_ref); + } + if (active_bytes_ref < total_bytes) { + // Preserve the "untouched-slot == zero" invariant the compute graph relies on: + // build_attn_v4 reads compressed/indexer prefixes by current batch end, which can + // include rows beyond the restored prefix on the first batch after restore. + ggml_backend_tensor_memset(tensor, 0, active_bytes_ref, total_bytes - active_bytes_ref); + } +} + +} // namespace + +llama_memory_deepseek4::llama_memory_deepseek4( + const llama_model & model, + ggml_type type_k, + bool offload, + uint32_t n_ctx_seq, + uint32_t n_seq_max) : + model(model), + n_ctx_seq(n_ctx_seq), + n_seq_max(n_seq_max), + layers(model.hparams.n_layer), + seq_pos_min_v(n_seq_max, -1), + seq_pos_max_v(n_seq_max, -1) { + struct ggml_backend_buft_comparator { + bool operator()(const ggml_backend_buffer_type_t & lhs, const ggml_backend_buffer_type_t & rhs) const { + return strcmp(ggml_backend_buft_name(lhs), ggml_backend_buft_name(rhs)) < 0; + } + }; + + std::map ctx_map; + + auto ctx_for_buft = [&](ggml_backend_buffer_type_t buft) -> ggml_context * { + auto it = ctx_map.find(buft); + if (it != ctx_map.end()) { + return it->second.get(); + } + + ggml_init_params params = { + /*.mem_size =*/ size_t(16u * model.hparams.n_layer * ggml_tensor_overhead()), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + + ggml_context * ctx = ggml_init(params); + if (!ctx) { + return nullptr; + } + + ctx_map.emplace(buft, ctx); + return ctx; + }; + + for (int32_t il = 0; il < (int32_t) model.hparams.n_layer; ++il) { + const auto & layer_model = model.layers[il]; + auto & layer = layers[il]; + + const uint32_t head_dim = model.hparams.n_embd_head_k(il); + const uint32_t ratio = deepseek4_compress_ratio(layer_model); + const uint32_t kv_size = model.hparams.n_swa + (ratio ? n_ctx_seq / ratio : 0); + + ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); + if (offload) { + buft = ggml_backend_dev_buffer_type(model.dev_layer(il)); + } + + ggml_context * ctx = ctx_for_buft(buft); + if (!ctx) { + throw std::runtime_error("failed to create DeepSeek4 state context"); + } + + layer.attn_kv = ggml_new_tensor_2d(ctx, type_k, head_dim, kv_size); + ggml_format_name(layer.attn_kv, "deepseek4_attn_kv_l%d", il); + + if (ratio > 0) { + const uint32_t attn_comp_slots = deepseek4_comp_slots(layer_model.attn_compress_ape, head_dim); + layer.attn_comp_kv_state = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, layer_model.attn_compress_ape->ne[0], attn_comp_slots * ratio); + ggml_format_name(layer.attn_comp_kv_state, "deepseek4_attn_comp_kv_state_l%d", il); + layer.attn_comp_score_state = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, layer_model.attn_compress_ape->ne[0], attn_comp_slots * ratio); + ggml_format_name(layer.attn_comp_score_state, "deepseek4_attn_comp_score_state_l%d", il); + } + + if (layer_model.indexer_proj && layer_model.indexer_attn_q_b && layer_model.indexer_compress_ape) { + const uint32_t idx_ratio = static_cast(layer_model.indexer_compress_ape->ne[1]); + const uint32_t idx_head_dim = model.hparams.indexer_head_size; + const uint32_t idx_kv_size = idx_ratio ? n_ctx_seq / idx_ratio : 0; + const uint32_t idx_comp_slots = deepseek4_comp_slots(layer_model.indexer_compress_ape, idx_head_dim); + + layer.indexer_kv = ggml_new_tensor_2d(ctx, type_k, idx_head_dim, idx_kv_size); + ggml_format_name(layer.indexer_kv, "deepseek4_indexer_kv_l%d", il); + + layer.indexer_comp_kv_state = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, layer_model.indexer_compress_ape->ne[0], idx_comp_slots * idx_ratio); + ggml_format_name(layer.indexer_comp_kv_state, "deepseek4_indexer_comp_kv_state_l%d", il); + layer.indexer_comp_score_state = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, layer_model.indexer_compress_ape->ne[0], idx_comp_slots * idx_ratio); + ggml_format_name(layer.indexer_comp_score_state, "deepseek4_indexer_comp_score_state_l%d", il); + } + } + + for (auto & [buft, ctx] : ctx_map) { + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); + if (!buf) { + throw std::runtime_error("failed to allocate DeepSeek4 state buffer"); + } + ggml_backend_buffer_clear(buf, 0); + ctxs_bufs.emplace_back(std::move(ctx), buf); + } + + for (auto & layer : layers) { + deepseek4_fill_f32_tensor(layer.attn_comp_score_state, -std::numeric_limits::infinity()); + deepseek4_fill_f32_tensor(layer.indexer_comp_score_state, -std::numeric_limits::infinity()); + } +} + +llama_memory_context_ptr llama_memory_deepseek4::init_batch( + llama_batch_allocr & balloc, + uint32_t n_ubatch, + bool embd_all) { + GGML_UNUSED(embd_all); + GGML_UNUSED(n_ubatch); + + balloc.split_reset(); + + std::vector ubatches; + while (true) { + // The compression-window indexer state is updated one position at a + // time, so each ubatch carries exactly one token from one sequence. + llama_ubatch ubatch = balloc.split_seq(1); + if (ubatch.n_tokens == 0) { + break; + } + + if (ubatch.n_tokens != 1 || ubatch.n_seqs_unq != 1) { + LLAMA_LOG_ERROR("%s: DeepSeek4 runtime currently supports a single token from a single sequence per ubatch\n", + __func__); + return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); + } + + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + if (ubatch.pos[i] < 0 || (uint32_t) ubatch.pos[i] >= n_ctx_seq) { + LLAMA_LOG_ERROR("%s: DeepSeek4 runtime position %d exceeds the configured context length %u\n", + __func__, ubatch.pos[i], n_ctx_seq); + return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); + } + } + + ubatches.push_back(std::move(ubatch)); + } + + if (balloc.get_n_used() < balloc.get_n_tokens()) { + return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); + } + + return std::make_unique(this, std::move(ubatches)); +} + +llama_memory_context_ptr llama_memory_deepseek4::init_full() { + std::vector ubatches = { make_dummy_ubatch() }; + return std::make_unique(this, std::move(ubatches)); +} + +llama_memory_context_ptr llama_memory_deepseek4::init_update(llama_context * lctx, bool optimize) { + GGML_UNUSED(lctx); + GGML_UNUSED(optimize); + return std::make_unique(LLAMA_MEMORY_STATUS_NO_UPDATE); +} + +bool llama_memory_deepseek4::get_can_shift() const { + return false; +} + +void llama_memory_deepseek4::clear(bool data) { + std::fill(seq_pos_min_v.begin(), seq_pos_min_v.end(), -1); + std::fill(seq_pos_max_v.begin(), seq_pos_max_v.end(), -1); + + if (data) { + for (auto & [_, buf] : ctxs_bufs) { + ggml_backend_buffer_clear(buf.get(), 0); + } + for (auto & layer : layers) { + deepseek4_fill_f32_tensor(layer.attn_comp_score_state, -std::numeric_limits::infinity()); + deepseek4_fill_f32_tensor(layer.indexer_comp_score_state, -std::numeric_limits::infinity()); + } + } +} + +bool llama_memory_deepseek4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { + const llama_pos r0 = p0 < 0 ? 0 : p0; + const llama_pos r1 = p1 < 0 ? std::numeric_limits::max() : p1; + + if (r0 >= r1) { + return true; + } + + llama_pos pos_min = -1; + llama_pos pos_max = -1; + if (seq_id < 0) { + for (size_t i = 0; i < seq_pos_min_v.size(); ++i) { + if (seq_pos_min_v[i] < 0) { + continue; + } + pos_min = pos_min < 0 ? seq_pos_min_v[i] : std::min(pos_min, seq_pos_min_v[i]); + pos_max = std::max(pos_max, seq_pos_max_v[i]); + } + } else { + if (static_cast(seq_id) >= seq_pos_min_v.size()) { + return false; + } + pos_min = seq_pos_min_v[seq_id]; + pos_max = seq_pos_max_v[seq_id]; + } + + if (pos_min < 0 || r1 <= pos_min || r0 > pos_max) { + return true; + } + + if (r0 <= pos_min && r1 > pos_max) { + clear(true); + return true; + } + + return false; +} + +void llama_memory_deepseek4::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { + GGML_UNUSED(seq_id_src); + GGML_UNUSED(seq_id_dst); + GGML_UNUSED(p0); + GGML_UNUSED(p1); +} + +void llama_memory_deepseek4::seq_keep(llama_seq_id seq_id) { + GGML_UNUSED(seq_id); +} + +void llama_memory_deepseek4::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { + GGML_UNUSED(seq_id); + GGML_UNUSED(p0); + GGML_UNUSED(p1); + GGML_UNUSED(shift); +} + +void llama_memory_deepseek4::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { + GGML_UNUSED(seq_id); + GGML_UNUSED(p0); + GGML_UNUSED(p1); + GGML_UNUSED(d); +} + +llama_pos llama_memory_deepseek4::seq_pos_min(llama_seq_id seq_id) const { + if (seq_id < 0 || (size_t) seq_id >= seq_pos_min_v.size()) { + return -1; + } + return seq_pos_min_v[seq_id]; +} + +llama_pos llama_memory_deepseek4::seq_pos_max(llama_seq_id seq_id) const { + if (seq_id < 0 || (size_t) seq_id >= seq_pos_max_v.size()) { + return -1; + } + return seq_pos_max_v[seq_id]; +} + +std::map llama_memory_deepseek4::memory_breakdown() const { + std::map mb; + for (const auto & [_, buf] : ctxs_bufs) { + mb[ggml_backend_buffer_get_type(buf.get())] += ggml_backend_buffer_get_size(buf.get()); + } + return mb; +} + +void llama_memory_deepseek4::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { + GGML_UNUSED(flags); + + const bool seq_specific = seq_id != -1; + const bool seq_valid = seq_id >= 0 && static_cast(seq_id) < seq_pos_min_v.size(); + const bool seq_active = !seq_specific || (seq_valid && seq_pos_min_v[seq_id] >= 0); + + const uint32_t version = DEEPSEEK4_STATE_VERSION; + const uint32_t n_layer = layers.size(); + const uint32_t seq_mode = seq_specific ? 1 : 0; + const uint32_t has_data = seq_active ? 1 : 0; + const uint32_t seq_count = seq_specific ? 1 : n_seq_max; + + io.write(&version, sizeof(version)); + io.write(&n_ctx_seq, sizeof(n_ctx_seq)); + io.write(&n_seq_max, sizeof(n_seq_max)); + io.write(&n_layer, sizeof(n_layer)); + io.write(&seq_mode, sizeof(seq_mode)); + io.write(&has_data, sizeof(has_data)); + io.write(&seq_count, sizeof(seq_count)); + + if (seq_specific) { + const llama_pos pos_min = seq_valid ? seq_pos_min_v[seq_id] : -1; + const llama_pos pos_max = seq_valid ? seq_pos_max_v[seq_id] : -1; + io.write(&pos_min, sizeof(pos_min)); + io.write(&pos_max, sizeof(pos_max)); + } else { + for (uint32_t i = 0; i < n_seq_max; ++i) { + const llama_pos pos_min = i < seq_pos_min_v.size() ? seq_pos_min_v[i] : -1; + const llama_pos pos_max = i < seq_pos_max_v.size() ? seq_pos_max_v[i] : -1; + io.write(&pos_min, sizeof(pos_min)); + io.write(&pos_max, sizeof(pos_max)); + } + } + + if (!has_data) { + return; + } + + // Compute the highest populated position over the seqs we are about to serialize so + // that n_ctx-scaling tensors can be trimmed to their active prefix. The model only + // supports n_seq_max == 1 in practice; for the broader (-1) save case we take the + // union of all seqs to stay correct if that ever changes. + llama_pos pos_max_global = -1; + if (seq_specific) { + if (seq_valid) { + pos_max_global = seq_pos_max_v[seq_id]; + } + } else { + for (size_t i = 0; i < seq_pos_max_v.size(); ++i) { + if (seq_pos_min_v[i] >= 0) { + pos_max_global = std::max(pos_max_global, seq_pos_max_v[i]); + } + } + } + + const uint32_t n_swa = model.hparams.n_swa; + + for (size_t il = 0; il < layers.size(); ++il) { + const auto & layer = layers[il]; + const auto & layer_model = model.layers[il]; + + // attn_kv: shape [head_dim, n_swa + n_ctx_seq/ratio]; rows used are + // [0, n_swa) (SWA circular slots) plus [n_swa, n_swa + ceil((pos_max+1)/ratio)). + // For ratio == 0 there is no compressed region and the tensor is sized for n_swa. + uint64_t attn_active_bytes = UINT64_MAX; + if (layer.attn_kv != nullptr) { + const uint32_t ratio = deepseek4_compress_ratio(layer_model); + const uint64_t row_size = layer.attn_kv->nb[1]; + const uint64_t total_rows = layer.attn_kv->ne[1]; + uint64_t active_rows = std::min(n_swa, total_rows); + if (ratio > 0 && pos_max_global >= 0) { + const uint64_t comp_rows = (uint64_t(pos_max_global) + ratio) / ratio; // ceil((pos_max+1)/ratio) + active_rows = std::min(uint64_t(n_swa) + comp_rows, total_rows); + } + attn_active_bytes = active_rows * row_size; + } + + // indexer_kv: shape [idx_head_dim, n_ctx_seq/idx_ratio]; rows used are + // [0, ceil((pos_max+1)/idx_ratio)). No n_swa offset for the indexer. + uint64_t indexer_active_bytes = UINT64_MAX; + if (layer.indexer_kv != nullptr && layer_model.indexer_compress_ape != nullptr) { + const uint32_t idx_ratio = static_cast(layer_model.indexer_compress_ape->ne[1]); + const uint64_t row_size = layer.indexer_kv->nb[1]; + const uint64_t total_rows = layer.indexer_kv->ne[1]; + uint64_t active_rows = 0; + if (idx_ratio > 0 && pos_max_global >= 0) { + active_rows = std::min((uint64_t(pos_max_global) + idx_ratio) / idx_ratio, total_rows); + } + indexer_active_bytes = active_rows * row_size; + } + + deepseek4_write_tensor(io, layer.attn_kv, attn_active_bytes); + // attn_comp_*/indexer_comp_* are fixed-size compression state and must be + // restored byte-for-byte (they encode incremental sums that the next batch + // continues from). Pass UINT64_MAX to keep the full-size write path. + deepseek4_write_tensor(io, layer.attn_comp_kv_state); + deepseek4_write_tensor(io, layer.attn_comp_score_state); + deepseek4_write_tensor(io, layer.indexer_kv, indexer_active_bytes); + deepseek4_write_tensor(io, layer.indexer_comp_kv_state); + deepseek4_write_tensor(io, layer.indexer_comp_score_state); + } +} + +void llama_memory_deepseek4::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { + GGML_UNUSED(flags); + + uint32_t version; + uint32_t n_ctx_seq_ref; + uint32_t n_seq_max_ref; + uint32_t n_layer_ref; + uint32_t seq_mode; + uint32_t has_data; + uint32_t seq_count; + + io.read_to(&version, sizeof(version)); + io.read_to(&n_ctx_seq_ref, sizeof(n_ctx_seq_ref)); + io.read_to(&n_seq_max_ref, sizeof(n_seq_max_ref)); + io.read_to(&n_layer_ref, sizeof(n_layer_ref)); + io.read_to(&seq_mode, sizeof(seq_mode)); + io.read_to(&has_data, sizeof(has_data)); + io.read_to(&seq_count, sizeof(seq_count)); + + if (version != DEEPSEEK4_STATE_VERSION) { + throw std::runtime_error("DeepSeek4 state version mismatch"); + } + if (n_ctx_seq_ref != n_ctx_seq) { + throw std::runtime_error("DeepSeek4 state context length mismatch"); + } + if (n_layer_ref != layers.size()) { + throw std::runtime_error("DeepSeek4 state layer count mismatch"); + } + + if (seq_mode == 1) { + if (seq_count != 1) { + throw std::runtime_error("DeepSeek4 sequence state metadata mismatch"); + } + + llama_pos pos_min; + llama_pos pos_max; + io.read_to(&pos_min, sizeof(pos_min)); + io.read_to(&pos_max, sizeof(pos_max)); + + if (seq_id < 0 || static_cast(seq_id) >= seq_pos_min_v.size()) { + throw std::runtime_error("DeepSeek4 sequence state destination is out of range"); + } + + seq_pos_min_v[seq_id] = has_data ? pos_min : -1; + seq_pos_max_v[seq_id] = has_data ? pos_max : -1; + } else if (seq_mode == 0) { + const uint32_t n_read = std::min(seq_count, n_seq_max); + for (uint32_t i = 0; i < seq_count; ++i) { + llama_pos pos_min; + llama_pos pos_max; + io.read_to(&pos_min, sizeof(pos_min)); + io.read_to(&pos_max, sizeof(pos_max)); + + if (i < n_read) { + seq_pos_min_v[i] = pos_min; + seq_pos_max_v[i] = pos_max; + } + } + for (uint32_t i = n_read; i < n_seq_max; ++i) { + seq_pos_min_v[i] = -1; + seq_pos_max_v[i] = -1; + } + } else { + throw std::runtime_error("DeepSeek4 state sequence mode mismatch"); + } + + GGML_UNUSED(n_seq_max_ref); + + if (!has_data) { + return; + } + + for (auto & layer : layers) { + deepseek4_read_tensor(io, layer.attn_kv); + deepseek4_read_tensor(io, layer.attn_comp_kv_state); + deepseek4_read_tensor(io, layer.attn_comp_score_state); + deepseek4_read_tensor(io, layer.indexer_kv); + deepseek4_read_tensor(io, layer.indexer_comp_kv_state); + deepseek4_read_tensor(io, layer.indexer_comp_score_state); + } +} + +const llama_memory_deepseek4::layer_state & llama_memory_deepseek4::get_layer(int32_t il) const { + return layers.at(il); +} + +uint32_t llama_memory_deepseek4::get_n_ctx_seq() const { + return n_ctx_seq; +} + +llama_memory_deepseek4_context::llama_memory_deepseek4_context(llama_memory_status status) : + status(status) { +} + +llama_memory_deepseek4_context::llama_memory_deepseek4_context( + llama_memory_deepseek4 * mem, + std::vector ubatches) : + status(LLAMA_MEMORY_STATUS_SUCCESS), + mem(mem), + ubatches(std::move(ubatches)) { +} + +bool llama_memory_deepseek4_context::next() { + if (status != LLAMA_MEMORY_STATUS_SUCCESS) { + return false; + } + + if (++i_next >= ubatches.size()) { + return false; + } + + return true; +} + +bool llama_memory_deepseek4_context::apply() { + if (status != LLAMA_MEMORY_STATUS_SUCCESS || mem == nullptr || ubatches.empty()) { + return status != LLAMA_MEMORY_STATUS_FAILED_PREPARE; + } + + const auto & ubatch = ubatches[i_next]; + const llama_seq_id seq_id = ubatch.seq_id[0][0]; + if (seq_id < 0 || (size_t) seq_id >= mem->seq_pos_min_v.size()) { + return false; + } + + auto & pos_min = mem->seq_pos_min_v[seq_id]; + auto & pos_max = mem->seq_pos_max_v[seq_id]; + + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + if (ubatch.seq_id[i][0] != seq_id) { + return false; + } + + const llama_pos pos = ubatch.pos[i]; + pos_min = pos_min < 0 ? pos : std::min(pos_min, pos); + pos_max = std::max(pos_max, pos); + } + + return true; +} + +const llama_ubatch & llama_memory_deepseek4_context::get_ubatch() const { + return ubatches.at(i_next); +} + +llama_memory_status llama_memory_deepseek4_context::get_status() const { + return status; +} + +const llama_memory_deepseek4::layer_state & llama_memory_deepseek4_context::get_layer(int32_t il) const { + return mem->get_layer(il); +} + +uint32_t llama_memory_deepseek4_context::get_n_ctx_seq() const { + return mem->get_n_ctx_seq(); +} diff --git a/src/llama-memory-deepseek4.h b/src/llama-memory-deepseek4.h new file mode 100644 index 000000000000..af73c9cfe09e --- /dev/null +++ b/src/llama-memory-deepseek4.h @@ -0,0 +1,109 @@ +#pragma once + +#include "llama-batch.h" +#include "llama-memory.h" +#include "ggml-cpp.h" + +#include + +struct ggml_context; +struct ggml_tensor; + +struct llama_model; +struct llama_context; + +class llama_memory_deepseek4 : public llama_memory_i { +public: + struct layer_state { + ggml_tensor * attn_kv = nullptr; + + ggml_tensor * attn_comp_kv_state = nullptr; + ggml_tensor * attn_comp_score_state = nullptr; + + ggml_tensor * indexer_kv = nullptr; + + ggml_tensor * indexer_comp_kv_state = nullptr; + ggml_tensor * indexer_comp_score_state = nullptr; + }; + + llama_memory_deepseek4( + const llama_model & model, + ggml_type type_k, + bool offload, + uint32_t n_ctx_seq, + uint32_t n_seq_max); + + ~llama_memory_deepseek4() override = default; + + llama_memory_context_ptr init_batch( + llama_batch_allocr & balloc, + uint32_t n_ubatch, + bool embd_all) override; + + llama_memory_context_ptr init_full() override; + + llama_memory_context_ptr init_update(llama_context * lctx, bool optimize) override; + + bool get_can_shift() const override; + + void clear(bool data) override; + + bool seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) override; + void seq_cp (llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) override; + void seq_keep(llama_seq_id seq_id) override; + void seq_add (llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) override; + void seq_div (llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) override; + + llama_pos seq_pos_min(llama_seq_id seq_id) const override; + llama_pos seq_pos_max(llama_seq_id seq_id) const override; + + std::map memory_breakdown() const override; + + void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override; + void state_read (llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override; + + const layer_state & get_layer(int32_t il) const; + uint32_t get_n_ctx_seq() const; + +private: + friend class llama_memory_deepseek4_context; + + const llama_model & model; + + const uint32_t n_ctx_seq; + const uint32_t n_seq_max; + + std::vector layers; + std::vector seq_pos_min_v; + std::vector seq_pos_max_v; + + std::vector> ctxs_bufs; +}; + +class llama_memory_deepseek4_context : public llama_memory_context_i { +public: + llama_memory_deepseek4_context(llama_memory_status status); + + llama_memory_deepseek4_context( + llama_memory_deepseek4 * mem, + std::vector ubatches); + + ~llama_memory_deepseek4_context() override = default; + + bool next() override; + bool apply() override; + + const llama_ubatch & get_ubatch() const override; + llama_memory_status get_status() const override; + + const llama_memory_deepseek4::layer_state & get_layer(int32_t il) const; + uint32_t get_n_ctx_seq() const; + +private: + const llama_memory_status status; + + llama_memory_deepseek4 * mem = nullptr; + + size_t i_next = 0; + std::vector ubatches; +}; diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 26864c18e973..5f28e31b47b9 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -212,8 +212,8 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp); add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_chexp); - add_kv(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp); - add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp); + add_kv(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, true); + add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, true); add_kv(LLM_KV_USE_PARALLEL_RESIDUAL, hparams.use_par_res); // add_kv(LLM_KV_TENSOR_DATA_LAYOUT, ???); add_kv(LLM_KV_EXPERT_COUNT, hparams.n_expert); @@ -397,6 +397,9 @@ void llama_model_saver::add_tensors_from_model() { add_tensor(model->cls_out); add_tensor(model->cls_out_b); add_tensor(model->cls_norm); + add_tensor(model->hc_head_base); + add_tensor(model->hc_head_fn); + add_tensor(model->hc_head_scale); for (const struct llama_layer & layer : model->layers) { for (size_t i = 0; i < sizeof(layer)/sizeof(struct ggml_tensor *); ++i) { diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 9e2a13cbd43e..c17e8c57a29a 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -10,6 +10,7 @@ #include "llama-kv-cache.h" #include "llama-kv-cache-iswa.h" +#include "llama-memory-deepseek4.h" #include "llama-memory-hybrid.h" #include "llama-memory-hybrid-iswa.h" #include "llama-memory-recurrent.h" @@ -2034,6 +2035,29 @@ void llama_model::load_hparams(llama_model_loader & ml) { default: type = LLM_TYPE_UNKNOWN; } } break; + case LLM_ARCH_DEEPSEEK4: + { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q); + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false); + ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head, false); + ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size, false); + ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k, false); + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); + hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train; + hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train; + ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false); + if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer, false)) { + std::fill_n(hparams.swiglu_clamp_exp.begin(), hparams.n_layer, 10.0f); + } + ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer, false); + + type = LLM_TYPE_UNKNOWN; + } break; case LLM_ARCH_PLM: { ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); @@ -5340,7 +5364,7 @@ bool llama_model::load_tensors(llama_model_loader & ml) { layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); } else { layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); - layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, i), {n_expert}, TENSOR_NOT_REQUIRED); if (n_expert == 0) { throw std::runtime_error("n_expert must be > 0"); @@ -5394,7 +5418,7 @@ bool llama_model::load_tensors(llama_model_loader & ml) { layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); } else { layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); - layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, i), {n_expert}, TENSOR_NOT_REQUIRED); if (n_expert == 0) { throw std::runtime_error("n_expert must be > 0"); @@ -5414,6 +5438,133 @@ bool llama_model::load_tensors(llama_model_loader & ml) { } } } break; + case LLM_ARCH_DEEPSEEK4: + { + const int64_t q_lora_rank = hparams.n_lora_q; + const int64_t n_ff_exp = hparams.n_ff_exp; + const int64_t n_expert_shared = hparams.n_expert_shared; + const int64_t n_embd_head = hparams.n_embd_head_k(); + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0); + + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED); + if (!output) { + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED); + } + + int64_t hc_mult = 0; + { + const auto * meta_base = ml.get_tensor_meta(tn(LLM_TENSOR_HC_HEAD_BASE).str().c_str()); + const auto * meta_fn = ml.get_tensor_meta(tn(LLM_TENSOR_HC_HEAD_FN).str().c_str()); + const auto * meta_scale = ml.get_tensor_meta(tn(LLM_TENSOR_HC_HEAD_SCALE).str().c_str()); + + hc_mult = meta_base ? meta_base->ne[0] : (meta_fn ? meta_fn->ne[1] : 4); + const int64_t hc_fn_in = meta_fn ? meta_fn->ne[0] : n_embd * hc_mult; + const int64_t hc_fn_out = meta_fn ? meta_fn->ne[1] : hc_mult; + const int64_t hc_scale_len = meta_scale ? meta_scale->ne[0] : 1; + + hc_head_base = create_tensor(tn(LLM_TENSOR_HC_HEAD_BASE), { hc_mult }, 0); + hc_head_fn = create_tensor(tn(LLM_TENSOR_HC_HEAD_FN), { hc_fn_in, hc_fn_out }, 0); + hc_head_scale = create_tensor(tn(LLM_TENSOR_HC_HEAD_SCALE), { hc_scale_len }, 0); + } + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, 0); + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), { q_lora_rank }, 0); + layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), { n_embd_head }, 0); + + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), { n_embd, q_lora_rank }, 0); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), { q_lora_rank, n_head * n_embd_head }, 0); + layer.attn_kv_latent = create_tensor(tn(LLM_TENSOR_ATTN_KV_LATENT, "weight", i), { n_embd, n_embd_head }, 0); + + { + const auto * meta_wo_a = ml.get_tensor_meta(tn(LLM_TENSOR_ATTN_OUT_A, "weight", i).str().c_str()); + const auto * meta_wo_b = ml.get_tensor_meta(tn(LLM_TENSOR_ATTN_OUT_B, "weight", i).str().c_str()); + const int64_t group_dim = n_embd_head; + const int64_t n_groups = n_head * n_embd_head / group_dim; + const int64_t o_rank = std::max(1, group_dim / 2); + const int64_t wo_a_ne0 = meta_wo_a ? meta_wo_a->ne[0] : group_dim; + const int64_t wo_a_ne1 = meta_wo_a ? meta_wo_a->ne[1] : n_groups * o_rank; + const int64_t wo_b_ne0 = meta_wo_b ? meta_wo_b->ne[0] : n_groups * o_rank; + const int64_t wo_b_ne1 = meta_wo_b ? meta_wo_b->ne[1] : n_embd; + layer.attn_out_a = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A, "weight", i), { wo_a_ne0, wo_a_ne1 }, 0); + layer.attn_out_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B, "weight", i), { wo_b_ne0, wo_b_ne1 }, 0); + } + + layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, i), { n_head }, 0); + + if (const auto * meta_ape = ml.get_tensor_meta(tn(LLM_TENSOR_ATTN_COMPRESS_APE, i).str().c_str())) { + layer.attn_compress_ape = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESS_APE, i), { meta_ape->ne[0], meta_ape->ne[1] }, 0); + layer.attn_compress_norm = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESS_NORM, "weight", i), { n_embd_head }, 0); + layer.attn_compress_kv = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESS_KV, "weight", i), { n_embd, meta_ape->ne[0] }, 0); + layer.attn_compress_gate = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESS_GATE, "weight", i), { n_embd, meta_ape->ne[0] }, 0); + } + + if (const auto * meta_indexer_proj = ml.get_tensor_meta(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i).str().c_str())) { + layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), { meta_indexer_proj->ne[0], meta_indexer_proj->ne[1] }, 0); + layer.indexer_attn_q_b = create_tensor( + tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), + { q_lora_rank, hparams.indexer_n_head * hparams.indexer_head_size }, + 0); + + const auto * meta_ape = ml.require_tensor_meta(tn(LLM_TENSOR_INDEXER_COMPRESS_APE, i).str()); + layer.indexer_compress_ape = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESS_APE, i), { meta_ape->ne[0], meta_ape->ne[1] }, 0); + layer.indexer_compress_norm = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESS_NORM, "weight", i), { hparams.indexer_head_size }, 0); + layer.indexer_compress_kv = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESS_KV, "weight", i), { n_embd, meta_ape->ne[0] }, 0); + layer.indexer_compress_gate = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESS_GATE, "weight", i), { n_embd, meta_ape->ne[0] }, 0); + } + + { + const auto * meta_hc_attn_base = ml.get_tensor_meta(tn(LLM_TENSOR_HC_ATTN_BASE, i).str().c_str()); + const auto * meta_hc_attn_fn = ml.get_tensor_meta(tn(LLM_TENSOR_HC_ATTN_FN, i).str().c_str()); + const auto * meta_hc_attn_scale = ml.get_tensor_meta(tn(LLM_TENSOR_HC_ATTN_SCALE, i).str().c_str()); + const auto * meta_hc_ffn_base = ml.get_tensor_meta(tn(LLM_TENSOR_HC_FFN_BASE, i).str().c_str()); + const auto * meta_hc_ffn_fn = ml.get_tensor_meta(tn(LLM_TENSOR_HC_FFN_FN, i).str().c_str()); + const auto * meta_hc_ffn_scale = ml.get_tensor_meta(tn(LLM_TENSOR_HC_FFN_SCALE, i).str().c_str()); + const int64_t hc_pre_out = 2 * hc_mult + hc_mult * hc_mult; + const int64_t hc_attn_base_ne = meta_hc_attn_base ? meta_hc_attn_base->ne[0] : hc_pre_out; + const int64_t hc_attn_fn_ne0 = meta_hc_attn_fn ? meta_hc_attn_fn->ne[0] : n_embd * hc_mult; + const int64_t hc_attn_fn_ne1 = meta_hc_attn_fn ? meta_hc_attn_fn->ne[1] : hc_pre_out; + const int64_t hc_attn_scale_ne = meta_hc_attn_scale ? meta_hc_attn_scale->ne[0] : 3; + const int64_t hc_ffn_base_ne = meta_hc_ffn_base ? meta_hc_ffn_base->ne[0] : hc_pre_out; + const int64_t hc_ffn_fn_ne0 = meta_hc_ffn_fn ? meta_hc_ffn_fn->ne[0] : n_embd * hc_mult; + const int64_t hc_ffn_fn_ne1 = meta_hc_ffn_fn ? meta_hc_ffn_fn->ne[1] : hc_pre_out; + const int64_t hc_ffn_scale_ne = meta_hc_ffn_scale ? meta_hc_ffn_scale->ne[0] : 3; + + layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, i), { hc_attn_base_ne }, 0); + layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, i), { hc_attn_fn_ne0, hc_attn_fn_ne1 }, 0); + layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, i), { hc_attn_scale_ne }, 0); + layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, i), { hc_ffn_base_ne }, 0); + layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, i), { hc_ffn_fn_ne0, hc_ffn_fn_ne1 }, 0); + layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, i), { hc_ffn_scale_ne }, 0); + } + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), { n_embd }, 0); + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert }, 0); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, i), { n_expert }, TENSOR_NOT_REQUIRED); + + if (const auto * meta_tid2eid = ml.get_tensor_meta(tn(LLM_TENSOR_FFN_GATE_TID2EID, i).str().c_str())) { + layer.ffn_gate_tid2eid = create_tensor(tn(LLM_TENSOR_FFN_GATE_TID2EID, i), { meta_tid2eid->ne[0], meta_tid2eid->ne[1] }, 0); + } + + if (n_expert == 0) { + throw std::runtime_error("n_expert must be > 0"); + } + if (n_expert_used == 0) { + throw std::runtime_error("n_expert_used must be > 0"); + } + + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), { n_ff_exp, n_embd, n_expert }, 0); + create_tensor_gate_up_exps(layer, i, n_embd, n_ff_exp, n_expert, 0); + + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), { n_embd, n_ff_exp * n_expert_shared }, 0); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_exp * n_expert_shared, n_embd }, 0); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), { n_embd, n_ff_exp * n_expert_shared }, 0); + } + } break; case LLM_ARCH_PLM: { const int64_t n_embd_head_qk_rope = hparams.n_rot(); @@ -5776,7 +5927,7 @@ bool llama_model::load_tensors(llama_model_loader & ml) { // MoE layers layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert }, flags); - layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), { n_expert }, flags); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, i), { n_expert }, flags); // MoE branch const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used; @@ -5888,7 +6039,7 @@ bool llama_model::load_tensors(llama_model_loader & ml) { layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, flags); } else { layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags); - layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, i), {n_expert}, TENSOR_NOT_REQUIRED); if (n_expert == 0) { throw std::runtime_error("n_expert must be > 0"); @@ -6016,7 +6167,7 @@ bool llama_model::load_tensors(llama_model_loader & ml) { const int64_t n_ff_shexp = hparams.n_ff_shexp; layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert}, 0); - layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert }, 0); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, i), {n_expert }, 0); // MoE branch layer.ffn_latent_down = create_tensor(tn(LLM_TENSOR_FFN_LATENT_DOWN, "weight", i), {n_embd, moe_n_embd}, TENSOR_NOT_REQUIRED); @@ -6144,7 +6295,7 @@ bool llama_model::load_tensors(llama_model_loader & ml) { layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, flags); } else { layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags); - layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED | flags); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, i), {n_expert}, TENSOR_NOT_REQUIRED | flags); if (n_expert == 0) { throw std::runtime_error("n_expert must be > 0"); @@ -6638,7 +6789,7 @@ bool llama_model::load_tensors(llama_model_loader & ml) { const int64_t n_ff_shexp = (hparams.n_ff_shexp ? hparams.n_ff_shexp : n_ff_exp) * n_expert_shared; layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags); - layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED | flags); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, i), {n_expert}, TENSOR_NOT_REQUIRED | flags); layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags); layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, flags); @@ -6694,7 +6845,7 @@ bool llama_model::load_tensors(llama_model_loader & ml) { layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); } else { layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); - layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, i), {n_expert}, TENSOR_NOT_REQUIRED); if (n_expert == 0) { throw std::runtime_error("n_expert must be > 0"); @@ -6785,7 +6936,7 @@ bool llama_model::load_tensors(llama_model_loader & ml) { if (static_cast(i) >= hparams.n_layer_dense_lead) { // MoE layers layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); - layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, i), {n_expert}, 0); // grouped expert weights layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); @@ -6838,7 +6989,7 @@ bool llama_model::load_tensors(llama_model_loader & ml) { int n_ff_exp = hparams.n_ff_exp; layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); - layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, i), {n_expert}, TENSOR_NOT_REQUIRED); layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED); layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), { n_ff_exp, n_embd, n_expert}, 0); layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); @@ -7082,7 +7233,7 @@ bool llama_model::load_tensors(llama_model_loader & ml) { layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, hparams.n_ff_exp, n_expert}, 0); layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {hparams.n_ff_exp, n_embd, n_expert}, 0); layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, hparams.n_ff_exp, n_expert}, 0); - layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, i), {n_expert}, 0); } else { // dense layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0); @@ -7251,7 +7402,7 @@ bool llama_model::load_tensors(llama_model_loader & ml) { layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff, n_expert}, 0); layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff, n_embd, n_expert}, 0); layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff, n_expert}, 0); - layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, i), {n_expert}, 0); } } break; case LLM_ARCH_KIMI_LINEAR: @@ -7383,7 +7534,7 @@ bool llama_model::load_tensors(llama_model_loader & ml) { layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp_actual, n_embd}, TENSOR_NOT_REQUIRED); layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp_actual}, TENSOR_NOT_REQUIRED); - layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, i), {n_expert}, 0); } } } break; @@ -7687,7 +7838,7 @@ bool llama_model::load_tensors(llama_model_loader & ml) { layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED); layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, TENSOR_NOT_REQUIRED); layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED); - layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, i), {n_expert}, TENSOR_NOT_REQUIRED); } } break; case LLM_ARCH_STEP35: @@ -7746,7 +7897,7 @@ bool llama_model::load_tensors(llama_model_loader & ml) { layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED); layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, TENSOR_NOT_REQUIRED); layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED); - layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, i), {n_expert}, TENSOR_NOT_REQUIRED); // shared expert MLP layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, hparams.n_ff_shexp}, TENSOR_NOT_REQUIRED); @@ -8444,6 +8595,15 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, { res = nullptr; } break; + case LLM_ARCH_DEEPSEEK4: + { + res = new llama_memory_deepseek4( + *this, + params.type_k, + cparams.offload_kqv, + cparams.n_ctx_seq, + cparams.n_seq_max); + } break; // Models that need standard caching should rely on recurrent/hybrid // checks default: @@ -8840,6 +9000,10 @@ ggml_cgraph * llama_model::build_graph(const llm_graph_params & params) const { { llm = std::make_unique(*this, params); } break; + case LLM_ARCH_DEEPSEEK4: + { + llm = std::make_unique(*this, params); + } break; case LLM_ARCH_CHATGLM: { llm = std::make_unique(*this, params); @@ -9236,6 +9400,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_DEEPSEEK: case LLM_ARCH_DEEPSEEK2: case LLM_ARCH_DEEPSEEK2OCR: + case LLM_ARCH_DEEPSEEK4: case LLM_ARCH_PLM: case LLM_ARCH_CHATGLM: case LLM_ARCH_GRANITE: diff --git a/src/llama-model.h b/src/llama-model.h index 5f101bd63745..120068e44830 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -484,6 +484,26 @@ 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 + // DeepSeek V4 + struct ggml_tensor * attn_kv_latent = nullptr; + struct ggml_tensor * attn_out_a = nullptr; + struct ggml_tensor * attn_out_b = nullptr; + struct ggml_tensor * attn_compress_ape = nullptr; + struct ggml_tensor * attn_compress_norm = nullptr; + struct ggml_tensor * attn_compress_kv = nullptr; + struct ggml_tensor * attn_compress_gate = nullptr; + struct ggml_tensor * indexer_compress_ape = nullptr; + struct ggml_tensor * indexer_compress_norm = nullptr; + struct ggml_tensor * indexer_compress_kv = nullptr; + struct ggml_tensor * indexer_compress_gate = nullptr; + struct ggml_tensor * hc_attn_base = nullptr; + struct ggml_tensor * hc_attn_fn = nullptr; + struct ggml_tensor * hc_attn_scale = nullptr; + struct ggml_tensor * hc_ffn_base = nullptr; + struct ggml_tensor * hc_ffn_fn = nullptr; + struct ggml_tensor * hc_ffn_scale = nullptr; + struct ggml_tensor * ffn_gate_tid2eid = nullptr; + // gemma4 layer output scale struct ggml_tensor * out_scale = nullptr; @@ -550,6 +570,11 @@ struct llama_model { struct ggml_tensor * per_layer_model_proj = nullptr; struct ggml_tensor * per_layer_proj_norm = nullptr; + // DeepSeek V4 hyper-connection head + struct ggml_tensor * hc_head_base = nullptr; + struct ggml_tensor * hc_head_fn = nullptr; + struct ggml_tensor * hc_head_scale = nullptr; + std::vector layers; //Dense linear projections for SentenceTransformers models like embeddinggemma diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp new file mode 100644 index 000000000000..b1e41bc65b81 --- /dev/null +++ b/src/models/deepseek4.cpp @@ -0,0 +1,957 @@ +#include "models.h" + +#include "llama-impl.h" +#include "llama-memory-deepseek4.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +static bool deepseek4_is_power_of_2(int64_t n) { + return n > 0 && (n & (n - 1)) == 0; +} + +static void deepseek4_fill_hadamard(std::vector & data, int64_t n) { + GGML_ASSERT(deepseek4_is_power_of_2(n)); + + data.assign(n*n, 0.0f); + data[0] = 1.0f / std::sqrt(float(n)); + + for (int64_t s = 1; s < n; s *= 2) { + for (int64_t i = 0; i < s; ++i) { + for (int64_t j = 0; j < s; ++j) { + const float v = data[i*n + j]; + data[(i + s)*n + j ] = v; + data[i*n + j + s] = v; + data[(i + s)*n + j + s] = -v; + } + } + } +} + +class llm_build_deepseek4_inputs : public llm_graph_input_i { +public: + explicit llm_build_deepseek4_inputs(uint32_t n_swa) : n_swa(n_swa) {} + + void set_input(const llama_ubatch * ubatch) override { + GGML_ASSERT(ubatch->n_tokens >= 1); + const uint32_t n_tokens = ubatch->n_tokens; + + auto set_i32_input = [&](ggml_tensor * tensor, auto fn) { + if (!tensor || !tensor->buffer) { + return; + } + + i32_data.resize(tensor->ne[0]); + for (int64_t i = 0; i < tensor->ne[0]; ++i) { + const int32_t p = ubatch->pos ? ubatch->pos[std::min(i, n_tokens - 1)] : 0; + i32_data[i] = fn(p); + } + ggml_backend_tensor_set(tensor, i32_data.data(), 0, ggml_nbytes(tensor)); + }; + + set_i32_input(attn_cache_idx, [&](int32_t p) { return p % (int32_t) n_swa; }); + + set_i32_input(comp_pos_r4, [](int32_t p) { return std::max(0, p + 1 - 4); }); + + set_i32_input(comp_pos_r128, [](int32_t p) { return std::max(0, p + 1 - 128); }); + + set_i32_input(comp_cache_idx_r4, [&](int32_t p) { return (int32_t) n_swa + p / 4; }); + + set_i32_input(indexer_cache_idx_r4, [](int32_t p) { return p / 4; }); + + set_i32_input(comp_cache_idx_r128, [&](int32_t p) { return (int32_t) n_swa + p / 128; }); + + set_i32_input(comp_slot_idx_r4, [](int32_t p) { return 4 + (p % 4); }); + + set_i32_input(comp_slot_idx_r128, [](int32_t p) { return p % 128; }); + + for (ggml_tensor * mask : kq_masks) { + if (!mask || !mask->buffer) { + continue; + } + + const int64_t n_kv = mask->ne[0]; + const int64_t n_q = mask->ne[1]; + f32_data.assign(ggml_nelements(mask), -INFINITY); + for (int64_t iq = 0; iq < n_q; ++iq) { + const int32_t q_pos = ubatch->pos ? ubatch->pos[std::min(iq, n_tokens - 1)] : 0; + for (int64_t ikv = 0; ikv < n_kv; ++ikv) { + if (ikv >= (int64_t) n_swa || ikv <= q_pos) { + f32_data[iq*n_kv + ikv] = 0.0f; + } + } + } + ggml_backend_tensor_set(mask, f32_data.data(), 0, ggml_nbytes(mask)); + } + + if (indexer_hadamard && indexer_hadamard->buffer) { + const int64_t n = indexer_hadamard->ne[0]; + GGML_ASSERT(indexer_hadamard->ne[1] == n); + if (indexer_hadamard_data.empty()) { + deepseek4_fill_hadamard(indexer_hadamard_data, n); + } + ggml_backend_tensor_set(indexer_hadamard, indexer_hadamard_data.data(), 0, ggml_nbytes(indexer_hadamard)); + } + } + + ggml_tensor * attn_cache_idx = nullptr; + ggml_tensor * comp_pos_r4 = nullptr; + ggml_tensor * comp_pos_r128 = nullptr; + ggml_tensor * comp_cache_idx_r4 = nullptr; + ggml_tensor * comp_cache_idx_r128 = nullptr; + ggml_tensor * indexer_cache_idx_r4 = nullptr; + ggml_tensor * comp_slot_idx_r4 = nullptr; + ggml_tensor * comp_slot_idx_r128 = nullptr; + ggml_tensor * indexer_hadamard = nullptr; + std::vector kq_masks; + + std::vector i32_data; + std::vector f32_data; + std::vector indexer_hadamard_data; + + const uint32_t n_swa; +}; + +} // namespace + +llm_build_deepseek4::llm_build_deepseek4(const llama_model & model, const llm_graph_params & params) : + llm_graph_context(params) { + GGML_ASSERT(model.arch == LLM_ARCH_DEEPSEEK4); + GGML_ASSERT(n_tokens >= 1); + + const auto * mctx_cur = dynamic_cast(mctx); + GGML_ASSERT(mctx_cur != nullptr); + GGML_ASSERT(hparams.n_swa > 0); + + // For prefill (n_tokens > 1) we only build a reservation graph on a + // single dummy token; the per-batch decode path runs single-token ubatches. + const bool reserve_only = n_tokens != 1; + const llama_pos start_pos = reserve_only ? 0 : ubatch.pos[0]; + const int64_t work_tokens = reserve_only ? 1 : n_tokens; + GGML_ASSERT(start_pos >= 0); + GGML_ASSERT((uint32_t) start_pos < mctx_cur->get_n_ctx_seq()); + + const int64_t head_dim = hparams.n_embd_head_k(); + const int64_t rope_dim = hparams.n_rot(); + const int64_t nope_dim = head_dim - rope_dim; + const int64_t total_q_dim = head_dim * n_head; + const int64_t hc_mult = model.hc_head_base ? model.hc_head_base->ne[0] : 0; + GGML_ASSERT(hc_mult > 0); + GGML_ASSERT(nope_dim >= 0); + + auto inp_ds4 = std::make_unique(hparams.n_swa); + inp_ds4->attn_cache_idx = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, work_tokens); + ggml_set_input(inp_ds4->attn_cache_idx); + ggml_set_name(inp_ds4->attn_cache_idx, "deepseek4_attn_cache_idx"); + inp_ds4->comp_pos_r4 = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, work_tokens); + ggml_set_input(inp_ds4->comp_pos_r4); + ggml_set_name(inp_ds4->comp_pos_r4, "deepseek4_comp_pos_r4"); + inp_ds4->comp_pos_r128 = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, work_tokens); + ggml_set_input(inp_ds4->comp_pos_r128); + ggml_set_name(inp_ds4->comp_pos_r128, "deepseek4_comp_pos_r128"); + inp_ds4->comp_cache_idx_r4 = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, work_tokens); + ggml_set_input(inp_ds4->comp_cache_idx_r4); + ggml_set_name(inp_ds4->comp_cache_idx_r4, "deepseek4_comp_cache_idx_r4"); + inp_ds4->comp_cache_idx_r128 = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, work_tokens); + ggml_set_input(inp_ds4->comp_cache_idx_r128); + ggml_set_name(inp_ds4->comp_cache_idx_r128, "deepseek4_comp_cache_idx_r128"); + inp_ds4->indexer_cache_idx_r4 = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, work_tokens); + ggml_set_input(inp_ds4->indexer_cache_idx_r4); + ggml_set_name(inp_ds4->indexer_cache_idx_r4, "deepseek4_indexer_cache_idx_r4"); + inp_ds4->comp_slot_idx_r4 = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, work_tokens); + ggml_set_input(inp_ds4->comp_slot_idx_r4); + ggml_set_name(inp_ds4->comp_slot_idx_r4, "deepseek4_comp_slot_idx_r4"); + inp_ds4->comp_slot_idx_r128 = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, work_tokens); + ggml_set_input(inp_ds4->comp_slot_idx_r128); + ggml_set_name(inp_ds4->comp_slot_idx_r128, "deepseek4_comp_slot_idx_r128"); + if (hparams.indexer_head_size > 0 && + hparams.indexer_top_k > 0 && + uint64_t(cparams.n_ctx_seq) > uint64_t(hparams.indexer_top_k) * 4u) { + inp_ds4->indexer_hadamard = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.indexer_head_size, hparams.indexer_head_size); + ggml_set_input(inp_ds4->indexer_hadamard); + ggml_set_name(inp_ds4->indexer_hadamard, "deepseek4_indexer_hadamard"); + } + auto * deepseek4_inputs = static_cast(res->add_input(std::move(inp_ds4))); + + auto scalar_view = [&](ggml_tensor * tensor, int64_t idx) -> ggml_tensor * { + return ggml_view_1d(ctx0, tensor, 1, idx * tensor->nb[0]); + }; + + auto vector_slice = [&](ggml_tensor * tensor, int64_t offset, int64_t len) -> ggml_tensor * { + return ggml_view_1d(ctx0, tensor, len, offset * tensor->nb[0]); + }; + + auto matrix_slice = [&](ggml_tensor * tensor, int64_t offset, int64_t rows, int64_t cols) -> ggml_tensor * { + return ggml_view_2d(ctx0, tensor, rows, cols, rows * tensor->nb[0], offset * tensor->nb[0]); + }; + + auto matrix_block = [&](ggml_tensor * tensor, int64_t row_offset, int64_t col_offset, int64_t rows, int64_t cols) -> ggml_tensor * { + return ggml_view_2d(ctx0, tensor, rows, cols, tensor->nb[1], row_offset * tensor->nb[0] + col_offset * tensor->nb[1]); + }; + + auto compression_ape_rows = [&](ggml_tensor * ape, int64_t comp_dim, int64_t comp_ratio) -> ggml_tensor * { + const int64_t start_mod = start_pos % comp_ratio; + if (start_mod + work_tokens <= comp_ratio) { + return matrix_block(ape, 0, start_mod, comp_dim, work_tokens); + } + if (start_mod != 0 || work_tokens % comp_ratio != 0) { + GGML_ABORT("deepseek4: unsupported multi-window APE slice pos=%d tokens=%" PRId64 " ratio=%" PRId64, + (int) start_pos, work_tokens, comp_ratio); + } + + ggml_tensor * out = nullptr; + const int64_t n_windows = work_tokens / comp_ratio; + for (int64_t iw = 0; iw < n_windows; ++iw) { + ggml_tensor * cur = matrix_block(ape, 0, 0, comp_dim, comp_ratio); + out = out ? ggml_concat(ctx0, out, cur, 1) : cur; + } + return out; + }; + + auto reshape_3d_checked = [&](ggml_tensor * tensor, int64_t ne0, int64_t ne1, int64_t ne2, const char * tag, int il = -1) -> ggml_tensor * { + const int64_t expected = ne0 * ne1 * ne2; + if (ggml_nelements(tensor) != expected) { + GGML_ABORT( + "deepseek4: reshape_3d mismatch in %s layer %d pos %d" + " ne=%" PRId64 " expected=%" PRId64 " target=(%" PRId64 ",%" PRId64 ",%" PRId64 ") tensor=%s", + tag, il, (int) start_pos, ggml_nelements(tensor), expected, ne0, ne1, ne2, + tensor->name[0] ? tensor->name : ""); + } + return ggml_reshape_3d(ctx0, tensor, ne0, ne1, ne2); + }; + + auto reshape_2d_checked = [&](ggml_tensor * tensor, int64_t ne0, int64_t ne1, const char * tag, int il = -1) -> ggml_tensor * { + const int64_t expected = ne0 * ne1; + if (ggml_nelements(tensor) != expected) { + GGML_ABORT( + "deepseek4: reshape_2d mismatch in %s layer %d pos %d" + " ne=%" PRId64 " expected=%" PRId64 " target=(%" PRId64 ",%" PRId64 ") tensor=%s" + " shape=(%" PRId64 ",%" PRId64 ",%" PRId64 ",%" PRId64 ")", + tag, il, (int) start_pos, ggml_nelements(tensor), expected, ne0, ne1, + tensor->name[0] ? tensor->name : "", + tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]); + } + return ggml_reshape_2d(ctx0, tensor, ne0, ne1); + }; + + auto add_eps = [&](ggml_tensor * tensor, float eps) -> ggml_tensor * { + return ggml_clamp(ctx0, tensor, eps, INFINITY); + }; + + auto cont_if_needed = [&](ggml_tensor * tensor) -> ggml_tensor * { + return ggml_is_contiguous(tensor) ? tensor : ggml_cont(ctx0, tensor); + }; + + auto mul_mat_checked = [&](ggml_tensor * a, ggml_tensor * b, const char * tag) -> ggml_tensor * { + if (ggml_is_transposed(a)) { + GGML_ABORT("deepseek4: transposed lhs in %s (%s)", tag, a->name[0] ? a->name : ""); + } + if (b->nb[0] != ggml_type_size(b->type)) { + GGML_ABORT( + "deepseek4: mul_mat rhs layout in %s (%s) nb0=%zu nb1=%zu", + tag, b->name[0] ? b->name : "", b->nb[0], b->nb[1]); + } + return ggml_mul_mat(ctx0, a, b); + }; + + auto repeat_checked = [&](ggml_tensor * src, ggml_tensor * dst, const char * tag) -> ggml_tensor * { + if (src->nb[0] != sizeof(float)) { + GGML_ABORT( + "deepseek4: repeat source layout in %s (%s) nb0=%zu nb1=%zu", + tag, src->name[0] ? src->name : "", src->nb[0], src->nb[1]); + } + if (dst->nb[0] != sizeof(float)) { + GGML_ABORT( + "deepseek4: repeat destination layout in %s (%s) nb0=%zu nb1=%zu", + tag, dst->name[0] ? dst->name : "", dst->nb[0], dst->nb[1]); + } + return ggml_repeat(ctx0, src, dst); + }; + + auto sum_rows_checked = [&](ggml_tensor * src, const char * tag) -> ggml_tensor * { + if (src->nb[0] != sizeof(float)) { + GGML_ABORT( + "deepseek4: sum_rows source layout in %s (%s) nb0=%zu nb1=%zu", + tag, src->name[0] ? src->name : "", src->nb[0], src->nb[1]); + } + return ggml_sum_rows(ctx0, src); + }; + + auto affine = [&](ggml_tensor * tensor, ggml_tensor * scale, ggml_tensor * bias) -> ggml_tensor * { + ggml_tensor * out = ggml_mul(ctx0, tensor, scale); + return ggml_add(ctx0, out, bias); + }; + + auto weighted_sum_hc = [&](ggml_tensor * x_hc, ggml_tensor * weights) -> ggml_tensor * { + // Weighted sum across the HC (Hadamard chunk) dimension via the standard + // mul_mat path: reshape [n_embd, hc_mult] -> transpose -> mul_mat with + // [hc_mult] weights -> [n_embd]. + ggml_tensor * x_mat = cont_if_needed(reshape_2d_checked(x_hc, n_embd, hc_mult, "weighted_sum_hc.x_hc")); + ggml_tensor * x_t = ggml_cont(ctx0, ggml_transpose(ctx0, x_mat)); + return mul_mat_checked(x_t, weights, "weighted_sum_hc"); + }; + + auto sinkhorn = [&](ggml_tensor * comb) -> ggml_tensor * { + // Iterative Sinkhorn normalization for the 4x4 expert combine matrix: + // alternating row/column normalization for a fixed number of iterations. + comb = ggml_soft_max(ctx0, comb); + comb = add_eps(comb, 1e-6f); + + ggml_tensor * col_sum = sum_rows_checked(ggml_cont(ctx0, ggml_transpose(ctx0, comb)), "sinkhorn.col_sum"); + col_sum = add_eps(col_sum, 1e-6f); + comb = ggml_div(ctx0, comb, repeat_checked(ggml_cont(ctx0, ggml_transpose(ctx0, col_sum)), comb, "sinkhorn.col_sum")); + + for (int i = 1; i < 20; ++i) { + ggml_tensor * row_sum = sum_rows_checked(comb, "sinkhorn.row_sum"); + row_sum = add_eps(row_sum, 1e-6f); + comb = ggml_div(ctx0, comb, repeat_checked(row_sum, comb, "sinkhorn.row_sum")); + + col_sum = sum_rows_checked(ggml_cont(ctx0, ggml_transpose(ctx0, comb)), "sinkhorn.col_sum_iter"); + col_sum = add_eps(col_sum, 1e-6f); + comb = ggml_div(ctx0, comb, repeat_checked(ggml_cont(ctx0, ggml_transpose(ctx0, col_sum)), comb, "sinkhorn.col_sum_iter")); + } + + return comb; + }; + + auto hc_pre = [&](ggml_tensor * x_hc, ggml_tensor * hc_fn, ggml_tensor * hc_scale, ggml_tensor * hc_base, int il) { + ggml_tensor * x_flat = cont_if_needed(reshape_2d_checked(x_hc, n_embd * hc_mult, work_tokens, "hc_pre.x_flat", il)); + ggml_tensor * x_norm = ggml_rms_norm(ctx0, x_flat, hparams.f_norm_rms_eps); + cb(x_norm, "hc_norm", il); + + ggml_tensor * mixes = mul_mat_checked(hc_fn, x_norm, "hc_pre.mixes"); + cb(mixes, "hc_mixes", il); + + ggml_tensor * pre = vector_slice(mixes, 0, hc_mult); + ggml_tensor * post = vector_slice(mixes, hc_mult, hc_mult); + ggml_tensor * comb = matrix_slice(mixes, 2 * hc_mult, hc_mult, hc_mult); + if (work_tokens > 1) { + pre = ggml_view_2d(ctx0, mixes, hc_mult, work_tokens, mixes->nb[1], 0); + post = ggml_view_2d(ctx0, mixes, hc_mult, work_tokens, mixes->nb[1], hc_mult * mixes->nb[0]); + comb = ggml_view_3d(ctx0, mixes, hc_mult, hc_mult, work_tokens, + hc_mult * mixes->nb[0], mixes->nb[1], 2 * hc_mult * mixes->nb[0]); + } + + pre = affine(pre, scalar_view(hc_scale, 0), vector_slice(hc_base, 0, hc_mult)); + pre = ggml_sigmoid(ctx0, pre); + pre = add_eps(pre, 1e-6f); + cb(pre, "hc_pre", il); + + post = affine(post, scalar_view(hc_scale, 1), vector_slice(hc_base, hc_mult, hc_mult)); + post = ggml_sigmoid(ctx0, post); + post = ggml_scale(ctx0, post, 2.0f); + cb(post, "hc_post_w", il); + + comb = affine(comb, scalar_view(hc_scale, 2), matrix_slice(hc_base, 2 * hc_mult, hc_mult, hc_mult)); + comb = sinkhorn(comb); + cb(comb, "hc_comb", il); + + ggml_tensor * y = weighted_sum_hc(x_hc, pre); + cb(y, "hc_reduce", il); + + return std::make_tuple(y, post, comb); + }; + + auto hc_post = [&](ggml_tensor * x_single, ggml_tensor * residual_hc, ggml_tensor * post, ggml_tensor * comb, int il) -> ggml_tensor * { + if (work_tokens > 1) { + ggml_tensor * residual_t = ggml_cont(ctx0, ggml_permute(ctx0, residual_hc, 1, 0, 2, 3)); + ggml_tensor * mixed_t = mul_mat_checked(comb, residual_t, "hc_post.mixed_batched"); + ggml_tensor * mixed = ggml_cont(ctx0, ggml_permute(ctx0, mixed_t, 1, 0, 2, 3)); + + ggml_tensor * x_repeat = repeat_checked(reshape_3d_checked(x_single, n_embd, 1, work_tokens, "hc_post.x_batched", il), + residual_hc, "hc_post.x_batched"); + ggml_tensor * post_repeat = repeat_checked(reshape_3d_checked(post, 1, hc_mult, work_tokens, "hc_post.post_batched", il), + residual_hc, "hc_post.post_batched"); + + ggml_tensor * out = ggml_add(ctx0, ggml_mul(ctx0, x_repeat, post_repeat), mixed); + cb(out, "hc_expand", il); + return out; + } + + ggml_tensor * residual = cont_if_needed(reshape_2d_checked(residual_hc, n_embd, hc_mult, "hc_post.residual", il)); + ggml_tensor * residual_t = ggml_cont(ctx0, ggml_transpose(ctx0, residual)); + ggml_tensor * mixed_t = mul_mat_checked(comb, residual_t, "hc_post.mixed"); + ggml_tensor * mixed = ggml_cont(ctx0, ggml_transpose(ctx0, mixed_t)); + + ggml_tensor * x_repeat = repeat_checked(x_single, residual, "hc_post.x"); + ggml_tensor * post_t = reshape_2d_checked(post, 1, hc_mult, "hc_post.post", il); + + ggml_tensor * out = ggml_add(ctx0, ggml_mul(ctx0, x_repeat, post_t), mixed); + cb(out, "hc_expand", il); + + return reshape_3d_checked(out, n_embd, hc_mult, work_tokens, "hc_post.out", il); + }; + + auto hc_head = [&](ggml_tensor * x_hc, ggml_tensor * hc_fn, ggml_tensor * hc_scale, ggml_tensor * hc_base) -> ggml_tensor * { + ggml_tensor * x_flat = cont_if_needed(reshape_2d_checked(x_hc, n_embd * hc_mult, work_tokens, "hc_head.x_flat")); + ggml_tensor * x_norm = ggml_rms_norm(ctx0, x_flat, hparams.f_norm_rms_eps); + ggml_tensor * mixes = mul_mat_checked(hc_fn, x_norm, "hc_head.mixes"); + ggml_tensor * pre = affine(mixes, scalar_view(hc_scale, 0), hc_base); + pre = ggml_sigmoid(ctx0, pre); + pre = add_eps(pre, 1e-6f); + return weighted_sum_hc(x_hc, pre); + }; + + auto build_grouped_out = [&](ggml_tensor * attn_out, const llama_layer & layer, int il) -> ggml_tensor * { + const int64_t group_dim = layer.attn_out_a->ne[0]; + const int64_t n_groups = total_q_dim / group_dim; + const int64_t o_rank = layer.attn_out_b->ne[0] / n_groups; + + GGML_ASSERT(group_dim > 0); + GGML_ASSERT(n_groups > 0); + GGML_ASSERT(layer.attn_out_b->ne[0] == n_groups * o_rank); + + ggml_tensor * grouped = nullptr; + for (int64_t g = 0; g < n_groups; ++g) { + ggml_tensor * xg = ggml_view_2d(ctx0, attn_out, group_dim, work_tokens, attn_out->nb[1], g * group_dim * attn_out->nb[0]); + ggml_tensor * wg = ggml_view_2d(ctx0, layer.attn_out_a, group_dim, o_rank, layer.attn_out_a->nb[1], g * o_rank * layer.attn_out_a->nb[1]); + ggml_tensor * og = mul_mat_checked(wg, xg, "build_grouped_out.group"); + cb(og, "attn_group_out", il); + grouped = grouped ? ggml_concat(ctx0, grouped, og, 0) : og; + } + + ggml_tensor * out = mul_mat_checked(layer.attn_out_b, grouped, "build_grouped_out.out"); + cb(out, "attn_out_proj", il); + return out; + }; + + auto build_expert_mix = [&](ggml_tensor * cur_ffn, ggml_tensor * selected_experts, ggml_tensor * weights, const llama_layer & layer, int il) -> ggml_tensor * { + const int64_t mix_tokens = cur_ffn->ne[1]; + ggml_tensor * cur_experts_in = reshape_3d_checked(cur_ffn, n_embd, 1, mix_tokens, "build_expert_mix.cur_ffn", il); + ggml_tensor * gate = nullptr; + ggml_tensor * up = nullptr; + + if (layer.ffn_gate_up_exps) { + ggml_tensor * gate_up = build_lora_mm_id(layer.ffn_gate_up_exps, cur_experts_in, selected_experts); + cb(gate_up, "ffn_moe_gate_up", il); + + const int64_t n_ff = gate_up->ne[0] / 2; + gate = ggml_view_3d(ctx0, gate_up, n_ff, gate_up->ne[1], gate_up->ne[2], gate_up->nb[1], gate_up->nb[2], 0); + up = ggml_view_3d(ctx0, gate_up, n_ff, gate_up->ne[1], gate_up->ne[2], gate_up->nb[1], gate_up->nb[2], n_ff * gate_up->nb[0]); + } else { + gate = build_lora_mm_id(layer.ffn_gate_exps, cur_experts_in, selected_experts); + up = build_lora_mm_id(layer.ffn_up_exps, cur_experts_in, selected_experts); + cb(gate, "ffn_moe_gate", il); + cb(up, "ffn_moe_up", il); + } + + const float swiglu_limit = hparams.swiglu_clamp_exp[il]; + if (swiglu_limit > 1e-6f) { + gate = ggml_clamp(ctx0, gate, -INFINITY, swiglu_limit); + up = ggml_clamp(ctx0, up, -swiglu_limit, swiglu_limit); + cb(gate, "ffn_moe_gate_clamped", il); + cb(up, "ffn_moe_up_clamped", il); + } + + ggml_tensor * act = ggml_swiglu_split(ctx0, gate, up); + cb(act, "ffn_moe_swiglu", il); + + ggml_tensor * experts = build_lora_mm_id(layer.ffn_down_exps, act, selected_experts); + experts = ggml_mul(ctx0, experts, weights); + cb(experts, "ffn_moe_down", il); + + ggml_tensor * experts_by_id = ggml_cont(ctx0, ggml_permute(ctx0, experts, 1, 0, 2, 3)); + ggml_tensor * out = sum_rows_checked(experts_by_id, "build_expert_mix.sum"); + out = reshape_3d_checked(out, 1, n_embd, mix_tokens, "build_expert_mix.sum_out", il); + out = reshape_2d_checked(out, n_embd, mix_tokens, "build_expert_mix.out", il); + + cb(out, "ffn_moe_out", il); + return out; + }; + + ggml_tensor * inpL = build_inp_embd(model.tok_embd); + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_tokens = res->get_inp_tokens(); + if (reserve_only) { + inpL = ggml_cont(ctx0, ggml_view_2d(ctx0, inpL, n_embd, 1, inpL->nb[1], 0)); + inp_pos = ggml_view_1d(ctx0, inp_pos, 1, 0); + if (inp_tokens) { + inp_tokens = ggml_view_1d(ctx0, inp_tokens, 1, 0); + } + } + GGML_UNUSED(inp_tokens); + + auto build_moe_v4 = [&](ggml_tensor * cur_ffn, ggml_tensor * inp_tokens_local, const llama_layer & layer, int il) -> ggml_tensor * { + const int64_t moe_tokens = cur_ffn->ne[1]; + ggml_tensor * scores = build_lora_mm(layer.ffn_gate_inp, cur_ffn); + scores = ggml_softplus(ctx0, scores); + scores = ggml_sqrt(ctx0, scores); + cb(scores, "ffn_scores", il); + + ggml_tensor * selection = scores; + if (layer.ffn_gate_tid2eid) { + ggml_tensor * hash_selected = ggml_get_rows(ctx0, layer.ffn_gate_tid2eid, inp_tokens_local); + ggml_tensor * score3d = reshape_3d_checked(scores, 1, n_expert, moe_tokens, "build_moe_v4.scores_hash", il); + ggml_tensor * selected_scores = ggml_get_rows(ctx0, score3d, hash_selected); + selection = ggml_set_rows(ctx0, ggml_fill(ctx0, score3d, -INFINITY), selected_scores, hash_selected); + selection = reshape_2d_checked(selection, n_expert, moe_tokens, "build_moe_v4.selection", il); + cb(selection, "ffn_hash_scores", il); + } else if (layer.ffn_exp_probs_b) { + selection = ggml_add(ctx0, scores, layer.ffn_exp_probs_b); + cb(selection, "ffn_biased_scores", il); + } + + ggml_tensor * selected_experts = ggml_top_k(ctx0, selection, n_expert_used); + cb(selected_experts, "ffn_topk", il); + + ggml_tensor * weights = ggml_get_rows(ctx0, reshape_3d_checked(scores, 1, n_expert, moe_tokens, "build_moe_v4.scores", il), selected_experts); + weights = reshape_2d_checked(weights, n_expert_used, moe_tokens, "build_moe_v4.weights_2d", il); + ggml_tensor * weights_sum = sum_rows_checked(weights, "build_moe_v4.weights_sum"); + weights_sum = ggml_clamp(ctx0, weights_sum, 6.103515625e-5f, INFINITY); + weights = ggml_div(ctx0, weights, weights_sum); + if (hparams.expert_weights_scale != 1.0f) { + weights = ggml_scale(ctx0, weights, hparams.expert_weights_scale); + } + weights = reshape_3d_checked(weights, 1, n_expert_used, moe_tokens, "build_moe_v4.weights", il); + cb(weights, "ffn_weights", il); + + return build_expert_mix(cur_ffn, selected_experts, weights, layer, il); + }; + + auto build_attn_v4 = [&](ggml_tensor * cur_attn, const llama_layer & layer, int il) -> ggml_tensor * { + const int64_t comp_ratio = layer.attn_compress_ape ? layer.attn_compress_ape->ne[1] : 0; + const float layer_freq_base = layer.attn_compress_ape ? hparams.rope_freq_base_train_swa : hparams.rope_freq_base_train; + const float layer_freq_scale = layer.attn_compress_ape ? hparams.rope_freq_scale_train_swa : 1.0f; + const float layer_ext_factor = layer.attn_compress_ape ? 1.0f : 0.0f; + const float layer_attn_factor = layer.attn_compress_ape && layer_freq_scale != 1.0f ? + 1.0f / (1.0f + 0.1f * std::log(1.0f / layer_freq_scale)) : 1.0f; + const float layer_beta_fast = layer.attn_compress_ape ? hparams.yarn_beta_fast : 0.0f; + const float layer_beta_slow = layer.attn_compress_ape ? hparams.yarn_beta_slow : 0.0f; + const int32_t layer_n_ctx_orig = layer.attn_compress_ape ? hparams.n_ctx_orig_yarn : 0; + + ggml_tensor * q_base = mul_mat_checked(layer.wq_a, cur_attn, "build_attn_v4.wq_a"); + q_base = build_norm(q_base, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + ggml_tensor * q = mul_mat_checked(layer.wq_b, q_base, "build_attn_v4.wq_b"); + q = reshape_3d_checked(q, head_dim, n_head, work_tokens, "build_attn_v4.q", il); + q = ggml_rms_norm(ctx0, q, hparams.f_norm_rms_eps); + cb(q, "q_proj", il); + + ggml_tensor * q_nope = ggml_view_3d(ctx0, q, nope_dim, n_head, work_tokens, q->nb[1], q->nb[2], 0); + ggml_tensor * q_pe = ggml_view_3d(ctx0, q, rope_dim, n_head, work_tokens, q->nb[1], q->nb[2], nope_dim * q->nb[0]); + q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, rope_dim, rope_type, layer_n_ctx_orig, layer_freq_base, layer_freq_scale, + layer_ext_factor, layer_attn_factor, layer_beta_fast, layer_beta_slow); + ggml_tensor * q_states = ggml_concat(ctx0, q_nope, q_pe, 0); + cb(q_states, "q_states", il); + + ggml_tensor * kv = mul_mat_checked(layer.attn_kv_latent, cur_attn, "build_attn_v4.kv_latent"); + kv = build_norm(kv, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + kv = reshape_3d_checked(kv, head_dim, 1, work_tokens, "build_attn_v4.kv", il); + cb(kv, "kv_latent", il); + + ggml_tensor * k_nope = ggml_view_3d(ctx0, kv, nope_dim, 1, work_tokens, kv->nb[1], kv->nb[2], 0); + ggml_tensor * k_pe = ggml_view_3d(ctx0, kv, rope_dim, 1, work_tokens, kv->nb[1], kv->nb[2], nope_dim * kv->nb[0]); + k_nope = cont_if_needed(k_nope); + k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, rope_dim, rope_type, layer_n_ctx_orig, layer_freq_base, layer_freq_scale, + layer_ext_factor, layer_attn_factor, layer_beta_fast, layer_beta_slow); + ggml_tensor * k_states = ggml_concat(ctx0, k_nope, k_pe, 0); + ggml_tensor * k_flat = cont_if_needed(reshape_2d_checked(k_states, head_dim, work_tokens, "build_attn_v4.k_flat", il)); + + const auto & state = mctx_cur->get_layer(il); + ggml_tensor * updated_cache = ggml_set_rows(ctx0, state.attn_kv, k_flat, deepseek4_inputs->attn_cache_idx); + ggml_tensor * updated_attn_comp_kv_state = state.attn_comp_kv_state; + ggml_tensor * updated_attn_comp_score_state = state.attn_comp_score_state; + ggml_tensor * updated_indexer_kv = state.indexer_kv; + ggml_tensor * updated_indexer_comp_kv_state = state.indexer_comp_kv_state; + ggml_tensor * updated_indexer_comp_score_state = state.indexer_comp_score_state; + + if (comp_ratio > 0) { + GGML_ASSERT(state.attn_comp_kv_state != nullptr); + GGML_ASSERT(state.attn_comp_score_state != nullptr); + + const int64_t comp_dim = layer.attn_compress_ape->ne[0]; + const int64_t comp_slots = comp_dim / head_dim; + const bool overlap = comp_slots > 1; + const bool should_compress = ((start_pos + work_tokens) % comp_ratio) == 0; + const bool multiwindow_r4 = + comp_ratio == 4 && overlap && work_tokens > comp_ratio && + (start_pos % comp_ratio) == 0 && (work_tokens % comp_ratio) == 0; + + ggml_tensor * comp_kv = mul_mat_checked(layer.attn_compress_kv, cur_attn, "build_attn_v4.comp_kv"); + ggml_tensor * comp_score = mul_mat_checked(layer.attn_compress_gate, cur_attn, "build_attn_v4.comp_score"); + comp_kv = ggml_cont(ctx0, ggml_cast(ctx0, comp_kv, GGML_TYPE_F32)); + comp_score = ggml_cont(ctx0, ggml_cast(ctx0, comp_score, GGML_TYPE_F32)); + + ggml_tensor * ape_row = compression_ape_rows(layer.attn_compress_ape, comp_dim, comp_ratio); + comp_score = ggml_cont(ctx0, ggml_add(ctx0, comp_score, ape_row)); + cb(comp_score, "attn_comp_score", il); + + ggml_tensor * comp_slot_idx = nullptr; + if (comp_ratio == 4) { + comp_slot_idx = deepseek4_inputs->comp_slot_idx_r4; + } else if (comp_ratio == 128) { + comp_slot_idx = deepseek4_inputs->comp_slot_idx_r128; + } else { + GGML_ABORT("deepseek4: unsupported compress ratio %" PRId64, comp_ratio); + } + + if (!multiwindow_r4) { + updated_attn_comp_kv_state = ggml_set_rows(ctx0, state.attn_comp_kv_state, comp_kv, comp_slot_idx); + updated_attn_comp_score_state = ggml_set_rows(ctx0, state.attn_comp_score_state, comp_score, comp_slot_idx); + } + + if (should_compress) { + ggml_tensor * comp_pos = nullptr; + ggml_tensor * comp_cache_idx = nullptr; + if (comp_ratio == 4) { + comp_pos = deepseek4_inputs->comp_pos_r4; + comp_cache_idx = deepseek4_inputs->comp_cache_idx_r4; + } else if (comp_ratio == 128) { + comp_pos = deepseek4_inputs->comp_pos_r128; + comp_cache_idx = deepseek4_inputs->comp_cache_idx_r128; + } else { + GGML_ABORT("deepseek4: unsupported compress ratio %" PRId64, comp_ratio); + } + + const int64_t n_comp_windows = multiwindow_r4 ? work_tokens / comp_ratio : 1; + ggml_tensor * final_carry_kv = nullptr; + ggml_tensor * final_carry_score = nullptr; + for (int64_t iw = 0; iw < n_comp_windows; ++iw) { + ggml_tensor * comp_kv_slots = nullptr; + ggml_tensor * comp_score_slots = nullptr; + + if (multiwindow_r4) { + ggml_tensor * kv_prev = iw == 0 ? + matrix_block(state.attn_comp_kv_state, 0, 0, head_dim, comp_ratio) : + matrix_block(comp_kv, 0, (iw - 1) * comp_ratio, head_dim, comp_ratio); + ggml_tensor * kv_cur = matrix_block(comp_kv, head_dim, iw * comp_ratio, head_dim, comp_ratio); + ggml_tensor * score_prev = iw == 0 ? + matrix_block(state.attn_comp_score_state, 0, 0, head_dim, comp_ratio) : + matrix_block(comp_score, 0, (iw - 1) * comp_ratio, head_dim, comp_ratio); + ggml_tensor * score_cur = matrix_block(comp_score, head_dim, iw * comp_ratio, head_dim, comp_ratio); + + comp_kv_slots = ggml_concat(ctx0, kv_prev, kv_cur, 1); + comp_score_slots = ggml_concat(ctx0, score_prev, score_cur, 1); + final_carry_kv = matrix_block(comp_kv, 0, iw * comp_ratio, comp_dim, comp_ratio); + final_carry_score = matrix_block(comp_score, 0, iw * comp_ratio, comp_dim, comp_ratio); + } else if (overlap) { + ggml_tensor * kv_prev = matrix_block(updated_attn_comp_kv_state, 0, 0, head_dim, comp_ratio); + ggml_tensor * kv_cur = matrix_block(updated_attn_comp_kv_state, head_dim, comp_ratio, head_dim, comp_ratio); + ggml_tensor * score_prev = matrix_block(updated_attn_comp_score_state, 0, 0, head_dim, comp_ratio); + ggml_tensor * score_cur = matrix_block(updated_attn_comp_score_state, head_dim, comp_ratio, head_dim, comp_ratio); + + comp_kv_slots = ggml_concat(ctx0, kv_prev, kv_cur, 1); + comp_score_slots = ggml_concat(ctx0, score_prev, score_cur, 1); + final_carry_kv = matrix_block(updated_attn_comp_kv_state, 0, comp_ratio, comp_dim, comp_ratio); + final_carry_score = matrix_block(updated_attn_comp_score_state, 0, comp_ratio, comp_dim, comp_ratio); + } else { + comp_kv_slots = updated_attn_comp_kv_state; + comp_score_slots = updated_attn_comp_score_state; + } + + ggml_tensor * comp_kv_seq = ggml_cont(ctx0, ggml_transpose(ctx0, comp_kv_slots)); + ggml_tensor * comp_score_seq = ggml_cont(ctx0, ggml_transpose(ctx0, comp_score_slots)); + ggml_tensor * comp_weights = ggml_soft_max(ctx0, comp_score_seq); + ggml_tensor * comp_weighted = ggml_mul(ctx0, comp_kv_seq, comp_weights); + ggml_tensor * comp_flat = sum_rows_checked(comp_weighted, "build_attn_v4.comp_sum"); + comp_flat = ggml_cont(ctx0, ggml_transpose(ctx0, comp_flat)); + comp_flat = build_norm(comp_flat, layer.attn_compress_norm, nullptr, LLM_NORM_RMS, il); + if (ggml_nelements(comp_flat) != head_dim) { + GGML_ABORT( + "deepseek4: comp_flat reshape mismatch at layer %d pos %d ratio %" PRId64 + " ne=%" PRId64 " expected=%" PRId64, + il, (int) start_pos, comp_ratio, ggml_nelements(comp_flat), head_dim); + } + + ggml_tensor * comp_states = reshape_3d_checked(comp_flat, head_dim, 1, 1, "build_attn_v4.comp_states", il); + ggml_tensor * comp_nope = ggml_view_3d(ctx0, comp_states, nope_dim, 1, 1, comp_states->nb[1], comp_states->nb[2], 0); + ggml_tensor * comp_pe = ggml_view_3d(ctx0, comp_states, rope_dim, 1, 1, comp_states->nb[1], comp_states->nb[2], nope_dim * comp_states->nb[0]); + comp_nope = cont_if_needed(comp_nope); + + const int64_t token_in_ubatch = multiwindow_r4 ? (iw + 1) * comp_ratio - 1 : work_tokens - 1; + ggml_tensor * comp_pos_i = ggml_view_1d(ctx0, comp_pos, 1, token_in_ubatch * comp_pos->nb[0]); + ggml_tensor * comp_cache_idx_i = ggml_view_1d(ctx0, comp_cache_idx, 1, token_in_ubatch * comp_cache_idx->nb[0]); + + comp_pe = ggml_rope_ext(ctx0, comp_pe, comp_pos_i, nullptr, rope_dim, rope_type, layer_n_ctx_orig, layer_freq_base, layer_freq_scale, + layer_ext_factor, layer_attn_factor, layer_beta_fast, layer_beta_slow); + comp_states = ggml_concat(ctx0, comp_nope, comp_pe, 0); + comp_flat = cont_if_needed(reshape_2d_checked(comp_states, head_dim, 1, "build_attn_v4.comp_flat", il)); + cb(comp_flat, "attn_comp_cache", il); + + updated_cache = ggml_set_rows(ctx0, updated_cache, comp_flat, comp_cache_idx_i); + } + + if (overlap) { + // HF seeds the next overlapping current window with the just-compressed window; new tokens overwrite it slot by slot. + updated_attn_comp_kv_state = ggml_concat(ctx0, final_carry_kv, final_carry_kv, 1); + updated_attn_comp_score_state = ggml_concat(ctx0, final_carry_score, final_carry_score, 1); + } + } + + ggml_build_forward_expand(gf, ggml_cpy(ctx0, updated_attn_comp_kv_state, state.attn_comp_kv_state)); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, updated_attn_comp_score_state, state.attn_comp_score_state)); + } + + const bool indexer_reaches_topk = + hparams.indexer_top_k > 0 && + comp_ratio > 0 && + uint64_t(cparams.n_ctx_seq) > uint64_t(hparams.indexer_top_k) * uint64_t(comp_ratio); + const bool has_indexer = + comp_ratio == 4 && + indexer_reaches_topk && + layer.indexer_proj != nullptr && + layer.indexer_attn_q_b != nullptr && + layer.indexer_compress_ape != nullptr && + layer.indexer_compress_norm != nullptr && + layer.indexer_compress_kv != nullptr && + layer.indexer_compress_gate != nullptr && + state.indexer_kv != nullptr && + state.indexer_comp_kv_state != nullptr && + state.indexer_comp_score_state != nullptr && + deepseek4_inputs->indexer_hadamard != nullptr; + + if (has_indexer) { + const int64_t indexer_head_dim = hparams.indexer_head_size; + const int64_t indexer_nope_dim = indexer_head_dim - rope_dim; + GGML_ASSERT(indexer_nope_dim >= 0); + + const int64_t indexer_comp_dim = layer.indexer_compress_ape->ne[0]; + const int64_t indexer_comp_slots = indexer_comp_dim / indexer_head_dim; + const bool indexer_overlap = indexer_comp_slots > 1; + const bool should_compress = ((start_pos + work_tokens) % comp_ratio) == 0; + const bool multiwindow_r4 = + indexer_overlap && work_tokens > comp_ratio && + (start_pos % comp_ratio) == 0 && (work_tokens % comp_ratio) == 0; + + ggml_tensor * indexer_comp_kv = mul_mat_checked(layer.indexer_compress_kv, cur_attn, "build_attn_v4.indexer_comp_kv"); + ggml_tensor * indexer_comp_score = mul_mat_checked(layer.indexer_compress_gate, cur_attn, "build_attn_v4.indexer_comp_score"); + indexer_comp_kv = ggml_cont(ctx0, ggml_cast(ctx0, indexer_comp_kv, GGML_TYPE_F32)); + indexer_comp_score = ggml_cont(ctx0, ggml_cast(ctx0, indexer_comp_score, GGML_TYPE_F32)); + + ggml_tensor * indexer_ape_row = compression_ape_rows(layer.indexer_compress_ape, indexer_comp_dim, comp_ratio); + indexer_comp_score = ggml_cont(ctx0, ggml_add(ctx0, indexer_comp_score, indexer_ape_row)); + cb(indexer_comp_score, "indexer_comp_score", il); + + if (!multiwindow_r4) { + updated_indexer_comp_kv_state = ggml_set_rows(ctx0, state.indexer_comp_kv_state, indexer_comp_kv, deepseek4_inputs->comp_slot_idx_r4); + updated_indexer_comp_score_state = ggml_set_rows(ctx0, state.indexer_comp_score_state, indexer_comp_score, deepseek4_inputs->comp_slot_idx_r4); + } + + if (should_compress) { + ggml_tensor * indexer_comp_pos = deepseek4_inputs->comp_pos_r4; + ggml_tensor * indexer_cache_idx = deepseek4_inputs->indexer_cache_idx_r4; + + const int64_t n_comp_windows = multiwindow_r4 ? work_tokens / comp_ratio : 1; + ggml_tensor * final_carry_kv = nullptr; + ggml_tensor * final_carry_score = nullptr; + for (int64_t iw = 0; iw < n_comp_windows; ++iw) { + ggml_tensor * indexer_comp_kv_slots = nullptr; + ggml_tensor * indexer_comp_score_slots = nullptr; + + if (multiwindow_r4) { + ggml_tensor * kv_prev = iw == 0 ? + matrix_block(state.indexer_comp_kv_state, 0, 0, indexer_head_dim, comp_ratio) : + matrix_block(indexer_comp_kv, 0, (iw - 1) * comp_ratio, indexer_head_dim, comp_ratio); + ggml_tensor * kv_cur = matrix_block(indexer_comp_kv, indexer_head_dim, iw * comp_ratio, indexer_head_dim, comp_ratio); + ggml_tensor * score_prev = iw == 0 ? + matrix_block(state.indexer_comp_score_state, 0, 0, indexer_head_dim, comp_ratio) : + matrix_block(indexer_comp_score, 0, (iw - 1) * comp_ratio, indexer_head_dim, comp_ratio); + ggml_tensor * score_cur = matrix_block(indexer_comp_score, indexer_head_dim, iw * comp_ratio, indexer_head_dim, comp_ratio); + + indexer_comp_kv_slots = ggml_concat(ctx0, kv_prev, kv_cur, 1); + indexer_comp_score_slots = ggml_concat(ctx0, score_prev, score_cur, 1); + final_carry_kv = matrix_block(indexer_comp_kv, 0, iw * comp_ratio, indexer_comp_dim, comp_ratio); + final_carry_score = matrix_block(indexer_comp_score, 0, iw * comp_ratio, indexer_comp_dim, comp_ratio); + } else if (indexer_overlap) { + ggml_tensor * kv_prev = matrix_block(updated_indexer_comp_kv_state, 0, 0, indexer_head_dim, comp_ratio); + ggml_tensor * kv_cur = matrix_block(updated_indexer_comp_kv_state, indexer_head_dim, comp_ratio, indexer_head_dim, comp_ratio); + ggml_tensor * score_prev = matrix_block(updated_indexer_comp_score_state, 0, 0, indexer_head_dim, comp_ratio); + ggml_tensor * score_cur = matrix_block(updated_indexer_comp_score_state, indexer_head_dim, comp_ratio, indexer_head_dim, comp_ratio); + + indexer_comp_kv_slots = ggml_concat(ctx0, kv_prev, kv_cur, 1); + indexer_comp_score_slots = ggml_concat(ctx0, score_prev, score_cur, 1); + final_carry_kv = matrix_block(updated_indexer_comp_kv_state, 0, comp_ratio, indexer_comp_dim, comp_ratio); + final_carry_score = matrix_block(updated_indexer_comp_score_state, 0, comp_ratio, indexer_comp_dim, comp_ratio); + } else { + indexer_comp_kv_slots = updated_indexer_comp_kv_state; + indexer_comp_score_slots = updated_indexer_comp_score_state; + } + + ggml_tensor * indexer_comp_kv_seq = ggml_cont(ctx0, ggml_transpose(ctx0, indexer_comp_kv_slots)); + ggml_tensor * indexer_comp_score_seq = ggml_cont(ctx0, ggml_transpose(ctx0, indexer_comp_score_slots)); + ggml_tensor * indexer_comp_weights = ggml_soft_max(ctx0, indexer_comp_score_seq); + ggml_tensor * indexer_comp_weighted = ggml_mul(ctx0, indexer_comp_kv_seq, indexer_comp_weights); + ggml_tensor * indexer_comp_flat = sum_rows_checked(indexer_comp_weighted, "build_attn_v4.indexer_comp_sum"); + indexer_comp_flat = ggml_cont(ctx0, ggml_transpose(ctx0, indexer_comp_flat)); + indexer_comp_flat = build_norm(indexer_comp_flat, layer.indexer_compress_norm, nullptr, LLM_NORM_RMS, il); + + ggml_tensor * indexer_comp_states = reshape_3d_checked(indexer_comp_flat, indexer_head_dim, 1, 1, "build_attn_v4.indexer_comp_states", il); + ggml_tensor * indexer_comp_nope = ggml_view_3d(ctx0, indexer_comp_states, indexer_nope_dim, 1, 1, indexer_comp_states->nb[1], indexer_comp_states->nb[2], 0); + ggml_tensor * indexer_comp_pe = ggml_view_3d(ctx0, indexer_comp_states, rope_dim, 1, 1, indexer_comp_states->nb[1], indexer_comp_states->nb[2], indexer_nope_dim * indexer_comp_states->nb[0]); + + const int64_t token_in_ubatch = multiwindow_r4 ? (iw + 1) * comp_ratio - 1 : work_tokens - 1; + ggml_tensor * indexer_comp_pos_i = ggml_view_1d(ctx0, indexer_comp_pos, 1, token_in_ubatch * indexer_comp_pos->nb[0]); + ggml_tensor * indexer_cache_idx_i = ggml_view_1d(ctx0, indexer_cache_idx, 1, token_in_ubatch * indexer_cache_idx->nb[0]); + + indexer_comp_pe = ggml_rope_ext(ctx0, indexer_comp_pe, indexer_comp_pos_i, nullptr, rope_dim, rope_type, + layer_n_ctx_orig, layer_freq_base, layer_freq_scale, layer_ext_factor, layer_attn_factor, layer_beta_fast, layer_beta_slow); + indexer_comp_states = ggml_concat(ctx0, indexer_comp_nope, indexer_comp_pe, 0); + indexer_comp_flat = cont_if_needed(reshape_2d_checked(indexer_comp_states, indexer_head_dim, 1, "build_attn_v4.indexer_comp_flat", il)); + indexer_comp_flat = ggml_mul_mat(ctx0, deepseek4_inputs->indexer_hadamard, indexer_comp_flat); + indexer_comp_flat = cont_if_needed(indexer_comp_flat); + cb(indexer_comp_flat, "indexer_comp_cache", il); + + updated_indexer_kv = ggml_set_rows(ctx0, updated_indexer_kv, indexer_comp_flat, indexer_cache_idx_i); + } + + if (indexer_overlap) { + // HF seeds the next overlapping current window with the just-compressed window; new tokens overwrite it slot by slot. + updated_indexer_comp_kv_state = ggml_concat(ctx0, final_carry_kv, final_carry_kv, 1); + updated_indexer_comp_score_state = ggml_concat(ctx0, final_carry_score, final_carry_score, 1); + } + } + + ggml_build_forward_expand(gf, ggml_cpy(ctx0, updated_indexer_comp_kv_state, state.indexer_comp_kv_state)); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, updated_indexer_comp_score_state, state.indexer_comp_score_state)); + if (should_compress) { + ggml_build_forward_expand(gf, ggml_cpy(ctx0, updated_indexer_kv, state.indexer_kv)); + } + } + + ggml_build_forward_expand(gf, ggml_cpy(ctx0, updated_cache, state.attn_kv)); + + const int64_t n_kv = std::min(start_pos + work_tokens, hparams.n_swa); + ggml_tensor * kv_prefix = ggml_view_2d(ctx0, updated_cache, head_dim, n_kv, updated_cache->nb[1], 0); + kv_prefix = ggml_cast(ctx0, kv_prefix, GGML_TYPE_F32); + int64_t n_comp_attn = comp_ratio > 0 ? (start_pos + work_tokens) / comp_ratio : 0; + if (comp_ratio > 0) { + const int64_t n_comp = (start_pos + work_tokens) / comp_ratio; + if (n_comp > 0) { + ggml_tensor * comp_prefix = ggml_view_2d(ctx0, updated_cache, head_dim, n_comp, updated_cache->nb[1], hparams.n_swa * updated_cache->nb[1]); + if (has_indexer && hparams.indexer_top_k > 0 && n_comp > hparams.indexer_top_k) { + const int64_t indexer_head_dim = hparams.indexer_head_size; + const int64_t indexer_nope_dim = indexer_head_dim - rope_dim; + + ggml_tensor * indexer_q = mul_mat_checked(layer.indexer_attn_q_b, q_base, "build_attn_v4.indexer_q"); + indexer_q = reshape_3d_checked(indexer_q, indexer_head_dim, hparams.indexer_n_head, work_tokens, "build_attn_v4.indexer_q_3d", il); + ggml_tensor * indexer_q_nope = ggml_view_3d(ctx0, indexer_q, indexer_nope_dim, hparams.indexer_n_head, work_tokens, indexer_q->nb[1], indexer_q->nb[2], 0); + ggml_tensor * indexer_q_pe = ggml_view_3d(ctx0, indexer_q, rope_dim, hparams.indexer_n_head, work_tokens, indexer_q->nb[1], indexer_q->nb[2], indexer_nope_dim * indexer_q->nb[0]); + indexer_q_pe = ggml_rope_ext(ctx0, indexer_q_pe, inp_pos, nullptr, rope_dim, rope_type, + layer_n_ctx_orig, layer_freq_base, layer_freq_scale, layer_ext_factor, layer_attn_factor, layer_beta_fast, layer_beta_slow); + indexer_q = ggml_concat(ctx0, indexer_q_nope, indexer_q_pe, 0); + indexer_q = cont_if_needed(reshape_2d_checked(indexer_q, indexer_head_dim, hparams.indexer_n_head, "build_attn_v4.indexer_q", il)); + indexer_q = ggml_mul_mat(ctx0, deepseek4_inputs->indexer_hadamard, indexer_q); + indexer_q = cont_if_needed(indexer_q); + cb(indexer_q, "indexer_q", il); + + ggml_tensor * indexer_kv_prefix = ggml_view_2d(ctx0, updated_indexer_kv, indexer_head_dim, n_comp, updated_indexer_kv->nb[1], 0); + ggml_tensor * index_scores = ggml_mul_mat(ctx0, indexer_kv_prefix, indexer_q); + index_scores = ggml_relu(ctx0, index_scores); + + ggml_tensor * index_weights = mul_mat_checked(layer.indexer_proj, cur_attn, "build_attn_v4.indexer_weights"); + const float index_scale = 1.0f / std::sqrt(float(indexer_head_dim)) / std::sqrt(float(hparams.indexer_n_head)); + index_weights = ggml_scale(ctx0, index_weights, index_scale); + index_weights = reshape_2d_checked(index_weights, 1, hparams.indexer_n_head, "build_attn_v4.index_weights", il); + index_scores = ggml_mul(ctx0, index_scores, index_weights); + index_scores = ggml_cont(ctx0, ggml_transpose(ctx0, index_scores)); + index_scores = sum_rows_checked(index_scores, "build_attn_v4.index_scores"); + index_scores = reshape_2d_checked(index_scores, n_comp, 1, "build_attn_v4.index_scores", il); + cb(index_scores, "index_scores", il); + + ggml_tensor * selected_comp = ggml_argsort_top_k(ctx0, index_scores, hparams.indexer_top_k); + cb(selected_comp, "index_topk", il); + comp_prefix = ggml_get_rows(ctx0, comp_prefix, selected_comp); + n_comp_attn = hparams.indexer_top_k; + } + comp_prefix = ggml_cast(ctx0, comp_prefix, GGML_TYPE_F32); + kv_prefix = ggml_concat(ctx0, kv_prefix, comp_prefix, 1); + } + } + const int64_t n_kv_total = n_kv + n_comp_attn; + ggml_tensor * kv_states = reshape_3d_checked(kv_prefix, head_dim, 1, n_kv_total, "build_attn_v4.kv_states", il); + ggml_tensor * kq_mask = nullptr; + if (work_tokens > 1) { + kq_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv_total, work_tokens); + ggml_set_input(kq_mask); + ggml_format_name(kq_mask, "deepseek4_kq_mask_l%d", il); + deepseek4_inputs->kq_masks.push_back(kq_mask); + } + + ggml_tensor * out = build_attn_mha( + q_states, + kv_states, + kv_states, + nullptr, + kq_mask, + layer.attn_sinks, + nullptr, + 1.0f / sqrtf(float(head_dim)), + il); + + out = reshape_3d_checked(out, head_dim, n_head, work_tokens, "build_attn_v4.out", il); + + ggml_tensor * o_nope = ggml_view_3d(ctx0, out, nope_dim, n_head, work_tokens, out->nb[1], out->nb[2], 0); + ggml_tensor * o_pe = ggml_view_3d(ctx0, out, rope_dim, n_head, work_tokens, out->nb[1], out->nb[2], nope_dim * out->nb[0]); + if (cparams.flash_attn) { + o_nope = ggml_cont(ctx0, o_nope); + o_pe = ggml_cont(ctx0, o_pe); + } + o_pe = ggml_rope_ext_back(ctx0, o_pe, inp_pos, nullptr, rope_dim, rope_type, layer_n_ctx_orig, layer_freq_base, layer_freq_scale, + layer_ext_factor, layer_attn_factor, layer_beta_fast, layer_beta_slow); + + out = ggml_concat(ctx0, o_nope, o_pe, 0); + out = cont_if_needed(reshape_2d_checked(out, total_q_dim, work_tokens, "build_attn_v4.out_2d", il)); + cb(out, "attn_out", il); + + return build_grouped_out(out, layer, il); + }; + + ggml_tensor * hc_target = ggml_new_tensor_3d(ctx0, inpL->type, n_embd, hc_mult, work_tokens); + ggml_tensor * inpL_hc = repeat_checked(reshape_3d_checked(inpL, n_embd, 1, work_tokens, "inpL_hc"), hc_target, "inpL_hc"); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + + ggml_tensor * residual = inpL_hc; + + auto [attn_in, attn_post_w, attn_comb] = hc_pre(inpL_hc, layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base, il); + attn_in = build_norm(attn_in, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + cb(attn_in, "attn_norm", il); + + ggml_tensor * attn_out = build_attn_v4(attn_in, layer, il); + inpL_hc = hc_post(attn_out, residual, attn_post_w, attn_comb, il); + + residual = inpL_hc; + + auto [ffn_in, ffn_post_w, ffn_comb] = hc_pre(inpL_hc, layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base, il); + ffn_in = build_norm(ffn_in, layer.ffn_norm, nullptr, LLM_NORM_RMS, il); + cb(ffn_in, "ffn_norm", il); + + ggml_tensor * moe_out = build_moe_v4(ffn_in, inp_tokens, layer, il); + ggml_tensor * shared_out = build_ffn(ffn_in, + layer.ffn_up_shexp, nullptr, nullptr, + layer.ffn_gate_shexp, nullptr, nullptr, + layer.ffn_down_shexp, nullptr, nullptr, + nullptr, + LLM_FFN_SILU, + LLM_FFN_PAR, + il); + cb(shared_out, "ffn_shared", il); + + ggml_tensor * ffn_out = ggml_add(ctx0, moe_out, shared_out); + cb(ffn_out, "ffn_out", il); + + inpL_hc = hc_post(ffn_out, residual, ffn_post_w, ffn_comb, il); + } + + ggml_tensor * cur = hc_head(inpL_hc, model.hc_head_fn, model.hc_head_scale, model.hc_head_base); + cb(cur, "hc_head", -1); + + cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = mul_mat_checked(model.output, cur, "output"); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} diff --git a/src/models/models.h b/src/models/models.h index 94991c55fe87..d8871d684ecd 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -190,6 +190,10 @@ struct llm_build_deepseek2 : public llm_graph_context { llm_build_deepseek2(const llama_model & model, const llm_graph_params & params); }; +struct llm_build_deepseek4 : public llm_graph_context { + llm_build_deepseek4(const llama_model & model, const llm_graph_params & params); +}; + struct llm_build_deepseek : public llm_graph_context { llm_build_deepseek(const llama_model & model, const llm_graph_params & params); }; diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 16af11a28623..0615c396bf9a 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -208,6 +208,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_EXPERT_USED_COUNT, uint32_t(1)); ms.add_kv(LLM_KV_EXPERT_SHARED_COUNT, uint32_t(1)); ms.add_kv(LLM_KV_EXPERT_GATING_FUNC, uint32_t(2)); // sigmoid + ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 1.0f); ms.add_kv(LLM_KV_EXPERT_GROUP_SCALE, 1.0f); ms.add_kv(LLM_KV_EXPERTS_PER_GROUP, uint32_t(1)); } @@ -331,6 +332,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_ARCTIC: case LLM_ARCH_DEEPSEEK: case LLM_ARCH_DEEPSEEK2: + case LLM_ARCH_DEEPSEEK4: case LLM_ARCH_GLM4_MOE: case LLM_ARCH_GLM_DSA: case LLM_ARCH_EXAONE_MOE: @@ -549,6 +551,9 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg std::string status_roundtrip = "\033[1;33mSKIP\033[0m"; char nmse_str[12] = {0}; bool skip = !arch_supported(arch) || (dc.split_mode == LLAMA_SPLIT_MODE_TENSOR && dc.devs.empty()); + if (arch == LLM_ARCH_DEEPSEEK4 && dc.split_mode == LLAMA_SPLIT_MODE_TENSOR) { + skip = true; // FIXME synthetic DeepSeek4 fixture needs dedicated tensor-split coverage. + } #if defined(GGML_USE_WEBGPU) skip = true; // FIXME #endif // GGML_USE_WEBGPU