diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index bf8af863a47..4ef03566fe6 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 @@ -777,7 +777,9 @@ 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_native_quant_tensor = name in getattr(self, "_preserve_native_quant_tensors", set()) + preserve_integer_tensor = name.endswith(".ffn.gate.tid2eid") or preserve_native_quant_tensor + 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 @@ -788,6 +790,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() @@ -865,6 +874,8 @@ def prepare_tensors(self): data_qtype = gguf.GGMLQuantizationType.TQ1_0 elif self.ftype == gguf.LlamaFileType.MOSTLY_TQ2_0: data_qtype = gguf.GGMLQuantizationType.TQ2_0 + elif self.ftype == gguf.LlamaFileType.MOSTLY_F8_E4M3_MXFP4: + data_qtype = gguf.GGMLQuantizationType.BF16 else: raise ValueError(f"Unknown file type: {self.ftype.name}") @@ -9182,6 +9193,428 @@ 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 = ( + "{{ '<｜begin▁of▁sentence｜>' }}" + "{% for message in messages %}" + "{% if message['role'] == 'system' %}" + "{{ message['content'] }}" + "{% elif message['role'] == 'user' %}" + "{{ '<｜User｜>' + message['content'] }}" + "{% elif message['role'] == 'assistant' %}" + "{{ message['content'] + '<｜end▁of▁sentence｜>' }}" + "{% endif %}" + "{% endfor %}" + "{% if add_generation_prompt %}" + "{{ '<｜Assistant｜>' }}" + "{% endif %}" + ) + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._expert_buffers: list[dict[str, Tensor]] | None = None + self._expert_seen: list[dict[str, set[int]]] | None = None + self._preserve_native_quant_tensors: set[str] = set() + self._native_quant_weight_types: dict[str, gguf.GGMLQuantizationType] = {} + self._native_quant_output_types: dict[str, gguf.GGMLQuantizationType] = {} + self._native_quant_scales: dict[str, Callable[[], Tensor]] = {} + + def set_vocab(self): + # transformers does not (yet) know about model_type=deepseek_v4, so the + # default AutoTokenizer.from_pretrained() in DeepseekV2Model.set_vocab + # fails inside AutoConfig before tokenizer files are touched. The V4 + # tokenizer is a vanilla PreTrainedTokenizerFast (model-agnostic), so + # try the parent path first and fall back to a direct load. + try: + super().set_vocab() + return + except (AttributeError, KeyError, ValueError) as e: + logger.info("DeepseekV4: AutoTokenizer path failed (%s); loading PreTrainedTokenizerFast directly", e) + + from transformers import PreTrainedTokenizerFast + tokenizer = PreTrainedTokenizerFast.from_pretrained(self.dir_model) + + tokens: list[str] = [] + toktypes: list[int] = [] + vocab_size = self.hparams.get("vocab_size", len(tokenizer.vocab)) + assert max(tokenizer.vocab.values()) < vocab_size + + tokpre = self.get_vocab_base_pre(tokenizer) + reverse_vocab = {id_: tok for tok, id_ in tokenizer.vocab.items()} + added_vocab = tokenizer.get_added_vocab() + added_tokens_decoder = tokenizer.added_tokens_decoder + + for i in range(vocab_size): + if i not in reverse_vocab: + tokens.append(f"[PAD{i}]") + toktypes.append(gguf.TokenType.UNUSED) + continue + token: str = reverse_vocab[i] + if token in added_vocab: + if not added_tokens_decoder[i].normalized: + token = tokenizer.decode(tokenizer.encode(token, add_special_tokens=False)) + if added_tokens_decoder[i].special or self.does_token_look_special(token): + toktypes.append(gguf.TokenType.CONTROL) + else: + token = token.replace(b"\xe2\x96\x81".decode("utf-8"), " ") + toktypes.append(gguf.TokenType.USER_DEFINED) + else: + toktypes.append(gguf.TokenType.NORMAL) + tokens.append(token) + + self.gguf_writer.add_tokenizer_model("gpt2") + self.gguf_writer.add_tokenizer_pre(tokpre) + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_types(toktypes) + + special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True) + special_vocab.add_to_gguf(self.gguf_writer) + + def dequant_model(self): + quant_method = (self.hparams.get("quantization_config") or {}).get("quant_method") + if quant_method == "fp8": + if self.ftype == gguf.LlamaFileType.MOSTLY_F8_E4M3_MXFP4: + for name, gen in list(self.model_tensors.items()): + if not name.endswith(".scale"): + continue + weight_name = name.removesuffix(".scale") + ".weight" + if weight_name not in self.model_tensors: + continue + + qtype = gguf.GGMLQuantizationType.MXFP4 if ".ffn.experts." in weight_name else gguf.GGMLQuantizationType.F8_E4M3_B128 + self._preserve_native_quant_tensors.add(weight_name) + self._native_quant_weight_types[weight_name] = qtype + self._native_quant_scales[weight_name] = gen + del self.model_tensors[name] + + return super().dequant_model() + + dequant_dtype = torch.float16 if self.ftype == gguf.LlamaFileType.MOSTLY_F16 else None + fp4_table = torch.tensor([ + 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + 0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, + ], dtype=torch.float32) + fp4_codes = torch.arange(256, dtype=torch.uint8) + fp4_pair_table = fp4_table[ + torch.stack((fp4_codes & 0x0F, (fp4_codes >> 4) & 0x0F), dim=1).long() + ] + + def finalize_dequant(data: Tensor) -> Tensor: + return data.to(dequant_dtype) if dequant_dtype is not None else data + + def dequant_with_scale(weight: Tensor, scale: Tensor) -> Tensor: + scale = scale.float() + + while scale.ndim < weight.ndim: + scale = scale.unsqueeze(-1) + + if scale.ndim != weight.ndim: + raise ValueError( + f"Unexpected DeepSeek V4 scale rank for weight {tuple(weight.shape)} and scale {tuple(scale.shape)}" + ) + + repeats: list[int] = [] + can_broadcast_blocks = True + for weight_dim, scale_dim in zip(weight.shape, scale.shape): + if scale_dim == weight_dim: + repeats.append(1) + continue + if scale_dim <= 0 or scale_dim > weight_dim: + raise ValueError( + f"Unexpected DeepSeek V4 scale shape {tuple(scale.shape)} for weight {tuple(weight.shape)}" + ) + if weight_dim % scale_dim != 0: + can_broadcast_blocks = False + break + repeats.append(weight_dim // scale_dim) + + if can_broadcast_blocks: + weight_shape: list[int] = [] + scale_shape: list[int] = [] + for scale_dim, repeat in zip(scale.shape, repeats): + weight_shape.extend((scale_dim, repeat)) + scale_shape.extend((scale_dim, 1)) + if dequant_dtype is not None: + return ( + weight.to(dequant_dtype).reshape(weight_shape) + * scale.to(dequant_dtype).reshape(scale_shape) + ).reshape(weight.shape) + + return (weight.float().reshape(weight_shape) * scale.reshape(scale_shape)).reshape(weight.shape) + + for dim, (weight_dim, scale_dim) in enumerate(zip(weight.shape, scale.shape)): + if scale_dim == weight_dim: + continue + repeat = (weight_dim + scale_dim - 1) // scale_dim + if repeat > 1: + scale = scale.repeat_interleave(repeat, dim) + + scale = scale[tuple(slice(0, size) for size in weight.shape)] + return finalize_dequant(weight.float() * scale) + + def dequant_packed_expert(weight: Tensor, scale: Tensor) -> Tensor: + weight = LazyTorchTensor.to_eager(weight) + scale = LazyTorchTensor.to_eager(scale).float() + + if weight.dtype != torch.int8 or weight.ndim != 2: + raise ValueError(f"Unexpected DeepSeek V4 expert weight {tuple(weight.shape)} {weight.dtype}") + + packed = weight.view(torch.uint8) + unpacked = fp4_pair_table[packed.long()].reshape(packed.shape[0], packed.shape[1] * 2) + + scale_groups = (unpacked.shape[1] + 31) // 32 + if scale.ndim != 2 or scale.shape[0] != unpacked.shape[0] or scale.shape[1] < scale_groups: + raise ValueError( + f"Unexpected DeepSeek V4 expert scale {tuple(scale.shape)} for weight {tuple(weight.shape)}" + ) + + scale = scale[:, :scale_groups] + if unpacked.shape[1] % 32 == 0: + data = unpacked.reshape(unpacked.shape[0], scale_groups, 32).mul_(scale.unsqueeze(-1)) + return finalize_dequant(data.reshape(unpacked.shape)) + + scale = scale.repeat_interleave(32, dim=1)[:, :unpacked.shape[1]] + return finalize_dequant(unpacked.mul_(scale)) + + for name, gen in list(self.model_tensors.items()): + if not name.endswith(".scale"): + continue + weight_name = name.removesuffix(".scale") + ".weight" + if weight_name not in self.model_tensors: + continue + + weight_gen = self.model_tensors[weight_name] + if ".ffn.experts." in weight_name: + self.model_tensors[weight_name] = ( + lambda weight_gen=weight_gen, scale_gen=gen: dequant_packed_expert(weight_gen(), scale_gen()) + ) + del self.model_tensors[name] + continue + + self.model_tensors[weight_name] = ( + lambda weight_gen=weight_gen, scale_gen=gen: dequant_with_scale(weight_gen(), scale_gen()) + ) + del self.model_tensors[name] + + return super().dequant_model() + + @staticmethod + def _pack_fp8_e4m3_b128(weight: Tensor, scale: Tensor, name: str) -> Tensor: + weight = LazyTorchTensor.to_eager(weight) + scale = LazyTorchTensor.to_eager(scale) + + if weight.dtype != torch.float8_e4m3fn or weight.ndim != 2: + raise ValueError(f"Unexpected DeepSeek V4 FP8 tensor {name}: {tuple(weight.shape)} {weight.dtype}") + + rows, cols = weight.shape + if rows % 128 != 0 or cols % 128 != 0: + raise ValueError(f"DeepSeek V4 FP8 tensor {name} shape {tuple(weight.shape)} is not divisible by 128x128") + + row_blocks = rows // 128 + col_blocks = cols // 128 + if scale.ndim != 2 or scale.shape != (row_blocks, col_blocks): + raise ValueError( + f"Unexpected DeepSeek V4 FP8 scale {tuple(scale.shape)} for tensor {name} with shape {tuple(weight.shape)}" + ) + + weight_u8 = weight.view(torch.uint8) + scale_u8 = scale.view(torch.uint8) + if scale_u8.shape != scale.shape: + raise ValueError(f"Unexpected DeepSeek V4 FP8 scale dtype {scale.dtype} for tensor {name}") + out = torch.empty((rows, col_blocks, 129), dtype=torch.uint8) + out.view(row_blocks, 128, col_blocks, 129)[:, :, :, 0].copy_(scale_u8[:, None, :]) + out[:, :, 1:].copy_(weight_u8.reshape(rows, col_blocks, 128)) + return out.reshape(rows, col_blocks * 129) + + @staticmethod + def _pack_mxfp4(weight: Tensor, scale: Tensor, name: str) -> Tensor: + weight = LazyTorchTensor.to_eager(weight) + scale = LazyTorchTensor.to_eager(scale) + + if weight.dtype != torch.int8 or weight.ndim != 2: + raise ValueError(f"Unexpected DeepSeek V4 packed expert tensor {name}: {tuple(weight.shape)} {weight.dtype}") + + rows, packed_cols = weight.shape + if packed_cols % 16 != 0: + raise ValueError(f"DeepSeek V4 packed expert tensor {name} has {packed_cols} bytes per row, not a multiple of 16") + + groups = packed_cols // 16 + if scale.ndim != 2 or scale.shape[0] != rows or scale.shape[1] < groups: + raise ValueError( + f"Unexpected DeepSeek V4 expert scale {tuple(scale.shape)} for tensor {name} with shape {tuple(weight.shape)}" + ) + + hf = weight.view(torch.uint8).reshape(rows, groups, 16) + scale_u8 = scale.view(torch.uint8) + if scale_u8.shape != scale.shape: + raise ValueError(f"Unexpected DeepSeek V4 expert scale dtype {scale.dtype} for tensor {name}") + out = torch.empty((rows, groups, 17), dtype=torch.uint8) + out[:, :, 0].copy_(scale_u8[:, :groups]) + lo = hf[:, :, :8] + hi = hf[:, :, 8:] + out[:, :, 1::2].copy_((lo & 0x0F) | ((hi & 0x0F) << 4)) + out[:, :, 2::2].copy_((lo >> 4) | (hi & 0xF0)) + return out.reshape(rows, groups * 17) + + def set_gguf_parameters(self): + self.hparams["num_key_value_heads"] = self.hparams.get("num_key_value_heads", 1) + self.hparams["rms_norm_eps"] = self.hparams.get("rms_norm_eps", self.hparams.get("norm_eps", 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) + + hparams = self.hparams + self.gguf_writer.add_vocab_size(hparams["vocab_size"]) + + if (q_lora_rank := hparams.get("q_lora_rank")) is not None: + self.gguf_writer.add_q_lora_rank(q_lora_rank) + + if (rope_dim := hparams.get("qk_rope_head_dim")) is not None: + self.gguf_writer.add_rope_dimension_count(rope_dim) + + if (sliding_window := hparams.get("sliding_window")) is not None: + self.gguf_writer.add_sliding_window(sliding_window) + + if (compress_rope_theta := hparams.get("compress_rope_theta")) is not None: + self.gguf_writer.add_rope_freq_base_swa(compress_rope_theta) + + self.gguf_writer.add_leading_dense_block_count(0) + + moe_intermediate_size = self.find_hparam(["moe_intermediate_size"], optional=False) + self.gguf_writer.add_expert_feed_forward_length(moe_intermediate_size) + + if (n_routed_experts := hparams.get("n_routed_experts")) is not None: + self.gguf_writer.add_expert_count(n_routed_experts) + + if (n_shared_experts := hparams.get("n_shared_experts")) is not None: + self.gguf_writer.add_expert_shared_count(n_shared_experts) + + if (routed_scaling_factor := hparams.get("routed_scaling_factor")) is not None: + self.gguf_writer.add_expert_weights_scale(routed_scaling_factor) + + if hparams.get("scoring_func") != "softmax": + self.gguf_writer.add_expert_weights_norm(True) + + if (swiglu_limit := hparams.get("swiglu_limit")) is not None: + self.gguf_writer.add_swiglu_clamp_exp([float(swiglu_limit)] * self.block_count) + + if (index_n_heads := hparams.get("index_n_heads")) is not None: + self.gguf_writer.add_indexer_head_count(index_n_heads) + + if (index_head_dim := hparams.get("index_head_dim")) is not None: + self.gguf_writer.add_indexer_key_length(index_head_dim) + + if (index_topk := hparams.get("index_topk")) 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 + + native_qtype = self._native_quant_weight_types.get(name) + if native_qtype is not None: + scale_gen = self._native_quant_scales[name] + if native_qtype == gguf.GGMLQuantizationType.F8_E4M3_B128: + data_torch = self._pack_fp8_e4m3_b128(data_torch, scale_gen(), name) + for new_name, data_torch in TextModel.modify_tensors(self, data_torch, name, bid): + self._native_quant_output_types[new_name] = native_qtype + yield new_name, data_torch + return + + if native_qtype == gguf.GGMLQuantizationType.MXFP4: + data_torch = self._pack_mxfp4(data_torch, scale_gen(), name) + else: + raise ValueError(f"Unsupported native quantization type for {name}: {native_qtype}") + + 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): + if native_qtype == gguf.GGMLQuantizationType.MXFP4: + self._native_quant_output_types[new_name] = native_qtype + yield new_name, data_torch + return + else: + return + + yield from TextModel.modify_tensors(self, data_torch, name, bid) + + def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool: + qtype = self._native_quant_output_types.get(new_name) + if qtype is not None: + return qtype + + return super().tensor_force_quant(name, new_name, bid, n_dims) + + def prepare_tensors(self): + super().prepare_tensors() + + if self._expert_seen is not None: + pending = [ + f"blk {bid} {w_name}: {len(xids)}/{self.hparams['n_routed_experts']}" + for bid, layer_seen in enumerate(self._expert_seen) + for w_name, xids in layer_seen.items() + if xids + ] + if pending: + raise ValueError(f"Unprocessed DeepSeek V4 experts: {pending}") + + @ModelBase.register( "Mistral3ForConditionalGeneration", "Ministral3ForCausalLM", @@ -13308,6 +13741,12 @@ 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_map[torch_float8_e8m0fnu] = np.uint8 + 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") @@ -13320,8 +13759,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", @@ -13348,6 +13787,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", @@ -13469,6 +13912,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 @@ -13496,6 +13945,7 @@ def main() -> None: "q8_0": gguf.LlamaFileType.MOSTLY_Q8_0, "tq1_0": gguf.LlamaFileType.MOSTLY_TQ1_0, "tq2_0": gguf.LlamaFileType.MOSTLY_TQ2_0, + "native": gguf.LlamaFileType.MOSTLY_F8_E4M3_MXFP4, "auto": gguf.LlamaFileType.GUESSED, } diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index a29dc707c3d..d563c81e84d 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -18,6 +18,7 @@ else() add_subdirectory(debug) add_subdirectory(embedding) add_subdirectory(eval-callback) + add_subdirectory(ds4-expert-profile) add_subdirectory(gguf-hash) add_subdirectory(gguf) diff --git a/examples/ds4-expert-profile/CMakeLists.txt b/examples/ds4-expert-profile/CMakeLists.txt new file mode 100644 index 00000000000..7f3eb219a37 --- /dev/null +++ b/examples/ds4-expert-profile/CMakeLists.txt @@ -0,0 +1,5 @@ +set(TARGET llama-ds4-expert-profile) +add_executable(${TARGET} ds4-expert-profile.cpp) +install(TARGETS ${TARGET} RUNTIME) +target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) +target_compile_features(${TARGET} PRIVATE cxx_std_17) diff --git a/examples/ds4-expert-profile/ds4-expert-profile.cpp b/examples/ds4-expert-profile/ds4-expert-profile.cpp new file mode 100644 index 00000000000..ab47bf4d1a4 --- /dev/null +++ b/examples/ds4-expert-profile/ds4-expert-profile.cpp @@ -0,0 +1,245 @@ +// Profile DeepSeek4 expert routing frequencies during inference. +// +// Captures the `ffn_topk` tensor output for each layer, builds per-layer +// expert-id histograms, and emits a JSON-ish report at the end. Use this to +// see whether routing is skewed enough to make hot-expert pinning worthwhile. + +#include "arg.h" +#include "common.h" +#include "log.h" +#include "llama.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +struct expert_profile { + int n_layer = 0; + int n_expert = 0; + std::map> per_layer; + uint64_t total_token_picks = 0; + std::vector scratch; + std::regex topk_re{"^ffn_topk(?:-([0-9]+))?$"}; +}; + +static bool ds4_cb_eval(struct ggml_tensor * t, bool ask, void * user_data) { + auto * prof = (expert_profile *) user_data; + if (!t || !t->name) { + return ask ? false : true; + } + + std::cmatch m; + if (!std::regex_match(t->name, m, prof->topk_re)) { + return ask ? false : true; + } + + if (ask) { + return true; + } + + int il = -1; + if (m.size() >= 2 && m[1].matched) { + il = std::atoi(m[1].str().c_str()); + } + + if (t->type != GGML_TYPE_I32) { + return true; + } + + auto nbytes = ggml_nbytes(t); + prof->scratch.resize(nbytes); + if (ggml_backend_buffer_is_host(t->buffer)) { + std::memcpy(prof->scratch.data(), t->data, nbytes); + } else { + ggml_backend_tensor_get(t, prof->scratch.data(), 0, nbytes); + } + + auto & hist = prof->per_layer[il]; + if ((int) hist.size() < prof->n_expert) { + hist.assign(prof->n_expert, 0); + } + + const int32_t * ids = (const int32_t *) prof->scratch.data(); + const size_t n_elems = nbytes / sizeof(int32_t); + for (size_t i = 0; i < n_elems; ++i) { + const int32_t e = ids[i]; + if (e >= 0 && e < prof->n_expert) { + hist[e]++; + prof->total_token_picks++; + } + } + + return true; +} + +int main(int argc, char ** argv) { + std::setlocale(LC_NUMERIC, "C"); + + common_params params; + common_init(); + + if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) { + return 1; + } + + llama_backend_init(); + llama_numa_init(params.numa); + + expert_profile prof; + params.cb_eval = ds4_cb_eval; + params.cb_eval_user_data = &prof; + params.warmup = false; + + auto llama_init = common_init_from_params(params); + auto * model = llama_init->model(); + auto * ctx = llama_init->context(); + if (!model || !ctx) { + LOG_ERR("failed to init\n"); + return 1; + } + + prof.n_layer = llama_model_n_layer(model); + prof.n_expert = 256; // hardcoded for DS4-Flash; could read from model metadata + + LOG_INF("\nds4-expert-profile: model has %d layers, %d experts\n", prof.n_layer, prof.n_expert); + LOG_INF("ds4-expert-profile: prompt length: %zu chars\n", params.prompt.size()); + + const llama_vocab * vocab = llama_model_get_vocab(model); + const bool add_bos = llama_vocab_get_add_bos(vocab); + auto tokens = common_tokenize(ctx, params.prompt, add_bos, true); + if (tokens.empty()) { + LOG_ERR("no tokens; provide a prompt with -p\n"); + return 1; + } + LOG_INF("ds4-expert-profile: tokenized to %zu tokens\n", tokens.size()); + + if (llama_decode(ctx, llama_batch_get_one(tokens.data(), tokens.size()))) { + LOG_ERR("decode failed\n"); + return 1; + } + + LOG_INF("\n=== expert routing report ===\n"); + LOG_INF("total expert picks observed: %" PRIu64 "\n", prof.total_token_picks); + + std::vector top_ks = {8, 16, 32, 64, 128}; + std::map hot_coverage_max; + std::map hot_coverage_avg_sum; + std::map hot_coverage_avg_count; + + LOG_INF("\nper-layer routing summary:\n"); + for (auto & [il, hist] : prof.per_layer) { + if (hist.empty()) continue; + + uint64_t total = 0; + for (uint64_t v : hist) total += v; + if (total == 0) continue; + + std::vector> sorted; + sorted.reserve(hist.size()); + for (size_t e = 0; e < hist.size(); ++e) { + if (hist[e] > 0) sorted.emplace_back((int) e, hist[e]); + } + std::sort(sorted.begin(), sorted.end(), [](auto & a, auto & b) { + return a.second > b.second; + }); + + const uint64_t hottest = sorted.empty() ? 0 : sorted.front().second; + const int unique_used = (int) sorted.size(); + + LOG_INF("layer %2d: total=%" PRIu64 " unique=%d hottest=%" PRIu64 "(%.1f%%)\n", + il, total, unique_used, hottest, 100.0 * hottest / total); + + for (int k : top_ks) { + uint64_t sum = 0; + for (int i = 0; i < k && i < (int) sorted.size(); ++i) { + sum += sorted[i].second; + } + const double frac = 100.0 * sum / total; + hot_coverage_max[k] = std::max(hot_coverage_max[k], frac); + hot_coverage_avg_sum[k] += frac; + hot_coverage_avg_count[k] += 1; + } + } + + LOG_INF("\n=== summary across layers ===\n"); + LOG_INF("top-K hot expert coverage:\n"); + for (int k : top_ks) { + if (hot_coverage_avg_count[k] == 0) continue; + const double avg = hot_coverage_avg_sum[k] / hot_coverage_avg_count[k]; + LOG_INF(" top-%-3d avg=%.1f%% max-layer=%.1f%%\n", + k, avg, hot_coverage_max[k]); + } + + LOG_INF("\nper-layer Pareto analysis (how many experts cover X%% of routings):\n"); + for (auto & [il, hist] : prof.per_layer) { + if (hist.empty()) continue; + std::vector sorted_h(hist); + std::sort(sorted_h.begin(), sorted_h.end(), std::greater()); + uint64_t total = 0; + for (uint64_t v : sorted_h) total += v; + if (total == 0) continue; + + uint64_t cum = 0; + int e50 = -1, e80 = -1, e90 = -1, e95 = -1, e99 = -1; + for (size_t i = 0; i < sorted_h.size(); ++i) { + cum += sorted_h[i]; + if (e50 < 0 && cum * 100 >= total * 50) e50 = (int)(i + 1); + if (e80 < 0 && cum * 100 >= total * 80) e80 = (int)(i + 1); + if (e90 < 0 && cum * 100 >= total * 90) e90 = (int)(i + 1); + if (e95 < 0 && cum * 100 >= total * 95) e95 = (int)(i + 1); + if (e99 < 0 && cum * 100 >= total * 99) e99 = (int)(i + 1); + } + LOG_INF("layer %2d: 50%%=top-%d 80%%=top-%d 90%%=top-%d 95%%=top-%d 99%%=top-%d\n", + il, e50, e80, e90, e95, e99); + } + + // Emit JSON profile to file (for runtime hot-expert pinning). + // Set DS4_PROFILE_JSON_OUT=path.json to enable. + if (const char * out_path = std::getenv("DS4_PROFILE_JSON_OUT")) { + FILE * fp = std::fopen(out_path, "w"); + if (fp) { + std::fprintf(fp, "{\n"); + std::fprintf(fp, " \"n_layer\": %d,\n", prof.n_layer); + std::fprintf(fp, " \"n_expert\": %d,\n", prof.n_expert); + std::fprintf(fp, " \"total_picks\": %" PRIu64 ",\n", prof.total_token_picks); + std::fprintf(fp, " \"layers\": {\n"); + bool first_layer = true; + for (auto & [il, hist] : prof.per_layer) { + if (hist.empty()) continue; + if (!first_layer) std::fprintf(fp, ",\n"); + first_layer = false; + // Sort experts by frequency descending; emit pairs. + std::vector> sorted; + sorted.reserve(hist.size()); + for (size_t e = 0; e < hist.size(); ++e) { + if (hist[e] > 0) sorted.emplace_back((int)e, hist[e]); + } + std::sort(sorted.begin(), sorted.end(), [](auto & a, auto & b) { + return a.second > b.second; + }); + std::fprintf(fp, " \"%d\": [", il); + for (size_t i = 0; i < sorted.size(); ++i) { + if (i) std::fprintf(fp, ","); + std::fprintf(fp, "[%d,%" PRIu64 "]", sorted[i].first, sorted[i].second); + } + std::fprintf(fp, "]"); + } + std::fprintf(fp, "\n }\n}\n"); + std::fclose(fp); + LOG_INF("\nds4-expert-profile: wrote JSON to %s\n", out_path); + } else { + LOG_ERR("ds4-expert-profile: could not open %s\n", out_path); + } + } + + llama_backend_free(); + return 0; +} diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index 6fcf5a43393..5ad121ae57f 100644 --- a/ggml/include/ggml-rpc.h +++ b/ggml/include/ggml-rpc.h @@ -8,10 +8,10 @@ extern "C" { #define RPC_PROTO_MAJOR_VERSION 4 #define RPC_PROTO_MINOR_VERSION 0 -#define RPC_PROTO_PATCH_VERSION 0 +#define RPC_PROTO_PATCH_VERSION 1 #ifdef __cplusplus -static_assert(GGML_OP_COUNT == 96, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); +static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); #endif #define GGML_RPC_MAX_SERVERS 16 diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 703e3783136..893ca815a60 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -429,7 +429,8 @@ extern "C" { GGML_TYPE_MXFP4 = 39, // MXFP4 (1 block) GGML_TYPE_NVFP4 = 40, // NVFP4 (4 blocks, E4M3 scale) GGML_TYPE_Q1_0 = 41, - GGML_TYPE_COUNT = 42, + GGML_TYPE_F8_E4M3_B128 = 42, // E4M3 FP8 values with one E8M0 scale per 128 values + GGML_TYPE_COUNT = 43, }; // precision @@ -467,6 +468,7 @@ extern "C" { GGML_FTYPE_MOSTLY_MXFP4 = 25, // except 1d tensors GGML_FTYPE_MOSTLY_NVFP4 = 26, // except 1d tensors GGML_FTYPE_MOSTLY_Q1_0 = 27, // except 1d tensors + GGML_FTYPE_MOSTLY_F8_E4M3_MXFP4 = 28, // except 1d tensors }; // available tensor operations: @@ -576,6 +578,7 @@ extern "C" { GGML_OP_OPT_STEP_SGD, GGML_OP_GLU, + GGML_OP_HC_WEIGHTED_SUM, GGML_OP_COUNT, }; @@ -603,6 +606,9 @@ extern "C" { GGML_UNARY_OP_CEIL, GGML_UNARY_OP_ROUND, GGML_UNARY_OP_TRUNC, + GGML_UNARY_OP_FP4_ACT_QUANT, + GGML_UNARY_OP_FP8_ACT_QUANT, + GGML_UNARY_OP_SINKHORN_4X4, GGML_UNARY_OP_COUNT, }; @@ -1246,7 +1252,18 @@ extern "C" { struct ggml_context * ctx, struct ggml_tensor * a); + // Blockwise activation quant-dequant simulation used by DeepSeek4 QAT paths. + GGML_API struct ggml_tensor * ggml_fp4_act_quant( + struct ggml_context * ctx, + struct ggml_tensor * a); + GGML_API struct ggml_tensor * ggml_fp8_act_quant( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_sinkhorn_4x4( + struct ggml_context * ctx, + struct ggml_tensor * a); // xIELU activation function // x = x * (c_a(alpha_n) + c_b(alpha_p, beta) * sigmoid(beta * x)) + eps * (x > 0) @@ -1413,6 +1430,13 @@ extern "C" { struct ggml_tensor * a, struct ggml_tensor * b); + // weighted sum over the HC dimension: + // a: [n_embd, hc_mult], b: [hc_mult] => result: [n_embd] + GGML_API struct ggml_tensor * ggml_hc_weighted_sum( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + // change the precision of a matrix multiplication // set to GGML_PREC_F32 for higher precision (useful for phi-2) GGML_API void ggml_mul_mat_set_prec( diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 41a61775bd6..0800063eac3 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -516,6 +516,12 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co return handle_generic(src_ss, /*scalar_only =*/ false); }; + auto handle_hc_weighted_sum = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_1); + return src_ss[0]; + }; + auto handle_concat = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { const ggml_backend_meta_split_axis concat_axis = ggml_backend_meta_split_axis(ggml_get_op_params_i32(tensor, 0)); if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[1].axis >= 0 && src_ss[1].axis < GGML_MAX_DIMS) { @@ -957,6 +963,9 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co case GGML_OP_GATED_DELTA_NET: { split_state = handle_gated_delta_net(src_ss); } break; + case GGML_OP_HC_WEIGHTED_SUM: { + split_state = handle_hc_weighted_sum(src_ss); + } break; case GGML_OP_UNARY: { split_state = handle_generic(src_ss, /*scalar_only =*/ false); } break; @@ -2123,4 +2132,3 @@ ggml_backend_t ggml_backend_meta_simple_backend(ggml_backend_t meta_backend, siz const ggml_backend_meta_context * backend_ctx = (const ggml_backend_meta_context *) meta_backend->context; return backend_ctx->backend_configs[index].backend; } - diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index d9f8aaec52f..af92bf80005 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -14,12 +14,16 @@ #include "ggml-impl.h" #include +#include #include #include #include #include #include #include +#include +#include +#include #include #ifdef __APPLE__ @@ -771,6 +775,59 @@ struct ggml_backend_sched_split { struct ggml_cgraph graph; }; +enum ggml_backend_sched_moe_prefetch_policy { + GGML_BACKEND_SCHED_MOE_PREFETCH_NONE, + GGML_BACKEND_SCHED_MOE_PREFETCH_SETMARKOV, +}; + +struct ggml_backend_sched_moe_transition { + std::vector from; + std::vector counts; +}; + +struct ggml_backend_sched_moe_cache { + const ggml_tensor * input; + int backend_id; + int n_expert; + int n_slots; + + size_t expert_size; + size_t slot_padding; + size_t slot_stride; + size_t weights_size; + size_t ids_nbytes; + + ggml_backend_buffer_t weights_buffer; + ggml_tensor weights_tensor; + + ggml_backend_buffer_t ids_buffer; + ggml_tensor ids_tensor; + + std::vector slot_of; + std::vector expert_in_slot; + std::vector slot_speculative; + std::vector lru_tick; + std::vector remapped_ids; + std::vector previous_experts; + std::vector transitions; + uint64_t now; + uint64_t hits; + uint64_t misses; + uint64_t bypasses; + uint64_t bytes_copied; + uint64_t speculative_hits; + uint64_t prefetches; + uint64_t prefetch_evictions; + uint64_t wrong_prefetches; + uint64_t bytes_prefetched; +}; + +struct ggml_backend_sched_moe_restore { + ggml_tensor * node; + ggml_tensor * src0; + ggml_tensor * src2; +}; + struct ggml_backend_sched { bool is_reset; // true if the scheduler has been reset since the last graph split bool is_alloc; @@ -817,6 +874,7 @@ struct ggml_backend_sched { size_t context_buffer_size; bool op_offload; + std::vector * moe_caches; int debug; @@ -874,6 +932,11 @@ static char causes[GGML_DEFAULT_GRAPH_SIZE*16 + GGML_SCHED_MAX_SPLITS_DEBUG*GGML #define GET_CAUSE(node) "" #endif +static int ggml_backend_sched_backend_from_non_weight_src( + ggml_backend_sched_t sched, + ggml_tensor * tensor, + int max_backend_id); + // returns the backend that should be used for the node based on the current locations static int ggml_backend_sched_backend_id_from_cur(ggml_backend_sched_t sched, struct ggml_tensor * tensor) { // assign pre-allocated nodes to their backend @@ -917,6 +980,11 @@ static int ggml_backend_sched_backend_id_from_cur(ggml_backend_sched_t sched, st int src_backend_id = ggml_backend_sched_backend_from_buffer(sched, src, tensor); // check if a backend with higher prio wants to offload the op if (sched->op_offload && src_backend_id == sched->n_backends - 1 && ggml_backend_buffer_is_host(src->buffer)) { + const int non_weight_src_backend_id = ggml_backend_sched_backend_from_non_weight_src(sched, tensor, src_backend_id); + if (non_weight_src_backend_id != -1) { + SET_CAUSE(tensor, "1.off-src%d", non_weight_src_backend_id); + return non_weight_src_backend_id; + } for (int b = 0; b < src_backend_id; b++) { if (ggml_backend_supports_op(sched->backends[b], tensor) && ggml_backend_offload_op(sched->backends[b], tensor)) { SET_CAUSE(tensor, "1.off"); @@ -982,6 +1050,35 @@ static void ggml_backend_sched_print_assignments(ggml_backend_sched_t sched, str } } +static int ggml_backend_sched_backend_from_non_weight_src( + ggml_backend_sched_t sched, + ggml_tensor * tensor, + int max_backend_id) { + for (int i = 0; i < GGML_MAX_SRC; ++i) { + ggml_tensor * src = tensor->src[i]; + if (src == nullptr) { + continue; + } + ggml_backend_buffer_t src_buffer = src->view_src != nullptr ? src->view_src->buffer : src->buffer; + if (src_buffer != nullptr && src_buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { + continue; + } + + int src_backend_id = tensor_backend_id(src); + if (src_backend_id == -1 && src->view_src != nullptr) { + src_backend_id = tensor_backend_id(src->view_src); + } + if (src_backend_id < 0 || src_backend_id >= max_backend_id) { + continue; + } + if (ggml_backend_supports_op(sched->backends[src_backend_id], tensor) && + ggml_backend_offload_op(sched->backends[src_backend_id], tensor)) { + return src_backend_id; + } + } + return -1; +} + static bool ggml_backend_sched_buffer_supported(ggml_backend_sched_t sched, struct ggml_tensor * t, int backend_id) { ggml_backend_buffer_t buf = t->view_src ? t->view_src->buffer : t->buffer; ggml_backend_buffer_type_t buft = NULL; @@ -1538,18 +1635,663 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { return true; } +static bool ggml_backend_sched_moe_log_enabled() { + static const bool enabled = []() { + const char * env = getenv("GGML_SCHED_MOE_LOG"); + return env != nullptr && env[0] != '\0' && strcmp(env, "0") != 0; + }(); + return enabled; +} + +static int ggml_backend_sched_moe_cache_slots() { + static const int slots = []() { + const char * env = getenv("GGML_SCHED_MOE_CACHE_SLOTS"); + if (env == nullptr || env[0] == '\0') { + return 0; + } + + errno = 0; + char * end = nullptr; + const long value = strtol(env, &end, 10); + if (errno != 0 || end == env || *end != '\0' || value < 0 || value > INT_MAX) { + GGML_LOG_WARN("%s: ignoring invalid GGML_SCHED_MOE_CACHE_SLOTS=%s\n", __func__, env); + return 0; + } + + return (int) value; + }(); + return slots; +} + +static enum ggml_backend_sched_moe_prefetch_policy ggml_backend_sched_moe_prefetch_policy() { + static const enum ggml_backend_sched_moe_prefetch_policy policy = []() { + const char * env = getenv("GGML_SCHED_MOE_CACHE_PREFETCH"); + if (env == nullptr || env[0] == '\0' || strcmp(env, "0") == 0 || strcmp(env, "none") == 0) { + return GGML_BACKEND_SCHED_MOE_PREFETCH_NONE; + } + if (strcmp(env, "setmarkov") == 0) { + return GGML_BACKEND_SCHED_MOE_PREFETCH_SETMARKOV; + } + + GGML_LOG_WARN("%s: ignoring invalid GGML_SCHED_MOE_CACHE_PREFETCH=%s\n", __func__, env); + return GGML_BACKEND_SCHED_MOE_PREFETCH_NONE; + }(); + return policy; +} + +static int ggml_backend_sched_moe_prefetch_limit() { + static const int limit = []() { + const char * env = getenv("GGML_SCHED_MOE_CACHE_PREFETCH_LIMIT"); + if (env == nullptr || env[0] == '\0') { + return 0; + } + + errno = 0; + char * end = nullptr; + const long value = strtol(env, &end, 10); + if (errno != 0 || end == env || *end != '\0' || value < 0 || value > INT_MAX) { + GGML_LOG_WARN("%s: ignoring invalid GGML_SCHED_MOE_CACHE_PREFETCH_LIMIT=%s\n", __func__, env); + return 0; + } + + return (int) value; + }(); + return limit; +} + +static bool ggml_backend_sched_moe_cache_prime_last_enabled() { + static const bool enabled = []() { + const char * env = getenv("GGML_SCHED_MOE_CACHE_PRIME"); + return env != nullptr && strcmp(env, "last") == 0; + }(); + return enabled; +} + +static const char * ggml_backend_sched_tensor_name(const ggml_tensor * tensor) { + return tensor->name[0] != '\0' ? tensor->name : ""; +} + +static bool ggml_backend_sched_same_layout(const ggml_tensor * a, const ggml_tensor * b) { + if (a->type != b->type) { + return false; + } + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (a->ne[i] != b->ne[i] || a->nb[i] != b->nb[i]) { + return false; + } + } + return true; +} + +static ggml_backend_sched_moe_cache * ggml_backend_sched_moe_cache_find( + ggml_backend_sched_t sched, + const ggml_tensor * input, + int backend_id) { + for (ggml_backend_sched_moe_cache * cache : *sched->moe_caches) { + if (cache->input == input && cache->backend_id == backend_id) { + return cache; + } + } + return nullptr; +} + +static size_t ggml_backend_sched_moe_cache_slot_padding(const ggml_tensor * input, size_t expert_size) { + const size_t type_size = ggml_type_size(input->type); + GGML_ASSERT(type_size > 0); + GGML_ASSERT(expert_size % type_size == 0); + + const size_t padding = std::min(expert_size, 512); + return ((padding + type_size - 1) / type_size) * type_size; +} + +static ggml_backend_sched_moe_cache * ggml_backend_sched_moe_cache_new( + ggml_backend_sched_t sched, + ggml_backend_t backend, + const ggml_tensor * input, + int backend_id, + int n_expert, + int n_slots, + size_t expert_size) { + GGML_ASSERT(n_slots > 0); + GGML_ASSERT(n_slots <= n_expert); + GGML_ASSERT(input->ne[3] == 1); + + ggml_backend_buffer_type_t buft = sched->bufts[backend_id]; + const size_t padding = ggml_backend_sched_moe_cache_slot_padding(input, expert_size); + + ggml_backend_sched_moe_cache * cache = new ggml_backend_sched_moe_cache(); + cache->input = input; + cache->backend_id = backend_id; + cache->n_expert = n_expert; + cache->n_slots = n_slots; + cache->expert_size = expert_size; + cache->slot_padding = padding; + cache->slot_stride = expert_size + padding; + cache->weights_size = 0; + + cache->weights_tensor = *input; + cache->weights_tensor.buffer = nullptr; + cache->weights_tensor.data = nullptr; + cache->weights_tensor.view_src = nullptr; + cache->weights_tensor.view_offs = 0; + cache->weights_tensor.extra = nullptr; + cache->weights_tensor.op = GGML_OP_NONE; + cache->weights_tensor.flags = 0; + cache->weights_tensor.ne[2] = n_slots + 1; // one dummy padding slot + cache->weights_tensor.ne[3] = 1; + cache->weights_tensor.nb[2] = cache->slot_stride; + cache->weights_tensor.nb[3] = cache->slot_stride * cache->weights_tensor.ne[2]; + for (int i = 0; i < GGML_MAX_SRC; ++i) { + cache->weights_tensor.src[i] = nullptr; + } + ggml_format_name(&cache->weights_tensor, "%s#moe-cache#%s", + ggml_backend_sched_tensor_name(input), ggml_backend_name(backend)); + + const size_t weights_size = ggml_backend_buft_get_alloc_size(buft, &cache->weights_tensor); + cache->weights_size = weights_size; + cache->weights_buffer = ggml_backend_buft_alloc_buffer(buft, weights_size); + if (cache->weights_buffer == nullptr) { + delete cache; + return nullptr; + } + ggml_backend_buffer_set_usage(cache->weights_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + if (ggml_backend_tensor_alloc(cache->weights_buffer, &cache->weights_tensor, ggml_backend_buffer_get_base(cache->weights_buffer)) != GGML_STATUS_SUCCESS) { + ggml_backend_buffer_free(cache->weights_buffer); + delete cache; + return nullptr; + } + ggml_backend_buffer_clear(cache->weights_buffer, 0); + + cache->slot_of.assign(n_expert, -1); + cache->expert_in_slot.assign(n_slots, -1); + cache->slot_speculative.assign(n_slots, 0); + cache->lru_tick.assign(n_slots, 0); + + sched->moe_caches->push_back(cache); + + GGML_LOG_INFO("%s: allocated MoE expert cache for %s on %s: slots=%d/%d, bytes=%zu\n", + __func__, ggml_backend_sched_tensor_name(input), ggml_backend_name(backend), + n_slots, n_expert, weights_size); + + return cache; +} + +static bool ggml_backend_sched_moe_cache_ensure_ids( + ggml_backend_sched_moe_cache * cache, + ggml_backend_buffer_type_t buft, + const ggml_tensor * ids_tensor) { + const size_t ids_nbytes = ggml_nbytes(ids_tensor); + if (cache->ids_buffer != nullptr && + cache->ids_nbytes == ids_nbytes && + ggml_backend_sched_same_layout(&cache->ids_tensor, ids_tensor)) { + return true; + } + + ggml_backend_buffer_free(cache->ids_buffer); + cache->ids_buffer = nullptr; + cache->ids_nbytes = 0; + + cache->ids_tensor = *ids_tensor; + cache->ids_tensor.buffer = nullptr; + cache->ids_tensor.data = nullptr; + cache->ids_tensor.view_src = nullptr; + cache->ids_tensor.view_offs = 0; + cache->ids_tensor.extra = nullptr; + cache->ids_tensor.op = GGML_OP_NONE; + cache->ids_tensor.flags = 0; + for (int i = 0; i < GGML_MAX_SRC; ++i) { + cache->ids_tensor.src[i] = nullptr; + } + ggml_format_name(&cache->ids_tensor, "%s#moe-cache-ids", + ggml_backend_sched_tensor_name(ids_tensor)); + + const size_t ids_alloc = ggml_backend_buft_get_alloc_size(buft, &cache->ids_tensor); + cache->ids_buffer = ggml_backend_buft_alloc_buffer(buft, ids_alloc); + if (cache->ids_buffer == nullptr) { + return false; + } + ggml_backend_buffer_set_usage(cache->ids_buffer, GGML_BACKEND_BUFFER_USAGE_COMPUTE); + if (ggml_backend_tensor_alloc(cache->ids_buffer, &cache->ids_tensor, ggml_backend_buffer_get_base(cache->ids_buffer)) != GGML_STATUS_SUCCESS) { + ggml_backend_buffer_free(cache->ids_buffer); + cache->ids_buffer = nullptr; + return false; + } + + cache->ids_nbytes = ids_nbytes; + return true; +} + +static bool ggml_backend_sched_moe_cache_prepare( + ggml_backend_sched_t sched, + ggml_backend_t split_backend, + int split_backend_id, + ggml_tensor * input, + ggml_tensor * node, + ggml_tensor * ids_tensor, + const std::vector & ids, + const std::vector & used_ids, + int64_t n_expert, + size_t expert_size, + int requested_slots, + enum ggml_backend_sched_moe_prefetch_policy prefetch_policy, + int prefetch_limit, + bool moe_log, + const char ** fail_reason, + std::vector & restores) { + if (fail_reason != nullptr) { + *fail_reason = nullptr; + } + auto fail = [fail_reason](const char * reason) { + if (fail_reason != nullptr) { + *fail_reason = reason; + } + return false; + }; + + if (requested_slots <= 0) { + return fail("disabled"); + } + if (input->ne[3] != 1) { + return fail("unsupported_shape"); + } + if (n_expert <= 0 || n_expert > INT_MAX) { + return fail("invalid_expert_count"); + } + + const int n_slots = std::min(requested_slots, (int) n_expert); + std::vector needed; + needed.reserve(n_slots); + for (int64_t i = 0; i < n_expert; ++i) { + if (ggml_bitset_get(used_ids.data(), i)) { + needed.push_back((int32_t) i); + } + } + if (needed.empty()) { + return fail("no_experts"); + } + + const bool too_many_experts = (int) needed.size() > n_slots; + const bool prime_last = too_many_experts && ggml_backend_sched_moe_cache_prime_last_enabled(); + if (too_many_experts && !prime_last) { + return fail("too_many_experts"); + } + + ggml_backend_sched_moe_cache * cache = ggml_backend_sched_moe_cache_find(sched, input, split_backend_id); + if (cache != nullptr && + (cache->n_expert != n_expert || cache->n_slots != n_slots || cache->expert_size != expert_size)) { + return fail("cache_metadata_mismatch"); + } + if (cache == nullptr) { + cache = ggml_backend_sched_moe_cache_new(sched, split_backend, input, split_backend_id, (int) n_expert, n_slots, expert_size); + if (cache == nullptr) { + return fail("cache_alloc_failed"); + } + } + + if (too_many_experts) { + std::vector prime_ids; + prime_ids.reserve((size_t) n_slots); + std::vector seen((size_t) n_expert, 0); + for (auto it = ids.rbegin(); it != ids.rend() && (int) prime_ids.size() < n_slots; ++it) { + const int32_t expert_id = *it; + if (expert_id < 0 || expert_id >= n_expert || seen[expert_id]) { + continue; + } + seen[expert_id] = 1; + prime_ids.push_back(expert_id); + } + std::reverse(prime_ids.begin(), prime_ids.end()); + + size_t primed = 0; + size_t primed_bytes = 0; + for (int32_t expert_id : prime_ids) { + int32_t slot = cache->slot_of[expert_id]; + if (slot >= 0) { + cache->slot_speculative[slot] = 1; + cache->lru_tick[slot] = ++cache->now; + continue; + } + + for (int32_t candidate = 0; candidate < cache->n_slots; ++candidate) { + if (cache->expert_in_slot[candidate] == -1) { + slot = candidate; + break; + } + } + + if (slot == -1) { + uint64_t best_tick = std::numeric_limits::max(); + for (int32_t candidate = 0; candidate < cache->n_slots; ++candidate) { + if (cache->lru_tick[candidate] < best_tick) { + best_tick = cache->lru_tick[candidate]; + slot = candidate; + } + } + } + if (slot == -1) { + continue; + } + + const int32_t old_expert = cache->expert_in_slot[slot]; + if (old_expert >= 0) { + if (cache->slot_speculative[slot]) { + cache->wrong_prefetches++; + } + cache->slot_of[old_expert] = -1; + } + + const size_t padding = expert_id < n_expert - 1 ? cache->slot_padding : 0; + const size_t copy_size = expert_size + padding; + ggml_backend_tensor_set_async(split_backend, + &cache->weights_tensor, + (const uint8_t *) input->data + (size_t) expert_id * expert_size, + (size_t) slot * cache->slot_stride, + copy_size); + + cache->expert_in_slot[slot] = expert_id; + cache->slot_of[expert_id] = slot; + cache->slot_speculative[slot] = 1; + cache->lru_tick[slot] = ++cache->now; + cache->prefetches++; + cache->bytes_prefetched += copy_size; + primed++; + primed_bytes += copy_size; + } + + if (moe_log) { + GGML_LOG_INFO("%s: moe_cache_prime tensor=%s backend=%s slots=%d primed=%zu primed_bytes=%zu used=%zu\n", + __func__, ggml_backend_sched_tensor_name(input), ggml_backend_name(split_backend), + n_slots, primed, primed_bytes, needed.size()); + } + return fail("too_many_experts_primed"); + } + + if (!ggml_backend_sched_moe_cache_ensure_ids(cache, sched->bufts[split_backend_id], ids_tensor)) { + return fail("ids_alloc_failed"); + } + + std::vector predicted_ids; + if (prefetch_policy == GGML_BACKEND_SCHED_MOE_PREFETCH_SETMARKOV && needed.size() <= 64) { + ggml_backend_sched_moe_transition * prediction = nullptr; + for (ggml_backend_sched_moe_transition & candidate : cache->transitions) { + if (candidate.from == needed) { + prediction = &candidate; + break; + } + } + + if (prediction != nullptr) { + std::vector> candidates; + candidates.reserve((size_t) n_expert); + for (int64_t expert_id = 0; expert_id < n_expert; ++expert_id) { + const uint32_t count = prediction->counts[expert_id]; + if (count > 0) { + candidates.push_back({ (int32_t) expert_id, count }); + } + } + std::sort(candidates.begin(), candidates.end(), + [](const std::pair & a, const std::pair & b) { + if (a.second != b.second) { + return a.second > b.second; + } + return a.first < b.first; + }); + + const size_t protect_limit = std::min((size_t) cache->n_slots, 6); + predicted_ids.resize(ggml_bitset_size(n_expert)); + for (size_t i = 0; i < std::min(protect_limit, candidates.size()); ++i) { + ggml_bitset_set(predicted_ids.data(), candidates[i].first); + } + } + } + + std::vector misses; + misses.reserve(needed.size()); + for (int32_t expert_id : needed) { + const int32_t slot = cache->slot_of[expert_id]; + if (slot >= 0) { + cache->hits++; + if (cache->slot_speculative[slot]) { + cache->speculative_hits++; + cache->slot_speculative[slot] = 0; + } + } else { + misses.push_back(expert_id); + cache->misses++; + } + } + + auto find_free_slot = [&]() -> int32_t { + for (int32_t slot = 0; slot < cache->n_slots; ++slot) { + if (cache->expert_in_slot[slot] == -1) { + return slot; + } + } + return -1; + }; + + auto find_victim_slot = [&](bool prefer_speculative) -> int32_t { + for (int pass = 0; pass < (prefer_speculative ? 4 : 2); ++pass) { + const bool speculative_only = prefer_speculative && (pass % 2 == 0); + const bool protect_predicted = pass < (prefer_speculative ? 2 : 1); + int32_t slot = -1; + uint64_t best_tick = std::numeric_limits::max(); + for (int32_t candidate = 0; candidate < cache->n_slots; ++candidate) { + const int32_t resident = cache->expert_in_slot[candidate]; + GGML_ASSERT(resident >= 0); + if (ggml_bitset_get(used_ids.data(), resident)) { + continue; + } + if (protect_predicted && !predicted_ids.empty() && ggml_bitset_get(predicted_ids.data(), resident)) { + continue; + } + if (speculative_only && !cache->slot_speculative[candidate]) { + continue; + } + if (cache->lru_tick[candidate] < best_tick) { + best_tick = cache->lru_tick[candidate]; + slot = candidate; + } + } + if (slot >= 0) { + return slot; + } + } + return -1; + }; + + auto copy_expert_to_slot = [&](int32_t expert_id, int32_t slot, bool speculative, bool prefetch) -> size_t { + const int32_t old_expert = cache->expert_in_slot[slot]; + if (old_expert >= 0) { + if (cache->slot_speculative[slot]) { + cache->wrong_prefetches++; + } + if (prefetch) { + cache->prefetch_evictions++; + } + cache->slot_of[old_expert] = -1; + } + + const size_t padding = expert_id < n_expert - 1 ? cache->slot_padding : 0; + const size_t copy_size = expert_size + padding; + ggml_backend_tensor_set_async(split_backend, + &cache->weights_tensor, + (const uint8_t *) input->data + (size_t) expert_id * expert_size, + (size_t) slot * cache->slot_stride, + copy_size); + + cache->expert_in_slot[slot] = expert_id; + cache->slot_of[expert_id] = slot; + cache->slot_speculative[slot] = speculative ? 1 : 0; + cache->lru_tick[slot] = ++cache->now; + return copy_size; + }; + + size_t copied_bytes = 0; + for (int32_t expert_id : misses) { + int32_t slot = find_free_slot(); + if (slot == -1) { + slot = find_victim_slot(true); + } + + if (slot == -1) { + cache->bypasses++; + return fail("no_evictable_slot"); + } + + const size_t copy_size = copy_expert_to_slot(expert_id, slot, false, false); + copied_bytes += copy_size; + cache->bytes_copied += copy_size; + } + + for (int32_t expert_id : needed) { + const int32_t slot = cache->slot_of[expert_id]; + GGML_ASSERT(slot >= 0); + cache->lru_tick[slot] = ++cache->now; + } + + cache->remapped_ids = ids; + for (int64_t i1 = 0; i1 < ids_tensor->ne[1]; i1++) { + for (int64_t i0 = 0; i0 < ids_tensor->ne[0]; i0++) { + const int64_t idx = i1 * ids_tensor->nb[1]/sizeof(int32_t) + i0 * ids_tensor->nb[0]/sizeof(int32_t); + const int32_t expert_id = ids[idx]; + const int32_t slot = cache->slot_of[expert_id]; + GGML_ASSERT(slot >= 0); + cache->remapped_ids[idx] = slot; + } + } + + ggml_backend_tensor_set_async(split_backend, &cache->ids_tensor, cache->remapped_ids.data(), 0, cache->ids_nbytes); + + size_t prefetched_bytes = 0; + size_t prefetch_count = 0; + if (prefetch_policy == GGML_BACKEND_SCHED_MOE_PREFETCH_SETMARKOV && needed.size() <= 64) { + if (!cache->previous_experts.empty() && cache->previous_experts.size() <= 64) { + ggml_backend_sched_moe_transition * transition = nullptr; + for (ggml_backend_sched_moe_transition & candidate : cache->transitions) { + if (candidate.from == cache->previous_experts) { + transition = &candidate; + break; + } + } + if (transition == nullptr) { + cache->transitions.push_back({ cache->previous_experts, std::vector((size_t) n_expert, 0) }); + transition = &cache->transitions.back(); + } + + for (int32_t expert_id : needed) { + uint32_t & count = transition->counts[expert_id]; + if (count < std::numeric_limits::max()) { + count++; + } + } + } + + ggml_backend_sched_moe_transition * prediction = nullptr; + for (ggml_backend_sched_moe_transition & candidate : cache->transitions) { + if (candidate.from == needed) { + prediction = &candidate; + break; + } + } + + if (prediction != nullptr) { + std::vector> candidates; + candidates.reserve((size_t) n_expert); + for (int64_t expert_id = 0; expert_id < n_expert; ++expert_id) { + const uint32_t count = prediction->counts[expert_id]; + if (count > 0) { + candidates.push_back({ (int32_t) expert_id, count }); + } + } + std::sort(candidates.begin(), candidates.end(), + [](const std::pair & a, const std::pair & b) { + if (a.second != b.second) { + return a.second > b.second; + } + return a.first < b.first; + }); + + if (prefetch_limit > 0) { + const size_t limit = std::min((size_t) std::min(prefetch_limit, cache->n_slots), candidates.size()); + for (size_t i = 0; i < limit; ++i) { + const int32_t expert_id = candidates[i].first; + int32_t slot = cache->slot_of[expert_id]; + if (slot >= 0) { + cache->lru_tick[slot] = ++cache->now; + continue; + } + + slot = find_free_slot(); + if (slot == -1) { + slot = find_victim_slot(true); + } + if (slot == -1) { + break; + } + + const size_t copy_size = copy_expert_to_slot(expert_id, slot, true, true); + prefetched_bytes += copy_size; + cache->bytes_prefetched += copy_size; + cache->prefetches++; + prefetch_count++; + } + } + } + + cache->previous_experts = needed; + } + + restores.push_back({ node, node->src[0], node->src[2] }); + node->src[0] = &cache->weights_tensor; + node->src[2] = &cache->ids_tensor; + + if (moe_log) { + GGML_LOG_INFO("%s: moe_cache tensor=%s backend=%s slots=%d expert_size=%zu cache_bytes=%zu used=%zu hits=%zu misses=%zu copied=%zu prefetches=%zu prefetched=%zu total_hits=%llu total_speculative_hits=%llu total_misses=%llu total_copied=%llu total_prefetches=%llu total_wrong_prefetches=%llu total_prefetch_evictions=%llu total_prefetched=%llu\n", + __func__, + ggml_backend_sched_tensor_name(input), + ggml_backend_name(split_backend), + cache->n_slots, + cache->expert_size, + cache->weights_size, + needed.size(), + needed.size() - misses.size(), + misses.size(), + copied_bytes, + prefetch_count, + prefetched_bytes, + (unsigned long long) cache->hits, + (unsigned long long) cache->speculative_hits, + (unsigned long long) cache->misses, + (unsigned long long) cache->bytes_copied, + (unsigned long long) cache->prefetches, + (unsigned long long) cache->wrong_prefetches, + (unsigned long long) cache->prefetch_evictions, + (unsigned long long) cache->bytes_prefetched); + } + + return true; +} + static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t sched) { GGML_ASSERT(sched); struct ggml_backend_sched_split * splits = sched->splits; ggml_tensor * prev_ids_tensor = nullptr; + int64_t prev_ids_n_expert = -1; std::vector ids; + std::vector id_counts; std::vector used_ids; + const bool moe_log = ggml_backend_sched_moe_log_enabled(); + const int moe_cache_slots = ggml_backend_sched_moe_cache_slots(); + const enum ggml_backend_sched_moe_prefetch_policy moe_prefetch_policy = ggml_backend_sched_moe_prefetch_policy(); + const int moe_prefetch_limit = ggml_backend_sched_moe_prefetch_limit(); for (int split_id = 0; split_id < sched->n_splits; split_id++) { struct ggml_backend_sched_split * split = &splits[split_id]; int split_backend_id = split->backend_id; ggml_backend_t split_backend = sched->backends[split_backend_id]; + std::vector moe_restores; // copy the input tensors to the split backend for (int input_id = 0; input_id < split->n_inputs; input_id++) { @@ -1574,13 +2316,18 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } // when offloading MoE weights, we can reduce the amount of data copied by copying only the experts that are used - ggml_tensor * node = split->graph.nodes[0]; - if (split->graph.n_nodes > 0 && - ggml_backend_buffer_get_usage(input->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS && - ggml_backend_buffer_is_host(input->buffer) && ( - (node->src[0] == input_cpy && node->op == GGML_OP_MUL_MAT_ID) - //|| (node->src[1] == input_cpy && node->op == GGML_OP_ADD_ID) /* GGML_OP_ADD_ID weights are small and not worth splitting */ - )) { + ggml_tensor * node = nullptr; + if (ggml_backend_buffer_get_usage(input->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS && + ggml_backend_buffer_is_host(input->buffer)) { + for (int node_id = 0; node_id < split->graph.n_nodes; ++node_id) { + ggml_tensor * candidate = split->graph.nodes[node_id]; + if (candidate->op == GGML_OP_MUL_MAT_ID && candidate->src[0] == input_cpy) { + node = candidate; + break; + } + } + } + if (node != nullptr) { const int64_t n_expert = node->op == GGML_OP_MUL_MAT_ID ? input->ne[2] : input->ne[1]; const size_t expert_size = node->op == GGML_OP_MUL_MAT_ID ? input->nb[2] : input->nb[1]; @@ -1601,7 +2348,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } } - if (ids_tensor != prev_ids_tensor) { + if (ids_tensor != prev_ids_tensor || n_expert != prev_ids_n_expert) { ids.resize(ggml_nbytes(ids_tensor) / sizeof(int32_t)); ggml_backend_tensor_get_async(ids_backend, ids_tensor, ids.data(), 0, ggml_nbytes(ids_tensor)); ggml_backend_synchronize(ids_backend); @@ -1609,55 +2356,132 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s // find the used experts used_ids.clear(); used_ids.resize(ggml_bitset_size(n_expert)); + id_counts.clear(); + id_counts.resize(n_expert); for (int64_t i1 = 0; i1 < ids_tensor->ne[1]; i1++) { for (int64_t i0 = 0; i0 < ids_tensor->ne[0]; i0++) { int32_t id = ids[i1 * ids_tensor->nb[1]/sizeof(int32_t) + i0 * ids_tensor->nb[0]/sizeof(int32_t)]; GGML_ASSERT(id >= 0 && id < n_expert); ggml_bitset_set(used_ids.data(), id); + id_counts[id]++; } } prev_ids_tensor = ids_tensor; + prev_ids_n_expert = n_expert; } - // group consecutive experts and copy them together - auto copy_experts = [&](int32_t first_id, int32_t last_id) { - const size_t expert_offset = first_id * expert_size; - const size_t expert_size_copy = (last_id - first_id + 1) * expert_size; - const size_t padding = std::min(expert_size, 512); - const size_t padding_end = last_id < n_expert - 1 ? padding : 0; - - ggml_backend_tensor_set_async(split_backend, - input_cpy, - (const uint8_t *)input->data + expert_offset, expert_offset, - // copy a bit extra at the to ensure there are no NaNs in the padding of the last expert - // this is necessary for MMQ in the CUDA backend - expert_size_copy + padding_end); - }; - - int id = 0; - while (!ggml_bitset_get(used_ids.data(), id)) { - id++; - } - int32_t first_id = id; - int32_t last_id = first_id; + const char * moe_cache_bypass_reason = nullptr; + const bool moe_cache_used = ggml_backend_sched_moe_cache_prepare( + sched, split_backend, split_backend_id, input, node, ids_tensor, ids, used_ids, + n_expert, expert_size, moe_cache_slots, moe_prefetch_policy, moe_prefetch_limit, + moe_log, &moe_cache_bypass_reason, moe_restores); + + if (!moe_cache_used) { + if (moe_log && moe_cache_slots > 0) { + GGML_LOG_INFO( + "%s: moe_cache_bypass tensor=%s node=%s ids=%s backend=%s slots=%d reason=%s n_expert=%lld expert_size=%zu\n", + __func__, + ggml_backend_sched_tensor_name(input), + ggml_backend_sched_tensor_name(node), + ggml_backend_sched_tensor_name(ids_tensor), + ggml_backend_name(split_backend), + moe_cache_slots, + moe_cache_bypass_reason != nullptr ? moe_cache_bypass_reason : "unknown", + (long long) n_expert, + expert_size); + } - for (++id; id < n_expert; ++id) { - if (!ggml_bitset_get(used_ids.data(), id)) { - continue; + // group consecutive experts and copy them together + size_t copy_bytes = 0; + int copy_ranges = 0; + auto copy_experts = [&](int32_t first_id, int32_t last_id) { + const size_t expert_offset = first_id * expert_size; + const size_t expert_size_copy = (last_id - first_id + 1) * expert_size; + const size_t padding = std::min(expert_size, 512); + const size_t padding_end = last_id < n_expert - 1 ? padding : 0; + const size_t bytes = expert_size_copy + padding_end; + + ggml_backend_tensor_set_async(split_backend, + input_cpy, + (const uint8_t *)input->data + expert_offset, expert_offset, + // copy a bit extra at the to ensure there are no NaNs in the padding of the last expert + // this is necessary for MMQ in the CUDA backend + bytes); + + if (moe_log) { + copy_bytes += bytes; + copy_ranges++; + } + }; + + int id = 0; + while (id < n_expert && !ggml_bitset_get(used_ids.data(), id)) { + id++; } + if (id < n_expert) { + int32_t first_id = id; + int32_t last_id = first_id; - if (id == last_id + 1) { - last_id = id; - continue; + for (++id; id < n_expert; ++id) { + if (!ggml_bitset_get(used_ids.data(), id)) { + continue; + } + + if (id == last_id + 1) { + last_id = id; + continue; + } + + copy_experts(first_id, last_id); + + first_id = id; + last_id = id; + } + copy_experts(first_id, last_id); } - copy_experts(first_id, last_id); + if (moe_log) { + std::string used_ids_str; + std::string used_id_counts_str; + size_t used_count = 0; + for (int64_t i = 0; i < n_expert; ++i) { + if (!ggml_bitset_get(used_ids.data(), i)) { + continue; + } + if (!used_ids_str.empty()) { + used_ids_str += ","; + } + used_ids_str += std::to_string(i); + if (!used_id_counts_str.empty()) { + used_id_counts_str += ","; + } + used_id_counts_str += std::to_string(i); + used_id_counts_str += ":"; + used_id_counts_str += std::to_string(id_counts[i]); + used_count++; + } - first_id = id; - last_id = id; + GGML_LOG_INFO( + "%s: moe_copy split=%d input=%d tensor=%s node=%s ids=%s src_backend=%s dst_backend=%s n_expert=%lld expert_size=%zu used=%zu used_bytes=%zu ranges=%d copy_bytes=%zu id_counts=[%s] ids=[%s]\n", + __func__, + split_id, + input_id, + ggml_backend_sched_tensor_name(input), + ggml_backend_sched_tensor_name(node), + ggml_backend_sched_tensor_name(ids_tensor), + ggml_backend_name(input_backend), + ggml_backend_name(split_backend), + (long long) n_expert, + expert_size, + used_count, + used_count * expert_size, + copy_ranges, + copy_bytes, + used_id_counts_str.c_str(), + used_ids_str.c_str()); + } } - copy_experts(first_id, last_id); } else { // try async copy, but if not possible, we can still use a sync copy without synchronizing the dst backend, since we handle the synchronization here with multiple copies and events // TODO: add public function to facilitate this, since applications do not have direct access to the backend interface @@ -1674,9 +2498,17 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } } + auto restore_moe_cache_nodes = [&]() { + for (ggml_backend_sched_moe_restore & restore : moe_restores) { + restore.node->src[0] = restore.src0; + restore.node->src[2] = restore.src2; + } + }; + if (!sched->callback_eval) { enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &split->graph); if (ec != GGML_STATUS_SUCCESS) { + restore_moe_cache_nodes(); return ec; } } else { @@ -1699,6 +2531,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &gv); if (ec != GGML_STATUS_SUCCESS) { + restore_moe_cache_nodes(); return ec; } @@ -1712,6 +2545,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s j0 = j1; } } + restore_moe_cache_nodes(); // record the event of this copy if (split->n_inputs > 0) { @@ -1787,6 +2621,7 @@ ggml_backend_sched_t ggml_backend_sched_new( sched->galloc = ggml_gallocr_new_n(sched->bufts, n_backends); sched->op_offload = op_offload; + sched->moe_caches = new std::vector(); ggml_backend_sched_reset(sched); @@ -1802,6 +2637,12 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) { ggml_backend_event_free(sched->events[b][c]); } } + for (ggml_backend_sched_moe_cache * cache : *sched->moe_caches) { + ggml_backend_buffer_free(cache->weights_buffer); + ggml_backend_buffer_free(cache->ids_buffer); + delete cache; + } + delete sched->moe_caches; ggml_gallocr_free(sched->galloc); ggml_free(sched->ctx); ggml_hash_set_free(&sched->hash_set); diff --git a/ggml/src/ggml-common.h b/ggml/src/ggml-common.h index f05683b44cd..8395b036292 100644 --- a/ggml/src/ggml-common.h +++ b/ggml/src/ggml-common.h @@ -109,6 +109,9 @@ typedef sycl::half2 ggml_half2; #define QI_NVFP4 (QK_NVFP4 / (4 * QR_NVFP4)) #define QR_NVFP4 2 +#define QI_F8_E4M3_B128 (QK_F8_E4M3_B128 / (4 * QR_F8_E4M3_B128)) +#define QR_F8_E4M3_B128 1 + #define QI5_0 (QK5_0 / (4 * QR5_0)) #define QR5_0 2 @@ -216,6 +219,13 @@ typedef struct { } block_nvfp4; static_assert(sizeof(block_nvfp4) == sizeof(uint8_t)*(QK_NVFP4/QK_NVFP4_SUB) + QK_NVFP4/2, "wrong nvfp4 block size/padding"); +#define QK_F8_E4M3_B128 128 +typedef struct { + uint8_t e; // E8M0 + uint8_t qs[QK_F8_E4M3_B128]; +} block_f8_e4m3_b128; +static_assert(sizeof(block_f8_e4m3_b128) == sizeof(uint8_t) + QK_F8_E4M3_B128, "wrong f8_e4m3_b128 block size/padding"); + #define QK5_0 32 typedef struct { ggml_half d; // delta diff --git a/ggml/src/ggml-cpu/arch/arm/quants.c b/ggml/src/ggml-cpu/arch/arm/quants.c index fe621332970..98ec4a42980 100644 --- a/ggml/src/ggml-cpu/arch/arm/quants.c +++ b/ggml/src/ggml-cpu/arch/arm/quants.c @@ -82,6 +82,10 @@ void quantize_row_q8_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, i #endif } +void ggml_vec_dot_f8_e4m3_b128_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + ggml_vec_dot_f8_e4m3_b128_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc); +} + void quantize_row_q8_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { assert(k % QK8_1 == 0); const int nb = k / QK8_1; @@ -4242,4 +4246,3 @@ void ggml_vec_dot_iq4_xs_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const v ggml_vec_dot_iq4_xs_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc); #endif } - diff --git a/ggml/src/ggml-cpu/arch/powerpc/quants.c b/ggml/src/ggml-cpu/arch/powerpc/quants.c index 644c380c738..1368474158c 100644 --- a/ggml/src/ggml-cpu/arch/powerpc/quants.c +++ b/ggml/src/ggml-cpu/arch/powerpc/quants.c @@ -2302,3 +2302,7 @@ void ggml_vec_dot_iq4_xs_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const v ggml_vec_dot_iq4_xs_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc); #endif } + +void ggml_vec_dot_f8_e4m3_b128_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + ggml_vec_dot_f8_e4m3_b128_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc); +} diff --git a/ggml/src/ggml-cpu/arch/riscv/quants.c b/ggml/src/ggml-cpu/arch/riscv/quants.c index d3278d6489f..079c387a6b5 100644 --- a/ggml/src/ggml-cpu/arch/riscv/quants.c +++ b/ggml/src/ggml-cpu/arch/riscv/quants.c @@ -4453,3 +4453,7 @@ void ggml_vec_dot_mxfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo ggml_vec_dot_mxfp4_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc); #endif } + +void ggml_vec_dot_f8_e4m3_b128_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + ggml_vec_dot_f8_e4m3_b128_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc); +} diff --git a/ggml/src/ggml-cpu/arch/s390/quants.c b/ggml/src/ggml-cpu/arch/s390/quants.c index 500857579a7..c75994c1857 100644 --- a/ggml/src/ggml-cpu/arch/s390/quants.c +++ b/ggml/src/ggml-cpu/arch/s390/quants.c @@ -1463,3 +1463,7 @@ void ggml_vec_dot_iq4_xs_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const v ggml_vec_dot_iq4_xs_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc); #endif } + +void ggml_vec_dot_f8_e4m3_b128_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + ggml_vec_dot_f8_e4m3_b128_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc); +} diff --git a/ggml/src/ggml-cpu/arch/x86/quants.c b/ggml/src/ggml-cpu/arch/x86/quants.c index 94b19b82bbc..ad5924d4174 100644 --- a/ggml/src/ggml-cpu/arch/x86/quants.c +++ b/ggml/src/ggml-cpu/arch/x86/quants.c @@ -936,7 +936,6 @@ void ggml_vec_dot_mxfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_mxfp4); const __m128i m4b = _mm_set1_epi8(0x0f); - const __m256i mone = _mm256_set1_epi16(1); __m256 accum1 = _mm256_setzero_ps(); __m256 accum2 = _mm256_setzero_ps(); @@ -950,14 +949,16 @@ void ggml_vec_dot_mxfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_1, m4b))); const __m256i q4b_2 = MM256_SET_M128I(_mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_2, 4), m4b)), _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_2, m4b))); - const __m256i p16_1 = mul_add_epi8(q4b_1, q8b_1); - const __m256i p16_2 = mul_add_epi8(q4b_2, q8b_2); - const __m256i p_1 = _mm256_madd_epi16(p16_1, mone); - const __m256i p_2 = _mm256_madd_epi16(p16_2, mone); + // mul_sum_i8_pairs_float lowers to a single VPDPBUSD on AVX-512 VNNI + // (and AVX-VNNI / AVX-VNNI-INT8) hosts, replacing the maddubs+madd_epi16 + // chain that the previous expansion produced. Falls back to the same + // sign+sign+maddubs sequence on plain AVX2. + const __m256 p_1 = mul_sum_i8_pairs_float(q4b_1, q8b_1); + const __m256 p_2 = mul_sum_i8_pairs_float(q4b_2, q8b_2); const __m256 scale0 = _mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y[ib + 0].d)*GGML_CPU_E8M0_TO_FP32_HALF(x[ib + 0].e)); const __m256 scale1 = _mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y[ib + 1].d)*GGML_CPU_E8M0_TO_FP32_HALF(x[ib + 1].e)); - accum1 = _mm256_fmadd_ps(scale0, _mm256_cvtepi32_ps(p_1), accum1); - accum2 = _mm256_fmadd_ps(scale1, _mm256_cvtepi32_ps(p_2), accum2); + accum1 = _mm256_fmadd_ps(scale0, p_1, accum1); + accum2 = _mm256_fmadd_ps(scale1, p_2, accum2); } sumf = hsum_float_8(_mm256_add_ps(accum1, accum2)); @@ -3968,3 +3969,7 @@ void ggml_vec_dot_iq4_xs_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const v ggml_vec_dot_iq4_xs_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc); #endif } + +void ggml_vec_dot_f8_e4m3_b128_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + ggml_vec_dot_f8_e4m3_b128_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc); +} diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 2b3eb5b5ce6..dbf2df85977 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -282,6 +282,12 @@ static const struct ggml_type_traits_cpu type_traits_cpu[GGML_TYPE_COUNT] = { .vec_dot_type = GGML_TYPE_Q8_0, .nrows = 1, }, + [GGML_TYPE_F8_E4M3_B128] = { + .from_float = quantize_row_f8_e4m3_b128, + .vec_dot = ggml_vec_dot_f8_e4m3_b128_q8_0, + .vec_dot_type = GGML_TYPE_Q8_0, + .nrows = 1, + }, [GGML_TYPE_Q2_K] = { .from_float = quantize_row_q2_K, .vec_dot = ggml_vec_dot_q2_K_q8_K, @@ -1822,6 +1828,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { ggml_compute_forward_mul_mat_id(params, tensor); } break; + case GGML_OP_HC_WEIGHTED_SUM: + { + ggml_compute_forward_hc_weighted_sum(params, tensor); + } break; case GGML_OP_OUT_PROD: { ggml_compute_forward_out_prod(params, tensor); @@ -2245,6 +2255,12 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_UNARY_OP_CEIL: case GGML_UNARY_OP_ROUND: case GGML_UNARY_OP_TRUNC: + case GGML_UNARY_OP_FP4_ACT_QUANT: + case GGML_UNARY_OP_FP8_ACT_QUANT: + { + n_tasks = n_threads; + } break; + case GGML_UNARY_OP_SINKHORN_4X4: { n_tasks = 1; } break; @@ -2287,6 +2303,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_CONCAT: case GGML_OP_MUL_MAT: case GGML_OP_MUL_MAT_ID: + case GGML_OP_HC_WEIGHTED_SUM: case GGML_OP_OUT_PROD: { n_tasks = n_threads; diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index a9bc21da6f0..1f796953df5 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -672,6 +672,7 @@ void ggml_compute_forward_add( case GGML_TYPE_Q8_0: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_F8_E4M3_B128: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -1123,6 +1124,7 @@ void ggml_compute_forward_add1( case GGML_TYPE_Q8_1: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_F8_E4M3_B128: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -1253,6 +1255,7 @@ void ggml_compute_forward_acc( case GGML_TYPE_Q8_1: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_F8_E4M3_B128: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -1505,6 +1508,57 @@ void ggml_compute_forward_sum_rows( } } +// ggml_compute_forward_hc_weighted_sum + +void ggml_compute_forward_hc_weighted_sum( + const ggml_compute_params * params, + ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT( dst->type == GGML_TYPE_F32); + // src0: [n_embd, hc_mult, n_batch], src1: [hc_mult, n_batch], + // dst: [n_embd, n_batch]; src0->ne[3] / src1->ne[2..3] all == 1. + GGML_ASSERT(src0->ne[1] == src1->ne[0]); + GGML_ASSERT(src0->ne[2] == src1->ne[1]); + GGML_ASSERT(src0->ne[3] == 1); + GGML_ASSERT(src1->ne[2] == 1 && src1->ne[3] == 1); + GGML_ASSERT(dst->ne[0] == src0->ne[0]); + GGML_ASSERT(dst->ne[1] == src0->ne[2]); + GGML_ASSERT(dst->ne[2] == 1 && dst->ne[3] == 1); + + const int64_t n_embd = src0->ne[0]; + const int64_t hc_mult = src0->ne[1]; + const int64_t n_batch = src0->ne[2]; + + const int ith = params->ith; + const int nth = params->nth; + + // Distribute work across (n_embd * n_batch) output elements so threads + // stay balanced even when n_batch == 1 (the legacy decode case). + const int64_t n_total = n_embd * n_batch; + const int64_t e_start = (n_total * ith) / nth; + const int64_t e_end = (n_total * (ith + 1)) / nth; + + const char * x = (const char *) src0->data; + const char * w = (const char *) src1->data; + float * out = (float *) dst->data; + + for (int64_t idx = e_start; idx < e_end; ++idx) { + const int64_t b = idx / n_embd; + const int64_t e = idx % n_embd; + float sum = 0.0f; + for (int64_t h = 0; h < hc_mult; ++h) { + const float xv = *(const float *) (x + e*src0->nb[0] + h*src0->nb[1] + b*src0->nb[2]); + const float wv = *(const float *) (w + h*src1->nb[0] + b*src1->nb[1]); + sum += xv * wv; + } + *(float *) ((char *) out + e*dst->nb[0] + b*dst->nb[1]) = sum; + } +} + // ggml_compute_forward_mean static void ggml_compute_forward_mean_f32( @@ -4342,6 +4396,7 @@ void ggml_compute_forward_out_prod( case GGML_TYPE_Q8_0: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_F8_E4M3_B128: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -4619,6 +4674,7 @@ void ggml_compute_forward_set( case GGML_TYPE_Q8_1: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_F8_E4M3_B128: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -4843,6 +4899,7 @@ void ggml_compute_forward_get_rows( case GGML_TYPE_Q8_1: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_F8_E4M3_B128: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -5569,6 +5626,7 @@ void ggml_compute_forward_clamp( case GGML_TYPE_Q8_1: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_F8_E4M3_B128: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -9756,6 +9814,18 @@ void ggml_compute_forward_unary( { ggml_compute_forward_trunc(params, dst); } break; + case GGML_UNARY_OP_FP4_ACT_QUANT: + { + ggml_compute_forward_fp4_act_quant(params, dst); + } break; + case GGML_UNARY_OP_FP8_ACT_QUANT: + { + ggml_compute_forward_fp8_act_quant(params, dst); + } break; + case GGML_UNARY_OP_SINKHORN_4X4: + { + ggml_compute_forward_sinkhorn_4x4(params, dst); + } break; case GGML_UNARY_OP_XIELU: { ggml_compute_forward_xielu(params, dst); diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index 3fa1443abc4..44f58fe9579 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -47,6 +47,7 @@ void ggml_compute_forward_rms_norm(const struct ggml_compute_params * params, st void ggml_compute_forward_rms_norm_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_group_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_l2_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_hc_weighted_sum(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_out_prod(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_scale(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_set(const struct ggml_compute_params * params, struct ggml_tensor * dst); diff --git a/ggml/src/ggml-cpu/quants.c b/ggml/src/ggml-cpu/quants.c index e5f9a4083f9..d3e287983ef 100644 --- a/ggml/src/ggml-cpu/quants.c +++ b/ggml/src/ggml-cpu/quants.c @@ -11,6 +11,7 @@ #include #include #include +#include #include // for qsort #include // for GGML_ASSERT @@ -58,6 +59,10 @@ void quantize_row_nvfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, i quantize_row_nvfp4_ref(x, y, k); } +void quantize_row_f8_e4m3_b128(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_f8_e4m3_b128_ref(x, y, k); +} + // // 2-6 bit quantization in super-blocks // @@ -311,6 +316,56 @@ void ggml_vec_dot_nvfp4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, *s = sumf; } +static inline float ggml_f8_e4m3fn_to_fp32_cpu(uint8_t x) { + if ((x & 0x7F) == 0) { + return 0.0f; + } + if ((x & 0x7F) == 0x7F) { + return NAN; + } + + const int sign = x >> 7; + const int exp = (x >> 3) & 0x0F; + const int man = x & 0x07; + const float val = exp == 0 ? ldexpf((float) man, -9) : ldexpf(1.0f + (float) man * 0.125f, exp - 7); + + return sign ? -val : val; +} + +void ggml_vec_dot_f8_e4m3_b128_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + assert(n % QK_F8_E4M3_B128 == 0); + + const block_f8_e4m3_b128 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + const int nb = n / QK_F8_E4M3_B128; + + float sumf = 0; + + for (int ib = 0; ib < nb; ++ib) { + const float dx = GGML_E8M0_TO_FP32(x[ib].e); + + for (int q8b = 0; q8b < QK_F8_E4M3_B128 / QK8_0; ++q8b) { + const block_q8_0 * yb = &y[ib * (QK_F8_E4M3_B128 / QK8_0) + q8b]; + const float dy = GGML_CPU_FP16_TO_FP32(yb->d); + float sumi = 0; + + for (int j = 0; j < QK8_0; ++j) { + sumi += ggml_f8_e4m3fn_to_fp32_cpu(x[ib].qs[q8b * QK8_0 + j]) * yb->qs[j]; + } + + sumf += dx * dy * sumi; + } + } + + *s = sumf; +} + void ggml_vec_dot_q5_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { const int qk = QK8_0; const int nb = n / qk; diff --git a/ggml/src/ggml-cpu/quants.h b/ggml/src/ggml-cpu/quants.h index d4bc87a1c05..3cc51ba7c41 100644 --- a/ggml/src/ggml-cpu/quants.h +++ b/ggml/src/ggml-cpu/quants.h @@ -22,6 +22,7 @@ void quantize_row_q8_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, in void quantize_row_mxfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); void quantize_row_nvfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_f8_e4m3_b128(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); void quantize_row_q2_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); void quantize_row_q3_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); @@ -46,6 +47,7 @@ void ggml_vec_dot_q8_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const voi void ggml_vec_dot_mxfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); void ggml_vec_dot_nvfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_f8_e4m3_b128_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); void ggml_vec_dot_q2_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); void ggml_vec_dot_q3_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); @@ -79,6 +81,7 @@ void ggml_vec_dot_q8_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, c void ggml_vec_dot_mxfp4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); void ggml_vec_dot_nvfp4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_f8_e4m3_b128_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); void ggml_vec_dot_tq1_0_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); void ggml_vec_dot_tq2_0_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); diff --git a/ggml/src/ggml-cpu/unary-ops.cpp b/ggml/src/ggml-cpu/unary-ops.cpp index 1d8344436f0..98d22fd2498 100644 --- a/ggml/src/ggml-cpu/unary-ops.cpp +++ b/ggml/src/ggml-cpu/unary-ops.cpp @@ -97,6 +97,119 @@ static inline float op_trunc(float x) { return truncf(x); } +static inline float act_quant_pow2_scale(float amax, float max_inv, float min_amax) { + const float scaled = fmaxf(amax, min_amax) * max_inv; + return exp2f(ceilf(log2f(scaled))); +} + +static inline uint8_t fp32_to_fp8_e4m3fn(float x) { + if (isnan(x)) { + return 0x7F; + } + + const uint8_t sign = signbit(x) ? 0x80 : 0x00; + const float ax = fabsf(x); + + if (ax == 0.0f) { + return sign; + } + + if (ax < 0x1p-6f) { + const int man = (int) roundf(ax * 512.0f); + if (man <= 0) { + return sign; + } + if (man >= 8) { + return sign | 0x08; + } + return sign | (uint8_t) man; + } + + int exp_unbiased; + const float fr = frexpf(ax, &exp_unbiased); + exp_unbiased -= 1; + + int exp = exp_unbiased + 7; + int man = (int) roundf((2.0f * fr - 1.0f) * 8.0f); + if (man == 8) { + man = 0; + exp++; + } + + if (exp > 15 || (exp == 15 && man > 6)) { + return sign | 0x7E; + } + + return sign | (uint8_t) ((exp << 3) | man); +} + +static inline float fp8_e4m3fn_to_fp32(uint8_t x) { + if ((x & 0x7F) == 0) { + return 0.0f; + } + if ((x & 0x7F) == 0x7F) { + return NAN; + } + + const int sign = x >> 7; + const int exp = (x >> 3) & 0x0F; + const int man = x & 0x07; + const float val = exp == 0 ? ldexpf((float) man, -9) : ldexpf(1.0f + (float) man * 0.125f, exp - 7); + + return sign ? -val : val; +} + +static inline float quant_dequant_fp8_e4m3(float x) { + return fp8_e4m3fn_to_fp32(fp32_to_fp8_e4m3fn(fminf(fmaxf(x, -448.0f), 448.0f))); +} + +static inline float quant_dequant_fp4_e2m1(float x) { + static const float values[16] = { + 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, + 0.0f,-0.5f,-1.0f,-1.5f,-2.0f,-3.0f,-4.0f,-6.0f, + }; + + const float xc = fminf(fmaxf(x, -6.0f), 6.0f); + int best = 0; + float best_err = fabsf(values[0] - xc); + for (int i = 1; i < 16; ++i) { + const float err = fabsf(values[i] - xc); + if (err < best_err) { + best = i; + best_err = err; + } + } + + return values[best]; +} + +template +static inline float act_quant_max_value() { + if constexpr (mode == 4) { + return 6.0f; + } else { + return 448.0f; + } +} + +template +static inline float act_quant_min_amax() { + if constexpr (mode == 4) { + return 0x1.8p-124f; + } else { + return 1.0e-4f; + } +} + +template +static inline float act_quant_dequant(float x) { + if constexpr (mode == 4) { + return quant_dequant_fp4_e2m1(x); + } else { + return quant_dequant_fp8_e4m3(x); + } +} + template static inline void vec_unary_op(int64_t n, dst_t * y, const src0_t * x) { constexpr auto src0_to_f32 = type_conversion_table::to_f32; @@ -107,6 +220,31 @@ static inline void vec_unary_op(int64_t n, dst_t * y, const src0_t * x) { } } +template +static inline void vec_act_quant_op(int64_t n, dst_t * y, const src0_t * x) { + constexpr auto src0_to_f32 = type_conversion_table::to_f32; + constexpr auto f32_to_dst = type_conversion_table::from_f32; + + GGML_ASSERT(n % block_size == 0); + + for (int64_t ib = 0; ib < n; ib += block_size) { + float amax = 0.0f; + for (int64_t i = 0; i < block_size; ++i) { + const float v = fabsf(src0_to_f32(x[ib + i])); + if (isfinite(v)) { + amax = fmaxf(amax, v); + } + } + + const float scale = act_quant_pow2_scale(amax, 1.0f / act_quant_max_value(), act_quant_min_amax()); + const float iscale = 1.0f / scale; + + for (int64_t i = 0; i < block_size; ++i) { + y[ib + i] = f32_to_dst(act_quant_dequant(src0_to_f32(x[ib + i]) * iscale) * scale); + } + } +} + template static void apply_unary_op(const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; @@ -132,6 +270,32 @@ static void apply_unary_op(const ggml_compute_params * params, ggml_tensor * dst } } +template +static void apply_act_quant_op(const ggml_compute_params * params, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + + GGML_ASSERT(ggml_is_contiguous_rows(src0) && ggml_is_contiguous_rows(dst) && ggml_are_same_shape(src0, dst)); + GGML_ASSERT(src0->ne[0] % block_size == 0); + + GGML_TENSOR_UNARY_OP_LOCALS + + GGML_ASSERT(nb0 == sizeof(dst_t)); + GGML_ASSERT(nb00 == sizeof(src0_t)); + + const auto [ir0, ir1] = get_thread_range(params, src0); + + for (int64_t ir = ir0; ir < ir1; ++ir) { + const int64_t i03 = ir/(ne02*ne01); + const int64_t i02 = (ir - i03*ne02*ne01)/ne01; + const int64_t i01 = (ir - i03*ne02*ne01 - i02*ne01); + + dst_t * dst_ptr = (dst_t *) ((char *) dst->data + i03*nb3 + i02*nb2 + i01*nb1 ); + const src0_t * src0_ptr = (const src0_t *) ((const char *) src0->data + i03*nb03 + i02*nb02 + i01*nb01); + + vec_act_quant_op(ne0, dst_ptr, src0_ptr); + } +} + // TODO: Use the 'traits' lookup table (for type conversion fns), instead of a mass of 'if' conditions with long templates template static void unary_op(const ggml_compute_params * params, ggml_tensor * dst) { @@ -154,6 +318,21 @@ static void unary_op(const ggml_compute_params * params, ggml_tensor * dst) { } } +template +static void act_quant_op(const ggml_compute_params * params, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + + /* */ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + apply_act_quant_op(params, dst); + } else if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { + apply_act_quant_op(params, dst); + } else { + fprintf(stderr, "%s: unsupported types: dst: %s, src0: %s\n", __func__, + ggml_type_name(dst->type), ggml_type_name(src0->type)); + GGML_ABORT("fatal error"); + } +} + template static void unary_op_params(const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; @@ -322,6 +501,98 @@ void ggml_compute_forward_trunc(const ggml_compute_params * params, ggml_tensor unary_op(params, dst); } +void ggml_compute_forward_fp4_act_quant(const ggml_compute_params * params, ggml_tensor * dst) { + act_quant_op<32, 4>(params, dst); +} + +void ggml_compute_forward_fp8_act_quant(const ggml_compute_params * params, ggml_tensor * dst) { + act_quant_op<64, 8>(params, dst); +} + +void ggml_compute_forward_sinkhorn_4x4(const ggml_compute_params * params, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + + GGML_ASSERT(src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(src0->ne[0] == 4 && src0->ne[1] == 4); + GGML_ASSERT(src0->ne[3] == 1); + GGML_ASSERT(ggml_is_contiguous(src0) && ggml_is_contiguous(dst)); + + // Distribute the per-batch 4x4 problems across worker threads. Each + // thread handles a slice of the batch dimension (src0->ne[2]). + const int64_t n_batch = src0->ne[2]; + const int ith = params->ith; + const int nth = params->nth; + + const int64_t b0 = (n_batch * ith) / nth; + const int64_t b1 = (n_batch * (ith + 1)) / nth; + + for (int64_t b = b0; b < b1; ++b) { + const float * src = (const float *) ((const char *) src0->data + b * src0->nb[2]); + float * out = (float *) ((char *) dst->data + b * dst->nb[2]); + float x[4][4]; + + for (int r = 0; r < 4; ++r) { + float maxv = src[4*r + 0]; + for (int c = 1; c < 4; ++c) { + maxv = fmaxf(maxv, src[4*r + c]); + } + + float sum = 0.0f; + for (int c = 0; c < 4; ++c) { + x[r][c] = expf(src[4*r + c] - maxv); + sum += x[r][c]; + } + + const float inv_sum = 1.0f / sum; + for (int c = 0; c < 4; ++c) { + x[r][c] = fmaxf(x[r][c] * inv_sum, 1e-6f); + } + } + + for (int c = 0; c < 4; ++c) { + float sum = 0.0f; + for (int r = 0; r < 4; ++r) { + sum += x[r][c]; + } + const float inv_sum = 1.0f / fmaxf(sum, 1e-6f); + for (int r = 0; r < 4; ++r) { + x[r][c] *= inv_sum; + } + } + + for (int it = 1; it < 20; ++it) { + for (int r = 0; r < 4; ++r) { + float sum = 0.0f; + for (int c = 0; c < 4; ++c) { + sum += x[r][c]; + } + const float inv_sum = 1.0f / fmaxf(sum, 1e-6f); + for (int c = 0; c < 4; ++c) { + x[r][c] *= inv_sum; + } + } + + for (int c = 0; c < 4; ++c) { + float sum = 0.0f; + for (int r = 0; r < 4; ++r) { + sum += x[r][c]; + } + const float inv_sum = 1.0f / fmaxf(sum, 1e-6f); + for (int r = 0; r < 4; ++r) { + x[r][c] *= inv_sum; + } + } + } + + for (int r = 0; r < 4; ++r) { + for (int c = 0; c < 4; ++c) { + out[4*r + c] = x[r][c]; + } + } + } +} + void ggml_compute_forward_xielu(const ggml_compute_params * params, ggml_tensor * dst) { const float alpha_n = ggml_get_op_params_f32(dst, 1); const float alpha_p = ggml_get_op_params_f32(dst, 2); @@ -334,4 +605,3 @@ void ggml_compute_forward_xielu(const ggml_compute_params * params, ggml_tensor unary_op_functor(params, dst, xielu_op_params); } - diff --git a/ggml/src/ggml-cpu/unary-ops.h b/ggml/src/ggml-cpu/unary-ops.h index bcad5a3af1a..8febdf791d9 100644 --- a/ggml/src/ggml-cpu/unary-ops.h +++ b/ggml/src/ggml-cpu/unary-ops.h @@ -28,6 +28,9 @@ void ggml_compute_forward_floor(const struct ggml_compute_params * params, struc void ggml_compute_forward_ceil(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_round(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_trunc(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_fp4_act_quant(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_fp8_act_quant(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_sinkhorn_4x4(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_xielu(const struct ggml_compute_params * params, struct ggml_tensor * dst); #ifdef __cplusplus diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 3aec1742ee1..bc265ec2e96 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -784,7 +784,10 @@ static __device__ __forceinline__ void ggml_cuda_memcpy_1(void * __restrict__ ds } static __device__ __forceinline__ float ggml_cuda_e8m0_to_fp32(uint8_t x) { -#if CUDART_VERSION >= 12080 +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + const uint32_t bits = x == 0 ? 0x00400000 : (uint32_t) x << 23; + return __uint_as_float(bits); +#elif CUDART_VERSION >= 12080 const nv_bfloat16 e = __nv_cvt_e8m0_to_bf16raw(x); return (float) e; #else @@ -830,6 +833,21 @@ static __device__ __forceinline__ float ggml_cuda_ue4m3_to_fp32(uint8_t x) { #endif // defined(GGML_USE_HIP) && defined(CDNA3) && defined(FP8_AVAILABLE) && HIP_VERSION >= 60200000 } +static __device__ __forceinline__ float ggml_cuda_f8_e4m3fn_to_fp32(uint8_t x) { + if ((x & 0x7F) == 0) { + return 0.0f; + } + if ((x & 0x7F) == 0x7F) { + return NAN; + } + + const int exp = (x >> 3) & 0x0F; + const int man = x & 0x07; + const float val = exp == 0 ? ldexpf((float) man, -9) : ldexpf(1.0f + (float) man * 0.125f, exp - 7); + + return (x & 0x80) ? -val : val; +} + __device__ __forceinline__ uint8_t ggml_cuda_float_to_fp4_e2m1(float x, float e) { const uint8_t sign_bit = (x < 0.0f) << 3; float ax = fabsf(x) * e; @@ -976,6 +994,13 @@ struct ggml_cuda_type_traits { static constexpr int qi = QI_NVFP4; }; +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_F8_E4M3_B128; + static constexpr int qr = QR_F8_E4M3_B128; + static constexpr int qi = QI_F8_E4M3_B128; +}; + template<> struct ggml_cuda_type_traits { static constexpr int qk = QK_K; diff --git a/ggml/src/ggml-cuda/convert.cu b/ggml/src/ggml-cuda/convert.cu index 61630a35a29..4d74b01ead2 100644 --- a/ggml/src/ggml-cuda/convert.cu +++ b/ggml/src/ggml-cuda/convert.cu @@ -486,6 +486,14 @@ static __global__ void dequantize_block_mxfp4(const void * __restrict__ vx, dst_ } } +static __device__ __forceinline__ void dequantize_f8_e4m3_b128(const void * __restrict__ vx, const int64_t ib, const int iqs, float2 & v) { + const block_f8_e4m3_b128 * x = (const block_f8_e4m3_b128 *) vx; + const float d = ggml_cuda_e8m0_to_fp32(x[ib].e); + + v.x = d * ggml_cuda_f8_e4m3fn_to_fp32(x[ib].qs[iqs + 0]); + v.y = d * ggml_cuda_f8_e4m3fn_to_fp32(x[ib].qs[iqs + 1]); +} + template static void dequantize_block_cuda(const void * vx, dst_t * y, const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03, @@ -758,6 +766,8 @@ to_fp16_cuda_t ggml_get_to_fp16_cuda(ggml_type type) { return dequantize_row_mxfp4_cuda; case GGML_TYPE_NVFP4: return dequantize_row_nvfp4_cuda; + case GGML_TYPE_F8_E4M3_B128: + return dequantize_block_cont_cuda; case GGML_TYPE_F32: return convert_unary_cont_cuda; case GGML_TYPE_BF16: @@ -813,6 +823,8 @@ to_fp32_cuda_t ggml_get_to_fp32_cuda(ggml_type type) { return dequantize_row_mxfp4_cuda; case GGML_TYPE_NVFP4: return dequantize_row_nvfp4_cuda; + case GGML_TYPE_F8_E4M3_B128: + return dequantize_block_cont_cuda; case GGML_TYPE_F16: return convert_unary_cont_cuda; case GGML_TYPE_BF16: @@ -838,6 +850,8 @@ to_fp16_nc_cuda_t ggml_get_to_fp16_nc_cuda(ggml_type type) { return dequantize_block_cuda; case GGML_TYPE_Q8_0: return dequantize_block_cuda; + case GGML_TYPE_F8_E4M3_B128: + return dequantize_block_cuda; case GGML_TYPE_BF16: return convert_unary_cuda; default: @@ -861,6 +875,8 @@ to_bf16_nc_cuda_t ggml_get_to_bf16_nc_cuda(ggml_type type) { return dequantize_block_cuda; case GGML_TYPE_Q8_0: return dequantize_block_cuda; + case GGML_TYPE_F8_E4M3_B128: + return dequantize_block_cuda; case GGML_TYPE_F16: return convert_unary_cuda; default: @@ -884,6 +900,8 @@ to_fp32_nc_cuda_t ggml_get_to_fp32_nc_cuda(ggml_type type) { return dequantize_block_cuda; case GGML_TYPE_Q8_0: return dequantize_block_cuda; + case GGML_TYPE_F8_E4M3_B128: + return dequantize_block_cuda; case GGML_TYPE_BF16: return convert_unary_cuda; default: diff --git a/ggml/src/ggml-cuda/cpy.cu b/ggml/src/ggml-cuda/cpy.cu index d208acf2d5f..95fdc7f5044 100644 --- a/ggml/src/ggml-cuda/cpy.cu +++ b/ggml/src/ggml-cuda/cpy.cu @@ -7,7 +7,7 @@ typedef void (*cpy_kernel_t)(const char * cx, char * cdst); -const int CUDA_CPY_TILE_DIM_2D = 32; // 2D tile dimension for transposed blocks +const int CUDA_CPY_TILE_DIM_2D = 16; // 2D tile dimension for transposed blocks const int CUDA_CPY_BLOCK_NM = 8; // block size of 3rd dimension if available const int CUDA_CPY_BLOCK_ROWS = 8; // block dimension for marching through rows diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 1c2c3b4ac69..957fa20c6c3 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -61,6 +61,7 @@ #include "ggml-cuda/tri.cuh" #include "ggml-cuda/cumsum.cuh" #include "ggml-cuda/fill.cuh" +#include "ggml-cuda/hc-weighted-sum.cuh" #include "ggml.h" #include @@ -2741,6 +2742,15 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_UNARY_OP_TRUNC: ggml_cuda_op_trunc(ctx, dst); break; + case GGML_UNARY_OP_FP4_ACT_QUANT: + ggml_cuda_op_fp4_act_quant(ctx, dst); + break; + case GGML_UNARY_OP_FP8_ACT_QUANT: + ggml_cuda_op_fp8_act_quant(ctx, dst); + break; + case GGML_UNARY_OP_SINKHORN_4X4: + ggml_cuda_op_sinkhorn_4x4(ctx, dst); + break; case GGML_UNARY_OP_EXPM1: ggml_cuda_op_expm1(ctx, dst); break; @@ -2820,6 +2830,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_MUL_MAT_ID: ggml_cuda_mul_mat_id(ctx, dst); break; + case GGML_OP_HC_WEIGHTED_SUM: + ggml_cuda_op_hc_weighted_sum(ctx, dst); + break; case GGML_OP_OUT_PROD: ggml_cuda_out_prod(ctx, dst); break; @@ -3089,6 +3102,30 @@ static void ggml_backend_cuda_synchronize(ggml_backend_t backend) { static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) { bool use_cuda_graph = true; + + // Wide-prefill graphs (e.g. DeepSeek4 batched prefill at ub>=768) + // exceed CUDA graph capture memory budgets. The most reliable signal + // for "this graph is processing many tokens at once" is MUL_MAT_ID's + // ne[2] dimension, which is exactly work_tokens. Regular MUL_MAT + // ne[1] is unreliable because some matmuls (e.g. V4's HC_POST + // batched mixer) have ne[1] = n_embd regardless of work_tokens. + int64_t max_mmid_tokens = 0; + for (int i = 0; i < cgraph->n_nodes; i++) { + ggml_tensor * node = cgraph->nodes[i]; + if (node->op == GGML_OP_MUL_MAT_ID) { + if (node->ne[2] > max_mmid_tokens) { + max_mmid_tokens = node->ne[2]; + } + } + } + if (max_mmid_tokens >= 384) { +#ifndef NDEBUG + GGML_LOG_DEBUG("%s: disabling CUDA graphs due to wide prefill mmid ne[2]=%lld\n", + __func__, (long long) max_mmid_tokens); +#endif + return false; + } + // Loop over nodes in GGML graph to obtain info needed for CUDA graph for (int i = 0; i < cgraph->n_nodes; i++) { @@ -4845,6 +4882,17 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g // TODO: should become: //return ggml_is_contiguous_rows(op->src[0]); return ggml_is_contiguous(op->src[0]); + case GGML_UNARY_OP_FP4_ACT_QUANT: + return op->src[0]->type == op->type && op->ne[0] % 32 == 0 && ggml_is_contiguous(op->src[0]) && + (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16); + case GGML_UNARY_OP_FP8_ACT_QUANT: + return op->src[0]->type == op->type && op->ne[0] % 64 == 0 && ggml_is_contiguous(op->src[0]) && + (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16); + case GGML_UNARY_OP_SINKHORN_4X4: + return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && + op->ne[0] == 4 && op->ne[1] == 4 && op->ne[3] == 1 && + ggml_are_same_shape(op->src[0], op) && + ggml_is_contiguous(op->src[0]); default: return false; } @@ -4908,6 +4956,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_TYPE_Q8_0: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_F8_E4M3_B128: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -5143,6 +5192,17 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_OP_MEAN: case GGML_OP_GROUP_NORM: return ggml_is_contiguous(op->src[0]); + case GGML_OP_HC_WEIGHTED_SUM: + return op->src[0]->type == GGML_TYPE_F32 && + op->src[1]->type == GGML_TYPE_F32 && + op->type == GGML_TYPE_F32 && + op->src[0]->ne[1] == op->src[1]->ne[0] && + op->src[0]->ne[2] == op->src[1]->ne[1] && + op->src[0]->ne[3] == 1 && + op->src[1]->ne[2] == 1 && op->src[1]->ne[3] == 1 && + op->ne[0] == op->src[0]->ne[0] && + op->ne[1] == op->src[0]->ne[2] && + op->ne[2] == 1 && op->ne[3] == 1; case GGML_OP_PAD: return true; case GGML_OP_UPSCALE: diff --git a/ggml/src/ggml-cuda/hc-weighted-sum.cu b/ggml/src/ggml-cuda/hc-weighted-sum.cu new file mode 100644 index 00000000000..29d1a13b747 --- /dev/null +++ b/ggml/src/ggml-cuda/hc-weighted-sum.cu @@ -0,0 +1,114 @@ +#include "hc-weighted-sum.cuh" + +// Per-batch n_embd-major layout. Each (block.y, thread block on x) pair +// owns one batch and a slice of n_embd. The h4 specialization keeps the +// 4 weights in registers. +static __global__ void hc_weighted_sum_h4_f32( + const char * __restrict__ x, + const char * __restrict__ w, + float * __restrict__ dst, + const int64_t n_embd, + const int64_t nbx0, + const int64_t nbx1, + const int64_t nbx2, + const int64_t nbw0, + const int64_t nbw1, + const int64_t nbd0, + const int64_t nbd1) { + const int64_t b = blockIdx.y; + const int64_t tid = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + const int64_t stride = (int64_t) blockDim.x * gridDim.x; + + const char * xb = x + b*nbx2; + const char * wb = w + b*nbw1; + char * db = ((char *) dst) + b*nbd1; + + const float w0 = *(const float *) (wb + 0*nbw0); + const float w1 = *(const float *) (wb + 1*nbw0); + const float w2 = *(const float *) (wb + 2*nbw0); + const float w3 = *(const float *) (wb + 3*nbw0); + + for (int64_t e = tid; e < n_embd; e += stride) { + const char * xe = xb + e*nbx0; + const float v = *(const float *) (xe + 0*nbx1) * w0 + + *(const float *) (xe + 1*nbx1) * w1 + + *(const float *) (xe + 2*nbx1) * w2 + + *(const float *) (xe + 3*nbx1) * w3; + *(float *) (db + e*nbd0) = v; + } +} + +static __global__ void hc_weighted_sum_f32( + const char * __restrict__ x, + const char * __restrict__ w, + float * __restrict__ dst, + const int64_t n_embd, + const int64_t hc_mult, + const int64_t nbx0, + const int64_t nbx1, + const int64_t nbx2, + const int64_t nbw0, + const int64_t nbw1, + const int64_t nbd0, + const int64_t nbd1) { + const int64_t b = blockIdx.y; + const int64_t tid = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + const int64_t stride = (int64_t) blockDim.x * gridDim.x; + + const char * xb = x + b*nbx2; + const char * wb = w + b*nbw1; + char * db = ((char *) dst) + b*nbd1; + + for (int64_t e = tid; e < n_embd; e += stride) { + const char * xe = xb + e*nbx0; + float sum = 0.0f; + for (int64_t h = 0; h < hc_mult; ++h) { + sum += *(const float *) (xe + h*nbx1) * *(const float *) (wb + h*nbw0); + } + *(float *) (db + e*nbd0) = sum; + } +} + +void ggml_cuda_op_hc_weighted_sum(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT( dst->type == GGML_TYPE_F32); + // src0: [n_embd, hc_mult, n_batch]; src1: [hc_mult, n_batch]; + // dst: [n_embd, n_batch]; src0->ne[3]/src1->ne[2..3] all == 1. + GGML_ASSERT(src0->ne[1] == src1->ne[0]); + GGML_ASSERT(src0->ne[2] == src1->ne[1]); + GGML_ASSERT(src0->ne[3] == 1); + GGML_ASSERT(src1->ne[2] == 1 && src1->ne[3] == 1); + GGML_ASSERT(dst->ne[0] == src0->ne[0]); + GGML_ASSERT(dst->ne[1] == src0->ne[2]); + GGML_ASSERT(dst->ne[2] == 1 && dst->ne[3] == 1); + + const int64_t n_embd = src0->ne[0]; + const int64_t hc_mult = src0->ne[1]; + const int64_t n_batch = src0->ne[2]; + + const int64_t num_blocks_x = (n_embd + CUDA_HC_WEIGHTED_SUM_BLOCK_SIZE - 1) / CUDA_HC_WEIGHTED_SUM_BLOCK_SIZE; + const dim3 block_nums((unsigned int) num_blocks_x, (unsigned int) n_batch, 1); + const dim3 block_dims(CUDA_HC_WEIGHTED_SUM_BLOCK_SIZE, 1, 1); + + const char * src0_d = (const char *) src0->data; + const char * src1_d = (const char *) src1->data; + float * dst_d = (float *) dst->data; + + if (hc_mult == 4) { + hc_weighted_sum_h4_f32<<>>( + src0_d, src1_d, dst_d, n_embd, + src0->nb[0], src0->nb[1], src0->nb[2], + src1->nb[0], src1->nb[1], + dst->nb[0], dst->nb[1]); + } else { + hc_weighted_sum_f32<<>>( + src0_d, src1_d, dst_d, n_embd, hc_mult, + src0->nb[0], src0->nb[1], src0->nb[2], + src1->nb[0], src1->nb[1], + dst->nb[0], dst->nb[1]); + } +} diff --git a/ggml/src/ggml-cuda/hc-weighted-sum.cuh b/ggml/src/ggml-cuda/hc-weighted-sum.cuh new file mode 100644 index 00000000000..ab1718300b6 --- /dev/null +++ b/ggml/src/ggml-cuda/hc-weighted-sum.cuh @@ -0,0 +1,5 @@ +#include "common.cuh" + +#define CUDA_HC_WEIGHTED_SUM_BLOCK_SIZE 256 + +void ggml_cuda_op_hc_weighted_sum(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 8f55cace1a1..181c065a4f9 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -4,9 +4,19 @@ #include "vecdotq.cuh" #include +#include +#include typedef float (*vec_dot_q_cuda_t)(const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs); +static bool ggml_cuda_f8_approx_dp4a_enabled() { + static const bool enabled = []() { + const char * env = std::getenv("GGML_CUDA_F8_APPROX_DP4A"); + return env != nullptr && std::strcmp(env, "0") != 0; + }(); + return enabled; +} + static constexpr __device__ vec_dot_q_cuda_t get_vec_dot_q_cuda(ggml_type type) { switch (type) { case GGML_TYPE_Q1_0: return vec_dot_q1_0_q8_1; @@ -17,6 +27,7 @@ static constexpr __device__ vec_dot_q_cuda_t get_vec_dot_q_cuda(ggml_type type) case GGML_TYPE_Q8_0: return vec_dot_q8_0_q8_1; case GGML_TYPE_MXFP4: return vec_dot_mxfp4_q8_1; case GGML_TYPE_NVFP4: return vec_dot_nvfp4_q8_1; + case GGML_TYPE_F8_E4M3_B128: return vec_dot_f8_e4m3_b128_q8_1; case GGML_TYPE_Q2_K: return vec_dot_q2_K_q8_1; case GGML_TYPE_Q3_K: return vec_dot_q3_K_q8_1; case GGML_TYPE_Q4_K: return vec_dot_q4_K_q8_1; @@ -45,6 +56,7 @@ static constexpr __host__ __device__ int get_vdr_mmvq(ggml_type type) { case GGML_TYPE_Q8_0: return VDR_Q8_0_Q8_1_MMVQ; case GGML_TYPE_MXFP4: return VDR_MXFP4_Q8_1_MMVQ; case GGML_TYPE_NVFP4: return VDR_NVFP4_Q8_1_MMVQ; + case GGML_TYPE_F8_E4M3_B128: return VDR_F8_E4M3_B128_Q8_1_MMVQ; case GGML_TYPE_Q2_K: return VDR_Q2_K_Q8_1_MMVQ; case GGML_TYPE_Q3_K: return VDR_Q3_K_Q8_1_MMVQ; case GGML_TYPE_Q4_K: return VDR_Q4_K_Q8_1_MMVQ; @@ -134,7 +146,7 @@ static constexpr __host__ __device__ int get_mmvq_mmid_max_batch_turing_plus(ggm case GGML_TYPE_IQ2_S: return 7; case GGML_TYPE_IQ3_S: return 6; case GGML_TYPE_IQ3_XXS: return 7; - case GGML_TYPE_MXFP4: return 7; + case GGML_TYPE_MXFP4: return 8; case GGML_TYPE_Q2_K: return 7; case GGML_TYPE_Q3_K: return 5; default: return MMVQ_MAX_BATCH_SIZE; @@ -368,10 +380,13 @@ static constexpr __host__ __device__ int calc_nwarps(ggml_type type, int ncols_d return 1; } -static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int table_id, bool small_k = false, int nwarps = 1) { +static constexpr __host__ __device__ int calc_rows_per_block(ggml_type type, int ncols_dst, int table_id, bool small_k = false, int nwarps = 1) { if (table_id == MMVQ_PARAMETERS_GENERIC || table_id == MMVQ_PARAMETERS_GCN) { switch (ncols_dst) { case 1: + if ((type == GGML_TYPE_F8_E4M3_B128 || type == GGML_TYPE_IQ4_XS) && !small_k) { + return 2; + } return small_k ? nwarps : 1; case 2: case 3: @@ -388,7 +403,7 @@ static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int return 1; } -template +template __launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id())*ggml_cuda_get_physical_warp_size(), 1) static __global__ void mul_mat_vec_q( const void * __restrict__ vx, const void * __restrict__ vy, const int32_t * __restrict__ ids, const ggml_cuda_mm_fusion_args_device fusion, float * __restrict__ dst, @@ -403,7 +418,7 @@ static __global__ void mul_mat_vec_q( constexpr int vdr = get_vdr_mmvq(type); constexpr mmvq_parameter_table_id table_id = get_device_table_id(); constexpr int nwarps = calc_nwarps(type, ncols_dst, table_id); - constexpr int rows_per_cuda_block = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps); + constexpr int rows_per_cuda_block = calc_rows_per_block(type, ncols_dst, table_id, small_k, nwarps); constexpr int warp_size = ggml_cuda_get_physical_warp_size(); constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type); @@ -413,6 +428,30 @@ static __global__ void mul_mat_vec_q( const int blocks_per_row_x = ncols_x / qk; constexpr int blocks_per_iter = vdr * nwarps*warp_size / qi; +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + constexpr bool use_f8_approx_dp4a = f8_approx_dp4a && type == GGML_TYPE_F8_E4M3_B128; + constexpr bool use_f8_approx_shared_lut = use_f8_approx_dp4a && ncols_dst <= 8; + constexpr bool use_f8_shared_lut = type == GGML_TYPE_F8_E4M3_B128 && ncols_dst == 1 && !small_k && !use_f8_approx_dp4a; +#else + constexpr bool use_f8_approx_dp4a = false; + constexpr bool use_f8_approx_shared_lut = false; + constexpr bool use_f8_shared_lut = false; +#endif + + __shared__ float f8_lut_shared[use_f8_shared_lut ? 256 : 1]; + __shared__ int8_t f8_i8_lut_shared[use_f8_approx_shared_lut ? 256 : 1]; + if constexpr (use_f8_shared_lut) { + for (int i = tid; i < 256; i += nwarps*warp_size) { + f8_lut_shared[i] = kvalues_f8_e4m3fn[i]; + } + __syncthreads(); + } else if constexpr (use_f8_approx_shared_lut) { + for (int i = tid; i < 256; i += nwarps*warp_size) { + f8_i8_lut_shared[i] = kvalues_f8_e4m3fn_i8_approx[i]; + } + __syncthreads(); + } + const uint32_t channel_dst = blockIdx.y; uint32_t channel_x; @@ -490,12 +529,38 @@ static __global__ void mul_mat_vec_q( for (int j = 0; j < ncols_dst; ++j) { #pragma unroll for (int i = 0; i < rows_per_cuda_block; ++i) { - tmp[j][i] += vec_dot_q_cuda( - vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + if constexpr (use_f8_approx_dp4a) { + if constexpr (use_f8_approx_shared_lut) { + tmp[j][i] += vec_dot_f8_e4m3_b128_q8_1_approx_dp4a_shared_lut( + vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs, f8_i8_lut_shared); + } else { + tmp[j][i] += vec_dot_f8_e4m3_b128_q8_1_approx_dp4a( + vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + } + } else if constexpr (use_f8_shared_lut) { + tmp[j][i] += vec_dot_f8_e4m3_b128_q8_1_shared_lut( + vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs, f8_lut_shared); + } else { + tmp[j][i] += vec_dot_q_cuda( + vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + } if constexpr (has_fusion) { if (use_gate) { - tmp_gate[j][i] += vec_dot_q_cuda( - vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + if constexpr (use_f8_approx_dp4a) { + if constexpr (use_f8_approx_shared_lut) { + tmp_gate[j][i] += vec_dot_f8_e4m3_b128_q8_1_approx_dp4a_shared_lut( + vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs, f8_i8_lut_shared); + } else { + tmp_gate[j][i] += vec_dot_f8_e4m3_b128_q8_1_approx_dp4a( + vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + } + } else if constexpr (use_f8_shared_lut) { + tmp_gate[j][i] += vec_dot_f8_e4m3_b128_q8_1_shared_lut( + vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs, f8_lut_shared); + } else { + tmp_gate[j][i] += vec_dot_q_cuda( + vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + } } } } @@ -658,14 +723,14 @@ static std::pair calc_launch_params( const int ncols_dst, const int nrows_x, const int nchannels_dst, const int nsamples_or_ntokens, const int warp_size, const mmvq_parameter_table_id table_id, const bool small_k = false) { const int nwarps = calc_nwarps(type, ncols_dst, table_id); - const int rpb = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps); + const int rpb = calc_rows_per_block(type, ncols_dst, table_id, small_k, nwarps); const int64_t nblocks = (nrows_x + rpb - 1) / rpb; const dim3 block_nums(nblocks, nchannels_dst, nsamples_or_ntokens); const dim3 block_dims(warp_size, nwarps, 1); return {block_nums, block_dims}; } -template +template static void mul_mat_vec_q_switch_fusion( const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t stride_row_x, const uint32_t stride_col_y, @@ -678,7 +743,7 @@ static void mul_mat_vec_q_switch_fusion( const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr; if constexpr (c_ncols_dst == 1) { if (has_fusion) { - mul_mat_vec_q<<>> + mul_mat_vec_q<<>> (vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride); @@ -688,12 +753,42 @@ static void mul_mat_vec_q_switch_fusion( GGML_ASSERT(!has_fusion && "fusion only supported for ncols_dst=1"); - mul_mat_vec_q<<>> + mul_mat_vec_q<<>> (vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride); } +template +static void mul_mat_vec_q_switch_fusion_runtime( + const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, + const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t stride_row_x, const uint32_t stride_col_y, + const uint32_t stride_col_dst, const uint3 channel_ratio, const uint32_t stride_channel_x, + const uint32_t stride_channel_y, const uint32_t stride_channel_dst, const uint3 sample_ratio, + const uint32_t stride_sample_x, const uint32_t stride_sample_y, const uint32_t stride_sample_dst, + const dim3 & block_nums, const dim3 & block_dims, const int nbytes_shared, + const uint32_t ids_stride, const bool f8_approx_dp4a, cudaStream_t stream) { + + if constexpr (type == GGML_TYPE_F8_E4M3_B128) { + if (f8_approx_dp4a) { + mul_mat_vec_q_switch_fusion( + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio, + stride_sample_x, stride_sample_y, stride_sample_dst, block_nums, block_dims, nbytes_shared, + ids_stride, stream); + return; + } + } else { + GGML_UNUSED(f8_approx_dp4a); + } + + mul_mat_vec_q_switch_fusion( + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio, + stride_sample_x, stride_sample_y, stride_sample_dst, block_nums, block_dims, nbytes_shared, + ids_stride, stream); +} + template static void mul_mat_vec_q_moe_launch( const void * vx, const void * vy, const int32_t * ids, float * dst, @@ -748,8 +843,10 @@ static void mul_mat_vec_q_switch_ncols_dst( constexpr int vdr = get_vdr_mmvq(type); const int blocks_per_row_x = ncols_x / qk; const int blocks_per_iter_1warp = vdr * warp_size / qi; + const int small_k_blocks_per_iter_1warp = + type == GGML_TYPE_F8_E4M3_B128 ? 4 * warp_size / qi : blocks_per_iter_1warp; const int nwarps = calc_nwarps(type, c_ncols_dst, table_id); - bool use = nwarps > 1 && blocks_per_row_x < nwarps * blocks_per_iter_1warp; + bool use = nwarps > 1 && blocks_per_row_x < nwarps * small_k_blocks_per_iter_1warp; constexpr std::array iq_slow_turing = { GGML_TYPE_IQ3_XXS, @@ -782,8 +879,8 @@ static void mul_mat_vec_q_switch_ncols_dst( return use; }; - if (has_ids && ncols_dst > 1) { - // Multi-token MUL_MAT_ID path - dedicated MoE kernel + if (has_ids) { + // MUL_MAT_ID path - dedicated MoE kernel mul_mat_vec_q_moe_launch( vx, vy, ids, dst, ncols_x, nchannels_y_fd, nrows_x, stride_row_x, stride_col_y, stride_col_dst, @@ -792,6 +889,8 @@ static void mul_mat_vec_q_switch_ncols_dst( return; } + const bool f8_approx_dp4a = type == GGML_TYPE_F8_E4M3_B128 && ggml_cuda_f8_approx_dp4a_enabled(); + switch (ncols_dst) { case 1: { constexpr int c_ncols_dst = 1; @@ -801,76 +900,76 @@ static void mul_mat_vec_q_switch_ncols_dst( if (use_small_k) { std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id, true); - mul_mat_vec_q_switch_fusion( + mul_mat_vec_q_switch_fusion_runtime( vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, dims.first, dims.second, 0, ids_stride, - stream); + f8_approx_dp4a, stream); } else { std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id); - mul_mat_vec_q_switch_fusion( + mul_mat_vec_q_switch_fusion_runtime( vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, dims.first, dims.second, 0, ids_stride, - stream); + f8_approx_dp4a, stream); } } break; case 2: { constexpr int c_ncols_dst = 2; std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id); - mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, - channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, - sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, - dims.first, dims.second, 0, ids_stride, stream); + mul_mat_vec_q_switch_fusion_runtime(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, + dims.first, dims.second, 0, ids_stride, f8_approx_dp4a, stream); } break; case 3: { constexpr int c_ncols_dst = 3; std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id); - mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, - channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, - sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, - dims.first, dims.second, 0, ids_stride, stream); + mul_mat_vec_q_switch_fusion_runtime(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, + dims.first, dims.second, 0, ids_stride, f8_approx_dp4a, stream); } break; case 4: { constexpr int c_ncols_dst = 4; std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id); - mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, - channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, - sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, - dims.first, dims.second, 0, ids_stride, stream); + mul_mat_vec_q_switch_fusion_runtime(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, + dims.first, dims.second, 0, ids_stride, f8_approx_dp4a, stream); } break; case 5: { constexpr int c_ncols_dst = 5; std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id); - mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, - channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, - sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, - dims.first, dims.second, 0, ids_stride, stream); + mul_mat_vec_q_switch_fusion_runtime(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, + dims.first, dims.second, 0, ids_stride, f8_approx_dp4a, stream); } break; case 6: { constexpr int c_ncols_dst = 6; std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id); - mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, - channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, - sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, - dims.first, dims.second, 0, ids_stride, stream); + mul_mat_vec_q_switch_fusion_runtime(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, + dims.first, dims.second, 0, ids_stride, f8_approx_dp4a, stream); } break; case 7: { constexpr int c_ncols_dst = 7; std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id); - mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, - channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, - sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, - dims.first, dims.second, 0, ids_stride, stream); + mul_mat_vec_q_switch_fusion_runtime(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, + dims.first, dims.second, 0, ids_stride, f8_approx_dp4a, stream); } break; case 8: { constexpr int c_ncols_dst = 8; std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id); - mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, - channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, - sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, - dims.first, dims.second, 0, ids_stride, stream); + mul_mat_vec_q_switch_fusion_runtime(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, + dims.first, dims.second, 0, ids_stride, f8_approx_dp4a, stream); } break; default: GGML_ABORT("fatal error"); @@ -936,6 +1035,12 @@ static void mul_mat_vec_q_switch_type( nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); break; + case GGML_TYPE_F8_E4M3_B128: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; case GGML_TYPE_Q2_K: mul_mat_vec_q_switch_ncols_dst (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index ef98f675aa7..5900fe0a15c 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -301,6 +301,9 @@ static void rms_norm_f32_cuda( if (ncols < 1024) { const dim3 block_dims(256, 1, 1); rms_norm_f32<256, false><< WARP_SIZE ? 32 * sizeof(float): 0, stream>>>(x, dst, ncols, stride_row, stride_channel, stride_sample, eps); + } else if (ncols == 1024) { + const dim3 block_dims(512, 1, 1); + rms_norm_f32<512, false><<>>(x, dst, ncols, stride_row, stride_channel, stride_sample, eps); } else { const dim3 block_dims(1024, 1, 1); rms_norm_f32<1024, false><< WARP_SIZE ? 32 * sizeof(float): 0, stream>>>(x, dst, ncols, stride_row, stride_channel, stride_sample, eps); @@ -349,6 +352,11 @@ static void rms_norm_mul_f32_cuda(const float * x, rms_norm_f32<256, true><< WARP_SIZE ? 32 * sizeof(float): 0, stream>>>( x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed); + } else if (ncols == 1024) { + const dim3 block_dims(512, 1, 1); + rms_norm_f32<512, true><<>>( + x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, + mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed); } else { const dim3 block_dims(1024, 1, 1); rms_norm_f32<1024, true><< WARP_SIZE ? 32 * sizeof(float): 0, stream>>>( @@ -372,6 +380,13 @@ static void rms_norm_mul_f32_cuda(const float * x, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add, add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed); + } else if (ncols == 1024) { + const dim3 block_dims(512, 1, 1); + rms_norm_f32<512, true, true><<>>( + x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, + mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add, + add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, + add_nchannels_packed, add_nsamples_packed); } else { const dim3 block_dims(1024, 1, 1); rms_norm_f32<1024, true, true><< WARP_SIZE ? 32 * sizeof(float): 0, stream>>>( diff --git a/ggml/src/ggml-cuda/quantize.cu b/ggml/src/ggml-cuda/quantize.cu index 4300ffc148c..3b48fb24bd1 100644 --- a/ggml/src/ggml-cuda/quantize.cu +++ b/ggml/src/ggml-cuda/quantize.cu @@ -1,6 +1,7 @@ #include "quantize.cuh" #include +template __launch_bounds__(CUDA_QUANTIZE_BLOCK_SIZE, 1) static __global__ void quantize_q8_1( const float * __restrict__ x, void * __restrict__ vy, @@ -30,13 +31,19 @@ static __global__ void quantize_q8_1( const float xi = i0 < ne00 ? x[i03*s03 + i02*s02 + i01*s01 + i00] : 0.0f; float amax = fabsf(xi); - float sum = xi; + float sum = 0.0f; + if constexpr (need_sum) { + sum = xi; + } amax = warp_reduce_max(amax); - sum = warp_reduce_sum(sum); + if constexpr (need_sum) { + sum = warp_reduce_sum(sum); + } - const float d = amax / 127.0f; - const int8_t q = amax == 0.0f ? 0 : roundf(xi / d); + const float d_inv = amax == 0.0f ? 0.0f : 127.0f / amax; + const float d = amax / 127.0f; + const int8_t q = roundf(xi * d_inv); y[ib].qs[iqs] = q; @@ -44,7 +51,11 @@ static __global__ void quantize_q8_1( return; } - y[ib].ds = make_half2(d, sum); + if constexpr (need_sum) { + y[ib].ds = make_half2(d, sum); + } else { + ((half *) &y[ib].ds)[0] = __float2half(d); + } } __device__ __forceinline__ uint8_t compute_e8m0_scale(float amax) { @@ -282,8 +293,12 @@ void quantize_row_q8_1_cuda( const int64_t block_num_x = (ne0 + CUDA_QUANTIZE_BLOCK_SIZE - 1) / CUDA_QUANTIZE_BLOCK_SIZE; const dim3 num_blocks(block_num_x, ne1, ne2*ne3); const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE, 1, 1); - quantize_q8_1<<>>(x, vy, ne00, s01, s02, s03, ne0, ne1, ne2_fastdiv); - GGML_UNUSED(type_src0); + if (type_src0 == GGML_TYPE_F8_E4M3_B128 || type_src0 == GGML_TYPE_MXFP4 || type_src0 == GGML_TYPE_NVFP4 || + type_src0 == GGML_TYPE_IQ4_XS) { + quantize_q8_1<<>>(x, vy, ne00, s01, s02, s03, ne0, ne1, ne2_fastdiv); + } else { + quantize_q8_1<<>>(x, vy, ne00, s01, s02, s03, ne0, ne1, ne2_fastdiv); + } } void quantize_mmq_q8_1_cuda( diff --git a/ggml/src/ggml-cuda/top-k.cu b/ggml/src/ggml-cuda/top-k.cu index 59ce36fb1c9..d420e0f2c5d 100644 --- a/ggml/src/ggml-cuda/top-k.cu +++ b/ggml/src/ggml-cuda/top-k.cu @@ -1,6 +1,10 @@ #include "argsort.cuh" #include "top-k.cuh" +#include +#include +#include + #ifdef GGML_CUDA_USE_CUB # include # if (CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 2) @@ -47,6 +51,93 @@ static int next_power_of_2(int x) { #endif // CUB_TOP_K_AVAILABLE +template +static __global__ void top_k_warp_f32_i32(const float * src, int * dst, const int k, const int nrows) { + constexpr int experts_per_thread = (ncols + WARP_SIZE - 1) / WARP_SIZE; + + const int row = blockIdx.x * blockDim.y + threadIdx.y; + if (row >= nrows) { + return; + } + + const int lane = threadIdx.x; + src += row * ncols; + dst += row * k; + + float vals[experts_per_thread]; + uint32_t active_mask = 0; + +#pragma unroll + for (int i = 0; i < experts_per_thread; ++i) { + const int idx = lane + i * WARP_SIZE; + const bool active = idx < ncols; + if (active) { + active_mask |= 1u << i; + } + float val = active ? src[idx] : -INFINITY; + vals[i] = __isnanf(val) ? -FLT_MAX : val; + } + + for (int out = 0; out < k; ++out) { + float max_val = -INFINITY; + int max_idx = INT_MAX; + +#pragma unroll + for (int i = 0; i < experts_per_thread; ++i) { + const int idx = lane + i * WARP_SIZE; + if (((active_mask >> i) & 1u) && (vals[i] > max_val || (vals[i] == max_val && idx < max_idx))) { + max_val = vals[i]; + max_idx = idx; + } + } + +#pragma unroll + for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) { + const float other_val = __shfl_xor_sync(0xFFFFFFFF, max_val, mask, WARP_SIZE); + const int other_idx = __shfl_xor_sync(0xFFFFFFFF, max_idx, mask, WARP_SIZE); + if (other_val > max_val || (other_val == max_val && other_idx < max_idx)) { + max_val = other_val; + max_idx = other_idx; + } + } + + if (lane == out) { + dst[out] = max_idx; + } + + if (max_idx < ncols && (max_idx & (WARP_SIZE - 1)) == lane) { + active_mask &= ~(1u << (max_idx / WARP_SIZE)); + } + } +} + +static bool top_k_warp(const float * src, int * dst, const int ncols, const int nrows, const int k, cudaStream_t stream) { + if (k <= 0 || k > WARP_SIZE) { + return false; + } + + constexpr int rows_per_block = 4; + const dim3 grid((nrows + rows_per_block - 1) / rows_per_block, 1, 1); + const dim3 block(WARP_SIZE, rows_per_block, 1); + + switch (ncols) { + case 128: + top_k_warp_f32_i32<128><<>>(src, dst, k, nrows); + return true; + case 256: + top_k_warp_f32_i32<256><<>>(src, dst, k, nrows); + return true; + case 512: + top_k_warp_f32_i32<512><<>>(src, dst, k, nrows); + return true; + case 576: + top_k_warp_f32_i32<576><<>>(src, dst, k, nrows); + return true; + default: + return false; + } +} + void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const float * src0_d = (const float *) src0->data; @@ -62,6 +153,9 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int64_t nrows = ggml_nrows(src0); const int64_t k = dst->ne[0]; ggml_cuda_pool & pool = ctx.pool(); + if (top_k_warp(src0_d, dst_d, ncols, nrows, k, stream)) { + return; + } #ifdef CUB_TOP_K_AVAILABLE // TODO: Switch to `DeviceSegmentedTopK` for multi-row TopK once implemented // https://github.com/NVIDIA/cccl/issues/6391 diff --git a/ggml/src/ggml-cuda/unary.cu b/ggml/src/ggml-cuda/unary.cu index 2aeba26f414..1fec733085d 100644 --- a/ggml/src/ggml-cuda/unary.cu +++ b/ggml/src/ggml-cuda/unary.cu @@ -114,6 +114,140 @@ static __device__ __forceinline__ float op_trunc(float x) { return trunc(x); } +static __device__ __forceinline__ float act_quant_pow2_scale(float amax, float max_inv, float min_amax) { + const float scaled = fmaxf(amax, min_amax) * max_inv; + return exp2f(ceilf(log2f(scaled))); +} + +static __device__ __forceinline__ uint8_t fp32_to_fp8_e4m3fn(float x) { + if (isnan(x)) { + return 0x7F; + } + + const uint8_t sign = signbit(x) ? 0x80 : 0x00; + const float ax = fabsf(x); + + if (ax == 0.0f) { + return sign; + } + + if (ax < 0x1p-6f) { + const int man = (int) roundf(ax * 512.0f); + if (man <= 0) { + return sign; + } + if (man >= 8) { + return sign | 0x08; + } + return sign | (uint8_t) man; + } + + int exp_unbiased; + const float fr = frexpf(ax, &exp_unbiased); + exp_unbiased -= 1; + + int exp = exp_unbiased + 7; + int man = (int) roundf((2.0f * fr - 1.0f) * 8.0f); + if (man == 8) { + man = 0; + exp++; + } + + if (exp > 15 || (exp == 15 && man > 6)) { + return sign | 0x7E; + } + + return sign | (uint8_t) ((exp << 3) | man); +} + +static __device__ __forceinline__ float fp8_e4m3fn_to_fp32(uint8_t x) { + if ((x & 0x7F) == 0) { + return 0.0f; + } + if ((x & 0x7F) == 0x7F) { + return NAN; + } + + const int sign = x >> 7; + const int exp = (x >> 3) & 0x0F; + const int man = x & 0x07; + const float val = exp == 0 ? ldexpf((float) man, -9) : ldexpf(1.0f + (float) man * 0.125f, exp - 7); + + return sign ? -val : val; +} + +static __device__ __forceinline__ float quant_dequant_fp8_e4m3(float x) { + return fp8_e4m3fn_to_fp32(fp32_to_fp8_e4m3fn(fminf(fmaxf(x, -448.0f), 448.0f))); +} + +static __device__ __forceinline__ float quant_dequant_fp4_e2m1(float x) { + const float xc = fminf(fmaxf(x, -6.0f), 6.0f); + const float values[16] = { + 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, + 0.0f,-0.5f,-1.0f,-1.5f,-2.0f,-3.0f,-4.0f,-6.0f, + }; + + int best = 0; + float best_err = fabsf(values[0] - xc); +#pragma unroll + for (int i = 1; i < 16; ++i) { + const float err = fabsf(values[i] - xc); + if (err < best_err) { + best = i; + best_err = err; + } + } + + return values[best]; +} + +template +static __device__ __forceinline__ float act_quant_max_value() { + if constexpr (mode == 4) { + return 6.0f; + } else { + return 448.0f; + } +} + +template +static __device__ __forceinline__ float act_quant_min_amax() { + if constexpr (mode == 4) { + return 0x1.8p-124f; + } else { + return 1.0e-4f; + } +} + +template +static __device__ __forceinline__ float act_quant_dequant(float x) { + if constexpr (mode == 4) { + return quant_dequant_fp4_e2m1(x); + } else { + return quant_dequant_fp8_e4m3(x); + } +} + +template +static __device__ __forceinline__ float act_quant_to_float(T x) { + return (float) x; +} + +template <> +__device__ __forceinline__ float act_quant_to_float(half x) { + return __half2float(x); +} + +template +static __device__ __forceinline__ T act_quant_from_float(float x) { + return (T) x; +} + +template <> +__device__ __forceinline__ half act_quant_from_float(float x) { + return __float2half(x); +} + template static __global__ void unary_op_kernel(const T * x, T * dst, const int k) { const int i = blockDim.x*blockIdx.x + threadIdx.x; @@ -125,12 +259,51 @@ static __global__ void unary_op_kernel(const T * x, T * dst, const int k) { dst[i] = (T)op((float)x[i]); } +template +static __global__ void act_quant_kernel(const T * x, T * dst, const int64_t ne0, const int64_t nrows) { + const int64_t groups_per_row = ne0 / block_size; + const int64_t group_idx = (int64_t) blockIdx.x; + const int64_t row = group_idx / groups_per_row; + const int64_t group = group_idx - row * groups_per_row; + const int64_t base = row * ne0 + group * block_size; + const int tid = threadIdx.x; + + __shared__ float amax_s[64]; + float amax = 0.0f; + if (tid < block_size && row < nrows) { + const float v = fabsf(act_quant_to_float(x[base + tid])); + amax = isfinite(v) ? v : 0.0f; + } + amax_s[tid] = amax; + __syncthreads(); + + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + amax_s[tid] = fmaxf(amax_s[tid], amax_s[tid + stride]); + } + __syncthreads(); + } + + const float scale = act_quant_pow2_scale(amax_s[0], 1.0f / act_quant_max_value(), act_quant_min_amax()); + const float iscale = 1.0f / scale; + if (tid < block_size && row < nrows) { + dst[base + tid] = act_quant_from_float(act_quant_dequant(act_quant_to_float(x[base + tid]) * iscale) * scale); + } +} + template static void unary_cuda(const T * x, T * dst, const int k, cudaStream_t stream) { const int num_blocks = (k + CUDA_NEG_BLOCK_SIZE - 1) / CUDA_NEG_BLOCK_SIZE; unary_op_kernel<<>>(x, dst, k); } +template +static void act_quant_cuda(const T * x, T * dst, const int64_t ne0, const int64_t nrows, cudaStream_t stream) { + GGML_ASSERT(ne0 % block_size == 0); + const int64_t num_groups = nrows * (ne0 / block_size); + act_quant_kernel<<>>(x, dst, ne0, nrows); +} + template void ggml_cuda_op_unary(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; @@ -151,6 +324,127 @@ void ggml_cuda_op_unary(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { } } +template +void ggml_cuda_op_act_quant(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const void * src0_d = src0->data; + void * dst_d = dst->data; + cudaStream_t stream = ctx.stream(); + + GGML_ASSERT(ggml_is_contiguous(src0)); + GGML_ASSERT(ggml_is_contiguous(dst)); + GGML_ASSERT(src0->ne[0] % block_size == 0); + + GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16); + GGML_ASSERT( dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_F16); + GGML_ASSERT(src0->type == dst->type); + + if (src0->type == GGML_TYPE_F16) { + act_quant_cuda((const half *)src0_d, (half *)dst_d, src0->ne[0], ggml_nrows(src0), stream); + } else { + act_quant_cuda((const float *)src0_d, (float *)dst_d, src0->ne[0], ggml_nrows(src0), stream); + } +} + +static __global__ void sinkhorn_4x4_kernel(const float * src, float * dst, const int64_t n_batch) { + const int64_t b = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + if (b >= n_batch) { + return; + } + src += 16 * b; + dst += 16 * b; + float x[4][4]; + + for (int r = 0; r < 4; ++r) { + float maxv = src[4*r + 0]; +#pragma unroll + for (int c = 1; c < 4; ++c) { + maxv = fmaxf(maxv, src[4*r + c]); + } + + float sum = 0.0f; +#pragma unroll + for (int c = 0; c < 4; ++c) { + x[r][c] = expf(src[4*r + c] - maxv); + sum += x[r][c]; + } + + const float inv_sum = 1.0f / sum; +#pragma unroll + for (int c = 0; c < 4; ++c) { + x[r][c] = fmaxf(x[r][c] * inv_sum, 1e-6f); + } + } + +#pragma unroll + for (int c = 0; c < 4; ++c) { + float sum = 0.0f; +#pragma unroll + for (int r = 0; r < 4; ++r) { + sum += x[r][c]; + } + const float inv_sum = 1.0f / fmaxf(sum, 1e-6f); +#pragma unroll + for (int r = 0; r < 4; ++r) { + x[r][c] *= inv_sum; + } + } + +#pragma unroll + for (int it = 1; it < 20; ++it) { +#pragma unroll + for (int r = 0; r < 4; ++r) { + float sum = 0.0f; +#pragma unroll + for (int c = 0; c < 4; ++c) { + sum += x[r][c]; + } + const float inv_sum = 1.0f / fmaxf(sum, 1e-6f); +#pragma unroll + for (int c = 0; c < 4; ++c) { + x[r][c] *= inv_sum; + } + } + +#pragma unroll + for (int c = 0; c < 4; ++c) { + float sum = 0.0f; +#pragma unroll + for (int r = 0; r < 4; ++r) { + sum += x[r][c]; + } + const float inv_sum = 1.0f / fmaxf(sum, 1e-6f); +#pragma unroll + for (int r = 0; r < 4; ++r) { + x[r][c] *= inv_sum; + } + } + } + +#pragma unroll + for (int r = 0; r < 4; ++r) { +#pragma unroll + for (int c = 0; c < 4; ++c) { + dst[4*r + c] = x[r][c]; + } + } +} + +void ggml_cuda_op_sinkhorn_4x4(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + + GGML_ASSERT(src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32); + GGML_ASSERT(src0->ne[0] == 4 && src0->ne[1] == 4 && src0->ne[3] == 1); + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(ggml_is_contiguous(src0) && ggml_is_contiguous(dst)); + + const int64_t n_batch = src0->ne[2]; + constexpr int block_size = 64; + const int64_t num_blocks = (n_batch + block_size - 1) / block_size; + sinkhorn_4x4_kernel<<<(unsigned int) num_blocks, block_size, 0, ctx.stream()>>>( + (const float *) src0->data, (float *) dst->data, n_batch); +} + void ggml_cuda_op_abs(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_op_unary(ctx, dst); } @@ -247,6 +541,14 @@ void ggml_cuda_op_trunc(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_op_unary(ctx, dst); } +void ggml_cuda_op_fp4_act_quant(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + ggml_cuda_op_act_quant<32, 4>(ctx, dst); +} + +void ggml_cuda_op_fp8_act_quant(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + ggml_cuda_op_act_quant<64, 8>(ctx, dst); +} + void ggml_cuda_op_expm1(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_op_unary(ctx, dst); } diff --git a/ggml/src/ggml-cuda/unary.cuh b/ggml/src/ggml-cuda/unary.cuh index 81ed873ecc3..c534a850571 100644 --- a/ggml/src/ggml-cuda/unary.cuh +++ b/ggml/src/ggml-cuda/unary.cuh @@ -75,6 +75,12 @@ void ggml_cuda_op_round(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_trunc(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +void ggml_cuda_op_fp4_act_quant(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_fp8_act_quant(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_sinkhorn_4x4(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + void ggml_cuda_op_reglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_geglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/vecdotq.cuh b/ggml/src/ggml-cuda/vecdotq.cuh index d1741cc8d7b..29e5c623dcc 100644 --- a/ggml/src/ggml-cuda/vecdotq.cuh +++ b/ggml/src/ggml-cuda/vecdotq.cuh @@ -106,12 +106,151 @@ static __device__ __forceinline__ uint32_t unpack_ksigns(const uint8_t v) { // VDR = vec dot ratio, how many contiguous integers each thread processes when the vec dot kernel is called // MMVQ = mul_mat_vec_q, MMQ = mul_mat_q +static const __device__ float kvalues_f8_e4m3fn[256] = { + 0.0f, 0.001953125f, 0.00390625f, 0.005859375f, 0.0078125f, 0.009765625f, 0.01171875f, 0.013671875f, + 0.015625f, 0.017578125f, 0.01953125f, 0.021484375f, 0.0234375f, 0.025390625f, 0.02734375f, 0.029296875f, + 0.03125f, 0.03515625f, 0.0390625f, 0.04296875f, 0.046875f, 0.05078125f, 0.0546875f, 0.05859375f, + 0.0625f, 0.0703125f, 0.078125f, 0.0859375f, 0.09375f, 0.1015625f, 0.109375f, 0.1171875f, + 0.125f, 0.140625f, 0.15625f, 0.171875f, 0.1875f, 0.203125f, 0.21875f, 0.234375f, + 0.25f, 0.28125f, 0.3125f, 0.34375f, 0.375f, 0.40625f, 0.4375f, 0.46875f, + 0.5f, 0.5625f, 0.625f, 0.6875f, 0.75f, 0.8125f, 0.875f, 0.9375f, + 1.0f, 1.125f, 1.25f, 1.375f, 1.5f, 1.625f, 1.75f, 1.875f, + 2.0f, 2.25f, 2.5f, 2.75f, 3.0f, 3.25f, 3.5f, 3.75f, + 4.0f, 4.5f, 5.0f, 5.5f, 6.0f, 6.5f, 7.0f, 7.5f, + 8.0f, 9.0f, 10.0f, 11.0f, 12.0f, 13.0f, 14.0f, 15.0f, + 16.0f, 18.0f, 20.0f, 22.0f, 24.0f, 26.0f, 28.0f, 30.0f, + 32.0f, 36.0f, 40.0f, 44.0f, 48.0f, 52.0f, 56.0f, 60.0f, + 64.0f, 72.0f, 80.0f, 88.0f, 96.0f, 104.0f, 112.0f, 120.0f, + 128.0f, 144.0f, 160.0f, 176.0f, 192.0f, 208.0f, 224.0f, 240.0f, + 256.0f, 288.0f, 320.0f, 352.0f, 384.0f, 416.0f, 448.0f, NAN, + 0.0f, -0.001953125f, -0.00390625f, -0.005859375f, -0.0078125f, -0.009765625f, -0.01171875f, -0.013671875f, + -0.015625f, -0.017578125f, -0.01953125f, -0.021484375f, -0.0234375f, -0.025390625f, -0.02734375f, -0.029296875f, + -0.03125f, -0.03515625f, -0.0390625f, -0.04296875f, -0.046875f, -0.05078125f, -0.0546875f, -0.05859375f, + -0.0625f, -0.0703125f, -0.078125f, -0.0859375f, -0.09375f, -0.1015625f, -0.109375f, -0.1171875f, + -0.125f, -0.140625f, -0.15625f, -0.171875f, -0.1875f, -0.203125f, -0.21875f, -0.234375f, + -0.25f, -0.28125f, -0.3125f, -0.34375f, -0.375f, -0.40625f, -0.4375f, -0.46875f, + -0.5f, -0.5625f, -0.625f, -0.6875f, -0.75f, -0.8125f, -0.875f, -0.9375f, + -1.0f, -1.125f, -1.25f, -1.375f, -1.5f, -1.625f, -1.75f, -1.875f, + -2.0f, -2.25f, -2.5f, -2.75f, -3.0f, -3.25f, -3.5f, -3.75f, + -4.0f, -4.5f, -5.0f, -5.5f, -6.0f, -6.5f, -7.0f, -7.5f, + -8.0f, -9.0f, -10.0f, -11.0f, -12.0f, -13.0f, -14.0f, -15.0f, + -16.0f, -18.0f, -20.0f, -22.0f, -24.0f, -26.0f, -28.0f, -30.0f, + -32.0f, -36.0f, -40.0f, -44.0f, -48.0f, -52.0f, -56.0f, -60.0f, + -64.0f, -72.0f, -80.0f, -88.0f, -96.0f, -104.0f, -112.0f, -120.0f, + -128.0f, -144.0f, -160.0f, -176.0f, -192.0f, -208.0f, -224.0f, -240.0f, + -256.0f, -288.0f, -320.0f, -352.0f, -384.0f, -416.0f, -448.0f, NAN, +}; + +static const __device__ int8_t kvalues_f8_e4m3fn_i8_approx[256] = { + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, + 2, 3, 3, 3, 3, 4, 4, 4, 5, 5, 6, 6, 7, 7, 8, 9, + 9, 10, 11, 12, 14, 15, 16, 17, 18, 20, 23, 25, 27, 29, 32, 34, + 36, 41, 45, 50, 54, 59, 64, 68, 73, 82, 91, 100, 109, 118, 127, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -2, -2, -2, -2, -2, + -2, -3, -3, -3, -3, -4, -4, -4, -5, -5, -6, -6, -7, -7, -8, -9, + -9, -10, -11, -12, -14, -15, -16, -17, -18, -20, -23, -25, -27, -29, -32, -34, + -36, -41, -45, -50, -54, -59, -64, -68, -73, -82, -91, -100, -109, -118, -127, 0, +}; + #define VDR_Q1_0_Q8_1_MMVQ 1 // Process one 32-element chunk at a time for parallelism #define VDR_Q1_0_Q8_1_MMQ 4 // Q1_0 has 128 bits (4 ints) per block #define VDR_Q4_0_Q8_1_MMVQ 2 #define VDR_Q4_0_Q8_1_MMQ 4 +template static __device__ __forceinline__ float vec_dot_f8_e4m3_b128_q8_1_impl( + const int * v, const int * u, const float & d8, const half & d_q8_1) { + + float sum = 0.0f; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { +#pragma unroll + for (int j = 0; j < 4; ++j) { + const uint8_t q = (uint32_t(v[i]) >> (8*j)) & 0xFF; + const int8_t y = (uint32_t(u[i]) >> (8*j)) & 0xFF; +#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) + const float x = kvalues_f8_e4m3fn[q]; +#else + const float x = __ldg(&kvalues_f8_e4m3fn[q]); +#endif + sum += x * y; + } + } + + return d8 * __half2float(d_q8_1) * sum; +} + +template static __device__ __forceinline__ float vec_dot_f8_e4m3_b128_q8_1_impl_approx_dp4a( + const int * v, const int * u, const float & d8, const half & d_q8_1) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + int x_i8 = 0; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const uint8_t q = (uint32_t(v[i]) >> (8*j)) & 0xFF; +#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) + const int8_t x = kvalues_f8_e4m3fn_i8_approx[q]; +#else + const int8_t x = __ldg(&kvalues_f8_e4m3fn_i8_approx[q]); +#endif + x_i8 |= (uint8_t) x << (8*j); + } + sumi = ggml_cuda_dp4a(x_i8, u[i], sumi); + } + + return (448.0f / 127.0f) * d8 * __half2float(d_q8_1) * sumi; +} + +template static __device__ __forceinline__ float vec_dot_f8_e4m3_b128_q8_1_impl_approx_dp4a_shared_lut( + const int * v, const int * u, const float & d8, const half & d_q8_1, const int8_t * __restrict__ values) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + int x_i8 = 0; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const uint8_t q = (uint32_t(v[i]) >> (8*j)) & 0xFF; + const int8_t x = values[q]; + x_i8 |= (uint8_t) x << (8*j); + } + sumi = ggml_cuda_dp4a(x_i8, u[i], sumi); + } + + return (448.0f / 127.0f) * d8 * __half2float(d_q8_1) * sumi; +} + +template static __device__ __forceinline__ float vec_dot_f8_e4m3_b128_q8_1_impl_shared_lut( + const int * v, const int * u, const float & d8, const half & d_q8_1, const float * __restrict__ values) { + + float sum = 0.0f; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { +#pragma unroll + for (int j = 0; j < 4; ++j) { + const uint8_t q = (uint32_t(v[i]) >> (8*j)) & 0xFF; + const int8_t y = (uint32_t(u[i]) >> (8*j)) & 0xFF; + sum += values[q] * y; + } + } + + return d8 * __half2float(d_q8_1) * sum; +} + template static __device__ __forceinline__ float vec_dot_q4_0_q8_1_impl( const int * v, const int * u, const float & d4, const half2 & ds8) { @@ -301,7 +440,7 @@ template static __device__ __forceinline__ float vec_dot_q8_0_16_q8_1_ return d8_1*sumf; } -#define VDR_MXFP4_Q8_1_MMVQ 2 +#define VDR_MXFP4_Q8_1_MMVQ 4 #define VDR_MXFP4_Q8_1_MMQ 4 static __device__ __forceinline__ float vec_dot_mxfp4_q8_1( @@ -811,6 +950,98 @@ static __device__ __forceinline__ float vec_dot_q8_0_q8_1( return vec_dot_q8_0_q8_1_impl(v, u, bq8_0->d, __low2half(bq8_1->ds)); } +#define VDR_F8_E4M3_B128_Q8_1_MMVQ 2 + +static __device__ __forceinline__ float vec_dot_f8_e4m3_b128_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_f8_e4m3_b128 * bq = (const block_f8_e4m3_b128 *) vbq + kbx; + + static_assert(VDR_F8_E4M3_B128_Q8_1_MMVQ <= QI8_1, "VDR must not span multiple Q8_1 blocks"); + const int y_block = iqs / QI8_1; + const int y_iqs = iqs % QI8_1; + + int v[VDR_F8_E4M3_B128_Q8_1_MMVQ]; + int u[VDR_F8_E4M3_B128_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_F8_E4M3_B128_Q8_1_MMVQ; ++i) { + v[i] = get_int_b1(bq->qs, iqs + i); + u[i] = get_int_b4(bq8_1[y_block].qs, y_iqs + i); + } + + return vec_dot_f8_e4m3_b128_q8_1_impl( + v, u, ggml_cuda_e8m0_to_fp32(bq->e), __low2half(bq8_1[y_block].ds)); +} + +static __device__ __forceinline__ float vec_dot_f8_e4m3_b128_q8_1_shared_lut( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs, + const float * __restrict__ values) { + + const block_f8_e4m3_b128 * bq = (const block_f8_e4m3_b128 *) vbq + kbx; + + static_assert(VDR_F8_E4M3_B128_Q8_1_MMVQ <= QI8_1, "VDR must not span multiple Q8_1 blocks"); + const int y_block = iqs / QI8_1; + const int y_iqs = iqs % QI8_1; + + int v[VDR_F8_E4M3_B128_Q8_1_MMVQ]; + int u[VDR_F8_E4M3_B128_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_F8_E4M3_B128_Q8_1_MMVQ; ++i) { + v[i] = get_int_b1(bq->qs, iqs + i); + u[i] = get_int_b4(bq8_1[y_block].qs, y_iqs + i); + } + + return vec_dot_f8_e4m3_b128_q8_1_impl_shared_lut( + v, u, ggml_cuda_e8m0_to_fp32(bq->e), __low2half(bq8_1[y_block].ds), values); +} + +static __device__ __forceinline__ float vec_dot_f8_e4m3_b128_q8_1_approx_dp4a( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_f8_e4m3_b128 * bq = (const block_f8_e4m3_b128 *) vbq + kbx; + + static_assert(VDR_F8_E4M3_B128_Q8_1_MMVQ <= QI8_1, "VDR must not span multiple Q8_1 blocks"); + const int y_block = iqs / QI8_1; + const int y_iqs = iqs % QI8_1; + + int v[VDR_F8_E4M3_B128_Q8_1_MMVQ]; + int u[VDR_F8_E4M3_B128_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_F8_E4M3_B128_Q8_1_MMVQ; ++i) { + v[i] = get_int_b1(bq->qs, iqs + i); + u[i] = get_int_b4(bq8_1[y_block].qs, y_iqs + i); + } + + return vec_dot_f8_e4m3_b128_q8_1_impl_approx_dp4a( + v, u, ggml_cuda_e8m0_to_fp32(bq->e), __low2half(bq8_1[y_block].ds)); +} + +static __device__ __forceinline__ float vec_dot_f8_e4m3_b128_q8_1_approx_dp4a_shared_lut( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs, + const int8_t * __restrict__ values) { + + const block_f8_e4m3_b128 * bq = (const block_f8_e4m3_b128 *) vbq + kbx; + + static_assert(VDR_F8_E4M3_B128_Q8_1_MMVQ <= QI8_1, "VDR must not span multiple Q8_1 blocks"); + const int y_block = iqs / QI8_1; + const int y_iqs = iqs % QI8_1; + + int v[VDR_F8_E4M3_B128_Q8_1_MMVQ]; + int u[VDR_F8_E4M3_B128_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_F8_E4M3_B128_Q8_1_MMVQ; ++i) { + v[i] = get_int_b1(bq->qs, iqs + i); + u[i] = get_int_b4(bq8_1[y_block].qs, y_iqs + i); + } + + return vec_dot_f8_e4m3_b128_q8_1_impl_approx_dp4a_shared_lut( + v, u, ggml_cuda_e8m0_to_fp32(bq->e), __low2half(bq8_1[y_block].ds), values); +} + static __device__ __forceinline__ float vec_dot_q2_K_q8_1( const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { diff --git a/ggml/src/ggml-quants.c b/ggml/src/ggml-quants.c index 15443aa554a..856b85790aa 100644 --- a/ggml/src/ggml-quants.c +++ b/ggml/src/ggml-quants.c @@ -549,6 +549,111 @@ void dequantize_row_nvfp4(const block_nvfp4 * GGML_RESTRICT x, float * GGML_REST } } +static inline float ggml_f8_e4m3fn_to_fp32(uint8_t x) { + if ((x & 0x7F) == 0) { + return 0.0f; + } + if ((x & 0x7F) == 0x7F) { + return NAN; + } + + const int sign = x >> 7; + const int exp = (x >> 3) & 0x0F; + const int man = x & 0x07; + const float val = exp == 0 ? ldexpf((float) man, -9) : ldexpf(1.0f + (float) man * 0.125f, exp - 7); + + return sign ? -val : val; +} + +static inline uint8_t ggml_fp32_to_f8_e4m3fn(float x) { + if (isnan(x)) { + return 0x7F; + } + + const uint8_t sign = signbit(x) ? 0x80 : 0x00; + const float ax = fabsf(x); + + if (ax == 0.0f) { + return sign; + } + + if (ax < 0x1p-6f) { + const int man = (int) roundf(ax * 512.0f); + if (man <= 0) { + return sign; + } + if (man >= 8) { + return sign | 0x08; + } + return sign | (uint8_t) man; + } + + int exp_unbiased; + const float fr = frexpf(ax, &exp_unbiased); + exp_unbiased -= 1; + + int exp = exp_unbiased + 7; + int man = (int) roundf((2.0f * fr - 1.0f) * 8.0f); + if (man == 8) { + man = 0; + exp++; + } + + if (exp > 15 || (exp == 15 && man > 6)) { + return sign | 0x7E; + } + + return sign | (uint8_t) ((exp << 3) | man); +} + +void quantize_row_f8_e4m3_b128_ref(const float * GGML_RESTRICT x, block_f8_e4m3_b128 * GGML_RESTRICT y, int64_t k) { + static const int qk = QK_F8_E4M3_B128; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + float amax = 0.0f; + + for (int j = 0; j < qk; j++) { + const float v = fabsf(x[i*qk + j]); + if (isfinite(v) && amax < v) { + amax = v; + } + } + + int e = 0; + if (amax > 0.0f) { + e = (int) ceilf(log2f(amax / ggml_f8_e4m3fn_to_fp32(0x7E))) + 127; + e = MAX(0, MIN(254, e)); + } + + y[i].e = (uint8_t) e; + + const float id = 1.0f / GGML_E8M0_TO_FP32(y[i].e); + for (int j = 0; j < qk; ++j) { + y[i].qs[j] = ggml_fp32_to_f8_e4m3fn(x[i*qk + j] * id); + } + } +} + +void dequantize_row_f8_e4m3_b128(const block_f8_e4m3_b128 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + static const int qk = QK_F8_E4M3_B128; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + const float d = GGML_E8M0_TO_FP32(x[i].e); + + for (int j = 0; j < qk; ++j) { + y[i*qk + j] = d * ggml_f8_e4m3fn_to_fp32(x[i].qs[j]); + } + } +} + // // 2-6 bit quantization in super-blocks // @@ -2235,6 +2340,12 @@ size_t quantize_nvfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, return nrow * ggml_row_size(GGML_TYPE_NVFP4, n_per_row); } +size_t quantize_f8_e4m3_b128(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_UNUSED(quant_weights); + quantize_row_f8_e4m3_b128_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * ggml_row_size(GGML_TYPE_F8_E4M3_B128, n_per_row); +} + // ====================== Ternary (de)-quantization (BitNet b1.58 and TriLMs) void quantize_row_tq1_0_ref(const float * GGML_RESTRICT x, block_tq1_0 * GGML_RESTRICT y, int64_t k) { @@ -5391,6 +5502,10 @@ bool ggml_validate_row_data(enum ggml_type type, const void * data, size_t nbyte GGML_UNUSED(data); GGML_UNUSED(nb); } break; + case GGML_TYPE_F8_E4M3_B128: + { + VALIDATE_ROW_DATA_E_E8M0_IMPL(block_f8_e4m3_b128, data, nb); + } break; case GGML_TYPE_Q2_K: { VALIDATE_ROW_DATA_DM_F16_IMPL(block_q2_K, data, nb, d, dmin); diff --git a/ggml/src/ggml-quants.h b/ggml/src/ggml-quants.h index d56c86da890..17fbb5613ea 100644 --- a/ggml/src/ggml-quants.h +++ b/ggml/src/ggml-quants.h @@ -24,6 +24,7 @@ GGML_API void quantize_row_q8_1_ref(const float * GGML_RESTRICT x, block_q8_1 * GGML_API void quantize_row_mxfp4_ref(const float * GGML_RESTRICT x, block_mxfp4 * GGML_RESTRICT y, int64_t k); GGML_API void quantize_row_nvfp4_ref(const float * GGML_RESTRICT x, block_nvfp4 * GGML_RESTRICT y, int64_t k); +GGML_API void quantize_row_f8_e4m3_b128_ref(const float * GGML_RESTRICT x, block_f8_e4m3_b128 * GGML_RESTRICT y, int64_t k); GGML_API void quantize_row_q2_K_ref(const float * GGML_RESTRICT x, block_q2_K * GGML_RESTRICT y, int64_t k); GGML_API void quantize_row_q3_K_ref(const float * GGML_RESTRICT x, block_q3_K * GGML_RESTRICT y, int64_t k); @@ -52,6 +53,7 @@ GGML_API void dequantize_row_q8_0(const block_q8_0 * GGML_RESTRICT x, float * GG GGML_API void dequantize_row_mxfp4(const block_mxfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); GGML_API void dequantize_row_nvfp4(const block_nvfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_f8_e4m3_b128(const block_f8_e4m3_b128 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); GGML_API void dequantize_row_q2_K(const block_q2_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); GGML_API void dequantize_row_q3_K(const block_q3_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); @@ -101,6 +103,7 @@ GGML_API size_t quantize_q8_0(const float * GGML_RESTRICT src, void * GGML_RESTR GGML_API size_t quantize_mxfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); GGML_API size_t quantize_nvfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_f8_e4m3_b128(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); GGML_API void iq2xs_init_impl(enum ggml_type type); GGML_API void iq2xs_free_impl(enum ggml_type type); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 54d3eae3e4d..0ed722cd649 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -43,6 +43,10 @@ #include #endif +#if defined(__linux__) +#include +#endif + #if defined(_WIN32) #define WIN32_LEAN_AND_MEAN #ifndef NOMINMAX @@ -375,6 +379,31 @@ void * ggml_aligned_malloc(size_t size) { GGML_LOG_ERROR("%s: %s (attempted to allocate %6.2f MB)\n", __func__, error_desc, size/(1024.0*1024.0)); return NULL; } +#if defined(__linux__) && !defined(GGML_USE_CPU_HBM) && !defined(TARGET_OS_OSX) + // For large allocations, hint the kernel to back this region with transparent + // huge pages. This dramatically reduces TLB pressure on memory-bandwidth-bound + // workloads such as large MoE expert matmuls where the working set is many GiB + // and the per-token weight read pattern walks millions of 4 KiB pages. + // + // The hint is best-effort: it only succeeds when the system THP policy is + // "always" or "madvise" and the allocation is mapped (large mallocs typically + // are), and silently does nothing otherwise. + // + // 2 MiB threshold avoids spending syscall time on small tensor metadata; + // madvise itself only operates at huge-page boundaries internally. + if (aligned_memory != NULL && size >= (2u << 20)) { + const uintptr_t hp_align = (1u << 21); // 2 MiB + uintptr_t addr_v = (uintptr_t) aligned_memory; + uintptr_t addr_a = (addr_v + hp_align - 1) & ~(hp_align - 1); + size_t off = (size_t) (addr_a - addr_v); + if (off < size) { + size_t hp_size = (size - off) & ~(hp_align - 1); + if (hp_size > 0) { + (void) madvise((void *) addr_a, hp_size, MADV_HUGEPAGE); + } + } + } +#endif return aligned_memory; #endif } @@ -744,6 +773,14 @@ static const struct ggml_type_traits type_traits[GGML_TYPE_COUNT] = { .to_float = (ggml_to_float_t) dequantize_row_nvfp4, .from_float_ref = (ggml_from_float_t)quantize_row_nvfp4_ref, }, + [GGML_TYPE_F8_E4M3_B128] = { + .type_name = "f8_e4m3_b128", + .blck_size = QK_F8_E4M3_B128, + .type_size = sizeof(block_f8_e4m3_b128), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_f8_e4m3_b128, + .from_float_ref = (ggml_from_float_t) quantize_row_f8_e4m3_b128_ref, + }, [GGML_TYPE_Q2_K] = { .type_name = "q2_K", .blck_size = QK_K, @@ -1073,9 +1110,10 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "OPT_STEP_SGD", "GLU", + "HC_WEIGHTED_SUM", }; -static_assert(GGML_OP_COUNT == 96, "GGML_OP_COUNT != 96"); +static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1183,9 +1221,10 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "sgd(x)", "glu(x)", + "hc_weighted_sum(x,w)", }; -static_assert(GGML_OP_COUNT == 96, "GGML_OP_COUNT != 96"); +static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -1212,9 +1251,12 @@ static const char * GGML_UNARY_OP_NAME[GGML_UNARY_OP_COUNT] = { "CEIL", "ROUND", "TRUNC", + "FP4_ACT_QUANT", + "FP8_ACT_QUANT", + "SINKHORN_4X4", }; -static_assert(GGML_UNARY_OP_COUNT == 22, "GGML_UNARY_OP_COUNT != 22"); +static_assert(GGML_UNARY_OP_COUNT == 25, "GGML_UNARY_OP_COUNT != 25"); static const char * GGML_GLU_OP_NAME[GGML_GLU_OP_COUNT] = { "REGLU", @@ -1408,6 +1450,7 @@ enum ggml_type ggml_ftype_to_ggml_type(enum ggml_ftype ftype) { case GGML_FTYPE_MOSTLY_Q8_0: wtype = GGML_TYPE_Q8_0; break; case GGML_FTYPE_MOSTLY_MXFP4: wtype = GGML_TYPE_MXFP4; break; case GGML_FTYPE_MOSTLY_NVFP4: wtype = GGML_TYPE_NVFP4; break; + case GGML_FTYPE_MOSTLY_F8_E4M3_MXFP4: wtype = GGML_TYPE_F8_E4M3_B128; break; case GGML_FTYPE_MOSTLY_Q2_K: wtype = GGML_TYPE_Q2_K; break; case GGML_FTYPE_MOSTLY_Q3_K: wtype = GGML_TYPE_Q3_K; break; case GGML_FTYPE_MOSTLY_Q4_K: wtype = GGML_TYPE_Q4_K; break; @@ -2942,6 +2985,28 @@ struct ggml_tensor * ggml_trunc_inplace( return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_TRUNC); } +struct ggml_tensor * ggml_fp4_act_quant( + struct ggml_context * ctx, + struct ggml_tensor * a) { + GGML_ASSERT(a->ne[0] % 32 == 0); + return ggml_unary(ctx, a, GGML_UNARY_OP_FP4_ACT_QUANT); +} + +struct ggml_tensor * ggml_fp8_act_quant( + struct ggml_context * ctx, + struct ggml_tensor * a) { + GGML_ASSERT(a->ne[0] % 64 == 0); + return ggml_unary(ctx, a, GGML_UNARY_OP_FP8_ACT_QUANT); +} + +struct ggml_tensor * ggml_sinkhorn_4x4( + struct ggml_context * ctx, + struct ggml_tensor * a) { + GGML_ASSERT(a->type == GGML_TYPE_F32); + GGML_ASSERT(a->ne[0] == 4 && a->ne[1] == 4); + return ggml_unary(ctx, a, GGML_UNARY_OP_SINKHORN_4X4); +} + struct ggml_tensor * ggml_glu( struct ggml_context * ctx, struct ggml_tensor * a, @@ -3249,6 +3314,28 @@ struct ggml_tensor * ggml_mul_mat( return result; } +// ggml_hc_weighted_sum + +struct ggml_tensor * ggml_hc_weighted_sum( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + GGML_ASSERT(a->type == GGML_TYPE_F32); + GGML_ASSERT(b->type == GGML_TYPE_F32); + + GGML_ASSERT(a->ne[1] == b->ne[0]); + GGML_ASSERT(a->ne[3] == 1); + GGML_ASSERT(b->ne[1] == a->ne[2] && b->ne[2] == 1 && b->ne[3] == 1); + + struct ggml_tensor * result = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, a->ne[0], a->ne[2]); + + result->op = GGML_OP_HC_WEIGHTED_SUM; + result->src[0] = a; + result->src[1] = b; + + return result; +} + void ggml_mul_mat_set_prec( struct ggml_tensor * a, enum ggml_prec prec) { @@ -6554,6 +6641,21 @@ static void ggml_compute_backward( grad))); // [m,p,qq,rr] } } break; + case GGML_OP_HC_WEIGHTED_SUM: { + if (src0_needs_grads || src1_needs_grads) { + struct ggml_tensor * grad_x = ggml_repeat(ctx, grad, src0); + + if (src0_needs_grads) { + struct ggml_tensor * src1_cont = ggml_is_contiguous(src1) ? src1 : ggml_cont(ctx, src1); + struct ggml_tensor * weights = ggml_reshape_2d(ctx, src1_cont, 1, src1->ne[0]); + ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, grad_x, weights)); + } + if (src1_needs_grads) { + struct ggml_tensor * weighted_grad = ggml_mul(ctx, src0, grad_x); + ggml_add_or_set(ctx, cgraph, isrc1, ggml_reshape(ctx, ggml_sum_rows(ctx, weighted_grad), src1)); + } + } + } break; case GGML_OP_SCALE: { if (src0_needs_grads) { float s; @@ -7681,6 +7783,7 @@ size_t ggml_quantize_chunk( case GGML_TYPE_Q8_0: result = quantize_q8_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_MXFP4: result = quantize_mxfp4 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_NVFP4: result = quantize_nvfp4 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_F8_E4M3_B128: result = quantize_f8_e4m3_b128(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_Q2_K: result = quantize_q2_K (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_Q3_K: result = quantize_q3_K (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_Q4_K: result = quantize_q4_K (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 83ae51ce9ce..624d92318f9 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, @@ -4025,6 +4112,7 @@ class GGMLQuantizationType(IntEnum): MXFP4 = 39 NVFP4 = 40 Q1_0 = 41 + F8_E4M3_B128 = 42 class ExpertGatingFuncType(IntEnum): @@ -4079,6 +4167,7 @@ class LlamaFileType(IntEnum): MOSTLY_MXFP4_MOE = 38 # except 1d tensors MOSTLY_NVFP4 = 39 # except 1d tensors MOSTLY_Q1_0 = 40 # except 1d tensors + MOSTLY_F8_E4M3_MXFP4 = 41 # except 1d tensors GUESSED = 1024 # not specified in the model file @@ -4197,6 +4286,7 @@ class VisionProjectorType: GGMLQuantizationType.MXFP4: (32, 1 + 16), GGMLQuantizationType.NVFP4: (64, 4 + 32), GGMLQuantizationType.Q1_0: (128, 2 + 16), + GGMLQuantizationType.F8_E4M3_B128: (128, 1 + 128), } diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 6a81ca37d8c..b2e915f31ab 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -10,7 +10,6 @@ from enum import Enum, auto from math import prod from pathlib import Path -from io import BufferedWriter from typing import IO, Any, Sequence, Mapping from string import ascii_letters, digits @@ -36,6 +35,7 @@ SHARD_NAME_FORMAT = "{:s}-{:05d}-of-{:05d}.gguf" +GGUF_WRITE_BUFFER_SIZE = 64 * 1024 * 1024 @dataclass @@ -63,7 +63,7 @@ class WriterState(Enum): class GGUFWriter: - fout: list[BufferedWriter] | None + fout: list[IO[bytes]] | None path: Path | None temp_file: tempfile.SpooledTemporaryFile[bytes] | None tensors: list[dict[str, TensorInfo]] @@ -179,7 +179,7 @@ def open_output_file(self, path: Path | None = None) -> None: if self.path is not None: filenames = self.print_plan() - self.fout = [open(filename, "wb") for filename in filenames] + self.fout = [open(filename, "wb", buffering=GGUF_WRITE_BUFFER_SIZE) for filename in filenames] self.state = WriterState.EMPTY def print_plan(self) -> list[Path]: @@ -384,7 +384,12 @@ def add_tensor( # Don't byteswap inplace since lazy copies cannot handle it tensor = tensor.byteswap(inplace=False) if self.use_temp_file and self.temp_file is None: - fp = tempfile.SpooledTemporaryFile(mode="w+b", max_size=256 * 1024 * 1024) + temp_dir = (self.path if self.path.is_dir() else self.path.parent) if self.path is not None else None + fp = tempfile.SpooledTemporaryFile( + mode="w+b", + max_size=256 * 1024 * 1024, + dir=str(temp_dir) if temp_dir is not None else None, + ) fp.seek(0) self.temp_file = fp @@ -401,7 +406,31 @@ def add_tensor( def write_padding(self, fp: IO[bytes], n: int, align: int | None = None) -> None: pad = GGUFWriter.ggml_pad(n, align if align is not None else self.data_alignment) - n if pad != 0: - fp.write(bytes([0] * pad)) + fp.write(b"\0" * pad) + + @staticmethod + def copy_file_range(src: IO[bytes], dst: IO[bytes], length: int = GGUF_WRITE_BUFFER_SIZE) -> None: + if not hasattr(os, "copy_file_range"): + shutil.copyfileobj(src, dst, length=length) + return + + try: + src.flush() + dst.flush() + src_fd = src.fileno() + dst_fd = dst.fileno() + except OSError: + shutil.copyfileobj(src, dst, length=length) + return + + while True: + try: + n = os.copy_file_range(src_fd, dst_fd, length) + except OSError: + shutil.copyfileobj(src, dst, length=length) + return + if n == 0: + return def write_tensor_data(self, tensor: np.ndarray[Any, Any], tensor_endianess: GGUFEndian | None = None) -> None: if self.state is not WriterState.TI_DATA and self.state is not WriterState.WEIGHTS: @@ -476,7 +505,7 @@ def write_tensors_to_file(self, *, progress: bool = False) -> None: else: self.temp_file.seek(0) - shutil.copyfileobj(self.temp_file, self.fout[0 if not self.small_first_shard else 1]) + self.copy_file_range(self.temp_file, self.fout[0 if not self.small_first_shard else 1]) self.flush() self.temp_file.close() diff --git a/gguf-py/gguf/quants.py b/gguf-py/gguf/quants.py index 1d9d9ab7d70..1090aaba34c 100644 --- a/gguf-py/gguf/quants.py +++ b/gguf-py/gguf/quants.py @@ -54,6 +54,8 @@ class QuantError(Exception): ... def quantize(data: np.ndarray, qtype: GGMLQuantizationType) -> np.ndarray: + if data.dtype == np.uint8 and qtype in (GGMLQuantizationType.MXFP4, GGMLQuantizationType.NVFP4, GGMLQuantizationType.F8_E4M3_B128): + return data if qtype == GGMLQuantizationType.F32: return data.astype(np.float32, copy=False) elif qtype == GGMLQuantizationType.F16: @@ -763,6 +765,31 @@ def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: return (d * vals.astype(np.float32)).reshape(n_super, 64) +class F8_E4M3_B128(__Quant, qtype=GGMLQuantizationType.F8_E4M3_B128): + @staticmethod + def e8m0_to_fp32(x: np.ndarray) -> np.ndarray: + bits = np.where(x == 0, np.uint32(0x00400000), x.astype(np.uint32) << np.uint32(23)) + return bits.view(np.float32) + + @staticmethod + def f8_e4m3fn_to_fp32(x: np.ndarray) -> np.ndarray: + sign = np.where((x & np.uint8(0x80)) == 0, np.float32(1.0), np.float32(-1.0)) + ax = x & np.uint8(0x7F) + exp = ((x >> np.uint8(3)) & np.uint8(0x0F)).astype(np.int32) + man = (x & np.uint8(0x07)).astype(np.float32) + val = np.where(exp == 0, man * np.float32(2.0 ** -9), (np.float32(1.0) + man * np.float32(0.125)) * (np.float32(2.0) ** (exp.astype(np.float32) - np.float32(7.0)))) + return np.where(ax == 0, np.float32(0.0), np.where(ax == 0x7F, np.float32(np.nan), sign * val)) + + @classmethod + def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + raise QuantError(f"{cls.qtype.name} is a native storage format and cannot be quantized from float data") + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + e, qs = np.hsplit(blocks, [1]) + return cls.e8m0_to_fp32(e) * cls.f8_e4m3fn_to_fp32(qs) + + class IQ2_XXS(__Quant, qtype=GGMLQuantizationType.IQ2_XXS): ksigns: bytes = ( b"\x00\x81\x82\x03\x84\x05\x06\x87\x88\x09\x0a\x8b\x0c\x8d\x8e\x0f" diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 01a9b236000..2ec2e6be17d 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/include/llama.h b/include/llama.h index eb869814097..f28c54b1d34 100644 --- a/include/llama.h +++ b/include/llama.h @@ -155,6 +155,7 @@ extern "C" { LLAMA_FTYPE_MOSTLY_MXFP4_MOE = 38, // except 1d tensors LLAMA_FTYPE_MOSTLY_NVFP4 = 39, // except 1d tensors LLAMA_FTYPE_MOSTLY_Q1_0 = 40, // except 1d tensors + LLAMA_FTYPE_MOSTLY_F8_E4M3_MXFP4 = 41, // except 1d tensors LLAMA_FTYPE_GUESSED = 1024, // not specified in the model file }; diff --git a/models/templates/deepseek-ai-DeepSeek-V4.jinja b/models/templates/deepseek-ai-DeepSeek-V4.jinja new file mode 100644 index 00000000000..103c60ba9a8 --- /dev/null +++ b/models/templates/deepseek-ai-DeepSeek-V4.jinja @@ -0,0 +1,188 @@ +{%- if not add_generation_prompt is defined -%} + {%- set add_generation_prompt = false -%} +{%- endif -%} +{%- if not thinking is defined -%} + {%- if enable_thinking is defined -%} + {%- set thinking = enable_thinking -%} + {%- else -%} + {%- set thinking = false -%} + {%- endif -%} +{%- endif -%} +{%- set dsml_token = '｜DSML｜' -%} +{%- set thinking_start_token = '' -%} +{%- set thinking_end_token = '' -%} +{%- set latest_reminder_token = '<｜latest_reminder｜>' -%} +{%- set task_tokens = { + 'action': '<｜action｜>', + 'query': '<｜query｜>', + 'authority': '<｜authority｜>', + 'domain': '<｜domain｜>', + 'title': '<｜title｜>', + 'read_url': '<｜read_url｜>' +} -%} +{%- set tools_header = '## Tools\n\nYou have access to a set of tools to help answer the user\'s question. You can invoke tools by writing a "<' + dsml_token + 'tool_calls>" block like the following:\n\n<' + dsml_token + 'tool_calls>\n<' + dsml_token + 'invoke name="$TOOL_NAME">\n<' + dsml_token + 'parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE\n...\n\n<' + dsml_token + 'invoke name="$TOOL_NAME2">\n...\n\n\n\nString parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.\n\nIf thinking_mode is enabled (triggered by ' + thinking_start_token + '), you MUST output your complete reasoning inside ' + thinking_start_token + '...' + thinking_end_token + ' BEFORE any tool calls or final response.\n\nOtherwise, output directly after ' + thinking_end_token + ' with tool calls or final response.\n\n### Available Tool Schemas\n\n' -%} +{%- set tools_footer = '\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n' -%} +{%- set response_format_header = '## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n' -%} +{%- set ns = namespace(system_prompt='', is_first_system=true, pending_assistant=false, has_tools=false, tools_text='', last_user_idx=-1) -%} +{%- if tools is defined and tools -%} + {%- set ns.has_tools = true -%} + {%- set ts = namespace(schemas='') -%} + {%- for tool in tools -%} + {%- if tool['type'] == 'function' -%} + {%- set ts.schemas = ts.schemas + (tool['function'] | tojson) + '\n' -%} + {%- endif -%} + {%- endfor -%} + {%- set ns.tools_text = tools_header + ts.schemas + tools_footer -%} +{%- endif -%} +{%- for message in messages -%} + {%- if message['role'] == 'system' -%} + {%- if ns.is_first_system -%} + {%- set ns.system_prompt = ns.system_prompt + (message['content'] or '') -%} + {%- set ns.is_first_system = false -%} + {%- else -%} + {%- set ns.system_prompt = ns.system_prompt + '\n\n' + (message['content'] or '') -%} + {%- endif -%} + {%- if message['tools'] is defined and message['tools'] -%} + {%- set ns.has_tools = true -%} + {%- set ts = namespace(schemas='') -%} + {%- for tool in message['tools'] -%} + {%- if tool['type'] == 'function' -%} + {%- set ts.schemas = ts.schemas + (tool['function'] | tojson) + '\n' -%} + {%- endif -%} + {%- endfor -%} + {%- set ns.tools_text = tools_header + ts.schemas + tools_footer -%} + {%- endif -%} + {%- if message['response_format'] is defined and message['response_format'] -%} + {%- set ns.system_prompt = ns.system_prompt + '\n\n' + response_format_header + (message['response_format'] | tojson) -%} + {%- endif -%} + {%- endif -%} + {%- if message['role'] == 'user' or message['role'] == 'developer' -%} + {%- set ns.last_user_idx = loop.index0 -%} + {%- endif -%} +{%- endfor -%} +{%- if ns.tools_text -%} + {%- if ns.system_prompt -%} + {%- set ns.system_prompt = ns.system_prompt + '\n\n' + ns.tools_text -%} + {%- else -%} + {%- set ns.system_prompt = ns.tools_text -%} + {%- endif -%} +{%- endif -%} +{{- bos_token -}} +{{- ns.system_prompt -}} +{%- for message in messages -%} + {%- if message['role'] == 'latest_reminder' -%} + {{- latest_reminder_token + (message['content'] or '') -}} + {%- elif message['role'] == 'developer' -%} + {{- '<｜User｜>' + (message['content'] or '') -}} + {%- if message['tools'] is defined and message['tools'] -%} + {%- set ts = namespace(schemas='') -%} + {%- for tool in message['tools'] -%} + {%- if tool['type'] == 'function' -%} + {%- set ts.schemas = ts.schemas + (tool['function'] | tojson) + '\n' -%} + {%- endif -%} + {%- endfor -%} + {{- '\n\n' + tools_header + ts.schemas + tools_footer -}} + {%- endif -%} + {%- if message['response_format'] is defined and message['response_format'] -%} + {{- '\n\n' + response_format_header + (message['response_format'] | tojson) -}} + {%- endif -%} + {%- if message['task'] is defined and message['task'] -%} + {%- if message['task'] == 'action' -%} + {{- '<｜Assistant｜>' -}} + {{- thinking_start_token if thinking else thinking_end_token -}} + {{- task_tokens[message['task']] -}} + {%- set ns.pending_assistant = false -%} + {%- else -%} + {{- task_tokens[message['task']] -}} + {%- set ns.pending_assistant = false -%} + {%- endif -%} + {%- else -%} + {%- set ns.pending_assistant = true -%} + {%- endif -%} + {%- elif message['role'] == 'user' -%} + {{- '<｜User｜>' -}} + {%- if message['content_blocks'] is defined and message['content_blocks'] -%} + {%- for block in message['content_blocks'] -%} + {%- if not loop.first -%}{{- '\n\n' -}}{%- endif -%} + {%- if block['type'] == 'tool_result' -%} + {{- '' -}} + {%- if block['content'] is iterable and block['content'] is not string -%} + {%- set parts = namespace(text='') -%} + {%- for part in block['content'] -%} + {%- if not loop.first -%}{%- set parts.text = parts.text + '\n\n' -%}{%- endif -%} + {%- set parts.text = parts.text + (part['text'] if part['type'] == 'text' else '[Unsupported ' + part['type'] + ']') -%} + {%- endfor -%} + {{- parts.text -}} + {%- else -%} + {{- block['content'] or '' -}} + {%- endif -%} + {{- '' -}} + {%- else -%} + {{- block['text'] if block['text'] is defined else '[Unsupported ' + block['type'] + ']' -}} + {%- endif -%} + {%- endfor -%} + {%- else -%} + {{- message['content'] or '' -}} + {%- endif -%} + {%- if message['task'] is defined and message['task'] -%} + {%- if message['task'] == 'action' -%} + {{- '<｜Assistant｜>' -}} + {{- thinking_start_token if thinking else thinking_end_token -}} + {{- task_tokens[message['task']] -}} + {%- set ns.pending_assistant = false -%} + {%- else -%} + {{- task_tokens[message['task']] -}} + {%- set ns.pending_assistant = false -%} + {%- endif -%} + {%- else -%} + {%- set ns.pending_assistant = true -%} + {%- endif -%} + {%- elif message['role'] == 'tool' -%} + {{- '<｜User｜>' + (message['content'] or '') + '' -}} + {%- set ns.pending_assistant = true -%} + {%- elif message['role'] == 'assistant' -%} + {%- if ns.pending_assistant -%} + {{- '<｜Assistant｜>' -}} + {%- endif -%} + {%- set prev_has_task = loop.index0 > 0 and messages[loop.index0 - 1]['task'] is defined and messages[loop.index0 - 1]['task'] -%} + {%- if thinking and not prev_has_task -%} + {%- if ns.has_tools or loop.index0 > ns.last_user_idx -%} + {{- thinking_start_token + (message['reasoning_content'] or '') + thinking_end_token -}} + {%- else -%} + {{- thinking_end_token -}} + {%- endif -%} + {%- elif not prev_has_task -%} + {{- thinking_end_token -}} + {%- endif -%} + {{- message['content'] or '' -}} + {%- if message['tool_calls'] -%} + {{- '\n\n<' + dsml_token + 'tool_calls>\n' -}} + {%- for tool in message['tool_calls'] -%} + {%- set func = tool['function'] -%} + {{- '<' + dsml_token + 'invoke name="' + func['name'] + '">\n' -}} + {%- set args = func['arguments'] -%} + {%- if args is string -%} + {%- set args = args | from_json -%} + {%- endif -%} + {%- for key, val in args.items() -%} + {%- if val is string -%} + {{- '<' + dsml_token + 'parameter name="' + key + '" string="true">' + val + '\n' -}} + {%- else -%} + {{- '<' + dsml_token + 'parameter name="' + key + '" string="false">' + (val | tojson) + '\n' -}} + {%- endif -%} + {%- endfor -%} + {{- '\n' -}} + {%- endfor -%} + {{- '' -}} + {%- endif -%} + {{- '<｜end▁of▁sentence｜>' -}} + {%- if message['task'] is defined and message['task'] -%} + {{- task_tokens[message['task']] -}} + {%- endif -%} + {%- set ns.pending_assistant = false -%} + {%- endif -%} +{%- endfor -%} +{%- if add_generation_prompt and ns.pending_assistant -%} + {{- '<｜Assistant｜>' -}} + {{- thinking_start_token if thinking else thinking_end_token -}} +{%- endif -%} diff --git a/scripts/moe-copy-lru-sim.py b/scripts/moe-copy-lru-sim.py new file mode 100755 index 00000000000..2048f87b6db --- /dev/null +++ b/scripts/moe-copy-lru-sim.py @@ -0,0 +1,868 @@ +#!/usr/bin/env python3 + +import argparse +import re +import sys +from collections import Counter, OrderedDict +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, Iterable, Iterator, List, Optional, Sequence, Tuple + + +MOE_COPY_RE = re.compile(r"\bmoe_copy\b(?P.*)\sids=\[(?P[^\]]*)\]") +MOE_CACHE_BYPASS_RE = re.compile(r"\bmoe_cache_bypass\b(?P.*)") +MOE_CACHE_RE = re.compile(r"\bmoe_cache\b(?P.*)") +FIELD_RE = re.compile(r"(\w+)=([^\s]+)") + + +@dataclass(frozen=True) +class MoeCopyEvent: + key: str + tensor: str + dst_backend: str + expert_size: int + used_bytes: int + copy_bytes: int + expert_ids: Tuple[int, ...] + expert_counts: Tuple[Tuple[int, int], ...] + + +@dataclass(frozen=True) +class MoeCacheEvent: + key: str + tensor: str + backend: str + slots: int + expert_size: int + cache_bytes: int + used: int + hits: int + misses: int + copied: int + total_hits: int + total_misses: int + total_copied: int + + +@dataclass(frozen=True) +class MoeCacheBypassEvent: + key: str + tensor: str + backend: str + slots: int + reason: str + n_expert: int + expert_size: int + + +@dataclass +class SimStats: + events: int = 0 + bypasses: int = 0 + cache_bytes: int = 0 + accesses: int = 0 + hits: int = 0 + misses: int = 0 + baseline_bytes: int = 0 + cache_copy_bytes: int = 0 + + +@dataclass +class PrefetchStats: + events: int = 0 + bypasses: int = 0 + cache_bytes: int = 0 + accesses: int = 0 + demand_hits: int = 0 + speculative_hits: int = 0 + misses: int = 0 + baseline_bytes: int = 0 + demand_copy_bytes: int = 0 + prefetch_copy_bytes: int = 0 + prefetches: int = 0 + wrong_prefetches: int = 0 + prefetch_evictions: int = 0 + + +@dataclass +class RuntimeStats: + slots: Optional[int] = None + expert_size: int = 0 + cache_bytes: int = 0 + events: int = 0 + accesses: int = 0 + hits: int = 0 + misses: int = 0 + copied: int = 0 + max_total_hits: int = 0 + max_total_misses: int = 0 + max_total_copied: int = 0 + + +def parse_slots(value: str) -> List[int]: + slots = [] + for item in value.split(","): + item = item.strip() + if not item: + continue + slot_count = int(item) + if slot_count < 0: + raise argparse.ArgumentTypeError("slot counts must be non-negative") + slots.append(slot_count) + if not slots: + raise argparse.ArgumentTypeError("at least one slot count is required") + return slots + + +def parse_positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("value must be positive") + return parsed + + +def parse_expert_counts(raw: Optional[str], expert_ids: Tuple[int, ...]) -> Tuple[Tuple[int, int], ...]: + if raw is None: + return tuple((expert_id, 1) for expert_id in expert_ids) + + raw = raw.strip() + if not raw.startswith("[") or not raw.endswith("]"): + raise ValueError(f"malformed id_counts field: {raw}") + + counts: Dict[int, int] = {} + body = raw[1:-1].strip() + if body: + for item in body.split(","): + if ":" not in item: + raise ValueError(f"malformed id_counts item: {item}") + expert_id_raw, count_raw = item.split(":", 1) + expert_id = int(expert_id_raw) + count = int(count_raw) + if count <= 0: + raise ValueError(f"id_counts entry for expert {expert_id} must be positive") + if expert_id in counts: + raise ValueError(f"id_counts has duplicate expert id: {expert_id}") + counts[expert_id] = count + + expert_id_set = set(expert_ids) + if set(counts) != expert_id_set: + raise ValueError("id_counts expert set does not match ids") + return tuple((expert_id, counts[expert_id]) for expert_id in expert_ids) + + +def parse_moe_copy_line(line: str) -> Optional[MoeCopyEvent]: + match = MOE_COPY_RE.search(line) + if match is None: + return None + + fields = dict(FIELD_RE.findall(match.group("fields"))) + try: + tensor = fields["tensor"] + dst_backend = fields["dst_backend"] + expert_size = int(fields["expert_size"]) + used_bytes = int(fields["used_bytes"]) + copy_bytes = int(fields["copy_bytes"]) + except KeyError as exc: + raise ValueError(f"missing moe_copy field: {exc.args[0]}") from exc + + expert_ids_raw = match.group("expert_ids").strip() + expert_ids = tuple(int(item) for item in expert_ids_raw.split(",") if item.strip()) + if len(expert_ids) != len(set(expert_ids)): + raise ValueError(f"moe_copy line has duplicate expert ids: {expert_ids_raw}") + expert_counts = parse_expert_counts(fields.get("id_counts"), expert_ids) + expected_used_bytes = len(expert_ids) * expert_size + if used_bytes != expected_used_bytes: + raise ValueError( + f"used_bytes={used_bytes} does not match " + f"{len(expert_ids)} expert ids * expert_size={expert_size}" + ) + if copy_bytes < used_bytes: + raise ValueError(f"copy_bytes={copy_bytes} is smaller than used_bytes={used_bytes}") + + return MoeCopyEvent( + key=f"{dst_backend}:{tensor}", + tensor=tensor, + dst_backend=dst_backend, + expert_size=expert_size, + used_bytes=used_bytes, + copy_bytes=copy_bytes, + expert_ids=expert_ids, + expert_counts=expert_counts, + ) + + +def parse_moe_cache_line(line: str) -> Optional[MoeCacheEvent]: + match = MOE_CACHE_RE.search(line) + if match is None: + return None + + fields = dict(FIELD_RE.findall(match.group("fields"))) + try: + tensor = fields["tensor"] + backend = fields["backend"] + slots = int(fields["slots"]) + expert_size = int(fields.get("expert_size", "0")) + cache_bytes = int(fields.get("cache_bytes", "0")) + used = int(fields["used"]) + hits = int(fields["hits"]) + misses = int(fields["misses"]) + copied = int(fields["copied"]) + total_hits = int(fields["total_hits"]) + total_misses = int(fields["total_misses"]) + total_copied = int(fields["total_copied"]) + except KeyError as exc: + raise ValueError(f"missing moe_cache field: {exc.args[0]}") from exc + + if slots < 0 or expert_size < 0 or cache_bytes < 0 or used < 0 or hits < 0 or misses < 0 or copied < 0: + raise ValueError("moe_cache counters must be non-negative") + if used != hits + misses: + raise ValueError(f"used={used} does not match hits={hits} + misses={misses}") + if total_hits < hits or total_misses < misses or total_copied < copied: + raise ValueError("moe_cache total counters are smaller than per-event counters") + + return MoeCacheEvent( + key=f"{backend}:{tensor}", + tensor=tensor, + backend=backend, + slots=slots, + expert_size=expert_size, + cache_bytes=cache_bytes, + used=used, + hits=hits, + misses=misses, + copied=copied, + total_hits=total_hits, + total_misses=total_misses, + total_copied=total_copied, + ) + + +def parse_moe_cache_bypass_line(line: str) -> Optional[MoeCacheBypassEvent]: + match = MOE_CACHE_BYPASS_RE.search(line) + if match is None: + return None + + fields = dict(FIELD_RE.findall(match.group("fields"))) + try: + tensor = fields["tensor"] + backend = fields["backend"] + slots = int(fields["slots"]) + reason = fields["reason"] + n_expert = int(fields["n_expert"]) + expert_size = int(fields["expert_size"]) + except KeyError as exc: + raise ValueError(f"missing moe_cache_bypass field: {exc.args[0]}") from exc + + if slots < 0 or n_expert < 0 or expert_size < 0: + raise ValueError("moe_cache_bypass numeric fields must be non-negative") + if not reason: + raise ValueError("moe_cache_bypass reason must be non-empty") + + return MoeCacheBypassEvent( + key=f"{backend}:{tensor}", + tensor=tensor, + backend=backend, + slots=slots, + reason=reason, + n_expert=n_expert, + expert_size=expert_size, + ) + + +def read_events(paths: Sequence[str]) -> Iterator[MoeCopyEvent]: + if not paths: + yield from read_events_from_lines(sys.stdin) + return + + for path_str in paths: + if path_str == "-": + yield from read_events_from_lines(sys.stdin) + else: + with Path(path_str).open("r", encoding="utf-8", errors="replace") as f: + yield from read_events_from_lines(f) + + +def read_events_from_lines(lines: Iterable[str]) -> Iterator[MoeCopyEvent]: + for line_no, line in enumerate(lines, 1): + try: + event = parse_moe_copy_line(line) + except ValueError as exc: + raise ValueError(f"line {line_no}: {exc}") from exc + if event is not None: + yield event + + +def read_runtime_events(paths: Sequence[str]) -> Tuple[List[MoeCacheEvent], List[MoeCacheBypassEvent]]: + cache_events: List[MoeCacheEvent] = [] + bypass_events: List[MoeCacheBypassEvent] = [] + + def read_lines(lines: Iterable[str]) -> None: + for line_no, line in enumerate(lines, 1): + try: + cache_event = parse_moe_cache_line(line) + bypass_event = parse_moe_cache_bypass_line(line) + except ValueError as exc: + raise ValueError(f"line {line_no}: {exc}") from exc + if cache_event is not None: + cache_events.append(cache_event) + if bypass_event is not None: + bypass_events.append(bypass_event) + + if not paths: + read_lines(sys.stdin) + return cache_events, bypass_events + + for path_str in paths: + if path_str == "-": + read_lines(sys.stdin) + else: + with Path(path_str).open("r", encoding="utf-8", errors="replace") as f: + read_lines(f) + + return cache_events, bypass_events + + +def read_cache_events(paths: Sequence[str]) -> Iterator[MoeCacheEvent]: + if not paths: + yield from read_cache_events_from_lines(sys.stdin) + return + + for path_str in paths: + if path_str == "-": + yield from read_cache_events_from_lines(sys.stdin) + else: + with Path(path_str).open("r", encoding="utf-8", errors="replace") as f: + yield from read_cache_events_from_lines(f) + + +def read_cache_events_from_lines(lines: Iterable[str]) -> Iterator[MoeCacheEvent]: + for line_no, line in enumerate(lines, 1): + try: + event = parse_moe_cache_line(line) + except ValueError as exc: + raise ValueError(f"line {line_no}: {exc}") from exc + if event is not None: + yield event + + +def read_cache_bypass_events_from_lines(lines: Iterable[str]) -> Iterator[MoeCacheBypassEvent]: + for line_no, line in enumerate(lines, 1): + try: + event = parse_moe_cache_bypass_line(line) + except ValueError as exc: + raise ValueError(f"line {line_no}: {exc}") from exc + if event is not None: + yield event + + +def simulate_lru(events: Sequence[MoeCopyEvent], slots: Sequence[int]) -> Dict[Tuple[int, str], SimStats]: + stats: Dict[Tuple[int, str], SimStats] = {} + caches: Dict[Tuple[int, str], OrderedDict[int, None]] = {} + expert_sizes: Dict[str, int] = {} + + for slot_count in slots: + for event in events: + previous_expert_size = expert_sizes.setdefault(event.key, event.expert_size) + if previous_expert_size != event.expert_size: + raise ValueError( + f"inconsistent expert_size for {event.key}: " + f"saw {event.expert_size}, expected {previous_expert_size}" + ) + + stat_key = (slot_count, event.key) + stat = stats.setdefault(stat_key, SimStats()) + cache = caches.setdefault(stat_key, OrderedDict()) + + needed = event.expert_ids + needed_set = set(needed) + + stat.events += 1 + stat.cache_bytes = max(stat.cache_bytes, slot_count * event.expert_size) + stat.accesses += len(needed) + stat.baseline_bytes += event.copy_bytes + + if slot_count == 0 or len(needed) > slot_count: + stat.bypasses += 1 + stat.misses += len(needed) + stat.cache_copy_bytes += event.copy_bytes + continue + + hits = [expert_id for expert_id in needed if expert_id in cache] + misses = [expert_id for expert_id in needed if expert_id not in cache] + + stat.hits += len(hits) + stat.misses += len(misses) + stat.cache_copy_bytes += len(misses) * event.expert_size + + while len(cache) + len(misses) > slot_count: + victim = next((expert_id for expert_id in cache if expert_id not in needed_set), None) + if victim is None: + victim = next(iter(cache)) + del cache[victim] + + for expert_id in misses: + cache[expert_id] = None + + for expert_id in needed: + if expert_id in cache: + cache.move_to_end(expert_id) + + return stats + + +def _validate_expert_size(expert_sizes: Dict[str, int], event: MoeCopyEvent) -> None: + previous_expert_size = expert_sizes.setdefault(event.key, event.expert_size) + if previous_expert_size != event.expert_size: + raise ValueError( + f"inconsistent expert_size for {event.key}: " + f"saw {event.expert_size}, expected {previous_expert_size}" + ) + + +def _evict_one_for_insert( + cache: OrderedDict[int, bool], + protected: set, + stat: PrefetchStats, + prefetch_eviction: bool) -> bool: + victim: Optional[int] = None + for expert_id, speculative in cache.items(): + if expert_id not in protected and speculative: + victim = expert_id + break + if victim is None: + for expert_id in cache: + if expert_id not in protected: + victim = expert_id + break + if victim is None: + return False + + if cache[victim]: + stat.wrong_prefetches += 1 + if prefetch_eviction: + stat.prefetch_evictions += 1 + del cache[victim] + return True + + +def _prefetch_candidates( + cache: OrderedDict[int, bool], + candidates: Iterable[int], + slot_count: int, + expert_size: int, + stat: PrefetchStats) -> None: + if slot_count <= 0: + return + + protected = set(candidates) + for expert_id in candidates: + if expert_id in cache: + cache.move_to_end(expert_id) + continue + + while len(cache) >= slot_count: + if not _evict_one_for_insert(cache, protected, stat, prefetch_eviction=True): + return + + cache[expert_id] = True + stat.prefetches += 1 + stat.prefetch_copy_bytes += expert_size + + +def _next_event_by_key(events: Sequence[MoeCopyEvent]) -> List[Optional[MoeCopyEvent]]: + next_events: List[Optional[MoeCopyEvent]] = [None] * len(events) + last_by_key: Dict[str, MoeCopyEvent] = {} + for index in range(len(events) - 1, -1, -1): + event = events[index] + next_events[index] = last_by_key.get(event.key) + last_by_key[event.key] = event + return next_events + + +def simulate_prefetch( + events: Sequence[MoeCopyEvent], + slots: Sequence[int], + policy: str, + prefetch_limit: Optional[int] = None) -> Dict[Tuple[str, int, str], PrefetchStats]: + if policy not in {"prompt", "freq", "markov", "setmarkov", "oracle"}: + raise ValueError(f"unsupported prefetch policy: {policy}") + + stats: Dict[Tuple[str, int, str], PrefetchStats] = {} + expert_sizes: Dict[str, int] = {} + next_events = _next_event_by_key(events) + + for slot_count in slots: + caches: Dict[Tuple[int, str], OrderedDict[int, bool]] = {} + frequencies: Dict[str, Counter] = {} + previous_ids: Dict[str, Tuple[int, ...]] = {} + transitions: Dict[str, Dict[int, Counter]] = {} + set_transitions: Dict[str, Dict[Tuple[int, ...], Counter]] = {} + + for event_index, event in enumerate(events): + _validate_expert_size(expert_sizes, event) + + stat_key = (policy, slot_count, event.key) + stat = stats.setdefault(stat_key, PrefetchStats()) + cache = caches.setdefault((slot_count, event.key), OrderedDict()) + frequency = frequencies.setdefault(event.key, Counter()) + + needed = event.expert_ids + needed_set = set(needed) + bypass = slot_count == 0 or len(needed) > slot_count + + stat.events += 1 + stat.cache_bytes = max(stat.cache_bytes, slot_count * event.expert_size) + stat.accesses += len(needed) + stat.baseline_bytes += event.copy_bytes + + if bypass: + stat.bypasses += 1 + stat.misses += len(needed) + stat.demand_copy_bytes += event.copy_bytes + else: + misses: List[int] = [] + for expert_id in needed: + if expert_id in cache: + if cache[expert_id]: + stat.speculative_hits += 1 + else: + stat.demand_hits += 1 + cache[expert_id] = False + cache.move_to_end(expert_id) + else: + misses.append(expert_id) + + stat.misses += len(misses) + stat.demand_copy_bytes += len(misses) * event.expert_size + + while len(cache) + len(misses) > slot_count: + if not _evict_one_for_insert(cache, needed_set, stat, prefetch_eviction=False): + break + + for expert_id in misses: + cache[expert_id] = False + cache.move_to_end(expert_id) + + frequency.update(dict(event.expert_counts)) + candidate_limit = slot_count if prefetch_limit is None else min(slot_count, prefetch_limit) + + if policy == "prompt": + if bypass: + candidates = [expert_id for expert_id, _ in frequency.most_common(candidate_limit)] + else: + candidates = [] + elif policy == "freq": + candidates = [expert_id for expert_id, _ in frequency.most_common(candidate_limit)] + elif policy == "markov": + previous = previous_ids.get(event.key) + if previous is not None and len(previous) <= 64 and len(needed) <= 64: + key_transitions = transitions.setdefault(event.key, {}) + for previous_id in previous: + key_transitions.setdefault(previous_id, Counter()).update(needed) + + scores = Counter() + for expert_id in needed: + scores.update(transitions.get(event.key, {}).get(expert_id, Counter())) + candidates = [expert_id for expert_id, _ in scores.most_common(candidate_limit)] + previous_ids[event.key] = needed + elif policy == "setmarkov": + previous = previous_ids.get(event.key) + if previous is not None and len(previous) <= 64 and len(needed) <= 64: + set_transitions.setdefault(event.key, {}).setdefault(previous, Counter()).update(needed) + + candidates = [ + expert_id + for expert_id, _ in set_transitions.get(event.key, {}).get(needed, Counter()).most_common(candidate_limit) + ] + previous_ids[event.key] = needed + else: + next_event = next_events[event_index] + if next_event is not None and len(next_event.expert_ids) <= slot_count: + candidates = list(next_event.expert_ids[:candidate_limit]) + else: + candidates = [] + + _prefetch_candidates(cache, candidates, slot_count, event.expert_size, stat) + + return stats + + +def summarize_runtime_cache(events: Sequence[MoeCacheEvent]) -> Dict[Tuple[int, str], RuntimeStats]: + stats: Dict[Tuple[int, str], RuntimeStats] = {} + for event in events: + stat = stats.setdefault((event.slots, event.key), RuntimeStats(slots=event.slots)) + if stat.slots is None: + stat.slots = event.slots + if stat.expert_size and event.expert_size and stat.expert_size != event.expert_size: + raise ValueError( + f"inconsistent expert_size for runtime cache {event.key} slots={event.slots}: " + f"saw {event.expert_size}, expected {stat.expert_size}" + ) + if stat.cache_bytes and event.cache_bytes and stat.cache_bytes != event.cache_bytes: + raise ValueError( + f"inconsistent cache_bytes for runtime cache {event.key} slots={event.slots}: " + f"saw {event.cache_bytes}, expected {stat.cache_bytes}" + ) + stat.expert_size = stat.expert_size or event.expert_size + stat.cache_bytes = stat.cache_bytes or event.cache_bytes + + stat.events += 1 + stat.accesses += event.used + stat.hits += event.hits + stat.misses += event.misses + stat.copied += event.copied + stat.max_total_hits = max(stat.max_total_hits, event.total_hits) + stat.max_total_misses = max(stat.max_total_misses, event.total_misses) + stat.max_total_copied = max(stat.max_total_copied, event.total_copied) + return stats + + +def summarize_runtime_bypasses(events: Sequence[MoeCacheBypassEvent]) -> Counter: + return Counter((event.key, event.slots, event.reason) for event in events) + + +def aggregate_stats(stats: Dict[Tuple[int, str], SimStats]) -> Dict[int, SimStats]: + aggregate: Dict[int, SimStats] = {} + for (slot_count, _), stat in stats.items(): + dst = aggregate.setdefault(slot_count, SimStats()) + dst.events += stat.events + dst.bypasses += stat.bypasses + dst.cache_bytes += stat.cache_bytes + dst.accesses += stat.accesses + dst.hits += stat.hits + dst.misses += stat.misses + dst.baseline_bytes += stat.baseline_bytes + dst.cache_copy_bytes += stat.cache_copy_bytes + return aggregate + + +def aggregate_prefetch_stats(stats: Dict[Tuple[str, int, str], PrefetchStats]) -> Dict[Tuple[str, int], PrefetchStats]: + aggregate: Dict[Tuple[str, int], PrefetchStats] = {} + for (policy, slot_count, _), stat in stats.items(): + dst = aggregate.setdefault((policy, slot_count), PrefetchStats()) + dst.events += stat.events + dst.bypasses += stat.bypasses + dst.cache_bytes += stat.cache_bytes + dst.accesses += stat.accesses + dst.demand_hits += stat.demand_hits + dst.speculative_hits += stat.speculative_hits + dst.misses += stat.misses + dst.baseline_bytes += stat.baseline_bytes + dst.demand_copy_bytes += stat.demand_copy_bytes + dst.prefetch_copy_bytes += stat.prefetch_copy_bytes + dst.prefetches += stat.prefetches + dst.wrong_prefetches += stat.wrong_prefetches + dst.prefetch_evictions += stat.prefetch_evictions + return aggregate + + +def aggregate_runtime_stats(stats: Dict[Tuple[int, str], RuntimeStats]) -> Dict[int, RuntimeStats]: + aggregate: Dict[int, RuntimeStats] = {} + for (slots, _), stat in stats.items(): + dst = aggregate.setdefault(slots, RuntimeStats(slots=slots)) + dst.cache_bytes += stat.cache_bytes + dst.events += stat.events + dst.accesses += stat.accesses + dst.hits += stat.hits + dst.misses += stat.misses + dst.copied += stat.copied + dst.max_total_hits += stat.max_total_hits + dst.max_total_misses += stat.max_total_misses + dst.max_total_copied += stat.max_total_copied + return aggregate + + +def stats_row(slot_count: int, key: str, stat: SimStats) -> str: + hit_rate = stat.hits / stat.accesses if stat.accesses else 0.0 + saved_bytes = stat.baseline_bytes - stat.cache_copy_bytes + saved_pct = saved_bytes / stat.baseline_bytes if stat.baseline_bytes else 0.0 + return "\t".join(( + str(slot_count), + key, + str(stat.cache_bytes), + str(stat.events), + str(stat.bypasses), + str(stat.accesses), + str(stat.hits), + str(stat.misses), + f"{hit_rate:.6f}", + str(stat.baseline_bytes), + str(stat.cache_copy_bytes), + str(saved_bytes), + f"{saved_pct:.6f}", + )) + + +def prefetch_stats_row(policy: str, slot_count: int, key: str, stat: PrefetchStats) -> str: + hits = stat.demand_hits + stat.speculative_hits + hit_rate = hits / stat.accesses if stat.accesses else 0.0 + critical_saved_bytes = stat.baseline_bytes - stat.demand_copy_bytes + critical_saved_pct = critical_saved_bytes / stat.baseline_bytes if stat.baseline_bytes else 0.0 + total_copy_bytes = stat.demand_copy_bytes + stat.prefetch_copy_bytes + net_saved_bytes = stat.baseline_bytes - total_copy_bytes + net_saved_pct = net_saved_bytes / stat.baseline_bytes if stat.baseline_bytes else 0.0 + return "\t".join(( + policy, + str(slot_count), + key, + str(stat.cache_bytes), + str(stat.events), + str(stat.bypasses), + str(stat.accesses), + str(hits), + str(stat.demand_hits), + str(stat.speculative_hits), + str(stat.misses), + f"{hit_rate:.6f}", + str(stat.baseline_bytes), + str(stat.demand_copy_bytes), + str(stat.prefetch_copy_bytes), + str(total_copy_bytes), + str(critical_saved_bytes), + f"{critical_saved_pct:.6f}", + str(net_saved_bytes), + f"{net_saved_pct:.6f}", + str(stat.prefetches), + str(stat.wrong_prefetches), + str(stat.prefetch_evictions), + )) + + +def runtime_stats_row(key: str, stat: RuntimeStats) -> str: + hit_rate = stat.hits / stat.accesses if stat.accesses else 0.0 + slots = "-" if stat.slots is None else str(stat.slots) + return "\t".join(( + key, + slots, + str(stat.cache_bytes), + str(stat.events), + str(stat.accesses), + str(stat.hits), + str(stat.misses), + f"{hit_rate:.6f}", + str(stat.copied), + str(stat.max_total_hits), + str(stat.max_total_misses), + str(stat.max_total_copied), + )) + + +def print_report(stats: Dict[Tuple[int, str], SimStats], show_details: bool) -> None: + print("slots\tkey\tcache_bytes\tevents\tbypasses\taccesses\thits\tmisses\thit_rate\tbaseline_bytes\tcache_copy_bytes\tsaved_bytes\tsaved_pct") + + for slot_count, stat in sorted(aggregate_stats(stats).items()): + print(stats_row(slot_count, "ALL", stat)) + + if not show_details: + return + + for (slot_count, key), stat in sorted(stats.items()): + print(stats_row(slot_count, key, stat)) + + +def print_prefetch_report(stats: Dict[Tuple[str, int, str], PrefetchStats], show_details: bool) -> None: + print( + "policy\tslots\tkey\tcache_bytes\tevents\tbypasses\taccesses\thits\t" + "demand_hits\tspeculative_hits\tmisses\thit_rate\tbaseline_bytes\t" + "demand_copy_bytes\tprefetch_copy_bytes\ttotal_copy_bytes\t" + "critical_saved_bytes\tcritical_saved_pct\tnet_saved_bytes\tnet_saved_pct\t" + "prefetches\twrong_prefetches\tprefetch_evictions" + ) + + for (policy, slot_count), stat in sorted(aggregate_prefetch_stats(stats).items()): + print(prefetch_stats_row(policy, slot_count, "ALL", stat)) + + if not show_details: + return + + for (policy, slot_count, key), stat in sorted(stats.items()): + print(prefetch_stats_row(policy, slot_count, key, stat)) + + +def print_runtime_report(stats: Dict[Tuple[int, str], RuntimeStats], show_details: bool) -> None: + print("key\tslots\tcache_bytes\tevents\taccesses\thits\tmisses\thit_rate\tcopied\tmax_total_hits\tmax_total_misses\tmax_total_copied") + for _, stat in sorted(aggregate_runtime_stats(stats).items()): + print(runtime_stats_row("ALL", stat)) + + if not show_details: + return + + for (_, key), stat in sorted(stats.items()): + print(runtime_stats_row(key, stat)) + + +def print_runtime_bypass_report(stats: Counter, show_details: bool) -> None: + print("bypass_key\tslots\treason\tevents") + + aggregate = Counter() + for (_, slots, reason), count in stats.items(): + aggregate[(slots, reason)] += count + for (slots, reason), count in sorted(aggregate.items()): + print(f"ALL\t{slots}\t{reason}\t{count}") + + if not show_details: + return + + for (key, slots, reason), count in sorted(stats.items()): + print(f"{key}\t{slots}\t{reason}\t{count}") + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = argparse.ArgumentParser( + description="Analyze GGML_SCHED_MOE_LOG output for MoE expert-copy and runtime-cache behavior.", + epilog=( + "Examples:\n" + " scripts/moe-copy-lru-sim.py --slots 32,64,128 trace.log\n" + " scripts/moe-copy-lru-sim.py --slots 48 --repeat 4 --policy oracle trace.log\n" + " scripts/moe-copy-lru-sim.py --slots 32 --policy prompt trace.log\n" + " scripts/moe-copy-lru-sim.py --runtime --details cache-enabled.log\n" + " # --runtime accepts moe_cache, moe_cache_bypass, or mixed logs" + ), + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("logs", nargs="*", help="log files to parse; omit or use '-' for stdin") + parser.add_argument("--slots", type=parse_slots, default=parse_slots("32,64,96,128"), help="comma-separated slot counts for moe_copy LRU simulation") + parser.add_argument("--repeat", type=parse_positive_int, default=1, help="repeat the parsed moe_copy event stream this many times with persistent simulated cache state") + parser.add_argument("--prefetch-limit", type=parse_positive_int, help="maximum experts to prefetch after each event for speculative policies; defaults to the slot count") + parser.add_argument( + "--policy", + choices=("lru", "prompt", "freq", "markov", "setmarkov", "oracle"), + default="lru", + help=( + "moe_copy simulation policy: lru is demand-only; prompt primes from bypass/prompt " + "events; freq keeps the most frequent experts hot; markov learns expert-to-expert " + "transitions; setmarkov learns expert-set transitions; oracle prefetches the next event " + "for an upper bound" + ), + ) + parser.add_argument("--details", action="store_true", help="also print per backend/tensor stats") + parser.add_argument("--runtime", action="store_true", help="summarize actual moe_cache/moe_cache_bypass runtime events instead of simulating moe_copy events") + args = parser.parse_args(argv) + + if args.runtime: + events, bypass_events = read_runtime_events(args.logs) + if not events and not bypass_events: + print("no moe_cache or moe_cache_bypass events found", file=sys.stderr) + return 1 + if events: + print_runtime_report(summarize_runtime_cache(events), args.details) + if bypass_events: + print_runtime_bypass_report(summarize_runtime_bypasses(bypass_events), args.details) + return 0 + + events = list(read_events(args.logs)) + if not events: + print("no moe_copy events found", file=sys.stderr) + return 1 + if args.repeat > 1: + events = events * args.repeat + + if args.policy == "lru": + print_report(simulate_lru(events, args.slots), args.details) + else: + print_prefetch_report(simulate_prefetch(events, args.slots, args.policy, args.prefetch_limit), args.details) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 7b1fcfca0ad..c9de3a2c118 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -25,6 +25,8 @@ add_library(llama llama-kv-cache.cpp llama-kv-cache-iswa.cpp llama-memory.cpp + llama-memory-deepseek4.cpp + llama-deepseek4-hot.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 633a66fc665..f2cf5e21f57 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 8f335f5c7b3..9438d9bb1d3 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 8126249e143..d495f77735a 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -4,6 +4,7 @@ #include "llama-arch.h" #include "llama-impl.h" #include "llama-batch.h" +#include "llama-deepseek4-hot.h" #include "llama-io.h" #include "llama-memory.h" #include "llama-mmap.h" @@ -346,6 +347,19 @@ llama_context::llama_context( LLAMA_LOG_INFO("%s: pipeline parallelism enabled\n", __func__); } + // DeepSeek4 hot-expert pinning: load profile and allocate per-layer + // hot subset tensors before the first sched_reserve so the graph + // builder can see them. No-op if DS4_HOT_PROFILE_JSON is unset. + if (model.arch == LLM_ARCH_DEEPSEEK4) { + auto & ds4_hot_mgr = ds4_hot::instance(); + if (ds4_hot_mgr.load_profile()) { + // Pass n_expert_used (P) so the manager can allocate the + // K + P + 1 dummy/padding slots needed by the dispatch graph. + ds4_hot_mgr.set_n_picks((int) model.hparams.n_expert_used); + ds4_hot_mgr.allocate(model); + } + } + sched_reserve(); if (!cparams.flash_attn) { @@ -469,6 +483,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 +2092,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 * 512, 256u * 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-deepseek4-hot.cpp b/src/llama-deepseek4-hot.cpp new file mode 100644 index 00000000000..2933cc34b9e --- /dev/null +++ b/src/llama-deepseek4-hot.cpp @@ -0,0 +1,538 @@ +#include "llama-deepseek4-hot.h" + +#include "llama.h" +#include "llama-impl.h" +#include "llama-model.h" +#include "ggml.h" +#include "ggml-backend.h" +#include "ggml-cpp.h" + +#include "../vendor/nlohmann/json.hpp" + +#include +#include +#include +#include +#include +#include +#include + +using nlohmann::json; + +namespace ds4_hot { + +struct hot_manager::ggml_buffers { + std::vector ctxs; + std::vector bufs; +}; + +hot_manager::~hot_manager() = default; + +const layer_hot_state * hot_manager::get(int il) const { + if (il < 0 || (size_t) il >= layers.size()) return nullptr; + return layers[il].get(); +} + +size_t hot_manager::total_gpu_bytes() const { + if (!bufs) return 0; + size_t total = 0; + for (const auto & b : bufs->bufs) { + if (b) total += ggml_backend_buffer_get_size(b.get()); + } + return total; +} + +bool hot_manager::load_profile(std::string path) { + if (active) return true; + + if (path.empty()) { + const char * env = std::getenv("DS4_HOT_PROFILE_JSON"); + if (!env || !*env) return false; + path = env; + } + + std::ifstream f(path); + if (!f.good()) { + LLAMA_LOG_ERROR("ds4-hot: failed to open profile %s\n", path.c_str()); + return false; + } + + json j; + try { + f >> j; + } catch (const std::exception & e) { + LLAMA_LOG_ERROR("ds4-hot: failed to parse %s: %s\n", path.c_str(), e.what()); + return false; + } + + if (!j.contains("hot") || !j.contains("k") || !j.contains("n_expert") || !j.contains("n_layer")) { + LLAMA_LOG_ERROR("ds4-hot: profile missing required fields (hot, k, n_expert, n_layer)\n"); + return false; + } + + n_layer = j.value("n_layer", 0); + n_expert = j.value("n_expert", 0); + k = j.value("k", 0); + category = j.value("category", std::string{}); + + if (k <= 0 || n_expert <= 0 || n_layer == 0) { + LLAMA_LOG_ERROR("ds4-hot: invalid profile dimensions: n_layer=%zu n_expert=%d k=%d\n", + n_layer, n_expert, k); + return false; + } + + layers.resize(n_layer); + + const auto & hot_obj = j["hot"]; + int loaded = 0; + for (auto it = hot_obj.begin(); it != hot_obj.end(); ++it) { + int il = std::atoi(it.key().c_str()); + if (il < 0 || (size_t) il >= n_layer) continue; + if (!it.value().is_array()) continue; + + auto state = std::make_unique(); + state->il = il; + state->hot_ids.reserve(k); + state->hot_set.reserve(k); + for (const auto & v : it.value()) { + int e = v.is_number_integer() ? v.get() : -1; + if (e < 0 || e >= n_expert) continue; + state->hot_ids.push_back(e); + state->hot_set.insert(e); + if ((int) state->hot_ids.size() >= k) break; + } + state->k = (int) state->hot_ids.size(); + if (state->k <= 0) continue; + + // Build cold set and remap tables. + state->remap_hot.assign(n_expert, -1); + state->remap_cold.assign(n_expert, -1); + for (int idx = 0; idx < state->k; ++idx) { + state->remap_hot[state->hot_ids[idx]] = idx; + } + + state->cold_ids.reserve(n_expert - state->k); + int cold_idx = 0; + for (int e = 0; e < n_expert; ++e) { + if (state->hot_set.count(e) == 0) { + state->cold_ids.push_back(e); + state->cold_set.insert(e); + state->remap_cold[e] = cold_idx++; + } + } + + layers[il] = std::move(state); + loaded++; + } + + LLAMA_LOG_INFO("ds4-hot: loaded profile %s category=%s k=%d n_layer=%zu n_expert=%d (entries=%d)\n", + path.c_str(), category.c_str(), k, n_layer, n_expert, loaded); + + active = (loaded > 0); + return active; +} + +namespace { + +// Track per-device allocations to avoid all hot tensors piling onto one GPU. +struct device_budget { + ggml_backend_buffer_type_t buft; + size_t reserved = 0; // bytes already targeted at this buft in current allocate() call + size_t free_at_start = 0; +}; + +static std::vector g_budgets; + +void init_budgets() { + g_budgets.clear(); + const int n_dev = ggml_backend_dev_count(); + for (int i = 0; i < n_dev; ++i) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_GPU) continue; + size_t free = 0, total = 0; + ggml_backend_dev_memory(dev, &free, &total); + device_budget b; + b.buft = ggml_backend_dev_buffer_type(dev); + b.free_at_start = free; + b.reserved = 0; + g_budgets.push_back(b); + } +} + +// Pick the GPU buffer type with the most remaining headroom that can fit +// `needed_bytes`. Reserves the bytes immediately so subsequent picks see +// the running total. +// +// If DS4_HOT_DEVICE is set in the environment, restrict picking to the +// matching CUDAN device (e.g. `DS4_HOT_DEVICE=CUDA0`). This is useful for +// debugging the dispatch path: pinning all hot tensors onto one GPU +// eliminates a class of multi-device scheduler interactions that have +// triggered illegal-memory-access crashes on certain prompts. +// +// Margin (default 1.5 GiB per device) is left untouched so prefill compute +// buffers can fit. DS4_HOT_MARGIN_MIB overrides this; use a larger value if +// you observe OOM errors during prefill of long prompts. +ggml_backend_buffer_type_t pick_gpu_buft(size_t needed_bytes) { + static const size_t margin = []() -> size_t { + const char * env = std::getenv("DS4_HOT_MARGIN_MIB"); + if (!env || !*env) return (size_t) 1536 * 1024 * 1024; // 1.5 GiB default + long v = std::strtol(env, nullptr, 10); + if (v <= 0) return (size_t) 1536 * 1024 * 1024; + return (size_t) v * 1024 * 1024; + }(); + + static const char * const force_device = std::getenv("DS4_HOT_DEVICE"); + + ggml_backend_buffer_type_t best = nullptr; + size_t best_remaining = 0; + for (auto & b : g_budgets) { + if (force_device && force_device[0]) { + if (std::strcmp(ggml_backend_buft_name(b.buft), force_device) != 0) { + continue; + } + } + size_t avail = b.free_at_start - std::min(b.free_at_start, b.reserved + margin); + if (avail < needed_bytes) continue; + size_t remaining_after = avail - needed_bytes; + if (remaining_after > best_remaining || best == nullptr) { + best_remaining = remaining_after; + best = b.buft; + } + } + if (best) { + for (auto & b : g_budgets) { + if (b.buft == best) { b.reserved += needed_bytes; break; } + } + } + return best; +} + +} // namespace + +bool hot_manager::allocate(const llama_model & model) { + if (!active) return false; + if (bufs && !bufs->bufs.empty()) return true; // already allocated + + init_budgets(); + + bufs = std::make_unique(); + + const auto & m_layers = model.layers; + if (m_layers.size() != n_layer) { + LLAMA_LOG_WARN("ds4-hot: profile n_layer=%zu but model has %zu layers; tolerating mismatch\n", + n_layer, m_layers.size()); + } + + // Pending uploads: a tensor pointer slot + the host bytes to copy into it. + struct pending_upload { + ggml_tensor ** slot; + std::vector data; + }; + + // Per-buft (i.e., per-GPU device) ggml_context that aggregates all hot + // tensors + lookup tables targeted at that device. We allocate one backing + // buffer per buft after the loop. + struct ctx_entry { + ggml_context_ptr ctx; + std::vector pending; + }; + std::map per_buft; + + auto get_ctx = [&](ggml_backend_buffer_type_t buft) -> ggml_context * { + auto it = per_buft.find(buft); + if (it != per_buft.end()) return it->second.ctx.get(); + // Reserve enough space for ~16 tensors per layer (3 weight + 4 lookup + headroom). + ggml_init_params p = { + /*.mem_size =*/ 16 * (size_t) ggml_tensor_overhead() * std::max(n_layer, 1), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context_ptr ctx_owner(ggml_init(p)); + ctx_entry e; + e.ctx = std::move(ctx_owner); + ggml_context * raw = e.ctx.get(); + per_buft.emplace(buft, std::move(e)); + return raw; + }; + + int n_alloc_layers = 0; + size_t total_bytes = 0; + + // Compute the total bytes one layer's hot tensors + lookup tables need so + // we can reserve all of them on the SAME device. This is essential — if + // gate_h, up_h, down_h end up on different GPUs the dual dispatch becomes + // a multi-backend mess and we lose the placement benefit. + auto layer_total_bytes = [&](int il, const llama_layer & lm) -> size_t { + size_t total = 0; + const layer_hot_state & st = *layers[il]; + const int64_t k_local = (int64_t) st.hot_ids.size(); + auto add_tensor = [&](const ggml_tensor * src) { + if (!src) return; + total += (ggml_nbytes(src) / src->ne[2]) * k_local; + }; + if (lm.ffn_gate_up_exps) { + add_tensor(lm.ffn_gate_up_exps); + } else { + add_tensor(lm.ffn_gate_exps); + add_tensor(lm.ffn_up_exps); + } + add_tensor(lm.ffn_down_exps); + // Lookup tables live in CPU buffer so they don't count against GPU budget. + return total; + }; + + auto extract_subset = [&](ggml_backend_buffer_type_t buft, const ggml_tensor * src, + const std::vector & hot_ids, + const std::string & dest_name, ggml_tensor ** out_tensor) -> bool { + if (!src) return false; + if (!src->buffer) return false; + + const int64_t ne0 = src->ne[0]; + const int64_t ne1 = src->ne[1]; + const int64_t n_expert_src = src->ne[2]; + if (n_expert_src != n_expert) { + LLAMA_LOG_WARN("ds4-hot: tensor %s has %ld experts, profile expects %d\n", + src->name, (long) n_expert_src, n_expert); + return false; + } + + const size_t per_expert_bytes = ggml_nbytes(src) / n_expert_src; + const int64_t k_local = (int64_t) hot_ids.size(); + // Allocate K + P + 1 experts: + // [0, K) - real hot experts + // [K, K+P) - per-pick dummy experts (zero-weighted; never collide + // with real expert IDs across picks of a single token, + // which fixes the CUDA mm_ids_helper dedup crash) + // [K+P] - trailing prefetch padding slot (kernel reads ahead) + const int64_t P = n_picks_; + const int64_t k_alloc = k_local + P + 1; + const size_t needed = per_expert_bytes * k_alloc; + + // Pull source data from CPU into a host buffer we can slice from. + std::vector host_data(ggml_nbytes(src)); + ggml_backend_tensor_get(src, host_data.data(), 0, host_data.size()); + + // Build the slice in a separate host buffer (zero-initialized so the + // dummy experts and trailing prefetch slot all hold zeros). + std::vector slice(needed, 0); + for (int64_t r = 0; r < k_local; ++r) { + const int32_t e = hot_ids[(size_t) r]; + const size_t src_off = per_expert_bytes * (size_t) e; + const size_t dst_off = per_expert_bytes * (size_t) r; + std::memcpy(slice.data() + dst_off, host_data.data() + src_off, per_expert_bytes); + } + + ggml_context * ctx = get_ctx(buft); + if (!ctx) return false; + ggml_tensor * dst = ggml_new_tensor_3d(ctx, src->type, ne0, ne1, k_alloc); + ggml_format_name(dst, "%s.hot", dest_name.c_str()); + + per_buft[buft].pending.push_back({ out_tensor, std::move(slice) }); + *out_tensor = dst; + total_bytes += needed; + return true; + }; + + auto add_lookup_f32 = [&](ggml_backend_buffer_type_t buft, const std::string & name, + const std::vector & values, int64_t ne0, int64_t ne1, + ggml_tensor ** out_tensor) -> bool { + ggml_context * ctx = get_ctx(buft); + if (!ctx) return false; + ggml_tensor * dst = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ne0, ne1); + ggml_format_name(dst, "%s", name.c_str()); + std::vector bytes(values.size() * sizeof(float)); + std::memcpy(bytes.data(), values.data(), bytes.size()); + per_buft[buft].pending.push_back({ out_tensor, std::move(bytes) }); + *out_tensor = dst; + total_bytes += values.size() * sizeof(float); + return true; + }; + + for (size_t il = 0; il < std::min(n_layer, m_layers.size()); ++il) { + if (!layers[il]) { + continue; + } + auto & state = *layers[il]; + const auto & lm = m_layers[il]; + + if (!lm.ffn_gate_up_exps && !(lm.ffn_gate_exps && lm.ffn_up_exps)) { + continue; + } + if (!lm.ffn_down_exps) { + continue; + } + + const bool has_combined = (lm.ffn_gate_up_exps != nullptr); + const ggml_tensor * probe = has_combined ? lm.ffn_gate_up_exps + : (lm.ffn_gate_exps ? lm.ffn_gate_exps : lm.ffn_up_exps); + + if (!probe || !probe->buffer) { + continue; + } + + bool buf_is_host = ggml_backend_buft_is_host(ggml_backend_buffer_get_type(probe->buffer)); + if (!buf_is_host) { + ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(probe->buffer)); + if (dev && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_GPU) { + layers[il].reset(); + continue; + } + } + + // Pick ONE GPU for all of this layer's hot tensors + lookup tables. + const size_t needed = layer_total_bytes((int) il, lm); + ggml_backend_buffer_type_t buft = pick_gpu_buft(needed); + if (!buft) { + LLAMA_LOG_WARN("ds4-hot: no GPU has %.1f MiB free for layer %zu hot pack; skipping\n", + needed / (1024.0 * 1024.0), il); + layers[il].reset(); + continue; + } + + // Hot-side lookup tables (hot_remap, is_hot) live on the SAME GPU as the + // hot weights so the get_rows + mul_mat_id chain can run entirely on + // that GPU without any cross-backend transfer of the per-pick IDs. + // Cold-side tables (cold_remap, is_cold) live on CPU so the cold + // mul_mat_id (CPU weights) consumes a CPU IDs tensor without sched + // having to bounce data between backends each step. + ggml_backend_buffer_type_t cpu_buft = ggml_backend_cpu_buffer_type(); + + // Build per-layer lookup tables. + // For the hot path we use float arithmetic in the graph to construct + // per-pick unique IDs: + // hot_id[k,t] = hot_remap_table[selected[k,t]] + is_cold[k,t] * hot_pick_arange[k] + // For hot picks the arange contribution is 0 -> hot_id in [0, K). + // For cold picks the base is K (sentinel) and the arange adds k -> id in [K, K+P). + // This guarantees all P picks within a token map to distinct expert IDs, + // which is required by the CUDA mm_ids_helper kernel (it dedups + // (token, expert) pairs and the downstream quantize kernel reads + // exactly P*T compact rows). + const int P = n_picks_; + + std::vector hot_remap_vals((size_t) n_expert, (float) state.k); // base sentinel = K + std::vector cold_remap_vals((size_t) n_expert, 0.0f); + std::vector is_hot_vals((size_t) n_expert, 0.0f); + std::vector is_cold_vals((size_t) n_expert, 1.0f); + for (int32_t e : state.hot_ids) { + hot_remap_vals[(size_t) e] = (float) state.remap_hot[(size_t) e]; + is_hot_vals[(size_t) e] = 1.0f; + is_cold_vals[(size_t) e] = 0.0f; + } + for (int32_t e : state.cold_ids) { + cold_remap_vals[(size_t) e] = (float) e; + } + + // Per-pick arange [0, 1, ..., P-1] + std::vector pick_arange_vals((size_t) P); + for (int i = 0; i < P; ++i) pick_arange_vals[(size_t) i] = (float) i; + + // Per-pick cold sentinel: cold_ids[k % n_cold] for k in [0, P). + // For the COLD path on CPU we actually want a SINGLE shared sentinel + // so that the CPU mul_mat_id's matrix_row_counts dedup collapses all + // hot picks within a token to one expert load (saves bandwidth). The + // CUDA mm_ids_helper bug that required per-pick uniqueness only + // affects the GPU hot path. So we set every entry to cold_ids[0]. + std::vector cold_sentinel_vals((size_t) P); + const int n_cold = (int) state.cold_ids.size(); + for (int i = 0; i < P; ++i) { + cold_sentinel_vals[(size_t) i] = n_cold > 0 ? (float) state.cold_ids[0] : 0.0f; + } + + bool ok_all = true; + if (has_combined) { + ok_all &= extract_subset(buft, lm.ffn_gate_up_exps, state.hot_ids, + "ds4_hot_gate_up_exps_l" + std::to_string(il), + &state.hot_gate_up_exps); + } else { + ok_all &= extract_subset(buft, lm.ffn_gate_exps, state.hot_ids, + "ds4_hot_gate_exps_l" + std::to_string(il), + &state.hot_gate_exps); + ok_all &= extract_subset(buft, lm.ffn_up_exps, state.hot_ids, + "ds4_hot_up_exps_l" + std::to_string(il), + &state.hot_up_exps); + } + ok_all &= extract_subset(buft, lm.ffn_down_exps, state.hot_ids, + "ds4_hot_down_exps_l" + std::to_string(il), + &state.hot_down_exps); + + // Track per-layer pick count for downstream graph builder access. + state.n_picks = P; + + ok_all &= add_lookup_f32(buft, "ds4_hot_remap_l" + std::to_string(il), + hot_remap_vals, 1, n_expert, &state.hot_remap_table); + ok_all &= add_lookup_f32(cpu_buft, "ds4_cold_remap_l" + std::to_string(il), + cold_remap_vals, 1, n_expert, &state.cold_remap_table); + ok_all &= add_lookup_f32(buft, "ds4_is_hot_l" + std::to_string(il), + is_hot_vals, 1, n_expert, &state.is_hot_mask); + ok_all &= add_lookup_f32(cpu_buft, "ds4_is_cold_l" + std::to_string(il), + is_cold_vals, 1, n_expert, &state.is_cold_mask); + // Per-pick constants live as [P, 1] tensors so they broadcast against + // [P, T] when multiplied. Hot side on GPU, cold side on CPU. + ok_all &= add_lookup_f32(buft, "ds4_hot_pick_arange_l" + std::to_string(il), + pick_arange_vals, P, 1, &state.hot_pick_arange); + ok_all &= add_lookup_f32(cpu_buft, "ds4_cold_pick_sentinel_l" + std::to_string(il), + cold_sentinel_vals, P, 1, &state.cold_pick_sentinel); + + if (!ok_all) { + state.hot_gate_up_exps = nullptr; + state.hot_gate_exps = nullptr; + state.hot_up_exps = nullptr; + state.hot_down_exps = nullptr; + state.hot_remap_table = nullptr; + state.cold_remap_table = nullptr; + state.is_hot_mask = nullptr; + state.is_cold_mask = nullptr; + state.hot_pick_arange = nullptr; + state.cold_pick_sentinel = nullptr; + layers[il].reset(); + continue; + } + n_alloc_layers++; + } + + // Now allocate backing buffers and upload the slices. + for (auto & [buft, e] : per_buft) { + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(e.ctx.get(), buft); + if (!buf) { + LLAMA_LOG_WARN("ds4-hot: could not allocate hot buffer for buft %s; skipping affected layers\n", + ggml_backend_buft_name(buft)); + for (auto & p : e.pending) { + if (p.slot) *p.slot = nullptr; + } + continue; + } + // Mark as model weights so the scheduler keeps the consuming ops on + // this device (matches normal model-weight placement semantics). + ggml_backend_buffer_set_usage(buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + for (auto & p : e.pending) { + if (!p.slot || !*p.slot) continue; + ggml_backend_tensor_set(*p.slot, p.data.data(), 0, p.data.size()); + } + bufs->bufs.emplace_back(buf); + bufs->ctxs.emplace_back(std::move(e.ctx)); + } + + // Re-validate: a layer is fully usable only if EVERY required tensor and + // lookup table is non-null after upload. + int n_usable = 0; + for (auto & lp : layers) { + if (!lp) continue; + if (lp->ready_for_dispatch()) { + n_usable++; + } else { + lp.reset(); + } + } + + LLAMA_LOG_INFO("ds4-hot: pinned hot experts for %d/%d CPU-MoE layers, ~%.1f MiB on GPU across %zu buffers (k=%d, category=%s)\n", + n_usable, n_alloc_layers, total_bytes / (1024.0 * 1024.0), bufs->bufs.size(), k, category.c_str()); + + return n_usable > 0; +} + +hot_manager & instance() { + static hot_manager mgr; + return mgr; +} + +} // namespace ds4_hot diff --git a/src/llama-deepseek4-hot.h b/src/llama-deepseek4-hot.h new file mode 100644 index 00000000000..b1ca9e547c9 --- /dev/null +++ b/src/llama-deepseek4-hot.h @@ -0,0 +1,142 @@ +// DeepSeek4 hot-expert pinning manager. +// +// Reads a per-layer hot-expert-ID profile (produced by ds4-expert-profile + +// ds4-hot-experts.py), extracts the K hot experts of each layer's +// `ffn_gate_up_exps` and `ffn_down_exps` tensors into a separate GPU buffer +// after model load, and exposes those subset tensors to the deepseek4 graph +// builder so build_moe_v4 / build_expert_mix can issue dual mul_mat_id +// dispatches (hot subset on GPU, cold subset on CPU). +// +// Activation: set DS4_HOT_PROFILE_JSON=path.json before starting llama-server +// or llama-cli. The JSON shape matches what ds4-hot-experts.py extract emits: +// { "n_layer": 43, "n_expert": 256, "k": 32, +// "category": "code", +// "hot": { "0": [12, 47, ...], ... } } +// +// This is Phase 1 (load-time extraction). Phase 2 (graph dispatch) lives in +// src/models/deepseek4.cpp. +#pragma once + +#include "ggml.h" + +#include +#include +#include +#include + +struct llama_model; +struct llama_context; + +namespace ds4_hot { + +struct layer_hot_state { + int il = -1; + int k = 0; + int n_picks = 0; // n_expert_used (P), e.g. 6 for DS4 + std::vector hot_ids; // size K, sorted by frequency desc + std::unordered_set hot_set; // for O(1) membership + std::vector cold_ids; // size n_expert - K + std::unordered_set cold_set; + std::vector remap_hot; // size n_expert: original -> 0..K-1 or -1 + std::vector remap_cold; // size n_expert: original -> 0..(n_expert-K)-1 or -1 + + // Pinned hot tensor data: extracted K hot expert rows + P zero-weighted + // dummy expert rows (one per pick index) + 1 trailing prefetch padding row. + // Total ne[2] = K + P + 1. The dummy experts at positions [K, K+P) let + // each pick within a token get a unique remapped ID even when most picks + // are cold, which is required by the CUDA mm_ids_helper kernel: it + // dedups (token, expert) pairs and produces fewer compacted rows when + // multiple picks share the same id, leaving the tail of ids_src1 + // uninitialized -> illegal memory access in quantize_mmq_mxfp4_cuda. + // Per-pick unique dummy experts (id = K + pick_idx) keep the helper + // emitting exactly P*T rows. + // For models with combined gate+up (DS-V3 style): hot_gate_up_exps is set, hot_gate_exps and hot_up_exps are null. + // For models with separate gate/up (DS4-Flash style): hot_gate_exps and hot_up_exps are set, hot_gate_up_exps is null. + ggml_tensor * hot_gate_up_exps = nullptr; + ggml_tensor * hot_gate_exps = nullptr; + ggml_tensor * hot_up_exps = nullptr; + ggml_tensor * hot_down_exps = nullptr; + + // Phase 2 graph-time lookup tables. + // + // hot_remap_table_f32[0, e] = remap_hot[e] (in [0, K)) if hot, K (base + // sentinel) if cold. Combined with a per-pick offset arange [0..P-1] + // in the graph: hot_ids = hot_remap + is_cold * arange so each cold + // pick gets a unique dummy id in [K, K+P). + // cold_remap_table_f32[0, e] = e if cold, 0 if hot. Combined with a + // per-pick cold sentinel arange [cold_ids[0]..cold_ids[P-1]] so each + // hot pick gets a different cold sentinel within the token (avoids + // the same dedup bug on the CPU mul_mat_id, defensively). + // is_hot_mask[0, e] / is_cold_mask[0, e] = 1.0 / 0.0. Used for the + // cold-path output mask (hot path no longer needs an output mask + // because the dummy experts produce zero output by construction). + // hot_pick_arange = [0, 1, ..., P-1] f32, length P. + // cold_pick_sentinel = [cold_ids[0], ..., cold_ids[P-1]] f32, length P. + ggml_tensor * hot_remap_table = nullptr; // f32 + ggml_tensor * cold_remap_table = nullptr; // f32 + ggml_tensor * is_hot_mask = nullptr; // f32 + ggml_tensor * is_cold_mask = nullptr; // f32 + ggml_tensor * hot_pick_arange = nullptr; // f32 [P] + ggml_tensor * cold_pick_sentinel = nullptr; // f32 [P] + + // Returns true if all tensors required for Phase 2 dual dispatch are non-null. + bool ready_for_dispatch() const { + const bool gate_up_ok = hot_gate_up_exps || (hot_gate_exps && hot_up_exps); + return gate_up_ok && hot_down_exps && hot_remap_table && cold_remap_table + && is_hot_mask && is_cold_mask + && hot_pick_arange && cold_pick_sentinel; + } +}; + +class hot_manager { +public: + hot_manager() = default; + ~hot_manager(); + + // Returns true if a profile path was provided and successfully loaded. + // Idempotent. Pulls path from DS4_HOT_PROFILE_JSON env var if path is empty. + bool load_profile(std::string path = {}); + + // Allocate per-layer hot subset tensors on the same device as the model's + // GPU split would prefer. Reads the original ffn_*_exps host data from + // each layer (which must already be loaded into CPU memory) and copies the + // K hot rows into a new GPU tensor. + // + // Must be called AFTER the model has been loaded and BEFORE inference. + bool allocate(const llama_model & model); + + bool is_active() const { return active; } + int k_per_layer() const { return k; } + size_t profile_n_layer() const { return n_layer; } + int profile_n_expert() const { return n_expert; } + int n_picks() const { return n_picks_; } + + void set_n_picks(int p) { n_picks_ = p; } + + // Per-layer accessors. il is the layer index. Returns nullptr if no hot + // state was allocated for that layer (e.g., layer is fully on GPU already + // and we skipped it). + const layer_hot_state * get(int il) const; + + // Total bytes pinned to GPU buffers across all layers (for reporting). + size_t total_gpu_bytes() const; + +private: + bool active = false; + std::string category = {}; + int k = 0; + int n_picks_ = 6; // n_expert_used (P); set from model hparams via set_n_picks + size_t n_layer = 0; + int n_expert = 0; + std::vector> layers; + + struct ggml_buffers; + std::unique_ptr bufs; +}; + +// Singleton accessor; convenient for plumbing through llama-context without +// changing the C API. The instance is created on first call and persists for +// the program lifetime. +hot_manager & instance(); + +} // namespace ds4_hot diff --git a/src/llama-memory-deepseek4.cpp b/src/llama-memory-deepseek4.cpp new file mode 100644 index 00000000000..9149ef7cafd --- /dev/null +++ b/src/llama-memory-deepseek4.cpp @@ -0,0 +1,722 @@ +#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 bool deepseek4_batch_log_enabled() { + const char * value = std::getenv("LLAMA_DEEPSEEK4_BATCH_LOG"); + return value != nullptr && std::strcmp(value, "0") != 0; +} + +static bool deepseek4_batch_prefill_enabled() { + // Default-on: batched prefill is ~7x faster than single-token at long + // context with no measurable correctness regression on the NMSE smoke + // tests. Set LLAMA_DEEPSEEK4_BATCH_PREFILL=0 to fall back to the + // single-token path. + const char * value = std::getenv("LLAMA_DEEPSEEK4_BATCH_PREFILL"); + return value == nullptr || std::strcmp(value, "0") != 0; +} + +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); + + const bool log_batch = deepseek4_batch_log_enabled(); + if (log_batch) { + std::fprintf(stderr, "%s: requested n_tokens=%u n_outputs=%u n_ubatch=%u embd_all=%d; current DeepSeek4 path splits to single-token ubatches\n", + __func__, balloc.get_n_tokens(), balloc.get_n_outputs(), n_ubatch, embd_all ? 1 : 0); + } + + balloc.split_reset(); + + // Only enable multi-token ubatches when batch_prefill is on AND the + // graph builder will agree (n_outputs != n_tokens, the prefill case). + // Otherwise the build sees work_tokens=1 (reserve_only) but the runtime + // would have given it a multi-token ubatch, and the mismatch corrupts + // logit reads. + const bool batch_prefill_active = + deepseek4_batch_prefill_enabled() && + balloc.get_n_outputs() != balloc.get_n_tokens(); + std::vector ubatches; + while (true) { + // Optional batched prefill (LLAMA_DEEPSEEK4_BATCH_PREFILL=1): split + // up to n_ubatch tokens per ubatch from a single sequence. Decoding + // and the unopt'ed path keep the legacy single-token semantics. + const uint32_t split_n = batch_prefill_active ? n_ubatch : 1; + llama_ubatch ubatch = balloc.split_seq(split_n); + if (ubatch.n_tokens == 0) { + break; + } + + if (ubatch.n_seqs_unq != 1 || (!batch_prefill_active && ubatch.n_tokens != 1)) { + LLAMA_LOG_ERROR("%s: DeepSeek4 runtime currently supports a single sequence per ubatch (got n_tokens=%u, n_seqs=%u, batch_prefill=%d)\n", + __func__, ubatch.n_tokens, ubatch.n_seqs_unq, batch_prefill_active ? 1 : 0); + 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 (log_batch) { + std::fprintf(stderr, "%s: prepared %zu %subatches\n", __func__, ubatches.size(), batch_prefill_active ? "" : "single-token "); + } + + 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 00000000000..af73c9cfe09 --- /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-loader.cpp b/src/llama-model-loader.cpp index 4e65a45a50d..1ca425cc627 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -44,6 +44,7 @@ static std::string llama_model_ftype_name(llama_ftype ftype) { case LLAMA_FTYPE_MOSTLY_Q8_0: return "Q8_0"; case LLAMA_FTYPE_MOSTLY_MXFP4_MOE: return "MXFP4 MoE"; case LLAMA_FTYPE_MOSTLY_NVFP4: return "NVFP4"; + case LLAMA_FTYPE_MOSTLY_F8_E4M3_MXFP4: return "F8_E4M3 + MXFP4"; case LLAMA_FTYPE_MOSTLY_Q2_K: return "Q2_K - Medium"; case LLAMA_FTYPE_MOSTLY_Q2_K_S: return "Q2_K - Small"; case LLAMA_FTYPE_MOSTLY_Q3_K_S: return "Q3_K - Small"; @@ -760,6 +761,7 @@ llama_model_loader::llama_model_loader( case GGML_TYPE_IQ3_S: ftype = LLAMA_FTYPE_MOSTLY_IQ3_S; break; case GGML_TYPE_NVFP4: ftype = LLAMA_FTYPE_MOSTLY_NVFP4; break; case GGML_TYPE_Q1_0: ftype = LLAMA_FTYPE_MOSTLY_Q1_0; break; + case GGML_TYPE_F8_E4M3_B128: ftype = LLAMA_FTYPE_MOSTLY_F8_E4M3_MXFP4; break; default: { LLAMA_LOG_WARN("%s: unknown type %s\n", __func__, ggml_type_name(type_max)); diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 26864c18e97..7b0041b6962 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 9e2a13cbd43..c17e8c57a29 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 5f101bd6374..120068e4483 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/llama-quant.cpp b/src/llama-quant.cpp index 25a333b4a7f..ba89571603b 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -800,6 +800,7 @@ ggml_type llama_ftype_get_default_type(llama_ftype ftype) { case LLAMA_FTYPE_MOSTLY_BF16: return GGML_TYPE_BF16; case LLAMA_FTYPE_ALL_F32: return GGML_TYPE_F32; case LLAMA_FTYPE_MOSTLY_Q1_0: return GGML_TYPE_Q1_0; + case LLAMA_FTYPE_MOSTLY_F8_E4M3_MXFP4: return GGML_TYPE_F8_E4M3_B128; case LLAMA_FTYPE_MOSTLY_MXFP4_MOE: return GGML_TYPE_MXFP4; diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp new file mode 100644 index 00000000000..daed6e22f80 --- /dev/null +++ b/src/models/deepseek4.cpp @@ -0,0 +1,1410 @@ +#include "models.h" + +#include "llama-impl.h" +#include "llama-memory-deepseek4.h" +#include "../llama-deepseek4-hot.h" + +#include +#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 bool deepseek4_batch_log_enabled() { + const char * value = std::getenv("LLAMA_DEEPSEEK4_BATCH_LOG"); + return value != nullptr && std::strcmp(value, "0") != 0; +} + +static bool deepseek4_batch_prefill_enabled() { + // Default-on: batched prefill is ~7x faster than single-token at long + // context with no measurable correctness regression on the NMSE smoke + // tests. Set LLAMA_DEEPSEEK4_BATCH_PREFILL=0 to fall back to the + // single-token path (used as an escape hatch if a downstream model + // shows quality regressions from the shared top-k indexer aggregation). + const char * value = std::getenv("LLAMA_DEEPSEEK4_BATCH_PREFILL"); + return value == nullptr || std::strcmp(value, "0") != 0; +} + +static bool deepseek4_indexer_collapse_q() { + // Default OFF: in batched prefill, the indexer keeps each query's score + // separate and only aggregates AT THE END (sum across n_head AND + // work_tokens) for a single ubatch-shared top-k. This is mathematically + // closer to the original per-token model and works correctly at long + // context (65K+ retrieval-style prompts). + // + // Set LLAMA_DEEPSEEK4_INDEXER_COLLAPSE_Q=1 to opt into the approximate + // path that sums indexer_q across queries BEFORE the score mul_mat. + // That cuts the score tensor from [n_comp, n_head, work_tokens] down + // to [n_comp, n_head] (fits in tight VRAM at large ub) and is faster + // per ubatch, but breaks long-context quality because the relu + // non-linearity in scoring means relu(sum_q (kv*q)) != sum_q + // relu(kv*q) -- the model attends to wrong KV slots. + static const bool enabled = []() { + const char * value = std::getenv("LLAMA_DEEPSEEK4_INDEXER_COLLAPSE_Q"); + return value != nullptr && std::strcmp(value, "0") != 0; + }(); + return enabled; +} + +static bool deepseek4_indexer_per_query() { + // Inverse of deepseek4_indexer_collapse_q (kept for code clarity at + // call sites). Per-query is the default; collapse-Q is opt-in. + return !deepseek4_indexer_collapse_q(); +} + +static bool deepseek4_hot_dispatch_enabled() { + // Default OFF until the prompt-content-sensitive crash on certain expert + // ID patterns is resolved. Set DS4_HOT_DISPATCH=1 to opt in. + static const bool enabled = []() { + const char * value = std::getenv("DS4_HOT_DISPATCH"); + if (value == nullptr) return false; + return std::strcmp(value, "0") != 0; + }(); + return enabled; +} + +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 (size_t mi = 0; mi < kq_masks.size(); ++mi) { + ggml_tensor * mask = kq_masks[mi]; + if (!mask || !mask->buffer) { + continue; + } + + const int64_t n_kv_padded = mask->ne[0]; + const int64_t n_q = mask->ne[1]; + // n_kv_total[mi] is the actual (unpadded) size; slots in + // [n_kv_total, n_kv_padded) are padding and stay at -INFINITY. + const int64_t n_kv_actual = (mi < kq_mask_n_kv_total.size()) ? kq_mask_n_kv_total[mi] : n_kv_padded; + 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_actual; ++ikv) { + if (ikv >= (int64_t) n_swa || ikv <= q_pos) { + f32_data[iq*n_kv_padded + 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 kq_mask_n_kv_total; + // Cache shared kq_mask tensors keyed by (n_kv_total, work_tokens) so all + // V4 layers with the same comp_ratio reuse a single graph input. Without + // this we hit GGML_SCHED_MAX_SPLIT_INPUTS (30) at >30 layers. + std::map, ggml_tensor *> kq_mask_by_shape; + + 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); + + const bool batch_prefill = deepseek4_batch_prefill_enabled() && n_outputs != n_tokens; + const bool reserve_only = n_tokens != 1 && !batch_prefill; + const llama_pos start_pos = reserve_only ? 0 : ubatch.pos[0]; + const int64_t work_tokens = reserve_only ? 1 : n_tokens; + if (deepseek4_batch_log_enabled()) { + std::fprintf(stderr, "%s: n_tokens=%" PRId64 " reserve_only=%d work_tokens=%" PRId64 " start_pos=%d\n", + __func__, n_tokens, reserve_only ? 1 : 0, work_tokens, (int) start_pos); + } + 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; + + // Fast path: the entire ubatch fits within one compression window. + if (start_mod + work_tokens <= comp_ratio) { + return matrix_block(ape, 0, start_mod, comp_dim, work_tokens); + } + + // General multi-window slice: decompose the ubatch into + // - start_remaining tokens from the current partial window + // - any number of complete windows + // - end_remaining tokens from the final partial window + // and concat the corresponding ape slices together. The number + // of concat ops is bounded by ceil(work_tokens / comp_ratio) + 1. + ggml_tensor * out = nullptr; + int64_t consumed = 0; + + if (start_mod > 0) { + const int64_t start_remaining = comp_ratio - start_mod; + out = matrix_block(ape, 0, start_mod, comp_dim, start_remaining); + consumed = start_remaining; + } + + while (consumed < work_tokens) { + const int64_t remaining = work_tokens - consumed; + const int64_t slice_len = std::min(remaining, comp_ratio); + ggml_tensor * cur = matrix_block(ape, 0, 0, comp_dim, slice_len); + out = out ? ggml_concat(ctx0, out, cur, 1) : cur; + consumed += slice_len; + } + + 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 * { + if (work_tokens > 1 && x_hc->ne[0] == n_embd && x_hc->ne[1] == hc_mult && x_hc->ne[2] == work_tokens && + weights->ne[0] == hc_mult && weights->ne[1] == work_tokens) { + return ggml_hc_weighted_sum(ctx0, x_hc, weights); + } + + if (x_hc->type == GGML_TYPE_F32 && weights->type == GGML_TYPE_F32 && + x_hc->ne[0] == n_embd && x_hc->ne[1] == hc_mult && x_hc->ne[2] == 1 && x_hc->ne[3] == 1 && + weights->ne[0] == hc_mult && weights->ne[1] == 1 && weights->ne[2] == 1 && weights->ne[3] == 1) { + return ggml_hc_weighted_sum(ctx0, x_hc, weights); + } + + 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 * { + if (comb->type == GGML_TYPE_F32 && + comb->ne[0] == 4 && comb->ne[1] == 4) { + return ggml_sinkhorn_4x4(ctx0, comb); + } + + 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; + + // Phase 2: hot-expert dual dispatch. + // If a hot-expert profile was loaded (DS4_HOT_PROFILE_JSON) and this + // layer has hot tensors pinned on GPU, route the K hot experts through + // the GPU-resident subset and only run the cold picks on the CPU + // tensor (with a sentinel cold expert in place of any hot picks). + const ds4_hot::layer_hot_state * hot = + ds4_hot::instance().is_active() ? ds4_hot::instance().get(il) : nullptr; + // Warmup/reserve graphs sometimes pass a selected_experts with + // ne[0] = n_expert instead of n_picks. In that case our per-pick + // arithmetic would assert in ggml_mul, so fall back to the single + // path code below. + const bool dispatch_dual = hot && hot->ready_for_dispatch() + && deepseek4_hot_dispatch_enabled() + && selected_experts->ne[0] == hot->n_picks; + + if (dispatch_dual) { + const int64_t n_picks_local = selected_experts->ne[0]; + const int64_t n_tokens_local = selected_experts->ne[1]; + + // Ensure selected_experts is contiguous before reshape (defensive). + ggml_tensor * sel_cont = ggml_cont(ctx0, selected_experts); + ggml_tensor * sel_flat = ggml_reshape_1d(ctx0, sel_cont, n_picks_local * n_tokens_local); + + // Lookup tables produce float values per pick (in [P*T] flat). + // Reshape each to [P, T] for the per-pick arithmetic and final + // mul_mat_id IDs cast. + ggml_tensor * hot_remap_flat = ggml_get_rows(ctx0, hot->hot_remap_table, sel_flat); + ggml_tensor * cold_remap_flat = ggml_get_rows(ctx0, hot->cold_remap_table, sel_flat); + ggml_tensor * is_hot_flat = ggml_get_rows(ctx0, hot->is_hot_mask, sel_flat); + ggml_tensor * is_cold_flat = ggml_get_rows(ctx0, hot->is_cold_mask, sel_flat); + + ggml_tensor * hot_remap = ggml_reshape_2d(ctx0, hot_remap_flat, n_picks_local, n_tokens_local); + ggml_tensor * cold_remap = ggml_reshape_2d(ctx0, cold_remap_flat, n_picks_local, n_tokens_local); + ggml_tensor * is_hot = ggml_reshape_2d(ctx0, is_hot_flat, n_picks_local, n_tokens_local); + ggml_tensor * is_cold = ggml_reshape_2d(ctx0, is_cold_flat, n_picks_local, n_tokens_local); + + // Construct per-pick unique IDs: + // hot_ids = hot_remap + is_cold * hot_pick_arange (broadcasts [P,1] -> [P,T]) + // cold_ids = cold_remap + is_hot * cold_pick_sentinel + // hot_pick_arange = [0, 1, ..., P-1] so cold picks land in [K, K+P) (the dummy zero-weighted experts). + // cold_pick_sentinel = [cold_ids[0], ..., cold_ids[P-1]] so hot picks each get a different cold sentinel. + ggml_tensor * hot_offset = ggml_mul(ctx0, is_cold, hot->hot_pick_arange); // [P, T] f32 + ggml_tensor * cold_offset = ggml_mul(ctx0, is_hot, hot->cold_pick_sentinel); // [P, T] f32 + + ggml_tensor * hot_ids_f = ggml_add(ctx0, hot_remap, hot_offset); + ggml_tensor * cold_ids_f = ggml_add(ctx0, cold_remap, cold_offset); + ggml_tensor * hot_ids = ggml_cast(ctx0, hot_ids_f, GGML_TYPE_I32); + ggml_tensor * cold_ids = ggml_cast(ctx0, cold_ids_f, GGML_TYPE_I32); + + // For the cold-path output mask, we still need [1, P, T] f32 to + // broadcast against the [n_embd, P, T] expert outputs. + ggml_tensor * is_cold_3d = ggml_reshape_3d(ctx0, is_cold, 1, n_picks_local, n_tokens_local); + + const float swiglu_limit = hparams.swiglu_clamp_exp[il]; + + // Diagnostic mode (DS4_HOT_DISPATCH=cold): only run cold path with + // cold_ids; no hot contribution. The mask still zeros out hot + // positions so the output is partial (hot picks contribute 0) but + // we can verify the cold-with-remap path doesn't crash. + const char * mode = std::getenv("DS4_HOT_DISPATCH_MODE"); + const bool cold_only = mode && std::strcmp(mode, "cold") == 0; + const bool hot_only = mode && std::strcmp(mode, "hot") == 0; + + ggml_tensor * out_h = nullptr; + ggml_tensor * out_c = nullptr; + + // === HOT path on GPU (K real hot experts + P dummy zero-weighted experts) === + // No output mask needed: cold-pick positions hit dummy experts + // (positions K..K+P-1) which are zero-initialized, so their + // contribution is naturally 0. For hot picks, hot_ids points at + // the right real expert in [0, K). + if (!cold_only) { + ggml_tensor * gate_h = nullptr; + ggml_tensor * up_h = nullptr; + // Diagnostic: DS4_HOT_USE_FULL_WEIGHTS=1 forces hot path to use the + // CPU-resident full-N tensor instead of the GPU-resident K-subset. + // hot_ids values in [0, K) are still valid for the full tensor. + const bool use_full = std::getenv("DS4_HOT_USE_FULL_WEIGHTS") != nullptr; + ggml_tensor * w_gate = use_full ? layer.ffn_gate_exps : hot->hot_gate_exps; + ggml_tensor * w_up = use_full ? layer.ffn_up_exps : hot->hot_up_exps; + ggml_tensor * w_down = use_full ? layer.ffn_down_exps : hot->hot_down_exps; + + if (hot->hot_gate_up_exps && !use_full) { + ggml_tensor * gate_up_h = build_lora_mm_id(hot->hot_gate_up_exps, cur_experts_in, hot_ids); + cb(gate_up_h, "ffn_moe_hot_gate_up", il); + const int64_t n_ff = gate_up_h->ne[0] / 2; + gate_h = ggml_view_3d(ctx0, gate_up_h, n_ff, gate_up_h->ne[1], gate_up_h->ne[2], + gate_up_h->nb[1], gate_up_h->nb[2], 0); + up_h = ggml_view_3d(ctx0, gate_up_h, n_ff, gate_up_h->ne[1], gate_up_h->ne[2], + gate_up_h->nb[1], gate_up_h->nb[2], n_ff * gate_up_h->nb[0]); + } else { + gate_h = build_lora_mm_id(w_gate, cur_experts_in, hot_ids); + up_h = build_lora_mm_id(w_up, cur_experts_in, hot_ids); + cb(gate_h, "ffn_moe_hot_gate", il); + cb(up_h, "ffn_moe_hot_up", il); + } + + if (swiglu_limit > 1e-6f) { + gate_h = ggml_clamp(ctx0, gate_h, -INFINITY, swiglu_limit); + up_h = ggml_clamp(ctx0, up_h, -swiglu_limit, swiglu_limit); + } + ggml_tensor * act_h = ggml_swiglu_split(ctx0, gate_h, up_h); + ggml_tensor * down_h = build_lora_mm_id(w_down, act_h, hot_ids); + out_h = ggml_mul(ctx0, down_h, weights); + cb(out_h, "ffn_moe_hot_out", il); + } + + // === COLD path on CPU (full original tensor with hot picks redirected to per-pick cold sentinels) === + // Per-pick cold sentinels avoid the same-expert-multiple-times + // problem on CPU mul_mat_id. The output mask zeros out hot-pick + // positions (we still need it because the cold sentinels are + // real cold experts producing real outputs). + if (!hot_only) { + ggml_tensor * gate_c = nullptr; + ggml_tensor * up_c = nullptr; + if (layer.ffn_gate_up_exps) { + ggml_tensor * gate_up_c = build_lora_mm_id(layer.ffn_gate_up_exps, cur_experts_in, cold_ids); + cb(gate_up_c, "ffn_moe_cold_gate_up", il); + const int64_t n_ff = gate_up_c->ne[0] / 2; + gate_c = ggml_view_3d(ctx0, gate_up_c, n_ff, gate_up_c->ne[1], gate_up_c->ne[2], + gate_up_c->nb[1], gate_up_c->nb[2], 0); + up_c = ggml_view_3d(ctx0, gate_up_c, n_ff, gate_up_c->ne[1], gate_up_c->ne[2], + gate_up_c->nb[1], gate_up_c->nb[2], n_ff * gate_up_c->nb[0]); + } else { + gate_c = build_lora_mm_id(layer.ffn_gate_exps, cur_experts_in, cold_ids); + up_c = build_lora_mm_id(layer.ffn_up_exps, cur_experts_in, cold_ids); + cb(gate_c, "ffn_moe_cold_gate", il); + cb(up_c, "ffn_moe_cold_up", il); + } + + if (swiglu_limit > 1e-6f) { + gate_c = ggml_clamp(ctx0, gate_c, -INFINITY, swiglu_limit); + up_c = ggml_clamp(ctx0, up_c, -swiglu_limit, swiglu_limit); + } + ggml_tensor * act_c = ggml_swiglu_split(ctx0, gate_c, up_c); + ggml_tensor * down_c = build_lora_mm_id(layer.ffn_down_exps, act_c, cold_ids); + out_c = ggml_mul(ctx0, down_c, weights); + out_c = ggml_mul(ctx0, out_c, is_cold_3d); + cb(out_c, "ffn_moe_cold_out", il); + } + + // === Combine === + ggml_tensor * experts; + if (out_h && out_c) { + experts = ggml_add(ctx0, out_h, out_c); + } else if (out_h) { + experts = out_h; + } else { + experts = out_c; + } + cb(experts, "ffn_moe_dual_combined", il); + + ggml_tensor * experts_by_id = ggml_cont(ctx0, ggml_permute(ctx0, experts, 1, 0, 2, 3)); + ggml_tensor * out_dual = sum_rows_checked(experts_by_id, "build_expert_mix.sum"); + out_dual = reshape_3d_checked(out_dual, 1, n_embd, mix_tokens, "build_expert_mix.sum_out", il); + out_dual = reshape_2d_checked(out_dual, n_embd, mix_tokens, "build_expert_mix.out", il); + cb(out_dual, "ffn_moe_out", il); + return out_dual; + } + + // === Default single-path (unchanged) === + 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 = ggml_fp8_act_quant(ctx0, 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"); + // mul_mat always produces a contiguous F32 output, so the + // cast/cont wrappers we used to use here are no-ops that just + // add CPY nodes to the prefill graph (43 layers x 2 nodes + // per ubatch). Drop them. + + ggml_tensor * ape_row = compression_ape_rows(layer.attn_compress_ape, comp_dim, comp_ratio); + comp_score = 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); + } + + if (multiwindow_r4) { + // Batched compression for prefill: instead of looping + // n_comp_windows = work_tokens/comp_ratio times and emitting + // O(n_comp_windows) graph nodes per layer per ubatch (~2.7K + // per layer at ub=512), do all windows in one set of ops. + // + // The original loop builds two head_dim-tall slabs per + // window (kv_prev and kv_cur) where comp_kv stacks two + // logical "slots" along dim 0 (comp_dim = 2*head_dim). + // We build each slab as a 3D batched tensor and then + // concat them along dim 1 to recover the [head_dim, 2r] + // per-window matrix. + const int64_t n = work_tokens / comp_ratio; + const int64_t r = comp_ratio; + const size_t type_size = ggml_type_size(GGML_TYPE_F32); + const size_t col_stride = comp_dim * type_size; + + // prev slab: iw=0 from state, iw>=1 from comp_kv first half + ggml_tensor * state_first_kv = ggml_view_3d(ctx0, state.attn_comp_kv_state, + head_dim, r, 1, col_stride, r * col_stride, 0); + ggml_tensor * state_first_score = ggml_view_3d(ctx0, state.attn_comp_score_state, + head_dim, r, 1, col_stride, r * col_stride, 0); + ggml_tensor * comp_kv_prev_strided = (n > 1) ? ggml_view_3d(ctx0, comp_kv, + head_dim, r, n - 1, col_stride, r * col_stride, 0) : nullptr; + ggml_tensor * comp_score_prev_strided = (n > 1) ? ggml_view_3d(ctx0, comp_score, + head_dim, r, n - 1, col_stride, r * col_stride, 0) : nullptr; + ggml_tensor * prev_kv_b = comp_kv_prev_strided ? ggml_concat(ctx0, state_first_kv, comp_kv_prev_strided, 2) : state_first_kv; + ggml_tensor * prev_score_b = comp_score_prev_strided ? ggml_concat(ctx0, state_first_score, comp_score_prev_strided, 2) : state_first_score; + + // cur slab: comp_kv second half across all n windows + ggml_tensor * cur_kv_b = ggml_view_3d(ctx0, comp_kv, + head_dim, r, n, col_stride, r * col_stride, head_dim * type_size); + ggml_tensor * cur_score_b = ggml_view_3d(ctx0, comp_score, + head_dim, r, n, col_stride, r * col_stride, head_dim * type_size); + + // [head_dim, 2r, n] + ggml_tensor * batched_kv_slots = ggml_concat(ctx0, prev_kv_b, cur_kv_b, 1); + ggml_tensor * batched_score_slots = ggml_concat(ctx0, prev_score_b, cur_score_b, 1); + + // permute (1, 0, 2, 3): [head_dim, 2r, n] -> [2r, head_dim, n] + ggml_tensor * batched_kv_seq = ggml_cont(ctx0, ggml_permute(ctx0, batched_kv_slots, 1, 0, 2, 3)); + ggml_tensor * batched_score_seq = ggml_cont(ctx0, ggml_permute(ctx0, batched_score_slots, 1, 0, 2, 3)); + + ggml_tensor * batched_weights = ggml_soft_max(ctx0, batched_score_seq); + ggml_tensor * batched_weighted = ggml_mul(ctx0, batched_kv_seq, batched_weights); + ggml_tensor * batched_flat = sum_rows_checked(batched_weighted, "build_attn_v4.comp_sum_b"); + // [1, head_dim, n] -> [head_dim, n] + batched_flat = cont_if_needed(reshape_2d_checked(batched_flat, head_dim, n, "build_attn_v4.comp_flat_b", il)); + batched_flat = build_norm(batched_flat, layer.attn_compress_norm, nullptr, LLM_NORM_RMS, il); + + // split nope/pe along dim 0 + ggml_tensor * batched_states = reshape_3d_checked(batched_flat, head_dim, 1, n, "build_attn_v4.comp_states_b", il); + ggml_tensor * batched_nope = ggml_view_3d(ctx0, batched_states, nope_dim, 1, n, + batched_states->nb[1], batched_states->nb[2], 0); + ggml_tensor * batched_pe = ggml_view_3d(ctx0, batched_states, rope_dim, 1, n, + batched_states->nb[1], batched_states->nb[2], nope_dim * batched_states->nb[0]); + batched_nope = ggml_fp8_act_quant(ctx0, cont_if_needed(batched_nope)); + + // positions / cache indices: stride-r view picks up + // the (r-1)-th token of each window. + const size_t i32 = ggml_type_size(GGML_TYPE_I32); + ggml_tensor * batched_pos = ggml_view_2d(ctx0, comp_pos, 1, n, r * i32, (r - 1) * i32); + batched_pos = ggml_reshape_1d(ctx0, ggml_cont(ctx0, batched_pos), n); + ggml_tensor * batched_cache_idx = ggml_view_2d(ctx0, comp_cache_idx, 1, n, r * i32, (r - 1) * i32); + batched_cache_idx = ggml_reshape_1d(ctx0, ggml_cont(ctx0, batched_cache_idx), n); + + batched_pe = ggml_rope_ext(ctx0, batched_pe, batched_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); + batched_states = ggml_concat(ctx0, batched_nope, batched_pe, 0); + batched_flat = cont_if_needed(reshape_2d_checked(batched_states, head_dim, n, "build_attn_v4.comp_flat_b2", il)); + cb(batched_flat, "attn_comp_cache_b", il); + + updated_cache = ggml_set_rows(ctx0, updated_cache, batched_flat, batched_cache_idx); + + // overlap state seeding: final window is comp_kv[:, (n-1)*r:n*r] + ggml_tensor * final_carry_kv = matrix_block(comp_kv, 0, (n - 1) * r, comp_dim, r); + ggml_tensor * final_carry_score = matrix_block(comp_score, 0, (n - 1) * r, comp_dim, r); + 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); + } else { + ggml_tensor * comp_kv_slots = nullptr; + ggml_tensor * comp_score_slots = nullptr; + ggml_tensor * final_carry_kv = nullptr; + ggml_tensor * final_carry_score = nullptr; + + 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 = ggml_fp8_act_quant(ctx0, cont_if_needed(comp_nope)); + + const int64_t token_in_ubatch = 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"); + + ggml_tensor * indexer_ape_row = compression_ape_rows(layer.indexer_compress_ape, indexer_comp_dim, comp_ratio); + indexer_comp_score = 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; + + if (multiwindow_r4) { + // See attn-side compression for the strided-view explanation. + const int64_t n = work_tokens / comp_ratio; + const int64_t r = comp_ratio; + const size_t type_size = ggml_type_size(GGML_TYPE_F32); + const size_t col_stride = indexer_comp_dim * type_size; + + ggml_tensor * state_first_kv = ggml_view_3d(ctx0, state.indexer_comp_kv_state, + indexer_head_dim, r, 1, col_stride, r * col_stride, 0); + ggml_tensor * state_first_score = ggml_view_3d(ctx0, state.indexer_comp_score_state, + indexer_head_dim, r, 1, col_stride, r * col_stride, 0); + ggml_tensor * comp_kv_prev_strided = (n > 1) ? ggml_view_3d(ctx0, indexer_comp_kv, + indexer_head_dim, r, n - 1, col_stride, r * col_stride, 0) : nullptr; + ggml_tensor * comp_score_prev_strided = (n > 1) ? ggml_view_3d(ctx0, indexer_comp_score, + indexer_head_dim, r, n - 1, col_stride, r * col_stride, 0) : nullptr; + ggml_tensor * prev_kv_b = comp_kv_prev_strided ? ggml_concat(ctx0, state_first_kv, comp_kv_prev_strided, 2) : state_first_kv; + ggml_tensor * prev_score_b = comp_score_prev_strided ? ggml_concat(ctx0, state_first_score, comp_score_prev_strided, 2) : state_first_score; + + ggml_tensor * cur_kv_b = ggml_view_3d(ctx0, indexer_comp_kv, + indexer_head_dim, r, n, col_stride, r * col_stride, indexer_head_dim * type_size); + ggml_tensor * cur_score_b = ggml_view_3d(ctx0, indexer_comp_score, + indexer_head_dim, r, n, col_stride, r * col_stride, indexer_head_dim * type_size); + + ggml_tensor * batched_kv_slots = ggml_concat(ctx0, prev_kv_b, cur_kv_b, 1); + ggml_tensor * batched_score_slots = ggml_concat(ctx0, prev_score_b, cur_score_b, 1); + + ggml_tensor * batched_kv_seq = ggml_cont(ctx0, ggml_permute(ctx0, batched_kv_slots, 1, 0, 2, 3)); + ggml_tensor * batched_score_seq = ggml_cont(ctx0, ggml_permute(ctx0, batched_score_slots, 1, 0, 2, 3)); + + ggml_tensor * batched_weights = ggml_soft_max(ctx0, batched_score_seq); + ggml_tensor * batched_weighted = ggml_mul(ctx0, batched_kv_seq, batched_weights); + ggml_tensor * batched_flat = sum_rows_checked(batched_weighted, "build_attn_v4.indexer_comp_sum_b"); + batched_flat = cont_if_needed(reshape_2d_checked(batched_flat, indexer_head_dim, n, "build_attn_v4.indexer_comp_flat_b", il)); + batched_flat = build_norm(batched_flat, layer.indexer_compress_norm, nullptr, LLM_NORM_RMS, il); + + ggml_tensor * batched_states = reshape_3d_checked(batched_flat, indexer_head_dim, 1, n, "build_attn_v4.indexer_comp_states_b", il); + ggml_tensor * batched_nope = ggml_view_3d(ctx0, batched_states, indexer_nope_dim, 1, n, + batched_states->nb[1], batched_states->nb[2], 0); + ggml_tensor * batched_pe = ggml_view_3d(ctx0, batched_states, rope_dim, 1, n, + batched_states->nb[1], batched_states->nb[2], indexer_nope_dim * batched_states->nb[0]); + + const size_t i32 = ggml_type_size(GGML_TYPE_I32); + ggml_tensor * batched_pos = ggml_view_2d(ctx0, indexer_comp_pos, 1, n, r * i32, (r - 1) * i32); + batched_pos = ggml_reshape_1d(ctx0, ggml_cont(ctx0, batched_pos), n); + ggml_tensor * batched_cache_idx = ggml_view_2d(ctx0, indexer_cache_idx, 1, n, r * i32, (r - 1) * i32); + batched_cache_idx = ggml_reshape_1d(ctx0, ggml_cont(ctx0, batched_cache_idx), n); + + batched_pe = ggml_rope_ext(ctx0, batched_pe, batched_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); + batched_states = ggml_concat(ctx0, batched_nope, batched_pe, 0); + batched_flat = cont_if_needed(reshape_2d_checked(batched_states, indexer_head_dim, n, "build_attn_v4.indexer_comp_flat_b2", il)); + batched_flat = ggml_mul_mat(ctx0, deepseek4_inputs->indexer_hadamard, batched_flat); + batched_flat = ggml_fp4_act_quant(ctx0, cont_if_needed(batched_flat)); + cb(batched_flat, "indexer_comp_cache_b", il); + + updated_indexer_kv = ggml_set_rows(ctx0, updated_indexer_kv, batched_flat, batched_cache_idx); + + ggml_tensor * final_carry_kv = matrix_block(indexer_comp_kv, 0, (n - 1) * r, indexer_comp_dim, r); + ggml_tensor * final_carry_score = matrix_block(indexer_comp_score, 0, (n - 1) * r, indexer_comp_dim, r); + 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); + } else { + ggml_tensor * indexer_comp_kv_slots = nullptr; + ggml_tensor * indexer_comp_score_slots = nullptr; + ggml_tensor * final_carry_kv = nullptr; + ggml_tensor * final_carry_score = nullptr; + + 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 = 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 = ggml_fp4_act_quant(ctx0, 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); + + if (work_tokens > 1 && !deepseek4_indexer_per_query()) { + // Batched prefill ubatch-shared top-k: collapse the + // work_tokens axis BEFORE the score mul_mat by summing + // indexer_q across queries. This avoids materializing + // a per-query [n_comp, n_head, work_tokens] score + // tensor (which OOMs at ub=512 + long context where + // n_comp * 64 * 512 * 4 bytes per layer x 21 r=4 + // layers blows past 10 GB on the GPUs). The scoring + // becomes mul_mat(kv [128, n_comp], sum_q [128, 64]) + // -> [n_comp, 64], which is the same shape as the + // existing decode (work_tokens=1) path. The + // approximation is small in practice for short + // prompts but at very long context (65K+ retrieval- + // style prompts) it can cause the model to attend + // to wrong KV slots because relu(sum_q kv*q) != + // sum_q relu(kv*q). Set + // LLAMA_DEEPSEEK4_INDEXER_PER_QUERY=1 to keep the + // exact per-query path (more accurate, more VRAM, + // requires a smaller -ub on tight VRAM hosts). + indexer_q = ggml_cont(ctx0, ggml_permute(ctx0, indexer_q, 1, 2, 0, 3)); + indexer_q = sum_rows_checked(indexer_q, "build_attn_v4.indexer_q_sum_b"); + indexer_q = cont_if_needed(reshape_2d_checked(indexer_q, indexer_head_dim, hparams.indexer_n_head, "build_attn_v4.indexer_q_collapsed", il)); + } else if (work_tokens > 1) { + // Per-query path: keep work_tokens as a separate + // dim through the score mul_mat. Used when + // LLAMA_DEEPSEEK4_INDEXER_PER_QUERY=1 is set. + indexer_q = cont_if_needed(reshape_3d_checked(indexer_q, indexer_head_dim, hparams.indexer_n_head, work_tokens, "build_attn_v4.indexer_q_b", il)); + } else { + 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 = ggml_fp4_act_quant(ctx0, 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); + // After the work_tokens collapse this is [n_comp, n_head] + // (same as decode). In per-query mode it is + // [n_comp, n_head, work_tokens]. + 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); + if (work_tokens > 1 && !deepseek4_indexer_per_query()) { + // index_weights starts as [indexer_n_head, work_tokens]; + // collapse the work_tokens axis the same way as the + // queries so the per-head weighting stays consistent + // with the collapsed scores. + index_weights = ggml_cont(ctx0, ggml_transpose(ctx0, index_weights)); + index_weights = sum_rows_checked(index_weights, "build_attn_v4.index_weights_sum_b"); + index_weights = reshape_2d_checked(index_weights, 1, hparams.indexer_n_head, "build_attn_v4.index_weights_collapsed", il); + } else if (work_tokens > 1) { + // Per-query: keep weights aligned with scores [.., n_head, work_tokens] + index_weights = reshape_3d_checked(index_weights, 1, hparams.indexer_n_head, work_tokens, "build_attn_v4.index_weights_b", il); + } else { + 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); + if (work_tokens > 1 && deepseek4_indexer_per_query()) { + // Aggregate per-query scores into a single ubatch-wide + // top-k. Sum across both indexer_n_head and the + // work_tokens axis so every query in the ubatch + // shares one selected prefix. + // Shape evolves [n_comp, n_head, work_tokens] -> + // [n_comp, n_head*work_tokens] -> + // [n_head*work_tokens, n_comp] -> + // [1, n_comp] -> [n_comp, 1] + index_scores = cont_if_needed(reshape_2d_checked(index_scores, n_comp, hparams.indexer_n_head * work_tokens, "build_attn_v4.index_scores_flat", il)); + index_scores = ggml_cont(ctx0, ggml_transpose(ctx0, index_scores)); + index_scores = sum_rows_checked(index_scores, "build_attn_v4.index_scores_sum"); + index_scores = reshape_2d_checked(index_scores, n_comp, 1, "build_attn_v4.index_scores_perq", il); + } else { + 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; + // FA kernels for K[0]=512 require K->ne[1] (= n_kv_total after permute) + // to be a multiple of FATTN_KQ_STRIDE, and require a non-null mask. + // Without this, the auto-FA reservation graph (which runs with + // work_tokens=1 and arbitrary n_kv_total) reports unsupported on CUDA, + // the scheduler places the FA tensor on CPU, and auto-FA disables + // FA globally for the entire context. Pad both kv_states and the + // mask to the next multiple of 256 so FA stays on GPU; the padded + // K/V slots are masked out with -INFINITY and contribute nothing. + constexpr int64_t kq_pad = 256; + const int64_t n_kv_total_padded = ((n_kv_total + kq_pad - 1) / kq_pad) * kq_pad; + const int64_t kv_pad = n_kv_total_padded - n_kv_total; + if (kv_pad > 0) { + kv_prefix = ggml_pad(ctx0, kv_prefix, 0, (int) kv_pad, 0, 0); + } + ggml_tensor * kv_states = reshape_3d_checked(kv_prefix, head_dim, 1, n_kv_total_padded, "build_attn_v4.kv_states", il); + + ggml_tensor * kq_mask = nullptr; + { + const auto key = std::make_pair(n_kv_total_padded, (int64_t) work_tokens); + auto it = deepseek4_inputs->kq_mask_by_shape.find(key); + if (it != deepseek4_inputs->kq_mask_by_shape.end()) { + kq_mask = it->second; + } else { + kq_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv_total_padded, work_tokens); + ggml_set_input(kq_mask); + ggml_format_name(kq_mask, "deepseek4_kq_mask_%lldx%lld", (long long) n_kv_total_padded, (long long) work_tokens); + deepseek4_inputs->kq_masks.push_back(kq_mask); + deepseek4_inputs->kq_mask_n_kv_total.push_back(n_kv_total); + deepseek4_inputs->kq_mask_by_shape[key] = kq_mask; + } + } + + // Flash attention requires the mask in F16; the dense path takes F32. + ggml_tensor * kq_mask_arg = cparams.flash_attn ? ggml_cast(ctx0, kq_mask, GGML_TYPE_F16) : kq_mask; + + ggml_tensor * out = build_attn_mha( + q_states, + kv_states, + kv_states, + nullptr, + kq_mask_arg, + 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 94991c55fe8..d8871d684ec 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/CMakeLists.txt b/tests/CMakeLists.txt index edb585b9f65..193e6638077 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -157,6 +157,9 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) llama_build_and_test(test-chat.cpp WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}) target_include_directories(test-chat PRIVATE ${PROJECT_SOURCE_DIR}/tools/server) target_link_libraries(test-chat PRIVATE server-context) + llama_build_and_test(test-server-prompt-cache.cpp) + target_include_directories(test-server-prompt-cache PRIVATE ${PROJECT_SOURCE_DIR}/tools/server) + target_link_libraries(test-server-prompt-cache PRIVATE server-context) # TODO: disabled on loongarch64 because the ggml-ci node lacks Python 3.8 if (NOT ${CMAKE_SYSTEM_PROCESSOR} MATCHES "loongarch64") llama_build_and_test(test-json-schema-to-grammar.cpp WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}) @@ -196,6 +199,19 @@ endif() llama_build_and_test(test-chat-peg-parser.cpp peg-parser/simple-tokenize.cpp) llama_build_and_test(test-jinja.cpp) llama_test(test-jinja NAME test-jinja-py ARGS -py LABEL python) +find_package(Python3 COMPONENTS Interpreter QUIET) +if (Python3_Interpreter_FOUND) + add_test( + NAME test-deepseek4-native-packers + WORKING_DIRECTORY ${PROJECT_SOURCE_DIR} + COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test-deepseek4-native-packers.py ${PROJECT_SOURCE_DIR}) + set_tests_properties(test-deepseek4-native-packers PROPERTIES LABELS python SKIP_RETURN_CODE 77) + add_test( + NAME test-moe-copy-lru-sim + WORKING_DIRECTORY ${PROJECT_SOURCE_DIR} + COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test-moe-copy-lru-sim.py ${PROJECT_SOURCE_DIR}) + set_tests_properties(test-moe-copy-lru-sim PROPERTIES LABELS python) +endif() llama_build_and_test(test-chat-auto-parser.cpp WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}) llama_build_and_test(test-chat-template.cpp) llama_build_and_test(test-json-partial.cpp) @@ -231,7 +247,9 @@ add_test(NAME test-download-model COMMAND ${CMAKE_COMMAND} set_tests_properties(test-download-model PROPERTIES FIXTURES_SETUP test-download-model) llama_build_and_test(test-thread-safety.cpp ARGS -m "${MODEL_DEST}" -ngl 99 -p "The meaning of life is" -n 128 -c 256 -ub 32 -np 4 -t 2) -set_tests_properties(test-thread-safety PROPERTIES FIXTURES_REQUIRED test-download-model) +set_tests_properties(test-thread-safety PROPERTIES + FIXTURES_REQUIRED test-download-model + ENVIRONMENT "CUDA_VISIBLE_DEVICES=0") llama_build_and_test(test-arg-parser.cpp) @@ -240,7 +258,16 @@ if (NOT LLAMA_SANITIZE_ADDRESS AND NOT GGML_SCHED_NO_REALLOC) llama_build_and_test(test-opt.cpp) endif() llama_build_and_test(test-gguf.cpp) -llama_build_and_test(test-backend-ops.cpp) + +set(LLAMA_BACKEND_OPS_SMOKE_FILTER + "HC_WEIGHTED_SUM(n_embd=64,hc_mult=4,slice_x=0,slice_w=0),\ +MUL_MAT(type_a=iq4_xs,type_b=f32,m=16,n=1,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],k_v=0,o=1),\ +MUL_MAT(type_a=f8_e4m3_b128,type_b=f32,m=1,n=64,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],k_v=0,o=1),\ +MUL_MAT_ID(type_a=iq4_xs,type_b=f32,n_mats=4,n_used=2,b=0,m=512,n=1,k=256),\ +MUL_MAT_ID(type_a=q3_K,type_b=f32,n_mats=4,n_used=2,b=0,m=512,n=1,k=256)") +llama_build(test-backend-ops.cpp get-model.cpp) +llama_test(test-backend-ops ARGS test -o "${LLAMA_BACKEND_OPS_SMOKE_FILTER}") +llama_test(test-backend-ops NAME test-backend-ops-full LABEL backend-full ARGS test) llama_build_and_test(test-model-load-cancel.cpp LABEL "model") llama_build_and_test(test-autorelease.cpp LABEL "model") diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 71601131671..b2a6e5fcae0 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -1496,6 +1496,7 @@ struct test_case { // build graph ggml_cgraph * gf = ggml_new_graph_custom(ctx.get(), graph_nodes, false); ggml_build_forward_expand(gf, out); + const int base_graph_nodes = ggml_graph_n_nodes(gf); // warmup run ggml_status status = ggml_backend_graph_compute(backend, gf); @@ -1504,47 +1505,80 @@ struct test_case { return false; } + auto tensor_op_size = [](ggml_tensor * t) { + size_t size = ggml_nbytes(t); + // add source tensors + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (t->src[i] != NULL) { + size += ggml_nbytes(t->src[i]); + } + } + return size; + }; + + auto graph_op_size = [&](int n_nodes) { + size_t size = 0; + for (int i = 0; i < n_nodes; ++i) { + ggml_tensor * node = ggml_graph_node(gf, i); + if (!ggml_is_view_op(node->op)) { + size += tensor_op_size(node); + } + } + return size; + }; + // determine number of runs int n_runs; bool is_cpu = ggml_backend_dev_type(ggml_backend_get_device(backend)) == GGML_BACKEND_DEVICE_TYPE_CPU; + const bool whole_graph = run_whole_graph(); + const int max_runs = whole_graph ? + std::max(1, ggml_graph_size(gf) / base_graph_nodes) : + std::max(1, ggml_graph_size(gf) - base_graph_nodes + 1); + const size_t size_per_run = whole_graph ? graph_op_size(base_graph_nodes) : op_size(out); if (op_flops(out) > 0) { // based on flops const uint64_t GFLOP = 1000 * 1000 * 1000; const uint64_t target_flops_cpu = 8ULL * GFLOP; const uint64_t target_flops_gpu = 100ULL * GFLOP; uint64_t target_flops = is_cpu ? target_flops_cpu : target_flops_gpu; - n_runs = (int)std::min(ggml_graph_size(gf) - ggml_graph_n_nodes(gf), target_flops / op_flops(out)) + 1; + n_runs = (int) std::min(max_runs, target_flops / op_flops(out) + 1); } else { // based on memory size const size_t GB = 1ULL << 30; const size_t target_size_cpu = 8 * GB; const size_t target_size_gpu = 32 * GB; size_t target_size = is_cpu ? target_size_cpu : target_size_gpu; - n_runs = (int)std::min(ggml_graph_size(gf) - ggml_graph_n_nodes(gf), target_size / op_size(out)) + 1; + n_runs = (int) std::min(max_runs, target_size / size_per_run + 1); } - // duplicate the op - for (int i = 1; i < n_runs; i++) { - ggml_graph_add_node(gf, out); + if (whole_graph) { + std::vector nodes; + nodes.reserve(base_graph_nodes); + for (int i = 0; i < base_graph_nodes; ++i) { + nodes.push_back(ggml_graph_node(gf, i)); + } + + for (int i = 1; i < n_runs; i++) { + for (ggml_tensor * node : nodes) { + ggml_graph_add_node(gf, node); + } + } + } else { + // duplicate the op + for (int i = 1; i < n_runs; i++) { + ggml_graph_add_node(gf, out); + } } // calculate memory - size_t mem = n_runs * op_size(out); - auto tensor_op_size = [](ggml_tensor * t) { - size_t size = ggml_nbytes(t); - // add source tensors - for (int i = 0; i < GGML_MAX_SRC; i++) { - if (t->src[i] != NULL) { - size += ggml_nbytes(t->src[i]); + size_t mem = n_runs * size_per_run; + if (!whole_graph) { + for (int i = 0; i < base_graph_nodes; ++i) { + if (ggml_is_view_op(ggml_graph_node(gf, i)->op) || ggml_graph_node(gf, i) == out) { + continue; } + mem += tensor_op_size(ggml_graph_node(gf, i)); } - return size; - }; - for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) { - if (ggml_is_view_op(ggml_graph_node(gf, i)->op) || ggml_graph_node(gf, i) == out) { - continue; - } - mem += tensor_op_size(ggml_graph_node(gf, i)); } // run @@ -1570,7 +1604,7 @@ struct test_case { double calculated_flops = (op_flops(out) > 0) ? (op_flops(out) * total_runs) / (total_time_us / 1e6) : 0.0; double calculated_bandwidth = (op_flops(out) == 0) ? total_mem / (total_time_us / 1e6) / 1024.0 / 1024.0 / 1024.0 : 0.0; - size_t calculated_memory_kb = op_size(out) / 1024; + size_t calculated_memory_kb = size_per_run / 1024; test_result result(ggml_backend_name(backend), current_op_name, vars(), "perf", true, true, "", avg_time_us, calculated_flops, calculated_bandwidth, calculated_memory_kb, total_runs); @@ -1932,6 +1966,11 @@ struct test_unary : public test_case { return VARS_TO_STR3(type, ne_a, v); } + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return ggml_unary_op_name(op); + } + test_unary(ggml_unary_op op, ggml_type type = GGML_TYPE_F32, std::array ne_a = {128, 2, 2, 2}, @@ -5534,6 +5573,13 @@ struct test_mul_mat_vec_fusion : public test_case { bool run_whole_graph() override { return true; } + uint64_t op_flops(ggml_tensor * t) override { + GGML_UNUSED(t); + const int64_t n_tokens = use_id ? n_used*m : m*batch_dims[0]*batch_dims[1]; + const int64_t n_matmuls = with_gate ? 2 : 1; + return 2ULL*n_matmuls*n_tokens*n*k; + } + ggml_tensor * build_gate(ggml_context * ctx, ggml_tensor * ffn_gate, ggml_tensor * ffn_up) { ggml_tensor * out = nullptr; if (with_gate) { @@ -5715,6 +5761,59 @@ struct test_sum_rows : public test_case { } }; +// GGML_OP_HC_WEIGHTED_SUM +struct test_hc_weighted_sum : public test_case { + const int64_t n_embd; + const int64_t hc_mult; + const bool slice_x; + const bool slice_w; + + std::string vars() override { + return VARS_TO_STR4(n_embd, hc_mult, slice_x, slice_w); + } + + test_hc_weighted_sum(int64_t n_embd = 64, int64_t hc_mult = 4, bool slice_x = false, bool slice_w = false) + : n_embd(n_embd), hc_mult(hc_mult), slice_x(slice_x), slice_w(slice_w) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * x = nullptr; + if (slice_x) { + int64_t ne_x_base[2] = { n_embd + 1, hc_mult }; + ggml_tensor * x_base = ggml_new_tensor(ctx, GGML_TYPE_F32, 2, ne_x_base); + ggml_set_param(x_base); + ggml_set_name(x_base, "x_base"); + x = ggml_view_2d(ctx, x_base, n_embd, hc_mult, x_base->nb[1], x_base->nb[0]); + } else { + int64_t ne_x[2] = { n_embd, hc_mult }; + x = ggml_new_tensor(ctx, GGML_TYPE_F32, 2, ne_x); + ggml_set_param(x); + } + ggml_set_name(x, "x"); + + ggml_tensor * w = nullptr; + if (slice_w) { + ggml_tensor * w_base = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hc_mult + 1); + ggml_set_param(w_base); + ggml_set_name(w_base, "w_base"); + w = ggml_view_1d(ctx, w_base, hc_mult, w_base->nb[0]); + } else { + w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hc_mult); + ggml_set_param(w); + } + ggml_set_name(w, "w"); + + ggml_tensor * out = ggml_hc_weighted_sum(ctx, x, w); + ggml_set_name(out, "out"); + + return out; + } + + uint64_t op_flops(ggml_tensor * t) override { + GGML_UNUSED(t); + return 2ull*n_embd*hc_mult; + } +}; + // GGML_OP_MEAN struct test_mean : public test_case { const ggml_type type; @@ -7339,12 +7438,27 @@ static std::vector> make_test_cases_eval() { if (op == GGML_UNARY_OP_XIELU) { continue; // need extra params, separate test } + if (op == GGML_UNARY_OP_FP4_ACT_QUANT || op == GGML_UNARY_OP_FP8_ACT_QUANT) { + continue; // require block-aligned row lengths; separate tests below + } + if (op == GGML_UNARY_OP_SINKHORN_4X4) { + continue; // requires a 4x4 F32 matrix; separate test below + } test_cases.emplace_back(new test_unary((ggml_unary_op) op, type, { 128, 2, 2, 2 }, v)); test_cases.emplace_back(new test_unary((ggml_unary_op) op, type, { 5, 7, 11, 13 }, v)); } } } + test_cases.emplace_back(new test_unary(GGML_UNARY_OP_SINKHORN_4X4, GGML_TYPE_F32, { 4, 4, 1, 1 }, 0)); + + for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) { + test_cases.emplace_back(new test_unary(GGML_UNARY_OP_FP4_ACT_QUANT, type, { 32, 5, 2, 1 }, 0)); + test_cases.emplace_back(new test_unary(GGML_UNARY_OP_FP4_ACT_QUANT, type, { 96, 3, 2, 1 }, 0)); + test_cases.emplace_back(new test_unary(GGML_UNARY_OP_FP8_ACT_QUANT, type, { 64, 5, 2, 1 }, 0)); + test_cases.emplace_back(new test_unary(GGML_UNARY_OP_FP8_ACT_QUANT, type, { 128, 3, 2, 1 }, 0)); + } + // fused relu + sqr (squared ReLU) for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) { test_cases.emplace_back(new test_relu_sqr(type, { 128, 2, 2, 2 })); @@ -8109,6 +8223,8 @@ static std::vector> make_test_cases_eval() { } test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 6, 4096, 5120, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F8_E4M3_B128, GGML_TYPE_F32, 1, 64, 256, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F8_E4M3_B128, GGML_TYPE_F32, 16, 1, 256, {1, 1}, {1, 1})); #if 0 // test the mat-mat path for Metal @@ -8508,6 +8624,12 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_sum_rows(GGML_TYPE_F32, { 33, 1, 1, 1 })); test_cases.emplace_back(new test_sum_rows(GGML_TYPE_F32, { 33, 1024, 1, 1 })); test_cases.emplace_back(new test_sum_rows(GGML_TYPE_F32, { 33, 256, 1, 1 })); + test_cases.emplace_back(new test_hc_weighted_sum(64, 4, false, false)); + test_cases.emplace_back(new test_hc_weighted_sum(4096, 4, false, false)); + test_cases.emplace_back(new test_hc_weighted_sum(127, 4, true, true)); + test_cases.emplace_back(new test_hc_weighted_sum(19, 1, false, false)); + test_cases.emplace_back(new test_hc_weighted_sum(65, 2, false, false)); + test_cases.emplace_back(new test_hc_weighted_sum(31, 7, true, true)); test_cases.emplace_back(new test_group_norm(GGML_TYPE_F32, {64, 64, 320, 1})); test_cases.emplace_back(new test_group_norm(GGML_TYPE_F32, {9, 9, 1280, 1})); test_cases.emplace_back(new test_group_norm_mul_add(GGML_TYPE_F32, {64, 64, 320, 1})); @@ -8682,6 +8804,8 @@ static std::vector> make_test_cases_eval() { } } } + test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_F8_E4M3_B128, GGML_GLU_OP_SWIGLU, 1, 32, 256, + false, 1, 1, false, false, true, {1, 1})); for (auto gate : {GATING_FUNC_SOFTMAX, GATING_FUNC_SIGMOID, GATING_FUNC_SOFTMAX_WEIGHT}) { for (bool with_norm : {false, true}) { @@ -8824,6 +8948,24 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 16416, 1, 128, {8, 1}, {4, 1}, {0, 2, 1, 3})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 128, 1, 16416, {8, 1}, {4, 1}, {0, 1, 2, 3}, 2*16416)); + // DeepSeek4 native FP8 projection shapes for focused CUDA MMVQ tuning. + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F8_E4M3_B128, GGML_TYPE_F32, 4096, 1, 2048, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F8_E4M3_B128, GGML_TYPE_F32, 4096, 1, 8192, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F8_E4M3_B128, GGML_TYPE_F32, 8192, 1, 4096, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F8_E4M3_B128, GGML_TYPE_F32, 1024, 1, 32768, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F8_E4M3_B128, GGML_TYPE_F32, 4096, 1, 512, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F8_E4M3_B128, GGML_TYPE_F32, 4096, 1, 1024, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F8_E4M3_B128, GGML_TYPE_F32, 4096, 8, 2048, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F8_E4M3_B128, GGML_TYPE_F32, 2048, 8, 4096, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F8_E4M3_B128, GGML_TYPE_F32, 4096, 16, 2048, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F8_E4M3_B128, GGML_TYPE_F32, 2048, 16, 4096, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F8_E4M3_B128, GGML_TYPE_F32, 4096, 32, 2048, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F8_E4M3_B128, GGML_TYPE_F32, 2048, 32, 4096, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_F8_E4M3_B128, GGML_GLU_OP_SWIGLU, 1, 4096, 2048, false, 1, 1, false, false, true, {1, 1})); + test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_F8_E4M3_B128, GGML_GLU_OP_SWIGLU, 1, 2048, 4096, false, 1, 1, false, false, true, {1, 1})); + test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_F8_E4M3_B128, GGML_GLU_OP_SWIGLU, 1, 4096, 512, false, 1, 1, false, false, true, {1, 1})); + test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_F8_E4M3_B128, GGML_GLU_OP_SWIGLU, 1, 4096, 8192, false, 1, 1, false, false, true, {1, 1})); + test_cases.emplace_back(new test_solve_tri(GGML_TYPE_F32, { 64, 64, 4, 4 }, { 32, 64, 4, 4 })); test_cases.emplace_back(new test_solve_tri(GGML_TYPE_F32, { 128, 128, 4, 2 }, { 32, 128, 4, 2 })); // qwen3next with CHUNK_SIZE 64 diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp index e6a5236645e..f3b0ce2135d 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp @@ -2758,6 +2758,28 @@ static void test_template_output_peg_parsers(bool detailed_debug) { { auto tst = peg_tester("models/templates/deepseek-ai-DeepSeek-V3.2.jinja", detailed_debug); + { + auto tmpls = read_templates("models/templates/deepseek-ai-DeepSeek-V3.2.jinja"); + common_chat_templates_inputs inputs; + inputs.messages = { message_user }; + inputs.add_generation_prompt = true; + inputs.reasoning_format = COMMON_REASONING_FORMAT_DEEPSEEK; + + inputs.enable_thinking = true; + auto thinking_params = common_chat_templates_apply(tmpls.get(), inputs); + assert_equals(true, thinking_params.supports_thinking); + if (!string_ends_with(thinking_params.prompt, "")) { + throw std::runtime_error("DeepSeek V3.2 thinking prompt must end with , got: " + thinking_params.prompt); + } + + inputs.enable_thinking = false; + auto no_thinking_params = common_chat_templates_apply(tmpls.get(), inputs); + assert_equals(true, no_thinking_params.supports_thinking); + if (!string_ends_with(no_thinking_params.prompt, "")) { + throw std::runtime_error("DeepSeek V3.2 non-thinking prompt must end with , got: " + no_thinking_params.prompt); + } + } + // Pure content (non-thinking mode) tst.test("Hello, world!\nWhat's up?") .enable_thinking(false) diff --git a/tests/test-deepseek4-native-packers.py b/tests/test-deepseek4-native-packers.py new file mode 100644 index 00000000000..6b6d464932b --- /dev/null +++ b/tests/test-deepseek4-native-packers.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 + +import importlib.util +import sys +from pathlib import Path +from typing import Tuple + + +def skip(message: str) -> None: + print(f"SKIP: {message}", file=sys.stderr) + sys.exit(77) + + +try: + import torch +except ModuleNotFoundError as exc: + skip(f"missing dependency: {exc.name}") + + +if not hasattr(torch, "float8_e4m3fn"): + skip("torch does not support float8_e4m3fn") + + +def load_converter(repo_root: Path): + sys.path.insert(0, str(repo_root)) + spec = importlib.util.spec_from_file_location("convert_hf_to_gguf", repo_root / "convert_hf_to_gguf.py") + if spec is None or spec.loader is None: + raise RuntimeError("failed to create convert_hf_to_gguf module spec") + + module = importlib.util.module_from_spec(spec) + try: + spec.loader.exec_module(module) + except ModuleNotFoundError as exc: + skip(f"missing dependency: {exc.name}") + return module + + +def make_u8(shape: Tuple[int, ...]) -> torch.Tensor: + n = 1 + for dim in shape: + n *= dim + return (torch.arange(n, dtype=torch.int32) % 256).to(torch.uint8).reshape(shape) + + +def assert_rejects_float_scale(fn, weight: torch.Tensor, scale: torch.Tensor, name: str) -> None: + try: + fn(weight, scale.float(), name) + except ValueError as exc: + assert "scale dtype" in str(exc), str(exc) + else: + raise AssertionError(f"{name} accepted a multi-byte float scale") + + +def reference_pack_fp8(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + rows, cols = weight.shape + col_blocks = cols // 128 + + weight_u8 = weight.view(torch.uint8) + scale_u8 = scale.view(torch.uint8) + out = torch.empty((rows, col_blocks, 129), dtype=torch.uint8) + out[:, :, 0].copy_(scale_u8.repeat_interleave(128, dim=0)) + out[:, :, 1:].copy_(weight_u8.reshape(rows, col_blocks, 128)) + return out.reshape(rows, col_blocks * 129) + + +def test_pack_fp8(pack_fp8) -> None: + rows, cols = 256, 384 + weight = make_u8((rows, cols)).view(torch.float8_e4m3fn) + scale = make_u8((rows // 128, cols // 128)) + + actual = pack_fp8(weight, scale, "fp8.weight") + expected = reference_pack_fp8(weight, scale) + assert torch.equal(actual, expected) + assert actual.shape == (rows, (cols // 128) * 129) + + if hasattr(torch, "float8_e8m0fnu"): + scale_e8 = scale.view(torch.float8_e8m0fnu) + assert torch.equal(pack_fp8(weight, scale_e8, "fp8.e8.weight"), expected) + + assert_rejects_float_scale(pack_fp8, weight, scale, "fp8.float-scale.weight") + + +def reference_pack_mxfp4(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + rows, packed_cols = weight.shape + groups = packed_cols // 16 + + hf = weight.view(torch.uint8).reshape(rows, groups, 16) + vals = torch.empty((rows, groups, 32), dtype=torch.uint8) + vals[:, :, 0::2].copy_(hf & 0x0F) + vals[:, :, 1::2].copy_(hf >> 4) + + out = torch.empty((rows, groups, 17), dtype=torch.uint8) + out[:, :, 0].copy_(scale.view(torch.uint8)[:, :groups]) + out[:, :, 1:].copy_(vals[:, :, :16] | (vals[:, :, 16:] << 4)) + return out.reshape(rows, groups * 17) + + +def test_pack_mxfp4(pack_mxfp4) -> None: + rows, packed_cols = 5, 48 + weight = make_u8((rows, packed_cols)).view(torch.int8) + scale = make_u8((rows, packed_cols // 16 + 2)) + + actual = pack_mxfp4(weight, scale, "experts.weight") + expected = reference_pack_mxfp4(weight, scale) + assert torch.equal(actual, expected) + assert actual.shape == (rows, (packed_cols // 16) * 17) + + if hasattr(torch, "float8_e8m0fnu"): + scale_e8 = scale.view(torch.float8_e8m0fnu) + assert torch.equal(pack_mxfp4(weight, scale_e8, "experts.e8.weight"), expected) + + assert_rejects_float_scale(pack_mxfp4, weight, scale, "experts.float-scale.weight") + + +def main() -> None: + repo_root = Path(sys.argv[1]) if len(sys.argv) > 1 else Path(__file__).resolve().parents[1] + converter = load_converter(repo_root) + + test_pack_fp8(converter.DeepseekV4Model._pack_fp8_e4m3_b128) + test_pack_mxfp4(converter.DeepseekV4Model._pack_mxfp4) + + +if __name__ == "__main__": + main() diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 16af11a2862..0615c396bf9 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 diff --git a/tests/test-moe-copy-lru-sim.py b/tests/test-moe-copy-lru-sim.py new file mode 100755 index 00000000000..350a931b042 --- /dev/null +++ b/tests/test-moe-copy-lru-sim.py @@ -0,0 +1,412 @@ +#!/usr/bin/env python3 + +import importlib.util +import subprocess +import sys +import tempfile +from pathlib import Path + + +SAMPLE_LOG = """\ +noise before +ggml_backend_sched_compute_splits: moe_copy split=1 input=0 tensor=blk.0.ffn_down_exps.weight node=ffn_down ids=topk src_backend=CPU dst_backend=CUDA0 n_expert=4 expert_size=100 used=2 used_bytes=200 ranges=1 copy_bytes=200 ids=[1,2] +ggml_backend_sched_compute_splits: moe_copy split=1 input=0 tensor=blk.0.ffn_down_exps.weight node=ffn_down ids=topk src_backend=CPU dst_backend=CUDA0 n_expert=4 expert_size=100 used=2 used_bytes=200 ranges=1 copy_bytes=200 ids=[2,3] +ggml_backend_sched_compute_splits: moe_copy split=1 input=0 tensor=blk.0.ffn_down_exps.weight node=ffn_down ids=topk src_backend=CPU dst_backend=CUDA0 n_expert=4 expert_size=100 used=2 used_bytes=200 ranges=1 copy_bytes=200 ids=[1,2] +ggml_backend_sched_compute_splits: moe_copy split=2 input=0 tensor=blk.1.ffn_down_exps.weight node=ffn_down ids=topk src_backend=CPU dst_backend=CUDA1 n_expert=4 expert_size=50 used=1 used_bytes=50 ranges=1 copy_bytes=50 ids=[0] +""" + +SAMPLE_RUNTIME_LOG = """\ +noise before +ggml_backend_sched_moe_cache_prepare: moe_cache tensor=blk.0.ffn_down_exps.weight backend=CUDA0 slots=2 expert_size=100 cache_bytes=300 used=2 hits=0 misses=2 copied=200 total_hits=0 total_misses=2 total_copied=200 +ggml_backend_sched_moe_cache_prepare: moe_cache tensor=blk.0.ffn_down_exps.weight backend=CUDA0 slots=2 expert_size=100 cache_bytes=300 used=2 hits=1 misses=1 copied=100 total_hits=1 total_misses=3 total_copied=300 +ggml_backend_sched_moe_cache_prepare: moe_cache tensor=blk.1.ffn_down_exps.weight backend=CUDA1 slots=1 expert_size=50 cache_bytes=50 used=1 hits=0 misses=1 copied=50 total_hits=0 total_misses=1 total_copied=50 +ggml_backend_sched_compute_splits: moe_cache_bypass tensor=blk.2.ffn_down_exps.weight node=ffn_down ids=topk backend=CUDA0 slots=2 reason=too_many_experts n_expert=4 expert_size=100 +ggml_backend_sched_compute_splits: moe_cache_bypass tensor=blk.2.ffn_down_exps.weight node=ffn_down ids=topk backend=CUDA0 slots=2 reason=ids_alloc_failed n_expert=4 expert_size=100 +ggml_backend_sched_compute_splits: moe_cache_bypass tensor=blk.3.ffn_down_exps.weight node=ffn_down ids=topk backend=CUDA1 slots=1 reason=too_many_experts n_expert=4 expert_size=50 +""" + + +SAMPLE_PROMPT_PRIME_LOG = """\ +ggml_backend_sched_compute_splits: moe_copy split=1 input=0 tensor=blk.0.ffn_down_exps.weight node=ffn_down ids=topk src_backend=CPU dst_backend=CUDA0 n_expert=4 expert_size=100 used=3 used_bytes=300 ranges=1 copy_bytes=300 id_counts=[0:5,1:4,2:1] ids=[0,1,2] +ggml_backend_sched_compute_splits: moe_copy split=1 input=0 tensor=blk.0.ffn_down_exps.weight node=ffn_down ids=topk src_backend=CPU dst_backend=CUDA0 n_expert=4 expert_size=100 used=2 used_bytes=200 ranges=1 copy_bytes=200 ids=[0,1] +""" + + +SAMPLE_ORACLE_LOG = """\ +ggml_backend_sched_compute_splits: moe_copy split=1 input=0 tensor=blk.0.ffn_down_exps.weight node=ffn_down ids=topk src_backend=CPU dst_backend=CUDA0 n_expert=4 expert_size=100 used=2 used_bytes=200 ranges=1 copy_bytes=200 ids=[0,1] +ggml_backend_sched_compute_splits: moe_copy split=1 input=0 tensor=blk.0.ffn_down_exps.weight node=ffn_down ids=topk src_backend=CPU dst_backend=CUDA0 n_expert=4 expert_size=100 used=2 used_bytes=200 ranges=1 copy_bytes=200 ids=[2,3] +""" + + +def load_sim(repo_root: Path): + script = repo_root / "scripts" / "moe-copy-lru-sim.py" + spec = importlib.util.spec_from_file_location("moe_copy_lru_sim", script) + if spec is None or spec.loader is None: + raise RuntimeError("failed to create simulator module spec") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_parser(sim) -> None: + events = list(sim.read_events_from_lines(SAMPLE_LOG.splitlines())) + assert len(events) == 4 + assert events[0].key == "CUDA0:blk.0.ffn_down_exps.weight" + assert events[0].expert_size == 100 + assert events[0].used_bytes == 200 + assert events[0].copy_bytes == 200 + assert events[0].expert_ids == (1, 2) + assert events[0].expert_counts == ((1, 1), (2, 1)) + + +def test_lru_batch_eviction(sim) -> None: + events = list(sim.read_events_from_lines(SAMPLE_LOG.splitlines())) + stats = sim.simulate_lru(events, [1, 2]) + + k1 = stats[(1, "CUDA0:blk.0.ffn_down_exps.weight")] + assert k1.events == 3 + assert k1.bypasses == 3 + assert k1.hits == 0 + assert k1.misses == 6 + assert k1.cache_bytes == 100 + assert k1.baseline_bytes == 600 + assert k1.cache_copy_bytes == 600 + + k2 = stats[(2, "CUDA0:blk.0.ffn_down_exps.weight")] + assert k2.events == 3 + assert k2.bypasses == 0 + assert k2.hits == 2 + assert k2.misses == 4 + assert k2.cache_bytes == 200 + assert k2.baseline_bytes == 600 + assert k2.cache_copy_bytes == 400 + + aggregate = sim.aggregate_stats(stats) + assert aggregate[2].cache_bytes == 300 + assert aggregate[2].baseline_bytes == 650 + assert aggregate[2].cache_copy_bytes == 450 + + +def test_prompt_prefetch_uses_bypass_hot_set(sim) -> None: + events = list(sim.read_events_from_lines(SAMPLE_PROMPT_PRIME_LOG.splitlines())) + assert events[0].expert_counts == ((0, 5), (1, 4), (2, 1)) + stats = sim.simulate_prefetch(events, [2], "prompt") + stat = stats[("prompt", 2, "CUDA0:blk.0.ffn_down_exps.weight")] + + assert stat.events == 2 + assert stat.bypasses == 1 + assert stat.accesses == 5 + assert stat.speculative_hits == 2 + assert stat.demand_hits == 0 + assert stat.misses == 3 + assert stat.baseline_bytes == 500 + assert stat.demand_copy_bytes == 300 + assert stat.prefetch_copy_bytes == 200 + assert stat.prefetches == 2 + + +def test_oracle_prefetch_bounds_next_event(sim) -> None: + events = list(sim.read_events_from_lines(SAMPLE_ORACLE_LOG.splitlines())) + stats = sim.simulate_prefetch(events, [2], "oracle") + stat = stats[("oracle", 2, "CUDA0:blk.0.ffn_down_exps.weight")] + + assert stat.events == 2 + assert stat.accesses == 4 + assert stat.speculative_hits == 2 + assert stat.misses == 2 + assert stat.baseline_bytes == 400 + assert stat.demand_copy_bytes == 200 + assert stat.prefetch_copy_bytes == 200 + assert stat.prefetches == 2 + assert stat.wrong_prefetches == 0 + + +def test_markov_prefetch_learns_repeated_sequence(sim) -> None: + events = list(sim.read_events_from_lines(SAMPLE_ORACLE_LOG.splitlines())) * 2 + stats = sim.simulate_prefetch(events, [2], "markov") + stat = stats[("markov", 2, "CUDA0:blk.0.ffn_down_exps.weight")] + + assert stat.events == 4 + assert stat.accesses == 8 + assert stat.speculative_hits == 2 + assert stat.misses == 6 + assert stat.prefetches == 4 + assert stat.demand_copy_bytes == 600 + assert stat.prefetch_copy_bytes == 400 + + set_stats = sim.simulate_prefetch(events, [2], "setmarkov") + set_stat = set_stats[("setmarkov", 2, "CUDA0:blk.0.ffn_down_exps.weight")] + assert set_stat.speculative_hits == 2 + assert set_stat.prefetches == 4 + + +def test_cli(repo_root: Path) -> None: + with tempfile.TemporaryDirectory() as tmp: + log_path = Path(tmp) / "moe.log" + log_path.write_text(SAMPLE_LOG, encoding="utf-8") + script = repo_root / "scripts" / "moe-copy-lru-sim.py" + result = subprocess.run( + [sys.executable, str(script), "--slots", "2", str(log_path)], + check=True, + capture_output=True, + text=True, + ) + assert "slots\tkey\tcache_bytes\tevents" in result.stdout + assert "2\tALL\t300\t4\t0\t7\t2\t5\t0.285714\t650\t450\t200\t0.307692" in result.stdout + + +def test_prefetch_cli(repo_root: Path) -> None: + with tempfile.TemporaryDirectory() as tmp: + log_path = Path(tmp) / "moe-prompt.log" + log_path.write_text(SAMPLE_PROMPT_PRIME_LOG, encoding="utf-8") + script = repo_root / "scripts" / "moe-copy-lru-sim.py" + result = subprocess.run( + [sys.executable, str(script), "--slots", "2", "--policy", "prompt", str(log_path)], + check=True, + capture_output=True, + text=True, + ) + assert "policy\tslots\tkey\tcache_bytes\tevents\tbypasses\taccesses" in result.stdout + assert "prompt\t2\tALL\t200\t2\t1\t5\t2\t0\t2\t3\t0.400000\t500\t300\t200\t500\t200\t0.400000\t0\t0.000000\t2\t0\t0" in result.stdout + + +def test_runtime_cache_parser_and_summary(sim) -> None: + events = list(sim.read_cache_events_from_lines(SAMPLE_RUNTIME_LOG.splitlines())) + assert len(events) == 3 + assert events[0].key == "CUDA0:blk.0.ffn_down_exps.weight" + assert events[0].slots == 2 + assert events[0].expert_size == 100 + assert events[0].cache_bytes == 300 + assert events[0].used == 2 + assert events[0].hits == 0 + assert events[0].misses == 2 + assert events[0].copied == 200 + + stats = sim.summarize_runtime_cache(events) + k0 = stats[(2, "CUDA0:blk.0.ffn_down_exps.weight")] + assert k0.slots == 2 + assert k0.expert_size == 100 + assert k0.cache_bytes == 300 + assert k0.events == 2 + assert k0.accesses == 4 + assert k0.hits == 1 + assert k0.misses == 3 + assert k0.copied == 300 + assert k0.max_total_hits == 1 + assert k0.max_total_misses == 3 + assert k0.max_total_copied == 300 + + aggregate = sim.aggregate_runtime_stats(stats) + assert aggregate[1].cache_bytes == 50 + assert aggregate[1].events == 1 + assert aggregate[1].accesses == 1 + assert aggregate[1].hits == 0 + assert aggregate[1].misses == 1 + assert aggregate[1].copied == 50 + assert aggregate[1].max_total_copied == 50 + assert aggregate[2].cache_bytes == 300 + assert aggregate[2].events == 2 + assert aggregate[2].accesses == 4 + assert aggregate[2].hits == 1 + assert aggregate[2].misses == 3 + assert aggregate[2].copied == 300 + assert aggregate[2].max_total_copied == 300 + + +def test_runtime_cache_cli(repo_root: Path) -> None: + with tempfile.TemporaryDirectory() as tmp: + log_path = Path(tmp) / "moe-runtime.log" + log_path.write_text(SAMPLE_RUNTIME_LOG, encoding="utf-8") + script = repo_root / "scripts" / "moe-copy-lru-sim.py" + result = subprocess.run( + [sys.executable, str(script), "--runtime", "--details", str(log_path)], + check=True, + capture_output=True, + text=True, + ) + assert "key\tslots\tcache_bytes\tevents\taccesses\thits\tmisses" in result.stdout + assert "ALL\t1\t50\t1\t1\t0\t1\t0.000000\t50\t0\t1\t50" in result.stdout + assert "ALL\t2\t300\t2\t4\t1\t3\t0.250000\t300\t1\t3\t300" in result.stdout + assert "CUDA0:blk.0.ffn_down_exps.weight\t2\t300\t2\t4\t1\t3\t0.250000\t300\t1\t3\t300" in result.stdout + assert "bypass_key\tslots\treason\tevents" in result.stdout + assert "ALL\t1\ttoo_many_experts\t1" in result.stdout + assert "ALL\t2\ttoo_many_experts\t1" in result.stdout + assert "CUDA0:blk.2.ffn_down_exps.weight\t2\tids_alloc_failed\t1" in result.stdout + + +def test_cli_tolerates_invalid_utf8(repo_root: Path) -> None: + with tempfile.TemporaryDirectory() as tmp: + log_path = Path(tmp) / "moe-invalid-utf8.log" + log_path.write_bytes( + b"\xef\xbf\x00spinner\n" + b"ggml_backend_sched_moe_cache_prepare: moe_cache tensor=t backend=CUDA0 slots=2 expert_size=100 " + b"cache_bytes=300 used=1 hits=0 misses=1 copied=100 total_hits=0 total_misses=1 total_copied=100\n" + ) + script = repo_root / "scripts" / "moe-copy-lru-sim.py" + result = subprocess.run( + [sys.executable, str(script), "--runtime", str(log_path)], + check=True, + capture_output=True, + text=True, + ) + assert "ALL\t2\t300\t1\t1\t0\t1\t0.000000\t100\t0\t1\t100" in result.stdout + + +def test_runtime_cache_bypass_parser_and_summary(sim) -> None: + events = list(sim.read_cache_bypass_events_from_lines(SAMPLE_RUNTIME_LOG.splitlines())) + assert len(events) == 3 + assert events[0].key == "CUDA0:blk.2.ffn_down_exps.weight" + assert events[0].slots == 2 + assert events[0].reason == "too_many_experts" + assert events[0].n_expert == 4 + assert events[0].expert_size == 100 + + stats = sim.summarize_runtime_bypasses(events) + assert stats[("CUDA0:blk.2.ffn_down_exps.weight", 2, "too_many_experts")] == 1 + assert stats[("CUDA0:blk.2.ffn_down_exps.weight", 2, "ids_alloc_failed")] == 1 + assert stats[("CUDA1:blk.3.ffn_down_exps.weight", 1, "too_many_experts")] == 1 + + +def test_runtime_bypass_only_cli(repo_root: Path) -> None: + log = """\ +ggml_backend_sched_compute_splits: moe_cache_bypass tensor=blk.2.ffn_down_exps.weight node=ffn_down ids=topk backend=CUDA0 slots=2 reason=too_many_experts n_expert=4 expert_size=100 +ggml_backend_sched_compute_splits: moe_cache_bypass tensor=blk.2.ffn_down_exps.weight node=ffn_down ids=topk backend=CUDA0 slots=2 reason=ids_alloc_failed n_expert=4 expert_size=100 +""" + with tempfile.TemporaryDirectory() as tmp: + log_path = Path(tmp) / "moe-runtime-bypass.log" + log_path.write_text(log, encoding="utf-8") + script = repo_root / "scripts" / "moe-copy-lru-sim.py" + result = subprocess.run( + [sys.executable, str(script), "--runtime", "--details", str(log_path)], + check=True, + capture_output=True, + text=True, + ) + assert "key\tslots\tcache_bytes\tevents\taccesses\thits\tmisses" not in result.stdout + assert "bypass_key\tslots\treason\tevents" in result.stdout + assert "ALL\t2\tids_alloc_failed\t1" in result.stdout + assert "ALL\t2\ttoo_many_experts\t1" in result.stdout + assert "CUDA0:blk.2.ffn_down_exps.weight\t2\ttoo_many_experts\t1" in result.stdout + + +def test_rejects_inconsistent_expert_size(sim) -> None: + log = """\ +ggml_backend_sched_compute_splits: moe_copy split=1 input=0 tensor=blk.0.ffn_down_exps.weight node=ffn_down ids=topk src_backend=CPU dst_backend=CUDA0 n_expert=4 expert_size=100 used=1 used_bytes=100 ranges=1 copy_bytes=100 ids=[1] +ggml_backend_sched_compute_splits: moe_copy split=1 input=0 tensor=blk.0.ffn_down_exps.weight node=ffn_down ids=topk src_backend=CPU dst_backend=CUDA0 n_expert=4 expert_size=200 used=1 used_bytes=200 ranges=1 copy_bytes=200 ids=[2] +""" + events = list(sim.read_events_from_lines(log.splitlines())) + try: + sim.simulate_lru(events, [2]) + except ValueError as exc: + assert "inconsistent expert_size" in str(exc) + else: + raise AssertionError("accepted inconsistent expert_size for a single cache key") + + +def test_rejects_inconsistent_copy_accounting(sim) -> None: + bad_used_bytes = "ggml_backend_sched_compute_splits: moe_copy split=1 input=0 tensor=blk.0.ffn_down_exps.weight node=ffn_down ids=topk src_backend=CPU dst_backend=CUDA0 n_expert=4 expert_size=100 used=2 used_bytes=100 ranges=1 copy_bytes=200 ids=[1,2]" + try: + list(sim.read_events_from_lines([bad_used_bytes])) + except ValueError as exc: + assert "used_bytes" in str(exc) + else: + raise AssertionError("accepted inconsistent used_bytes") + + bad_copy_bytes = "ggml_backend_sched_compute_splits: moe_copy split=1 input=0 tensor=blk.0.ffn_down_exps.weight node=ffn_down ids=topk src_backend=CPU dst_backend=CUDA0 n_expert=4 expert_size=100 used=2 used_bytes=200 ranges=1 copy_bytes=150 ids=[1,2]" + try: + list(sim.read_events_from_lines([bad_copy_bytes])) + except ValueError as exc: + assert "copy_bytes" in str(exc) + else: + raise AssertionError("accepted copy_bytes smaller than used_bytes") + + bad_id_counts = "ggml_backend_sched_compute_splits: moe_copy split=1 input=0 tensor=blk.0.ffn_down_exps.weight node=ffn_down ids=topk src_backend=CPU dst_backend=CUDA0 n_expert=4 expert_size=100 used=2 used_bytes=200 ranges=1 copy_bytes=200 id_counts=[1:2] ids=[1,2]" + try: + list(sim.read_events_from_lines([bad_id_counts])) + except ValueError as exc: + assert "id_counts" in str(exc) + else: + raise AssertionError("accepted id_counts that did not match ids") + + +def test_rejects_inconsistent_runtime_cache_accounting(sim) -> None: + bad_used = "ggml_backend_sched_moe_cache_prepare: moe_cache tensor=blk.0.ffn_down_exps.weight backend=CUDA0 slots=2 used=2 hits=2 misses=1 copied=100 total_hits=2 total_misses=1 total_copied=100" + try: + list(sim.read_cache_events_from_lines([bad_used])) + except ValueError as exc: + assert "used" in str(exc) + else: + raise AssertionError("accepted inconsistent runtime used/hit/miss counts") + + bad_total = "ggml_backend_sched_moe_cache_prepare: moe_cache tensor=blk.0.ffn_down_exps.weight backend=CUDA0 slots=2 used=1 hits=0 misses=1 copied=100 total_hits=0 total_misses=0 total_copied=100" + try: + list(sim.read_cache_events_from_lines([bad_total])) + except ValueError as exc: + assert "total" in str(exc) + else: + raise AssertionError("accepted runtime total counters below per-event counters") + + mixed_slots = """\ +ggml_backend_sched_moe_cache_prepare: moe_cache tensor=blk.0.ffn_down_exps.weight backend=CUDA0 slots=2 cache_bytes=200 used=1 hits=0 misses=1 copied=100 total_hits=0 total_misses=1 total_copied=100 +ggml_backend_sched_moe_cache_prepare: moe_cache tensor=blk.0.ffn_down_exps.weight backend=CUDA0 slots=3 cache_bytes=300 used=1 hits=0 misses=1 copied=100 total_hits=0 total_misses=1 total_copied=100 +""" + events = list(sim.read_cache_events_from_lines(mixed_slots.splitlines())) + stats = sim.summarize_runtime_cache(events) + assert stats[(2, "CUDA0:blk.0.ffn_down_exps.weight")].events == 1 + assert stats[(3, "CUDA0:blk.0.ffn_down_exps.weight")].events == 1 + assert stats[(2, "CUDA0:blk.0.ffn_down_exps.weight")].cache_bytes == 200 + assert stats[(3, "CUDA0:blk.0.ffn_down_exps.weight")].cache_bytes == 300 + + inconsistent_cache_bytes = """\ +ggml_backend_sched_moe_cache_prepare: moe_cache tensor=blk.0.ffn_down_exps.weight backend=CUDA0 slots=2 expert_size=100 cache_bytes=200 used=1 hits=0 misses=1 copied=100 total_hits=0 total_misses=1 total_copied=100 +ggml_backend_sched_moe_cache_prepare: moe_cache tensor=blk.0.ffn_down_exps.weight backend=CUDA0 slots=2 expert_size=100 cache_bytes=300 used=1 hits=0 misses=1 copied=100 total_hits=0 total_misses=1 total_copied=100 +""" + try: + sim.summarize_runtime_cache(list(sim.read_cache_events_from_lines(inconsistent_cache_bytes.splitlines()))) + except ValueError as exc: + assert "cache_bytes" in str(exc) + else: + raise AssertionError("accepted inconsistent runtime cache footprint") + + inconsistent_expert_size = """\ +ggml_backend_sched_moe_cache_prepare: moe_cache tensor=blk.0.ffn_down_exps.weight backend=CUDA0 slots=2 expert_size=100 cache_bytes=200 used=1 hits=0 misses=1 copied=100 total_hits=0 total_misses=1 total_copied=100 +ggml_backend_sched_moe_cache_prepare: moe_cache tensor=blk.0.ffn_down_exps.weight backend=CUDA0 slots=2 expert_size=101 cache_bytes=200 used=1 hits=0 misses=1 copied=100 total_hits=0 total_misses=1 total_copied=100 +""" + try: + sim.summarize_runtime_cache(list(sim.read_cache_events_from_lines(inconsistent_expert_size.splitlines()))) + except ValueError as exc: + assert "expert_size" in str(exc) + else: + raise AssertionError("accepted inconsistent runtime cache expert size") + + bad_bypass = "ggml_backend_sched_compute_splits: moe_cache_bypass tensor=blk.0.ffn_down_exps.weight node=ffn_down ids=topk backend=CUDA0 slots=-1 reason=too_many_experts n_expert=4 expert_size=100" + try: + list(sim.read_cache_bypass_events_from_lines([bad_bypass])) + except ValueError as exc: + assert "non-negative" in str(exc) + else: + raise AssertionError("accepted negative runtime bypass slots") + + +def main() -> None: + repo_root = Path(sys.argv[1]) if len(sys.argv) > 1 else Path(__file__).resolve().parents[1] + sim = load_sim(repo_root) + test_parser(sim) + test_lru_batch_eviction(sim) + test_prompt_prefetch_uses_bypass_hot_set(sim) + test_oracle_prefetch_bounds_next_event(sim) + test_markov_prefetch_learns_repeated_sequence(sim) + test_cli(repo_root) + test_prefetch_cli(repo_root) + test_runtime_cache_parser_and_summary(sim) + test_runtime_cache_cli(repo_root) + test_cli_tolerates_invalid_utf8(repo_root) + test_runtime_cache_bypass_parser_and_summary(sim) + test_runtime_bypass_only_cli(repo_root) + test_rejects_inconsistent_expert_size(sim) + test_rejects_inconsistent_copy_accounting(sim) + test_rejects_inconsistent_runtime_cache_accounting(sim) + + +if __name__ == "__main__": + main() diff --git a/tests/test-quant-type-selection.cpp b/tests/test-quant-type-selection.cpp index 3c8983360e2..6fc5f2a989d 100644 --- a/tests/test-quant-type-selection.cpp +++ b/tests/test-quant-type-selection.cpp @@ -56,6 +56,7 @@ static const ftype_name_entry ftype_name_table[] = { { "TQ1_0", LLAMA_FTYPE_MOSTLY_TQ1_0 }, { "TQ2_0", LLAMA_FTYPE_MOSTLY_TQ2_0 }, { "MXFP4_MOE", LLAMA_FTYPE_MOSTLY_MXFP4_MOE }, + { "F8_E4M3_MXFP4", LLAMA_FTYPE_MOSTLY_F8_E4M3_MXFP4 }, { "NVFP4", LLAMA_FTYPE_MOSTLY_NVFP4 }, }; diff --git a/tests/test-server-prompt-cache.cpp b/tests/test-server-prompt-cache.cpp new file mode 100644 index 00000000000..2187d3b962c --- /dev/null +++ b/tests/test-server-prompt-cache.cpp @@ -0,0 +1,179 @@ +#include "server-task.h" + +#include +#include +#include +#include + +static void require(bool condition, const char * message) { + if (!condition) { + std::fprintf(stderr, "%s\n", message); + std::exit(1); + } +} + +static server_prompt make_prompt(std::initializer_list tokens) { + server_prompt prompt; + prompt.tokens = server_tokens(llama_tokens(tokens), false); + return prompt; +} + +static void add_checkpoint(server_prompt & prompt, int64_t n_tokens) { + server_prompt_checkpoint checkpoint = {}; + checkpoint.pos_min = 0; + checkpoint.pos_max = n_tokens; + checkpoint.n_tokens = n_tokens; + checkpoint.data = { 1, 2, 3, 4 }; + prompt.checkpoints.push_back(std::move(checkpoint)); +} + +static void add_checkpoint_with_bounds(server_prompt & prompt, llama_pos pos_max, int64_t n_tokens, uint8_t marker) { + server_prompt_checkpoint checkpoint = {}; + checkpoint.pos_min = 0; + checkpoint.pos_max = pos_max; + checkpoint.n_tokens = n_tokens; + checkpoint.data = { marker }; + prompt.checkpoints.push_back(std::move(checkpoint)); +} + +static void test_find_checkpoint_before_tail_truncation_pos() { + server_prompt prompt = make_prompt({ 1, 2, 3, 4, 5, 6, 7, 8 }); + + add_checkpoint_with_bounds(prompt, 3, 4, 4); + add_checkpoint_with_bounds(prompt, 5, 6, 6); + add_checkpoint_with_bounds(prompt, 7, 12, 12); // invalid: more tokens than prompt + + const server_prompt_checkpoint * latest = server_prompt_find_checkpoint_before_pos(prompt, 7); + require(latest != nullptr, "expected a checkpoint before tail truncation position"); + require(latest->n_tokens == 6, "expected latest compatible checkpoint before p0"); + require(latest->data == std::vector{ 6 }, "expected latest compatible checkpoint data"); + + const server_prompt_checkpoint * earlier = server_prompt_find_checkpoint_before_pos(prompt, 5); + require(earlier != nullptr, "expected an earlier checkpoint before p0"); + require(earlier->n_tokens == 4, "expected checkpoint with pos_max strictly before p0"); + + require(server_prompt_find_checkpoint_before_pos(prompt, 3) == nullptr, + "checkpoint at pos_max >= p0 must not be used for tail truncation restore"); +} + +static void test_oaicompat_chat_streams_reasoning_delta() { + common_chat_parser_params parser_params; + parser_params.reasoning_format = COMMON_REASONING_FORMAT_DEEPSEEK; + parser_params.generation_prompt = ""; + + task_result_state state(parser_params); + + server_task_result_cmpl_partial partial = {}; + partial.content = "I am thinking"; + partial.n_decoded = 1; + partial.res_type = TASK_RESPONSE_TYPE_OAI_CHAT; + partial.oaicompat_model = "test-model"; + partial.oaicompat_cmpl_id = "chatcmpl-test"; + partial.update(state); + + json chunks = partial.to_json_oaicompat_chat(); + bool found_reasoning = false; + for (const auto & chunk : chunks) { + if (!chunk.contains("choices") || chunk.at("choices").empty()) { + continue; + } + const auto & delta = chunk.at("choices").at(0).at("delta"); + if (delta.contains("reasoning_content") && delta.at("reasoning_content") == "I am thinking") { + found_reasoning = true; + } + } + + require(found_reasoning, "streaming chat response should broadcast reasoning_content deltas"); +} + +static void test_oaicompat_chat_final_contains_reasoning() { + server_task_result_cmpl_final final = {}; + final.res_type = TASK_RESPONSE_TYPE_OAI_CHAT; + final.oaicompat_model = "test-model"; + final.oaicompat_cmpl_id = "chatcmpl-test"; + final.include_usage = true; + final.oaicompat_msg.role = "assistant"; + final.oaicompat_msg.reasoning_content = "I am thinking"; + final.oaicompat_msg.content = "The answer is 4."; + + json body = final.to_json_oaicompat_chat(); + const auto & message = body.at("choices").at(0).at("message"); + require(message.contains("reasoning_content"), "final chat response should contain reasoning_content"); + require(message.at("reasoning_content") == "I am thinking", + "final chat response should extract reasoning before "); + require(message.at("content") == "The answer is 4.", + "final chat response should keep post-thinking content separate"); +} + +static void test_full_removal_keeps_exact_shorter_without_checkpoint() { + server_prompt_cache cache(0, 0); + + server_prompt long_prompt = make_prompt({ 1, 2, 3, 4, 5, 6 }); + server_prompt short_prompt = make_prompt({ 1, 2, 3, 4 }); + + require(cache.alloc(long_prompt, 8, COMMON_CONTEXT_SEQ_RM_TYPE_FULL) != nullptr, + "failed to allocate initial long prompt"); + require(cache.alloc(short_prompt, 8, COMMON_CONTEXT_SEQ_RM_TYPE_FULL) != nullptr, + "short exact prompt should be cached when the longer state has no restorable prefix checkpoint"); + require(cache.states.size() == 2, + "cache should retain both long and short prompts without a restorable prefix checkpoint"); +} + +static void test_full_removal_reuses_longer_checkpoint_for_shorter_prompt() { + server_prompt_cache cache(0, 0); + + server_prompt long_prompt = make_prompt({ 1, 2, 3, 4, 5, 6 }); + server_prompt short_prompt = make_prompt({ 1, 2, 3, 4 }); + add_checkpoint(long_prompt, short_prompt.n_tokens()); + + require(cache.alloc(long_prompt, 8, COMMON_CONTEXT_SEQ_RM_TYPE_FULL) != nullptr, + "failed to allocate checkpointed long prompt"); + require(cache.alloc(short_prompt, 8, COMMON_CONTEXT_SEQ_RM_TYPE_FULL) == nullptr, + "short exact prompt should be skipped when a longer checkpoint can restore it"); + require(cache.states.size() == 1, + "checkpointed long prompt should make the exact shorter prompt redundant"); +} + +static void test_full_removal_only_removes_obsolete_shorter_with_checkpoint() { + { + server_prompt_cache cache(0, 0); + + server_prompt short_prompt = make_prompt({ 1, 2, 3, 4 }); + server_prompt long_prompt = make_prompt({ 1, 2, 3, 4, 5, 6 }); + + require(cache.alloc(short_prompt, 8, COMMON_CONTEXT_SEQ_RM_TYPE_FULL) != nullptr, + "failed to allocate initial short prompt"); + require(cache.alloc(long_prompt, 8, COMMON_CONTEXT_SEQ_RM_TYPE_FULL) != nullptr, + "long prompt without checkpoint should still be cached"); + require(cache.states.size() == 2, + "short prompt must not be removed when long prompt cannot restore that prefix"); + } + + { + server_prompt_cache cache(0, 0); + + server_prompt short_prompt = make_prompt({ 1, 2, 3, 4 }); + server_prompt long_prompt = make_prompt({ 1, 2, 3, 4, 5, 6 }); + add_checkpoint(long_prompt, short_prompt.n_tokens()); + + require(cache.alloc(short_prompt, 8, COMMON_CONTEXT_SEQ_RM_TYPE_FULL) != nullptr, + "failed to allocate initial short prompt"); + require(cache.alloc(long_prompt, 8, COMMON_CONTEXT_SEQ_RM_TYPE_FULL) != nullptr, + "failed to allocate checkpointed long prompt"); + require(cache.states.size() == 1, + "short prompt should be removed once long prompt has a restorable prefix checkpoint"); + require(cache.states.front().n_tokens() == long_prompt.n_tokens(), + "remaining cache entry should be the checkpointed long prompt"); + } +} + +int main() { + test_find_checkpoint_before_tail_truncation_pos(); + test_oaicompat_chat_streams_reasoning_delta(); + test_oaicompat_chat_final_contains_reasoning(); + test_full_removal_keeps_exact_shorter_without_checkpoint(); + test_full_removal_reuses_longer_checkpoint_for_shorter_prompt(); + test_full_removal_only_removes_obsolete_shorter_with_checkpoint(); + + return 0; +} diff --git a/tools/quantize/quantize.cpp b/tools/quantize/quantize.cpp index 3d33d47d98b..a575697aeae 100644 --- a/tools/quantize/quantize.cpp +++ b/tools/quantize/quantize.cpp @@ -36,6 +36,7 @@ static const std::vector QUANT_OPTIONS = { { "Q4_0", LLAMA_FTYPE_MOSTLY_Q4_0, " 4.34G, +0.4685 ppl @ Llama-3-8B", }, { "Q4_1", LLAMA_FTYPE_MOSTLY_Q4_1, " 4.78G, +0.4511 ppl @ Llama-3-8B", }, { "MXFP4_MOE",LLAMA_FTYPE_MOSTLY_MXFP4_MOE," MXFP4 MoE", }, + { "F8_E4M3_MXFP4", LLAMA_FTYPE_MOSTLY_F8_E4M3_MXFP4, " FP8 E4M3 dense + MXFP4 MoE", }, { "Q5_0", LLAMA_FTYPE_MOSTLY_Q5_0, " 5.21G, +0.1316 ppl @ Llama-3-8B", }, { "Q5_1", LLAMA_FTYPE_MOSTLY_Q5_1, " 5.65G, +0.1062 ppl @ Llama-3-8B", }, { "IQ2_XXS", LLAMA_FTYPE_MOSTLY_IQ2_XXS, " 2.06 bpw quantization", }, diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index e3822225bdb..921e16145fe 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -139,7 +139,7 @@ struct server_slot { SRV_WRN(" - saving prompt with length %d, total state size = %.3f MiB\n", (int) prompt.tokens.size(), cur_size / (1024.0 * 1024.0)); - auto * cur = prompt_cache.alloc(prompt, cur_size); + auto * cur = prompt_cache.alloc(prompt, cur_size, ctx_seq_rm_type); if (cur == nullptr) { return; } @@ -148,7 +148,7 @@ struct server_slot { } bool prompt_load(server_prompt_cache & prompt_cache, const server_tokens & tokens) { - bool res = prompt_cache.load(prompt, tokens, ctx, id); + bool res = prompt_cache.load(prompt, tokens, ctx, id, ctx_seq_rm_type); if (!res) { SLT_WRN(*this, "%s", "failed to load prompt from cache\n"); } @@ -933,6 +933,16 @@ struct server_context_impl { batch = llama_batch_init(std::max(n_batch, params_base.n_parallel), 0, 1); } + // Models that don't support partial sequence removal (e.g., DeepSeek V4 + // which has a fixed-size sliding-window + indexer KV state) crash later + // in update_slots() when the prompt cache tries to do prefix-matched + // reuse and llama_memory_seq_pos_min returns -1. Force the prompt + // cache off in that case to avoid a confusing GGML_ABORT. + if (params_base.cache_ram_mib != 0 && ctx_seq_rm_type != COMMON_CONTEXT_SEQ_RM_TYPE_PART) { + SRV_WRN("%s", "prompt cache disabled: model does not support partial sequence removal\n"); + params_base.cache_ram_mib = 0; + } + if (params_base.cache_ram_mib != 0) { if (params_base.cache_ram_mib < 0) { SRV_WRN("prompt cache is enabled, size limit: %s\n", "no limit"); @@ -2424,8 +2434,10 @@ struct server_context_impl { if (n_past > 0 && n_past < slot.prompt.n_tokens()) { const auto pos_min = llama_memory_seq_pos_min(llama_get_memory(ctx), slot.id); if (pos_min == -1) { - SLT_ERR(slot, "n_past = %d, slot.prompt.tokens.size() = %d, seq_id = %d, pos_min = %d\n", n_past, (int) slot.prompt.tokens.size(), slot.id, pos_min); - GGML_ABORT("pos_min == -1, but n_past > 0 - should not happen: https://github.com/ggml-org/llama.cpp/pull/13833#discussion_r2116181237"); + SLT_WRN(slot, "n_past = %d, slot.prompt.tokens.size() = %d, seq_id = %d, pos_min = %d - forcing full prompt re-evaluation (non-standard attention architecture cache mismatch)\n", n_past, (int) slot.prompt.tokens.size(), slot.id, pos_min); + // Non-standard attention (e.g. DeepSeek V4 CSA+HCA): cache state can be inconsistent. Safe fallback. + pos_next = 0; + n_past = 0; } // when the prompt prefix does not match, print the tokens around the mismatch @@ -2559,12 +2571,36 @@ struct server_context_impl { SLT_INF(slot, "n_tokens = %d, memory_seq_rm [%d, end)\n", slot.prompt.n_tokens(), p0); if (!llama_memory_seq_rm(llama_get_memory(ctx), slot.id, p0, -1)) { - SLT_WRN(slot, "failed to truncate tokens with position >= %d - clearing the memory\n", p0); + bool restored = false; + + if (slot.ctx_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL) { + const server_prompt_checkpoint * checkpoint = server_prompt_find_checkpoint_before_pos(slot.prompt, p0); + + if (checkpoint != nullptr) { + const size_t checkpoint_size = checkpoint->data.size(); + const size_t n = llama_state_seq_set_data_ext(ctx, checkpoint->data.data(), checkpoint_size, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + + if (n == checkpoint_size) { + slot.prompt.tokens.keep_first(checkpoint->n_tokens); + slot.n_prompt_tokens_cache = checkpoint->n_tokens; + restored = true; + SLT_WRN(slot, "restored context checkpoint after failed memory_seq_rm (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB)\n", + checkpoint->pos_min, checkpoint->pos_max, checkpoint->n_tokens, (float) checkpoint_size / 1024 / 1024); + } else { + SLT_ERR(slot, "failed to restore context checkpoint after failed memory_seq_rm (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB)\n", + checkpoint->pos_min, checkpoint->pos_max, checkpoint->n_tokens, (float) checkpoint_size / 1024 / 1024); + } + } + } + + if (!restored) { + SLT_WRN(slot, "failed to truncate tokens with position >= %d - clearing the memory\n", p0); - slot.prompt_clear(true); + slot.prompt_clear(true); - // there is no common part left - slot.n_prompt_tokens_cache = 0; + // there is no common part left + slot.n_prompt_tokens_cache = 0; + } } // If using an alora, there may be uncached tokens that come diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 45e5168fabe..559e6fc16c6 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -156,6 +156,20 @@ common_chat_msg task_result_state::update_chat_msg( generated_text, is_partial, chat_parser_params); + + if (is_partial && + chat_parser_params.reasoning_format == COMMON_REASONING_FORMAT_DEEPSEEK && + string_ends_with(chat_parser_params.generation_prompt, "") && + generated_text.find("") == std::string::npos) { + std::string reasoning = generated_text; + if (string_starts_with(reasoning, "")) { + reasoning.erase(0, std::string("").size()); + } + + new_msg.role = "assistant"; + new_msg.reasoning_content = std::move(reasoning); + } + if (!new_msg.empty()) { new_msg.set_tool_call_ids(generated_tool_call_ids, gen_tool_call_id); chat_msg = new_msg; @@ -1961,6 +1975,31 @@ json server_task_result_apply_lora::to_json() { // // server_prompt_cache // +static bool server_prompt_can_restore_prefix( + const server_prompt & prompt, + int64_t n_tokens, + common_context_seq_rm_type seq_rm_type) { + if (n_tokens < 0 || n_tokens > prompt.n_tokens()) { + return false; + } + + if (seq_rm_type != COMMON_CONTEXT_SEQ_RM_TYPE_FULL) { + return true; + } + + if (n_tokens == prompt.n_tokens()) { + return true; + } + + for (const auto & checkpoint : prompt.checkpoints) { + if (!checkpoint.empty() && checkpoint.n_tokens == n_tokens) { + return true; + } + } + + return false; +} + size_t server_prompt_cache::size() const { size_t res = 0; @@ -1981,12 +2020,16 @@ size_t server_prompt_cache::n_tokens() const { return res; } -server_prompt * server_prompt_cache::alloc(const server_prompt & prompt, size_t state_size) { +server_prompt * server_prompt_cache::alloc( + const server_prompt & prompt, + size_t state_size, + common_context_seq_rm_type seq_rm_type) { // first check if the current state is contained fully in the cache for (auto it = states.begin(); it != states.end(); ++it) { const int cur_lcp_len = it->tokens.get_common_prefix(prompt.tokens); - if (cur_lcp_len == (int) prompt.tokens.size()) { + if (cur_lcp_len == (int) prompt.tokens.size() && + server_prompt_can_restore_prefix(*it, prompt.n_tokens(), seq_rm_type)) { SRV_WRN("%s", " - prompt is already in the cache, skipping\n"); return nullptr; } @@ -1996,7 +2039,8 @@ server_prompt * server_prompt_cache::alloc(const server_prompt & prompt, size_t for (auto it = states.begin(); it != states.end();) { const int len = it->tokens.get_common_prefix(prompt.tokens); - if (len == (int) it->tokens.size()) { + if (len == (int) it->tokens.size() && + server_prompt_can_restore_prefix(prompt, it->n_tokens(), seq_rm_type)) { SRV_WRN(" - removing obsolete cached prompt with length %d\n", len); it = states.erase(it); @@ -2032,7 +2076,12 @@ server_prompt * server_prompt_cache::alloc(const server_prompt & prompt, size_t return &cur; } -bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx, int32_t id_slot) { +bool server_prompt_cache::load( + server_prompt & prompt, + const server_tokens & tokens_new, + llama_context * ctx, + int32_t id_slot, + common_context_seq_rm_type seq_rm_type) { const int lcp_best = prompt.tokens.get_common_prefix(tokens_new); float f_keep_best = prompt.tokens.size() > 0 ? float(lcp_best) / prompt.tokens.size() : -1.0f; // empty slot: any cache entry wins @@ -2041,13 +2090,32 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok SRV_WRN(" - looking for better prompt, base f_keep = %.3f, sim = %.3f\n", f_keep_best, sim_best); auto it_best = states.end(); + const server_prompt_checkpoint * checkpoint_best = nullptr; + int64_t n_tokens_best = -1; // find the most similar cached prompt, that would also preserve the most context for (auto it = states.begin(); it != states.end(); ++it) { const int lcp_cur = it->tokens.get_common_prefix(tokens_new); - const float f_keep_cur = float(lcp_cur) / it->tokens.size(); - const float sim_cur = float(lcp_cur) / tokens_new.size(); + int64_t n_tokens_cur = lcp_cur; + const server_prompt_checkpoint * checkpoint_cur = nullptr; + + if (seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL && lcp_cur < (int) it->tokens.size()) { + n_tokens_cur = -1; + for (const auto & checkpoint : it->checkpoints) { + if (!checkpoint.empty() && checkpoint.n_tokens <= lcp_cur && checkpoint.n_tokens > n_tokens_cur) { + checkpoint_cur = &checkpoint; + n_tokens_cur = checkpoint.n_tokens; + } + } + + if (checkpoint_cur == nullptr) { + continue; + } + } + + const float f_keep_cur = float(n_tokens_cur) / it->tokens.size(); + const float sim_cur = float(n_tokens_cur) / tokens_new.size(); // don't trash large prompts if (f_keep_cur < 0.25f) { @@ -2059,14 +2127,23 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok sim_best = sim_cur; it_best = it; + checkpoint_best = checkpoint_cur; + n_tokens_best = n_tokens_cur; } } if (it_best != states.end()) { - SRV_WRN(" - found better prompt with f_keep = %.3f, sim = %.3f\n", f_keep_best, sim_best); + if (checkpoint_best != nullptr) { + SRV_WRN(" - found better prompt checkpoint with f_keep = %.3f, sim = %.3f, n_tokens = %" PRId64 "\n", + f_keep_best, sim_best, n_tokens_best); + } else { + SRV_WRN(" - found better prompt with f_keep = %.3f, sim = %.3f\n", f_keep_best, sim_best); + } - const size_t size = it_best->data.size(); - const size_t n = llama_state_seq_set_data_ext(ctx, it_best->data.data(), size, id_slot, 0); + const std::vector & data = checkpoint_best != nullptr ? checkpoint_best->data : it_best->data; + const llama_state_seq_flags flags = checkpoint_best != nullptr ? LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY : 0; + const size_t size = data.size(); + const size_t n = llama_state_seq_set_data_ext(ctx, data.data(), size, id_slot, flags); if (n != size) { SRV_WRN("failed to restore state with size %zu\n", size); @@ -2077,6 +2154,16 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok it_best->data.shrink_to_fit(); prompt = std::move(*it_best); + if (checkpoint_best != nullptr) { + prompt.tokens.keep_first(n_tokens_best); + for (auto it = prompt.checkpoints.begin(); it != prompt.checkpoints.end();) { + if (it->n_tokens > n_tokens_best) { + it = prompt.checkpoints.erase(it); + } else { + ++it; + } + } + } states.erase(it_best); } diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 289e1fb8d24..69f0f9596c3 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -619,6 +619,18 @@ struct server_prompt { } }; +inline const server_prompt_checkpoint * server_prompt_find_checkpoint_before_pos( + const server_prompt & prompt, + llama_pos p0) { + for (auto it = prompt.checkpoints.rbegin(); it != prompt.checkpoints.rend(); ++it) { + if (it->pos_max < p0 && it->n_tokens >= 0 && (size_t) it->n_tokens <= prompt.tokens.size()) { + return &*it; + } + } + + return nullptr; +} + struct server_prompt_cache { server_prompt_cache(int32_t limit_size_mib, size_t limit_tokens) { this->limit_size = 1024ull*1024ull*(limit_size_mib < 0 ? 0 : limit_size_mib); @@ -637,9 +649,17 @@ struct server_prompt_cache { size_t n_tokens() const; - server_prompt * alloc(const server_prompt & prompt, size_t state_size); - - bool load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx, int32_t id_slot); + server_prompt * alloc( + const server_prompt & prompt, + size_t state_size, + common_context_seq_rm_type seq_rm_type); + + bool load( + server_prompt & prompt, + const server_tokens & tokens_new, + llama_context * ctx, + int32_t id_slot, + common_context_seq_rm_type seq_rm_type); void update(); };