diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index b0e1d03f..3b581cdb 100644 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -2333,6 +2333,9 @@ class DFlashDraftModel(Qwen3Model): _saw_token_embd = False _saw_output = False + def _causal_attention(self) -> bool: + return False + def _require_target_model_dir(self) -> Path: if self.target_model_dir is None: raise ValueError("DFlashDraftModel conversion requires --target-model-dir ") @@ -2426,7 +2429,7 @@ class DFlashDraftModel(Qwen3Model): def set_gguf_parameters(self): super().set_gguf_parameters() - self.gguf_writer.add_causal_attention(False) + self.gguf_writer.add_causal_attention(self._causal_attention()) # MiMo DFlash draft uses partial rotary (partial_rotary_factor=0.5): RoPE is applied to # only head_dim*partial_rotary_factor dims, the rest are NoPE. Honoring it is required; # otherwise the upper half of every head gets spurious position rotation it was never @@ -2613,6 +2616,133 @@ class DFlashDraftModel(Qwen3Model): return tensors +@Model.register("DFlashLagunaForCausalLM") +class DFlashLagunaModel(DFlashDraftModel): + model_arch = gguf.MODEL_ARCH.DFLASH_DRAFT + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._laguna_qkv_ids: set[int] = set() + self._laguna_gate_ids: set[int] = set() + self._laguna_aux_norm_ids: set[int] = set() + + def _causal_attention(self) -> bool: + return True + + def set_gguf_parameters(self): + dflash_cfg = self.hparams.get("dflash_config") + if not isinstance(dflash_cfg, dict) or dflash_cfg.get("causal") is not True: + raise ValueError("DFlashLagunaForCausalLM requires dflash_config.causal=true") + if self.hparams.get("gating") != "per-head": + raise ValueError("DFlashLagunaForCausalLM currently requires gating='per-head'") + target_hidden_size = self._get_target_hidden_size() + draft_hidden_size = int(self.hparams["hidden_size"]) + if target_hidden_size != draft_hidden_size: + raise ValueError( + "DFlashLagunaForCausalLM requires matching target and draft hidden sizes, " + f"got target={target_hidden_size} and draft={draft_hidden_size}" + ) + layer_types = self.hparams.get("layer_types") + if not isinstance(layer_types, list) or len(layer_types) != self.block_count: + raise ValueError( + "DFlashLagunaForCausalLM requires one layer_types entry per draft layer" + ) + if any(str(layer_type) != "sliding_attention" for layer_type in layer_types): + raise ValueError( + "DFlashLagunaForCausalLM currently requires every draft layer to use sliding_attention" + ) + if not self.hparams.get("sliding_window"): + raise ValueError("DFlashLagunaForCausalLM requires sliding_window metadata") + + self.hparams["use_sliding_window"] = True + super().set_gguf_parameters() + self.gguf_writer.add_bool(f"{self.gguf_writer.arch}.dflash.laguna", True) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + top_level_name = name[6:] if name.startswith("model.") else name + hidden_size = int(self.hparams["hidden_size"]) + + if top_level_name.startswith("aux_hidden_norms.") and top_level_name.endswith(".weight"): + parts = top_level_name.split(".") + if len(parts) != 3 or not parts[1].isdigit(): + raise ValueError(f"DFlashLagunaForCausalLM: invalid auxiliary norm name {name!r}") + aux_id = int(parts[1]) + if data_torch.ndim != 1 or data_torch.shape[0] != hidden_size: + raise ValueError( + f"DFlashLagunaForCausalLM: auxiliary norm {name!r} has shape " + f"{tuple(data_torch.shape)}, expected [{hidden_size}]" + ) + self._laguna_aux_norm_ids.add(aux_id) + tensor_name = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.DFLASH_AUX_HIDDEN_NORM].format(bid=aux_id) + return [(f"{tensor_name}.weight", data_torch)] + + if top_level_name.endswith(".self_attn.qkv_proj.weight"): + if bid is None: + raise ValueError(f"DFlashLagunaForCausalLM: can not infer block id for tensor {name!r}") + n_head = int(self.hparams["num_attention_heads"]) + n_head_kv = int(self.hparams["num_key_value_heads"]) + head_dim = int(self.hparams.get("head_dim", self.hparams["hidden_size"] // n_head)) + q_width = n_head * head_dim + k_width = n_head_kv * head_dim + v_width = n_head_kv * head_dim + expected_width = q_width + k_width + v_width + if data_torch.ndim != 2 or data_torch.shape != (expected_width, hidden_size): + raise ValueError( + f"DFlashLagunaForCausalLM: packed QKV tensor {name!r} has shape " + f"{tuple(data_torch.shape)}, expected [{expected_width}, {hidden_size}]" + ) + q_weight, k_weight, v_weight = data_torch.split([q_width, k_width, v_width], dim=0) + self._laguna_qkv_ids.add(bid) + result: list[tuple[str, Tensor]] = [] + for suffix, weight in (("q_proj", q_weight), ("k_proj", k_weight), ("v_proj", v_weight)): + split_name = name.replace("qkv_proj", suffix) + result.extend(super().modify_tensors(weight, split_name, bid)) + return result + + if top_level_name.endswith(".self_attn.g_proj.weight"): + if bid is None: + raise ValueError(f"DFlashLagunaForCausalLM: can not infer block id for tensor {name!r}") + gate = data_torch.squeeze().contiguous() + n_head = int(self.hparams["num_attention_heads"]) + if gate.ndim != 2 or gate.shape != (n_head, hidden_size): + raise ValueError( + f"DFlashLagunaForCausalLM: attention gate {name!r} has shape " + f"{tuple(gate.shape)}, expected [{n_head}, {hidden_size}]" + ) + self._laguna_gate_ids.add(bid) + tensor_name = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.ATTN_GATE].format(bid=bid) + return [(f"{tensor_name}.weight", gate)] + + return super().modify_tensors(data_torch, name, bid) + + def prepare_tensors(self): + super().prepare_tensors() + + expected_layers = set(range(self.block_count)) + dflash_cfg = self.hparams.get("dflash_config") + if not isinstance(dflash_cfg, dict): + raise ValueError("DFlashLagunaForCausalLM requires dflash_config metadata") + target_layer_ids = dflash_cfg.get("target_layer_ids", []) + if not isinstance(target_layer_ids, list) or not target_layer_ids: + raise ValueError("DFlashLagunaForCausalLM requires non-empty target_layer_ids metadata") + expected_aux = set(range(len(target_layer_ids))) + if self._laguna_qkv_ids != expected_layers: + raise ValueError( + f"DFlashLagunaForCausalLM: packed QKV layers {sorted(self._laguna_qkv_ids)} " + f"do not match expected {sorted(expected_layers)}" + ) + if self._laguna_gate_ids != expected_layers: + raise ValueError( + f"DFlashLagunaForCausalLM: attention gate layers {sorted(self._laguna_gate_ids)} " + f"do not match expected {sorted(expected_layers)}" + ) + if self._laguna_aux_norm_ids != expected_aux: + raise ValueError( + f"DFlashLagunaForCausalLM: auxiliary norm ids {sorted(self._laguna_aux_norm_ids)} " + f"do not match expected {sorted(expected_aux)}" + ) + + @Model.register("MellumForCausalLM") class MellumModel(Model): model_arch = gguf.MODEL_ARCH.MELLUM diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 413403b5..5cea96dd 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -398,6 +398,7 @@ class MODEL_TENSOR(IntEnum): MTP_CENTROIDS = auto() DFLASH_FC = auto() DFLASH_HIDDEN_NORM = auto() + DFLASH_AUX_HIDDEN_NORM = auto() # openPangu-2.0 (DSA lightning indexer) INDEXER_K_NORM = auto() INDEXER_PROJ = auto() # weights_proj @@ -608,6 +609,7 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = { MODEL_TENSOR.MTP_CENTROIDS: "mtp_centroids", MODEL_TENSOR.DFLASH_FC: "dflash_fc", MODEL_TENSOR.DFLASH_HIDDEN_NORM: "dflash_hidden_norm", + MODEL_TENSOR.DFLASH_AUX_HIDDEN_NORM: "dflash_aux_hidden_norm.{bid}", # openPangu-2.0 MODEL_TENSOR.INDEXER_K_NORM: "blk.{bid}.attn_indexer_k_norm", MODEL_TENSOR.INDEXER_PROJ: "blk.{bid}.attn_indexer_weights_proj", @@ -1498,6 +1500,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_K_NORM, MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_GATE, MODEL_TENSOR.ATTN_SINKS, MODEL_TENSOR.ATTN_OUT, MODEL_TENSOR.ATTN_POST_NORM, @@ -1506,6 +1509,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.FFN_UP, MODEL_TENSOR.DFLASH_FC, MODEL_TENSOR.DFLASH_HIDDEN_NORM, + MODEL_TENSOR.DFLASH_AUX_HIDDEN_NORM, ], MODEL_ARCH.BITNET: [ MODEL_TENSOR.ATTN_Q, diff --git a/src/graphs/build_dflash.cpp b/src/graphs/build_dflash.cpp index 38df86a7..66d9e611 100644 --- a/src/graphs/build_dflash.cpp +++ b/src/graphs/build_dflash.cpp @@ -30,7 +30,39 @@ ggml_cgraph * llm_build_context::build_dflash_kv_cache() { ggml_set_input(lctx.dflash.kv.cache_input_pos_ctx); cb(lctx.dflash.kv.cache_input_pos_ctx, "dflash_kv_input_pos_ctx", -1); - ggml_tensor * fused_target = llm_build_lora_mm(lctx, ctx0, model.dflash_fc, lctx.dflash.kv.cache_input_target_features); + ggml_tensor * target_features = lctx.dflash.kv.cache_input_target_features; + if (hparams.dflash_laguna) { + GGML_ASSERT(model.dflash_aux_hidden_norms.size() == hparams.dflash_n_target_layers); + const int64_t slice_width = n_target_features / hparams.dflash_n_target_layers; + ggml_tensor * normalized_features = nullptr; + for (uint32_t i = 0; i < hparams.dflash_n_target_layers; ++i) { + ggml_tensor * slice = ggml_view_2d( + ctx0, + target_features, + slice_width, + update_rows, + target_features->nb[1], + i * slice_width * target_features->nb[0]); + slice = llm_build_norm( + ctx0, + slice, + hparams, + model.dflash_aux_hidden_norms[i], + nullptr, + LLM_NORM_RMS, + cb, + -1); + cb(slice, "dflash_kv_aux_norm", (int) i); + normalized_features = normalized_features == nullptr + ? slice + : ggml_concat(ctx0, normalized_features, slice, 0); + } + GGML_ASSERT(normalized_features != nullptr); + target_features = normalized_features; + cb(target_features, "dflash_kv_normalized_target_features", -1); + } + + ggml_tensor * fused_target = llm_build_lora_mm(lctx, ctx0, model.dflash_fc, target_features); fused_target = llm_build_norm(ctx0, fused_target, hparams, model.dflash_hidden_norm, nullptr, LLM_NORM_RMS, cb, -1); cb(fused_target, "dflash_kv_fused_target", -1); @@ -38,7 +70,21 @@ ggml_cgraph * llm_build_context::build_dflash_kv_cache() { GGML_ASSERT(il < (int32_t) lctx.dflash.kv.k_ctx_cache.size()); GGML_ASSERT(il < (int32_t) lctx.dflash.kv.v_ctx_cache.size()); - ggml_tensor * Kcur_ctx_proj = llm_build_lora_mm(lctx, ctx0, model.layers[il].wk, fused_target); + ggml_tensor * layer_target = fused_target; + if (hparams.dflash_laguna) { + layer_target = llm_build_norm( + ctx0, + layer_target, + hparams, + model.layers[il].attn_norm, + nullptr, + LLM_NORM_RMS, + cb, + il); + cb(layer_target, "dflash_kv_attn_norm", il); + } + + ggml_tensor * Kcur_ctx_proj = llm_build_lora_mm(lctx, ctx0, model.layers[il].wk, layer_target); if (model.layers[il].bk) { Kcur_ctx_proj = ggml_add(ctx0, Kcur_ctx_proj, model.layers[il].bk); } cb(Kcur_ctx_proj, "dflash_kv_k_proj", il); @@ -54,7 +100,7 @@ ggml_cgraph * llm_build_context::build_dflash_kv_cache() { Kcur_ctx = ggml_cont(ctx0, ggml_permute(ctx0, Kcur_ctx, 0, 2, 1, 3)); cb(Kcur_ctx, "dflash_kv_k_physical", il); - ggml_tensor * Vcur_ctx = llm_build_lora_mm(lctx, ctx0, model.layers[il].wv, fused_target); + ggml_tensor * Vcur_ctx = llm_build_lora_mm(lctx, ctx0, model.layers[il].wv, layer_target); if (model.layers[il].bv) { Vcur_ctx = ggml_add(ctx0, Vcur_ctx, model.layers[il].bv); } cb(Vcur_ctx, "dflash_kv_v_proj", il); if (std::abs(hparams.f_attn_v_scale - 1.0f) > 1e-4f) { @@ -232,6 +278,7 @@ ggml_cgraph * llm_build_context::build_dflash() { ggml_tensor * cur = llm_build_norm(ctx0, inpL, hparams, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, cb, il); cb(cur, "attn_norm", il); + ggml_tensor * input_normed = cur; ggml_tensor * Qcur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wq, cur); ggml_tensor * Kcur_noise = llm_build_lora_mm(lctx, ctx0, model.layers[il].wk, cur); @@ -339,6 +386,17 @@ ggml_cgraph * llm_build_context::build_dflash() { // cur->op_params[4] = hparams.n_swa; //} + if (hparams.dflash_laguna) { + GGML_ASSERT(model.layers[il].wqkv_gate != nullptr); + ggml_tensor * gate = llm_build_lora_mm(lctx, ctx0, model.layers[il].wqkv_gate, input_normed); + gate = ggml_softplus(ctx0, gate); + cb(gate, "attn_gate", il); + GGML_ASSERT(gate->ne[0] == n_head); + gate = ggml_reshape_3d(ctx0, gate, 1, n_head, n_tokens); + cur = ggml_mul(ctx0, cur, gate); + cb(cur, "attn_gated", il); + } + cur = ggml_reshape_2d(ctx0, cur, model.layers[il].wo->ne[0], n_tokens); cb(cur, "flash_attn_reshaped", il); diff --git a/src/graphs/build_laguna.cpp b/src/graphs/build_laguna.cpp index 55d60902..02224c41 100644 --- a/src/graphs/build_laguna.cpp +++ b/src/graphs/build_laguna.cpp @@ -8,6 +8,7 @@ ggml_cgraph * llm_build_context::build_laguna() { ggml_tensor * inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb); ggml_tensor * inp_pos = build_inp_pos(); ggml_tensor * inp_out_ids = n_tokens > 1 ? build_inp_out_ids() : nullptr; + const bool needs_dflash_full_final_rows = lctx.dflash.capture != nullptr; ggml_tensor * KQ_mask = build_inp_KQ_mask(); // Laguna M.1 has only global-attention layers and leaves n_swa at zero; building // the SWA mask in that case trips the generic SWA precondition. @@ -23,8 +24,10 @@ ggml_cgraph * llm_build_context::build_laguna() { GGML_ASSERT(KQ_mask_l != nullptr); auto rope_factors = is_swa ? nullptr : build_rope_factors(il); + const bool is_final_layer = il == n_layer - 1; + ggml_tensor * attn_out_ids = is_final_layer && !needs_dflash_full_final_rows ? inp_out_ids : nullptr; auto cur = build_std_attention(gf, model.layers[il].attn_norm, inpL, - inp_pos, il == n_layer - 1 ? inp_out_ids : nullptr, rope_factors, + inp_pos, attn_out_ids, rope_factors, KQ_mask_l, nullptr, nullptr, 1.0f / sqrtf(float(n_embd_head_k)), 0.0f, n_swa_l, il, true, false, true); if (model.layers[il].ffn_gate_inp == nullptr) { @@ -53,6 +56,11 @@ ggml_cgraph * llm_build_context::build_laguna() { cur = lctx.cvec.apply_to(ctx0, cur, il); cb(cur, "l_out", il); + if (is_final_layer && needs_dflash_full_final_rows && inp_out_ids != nullptr) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + cb(cur, "l_out_selected", il); + } + inpL = cur; } diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index bf07387f..fcde0a95 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -168,6 +168,7 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_DFLASH_TARGET_LAYER_IDS, "%s.dflash.target_layer_ids" }, { LLM_KV_DFLASH_N_TARGET_FEATURES, "%s.dflash.n_target_features" }, { LLM_KV_DFLASH_BACKBONE_ROTARY_BASE, "%s.dflash.backbone_rotary_base" }, + { LLM_KV_DFLASH_LAGUNA, "%s.dflash.laguna" }, { LLM_KV_ATTENTION_HEAD_COUNT, "%s.attention.head_count" }, { LLM_KV_ATTENTION_HEAD_COUNT_KV, "%s.attention.head_count_kv" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index 211a3052..bdd9e2a3 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -151,6 +151,7 @@ enum llm_kv { LLM_KV_DFLASH_TARGET_LAYER_IDS, LLM_KV_DFLASH_N_TARGET_FEATURES, LLM_KV_DFLASH_BACKBONE_ROTARY_BASE, + LLM_KV_DFLASH_LAGUNA, LLM_KV_ATTENTION_HEAD_COUNT, LLM_KV_ATTENTION_HEAD_COUNT_KV, @@ -418,6 +419,7 @@ enum llm_tensor { LLM_TENSOR_MTP_CENTROIDS, LLM_TENSOR_DFLASH_FC, LLM_TENSOR_DFLASH_HIDDEN_NORM, + LLM_TENSOR_DFLASH_AUX_HIDDEN_NORM, // openPangu-2.0 LLM_TENSOR_ATTN_QA_CONV, // MoME causal conv on q-lora latent diff --git a/src/llama-context.h b/src/llama-context.h index 676e09d4..22ef1c33 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -362,10 +362,14 @@ struct llama_context { struct capture_state { std::vector layer_ids; std::vector> layer_rows; + std::vector layer_rows_written; int32_t row_count = 0; int32_t row_width = 0; + int32_t expected_rows = 0; uint64_t capture_batch_id = 0; std::vector layer_seen_batch_id; + bool readback_pending = false; + bool invalid = false; ggml_backend_sched_eval_callback prev_cb_eval = nullptr; void * prev_cb_eval_user_data = nullptr; }; diff --git a/src/llama-dflash.cpp b/src/llama-dflash.cpp index 38c024c5..b5f51bf1 100644 --- a/src/llama-dflash.cpp +++ b/src/llama-dflash.cpp @@ -230,6 +230,30 @@ static bool validate_dflash_graph_contract(const llama_context & lctx) { const auto & model = lctx.model; const auto & hparams = model.hparams; + if (hparams.dflash_laguna) { + if (!hparams.causal_attn || hparams.n_swa == 0) { + LLAMA_LOG_ERROR("%s: Laguna DFlash requires causal SWA metadata\n", __func__); + return false; + } + if (hparams.dflash_n_target_layers == 0 || + hparams.dflash_n_target_features % hparams.dflash_n_target_layers != 0 || + model.dflash_aux_hidden_norms.size() != hparams.dflash_n_target_layers) { + LLAMA_LOG_ERROR("%s: Laguna DFlash requires one auxiliary norm per target feature slice\n", __func__); + return false; + } + const int64_t aux_width = hparams.dflash_n_target_features / hparams.dflash_n_target_layers; + for (uint32_t i = 0; i < hparams.dflash_n_target_layers; ++i) { + const ggml_tensor * aux_norm = model.dflash_aux_hidden_norms[i]; + if (aux_norm == nullptr || aux_norm->ne[0] != aux_width) { + LLAMA_LOG_ERROR("%s: Laguna DFlash auxiliary norm %u has invalid width\n", __func__, i); + return false; + } + } + } else if (!model.dflash_aux_hidden_norms.empty()) { + LLAMA_LOG_ERROR("%s: generic DFlash must not carry Laguna auxiliary norms\n", __func__); + return false; + } + auto rope_dim_for_layer = [&hparams](int32_t il) -> uint32_t { if (hparams.rope_dim_per_layer[il] != 0) { return hparams.rope_dim_per_layer[il]; @@ -273,6 +297,20 @@ static bool validate_dflash_graph_contract(const llama_context & lctx) { return false; } + if (hparams.dflash_laguna) { + const ggml_tensor * gate = model.layers[il].wqkv_gate; + if (!hparams.swa_layers[il] || gate == nullptr || + gate->ne[0] != (int64_t) hparams.n_embd || + gate->ne[1] != (int64_t) hparams.n_head((uint32_t) il)) { + LLAMA_LOG_ERROR("%s: Laguna DFlash layer %d requires SWA and a head-wise attention gate\n", + __func__, il); + return false; + } + } else if (model.layers[il].wqkv_gate != nullptr) { + LLAMA_LOG_ERROR("%s: generic DFlash layer %d has an unsupported attention gate\n", __func__, il); + return false; + } + const uint32_t rope_dim = rope_dim_for_layer(il); const float rope_base = rope_base_for_layer(il); const float rope_scale = rope_scale_for_layer(il); diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index d5bad468..6ceb548a 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -907,6 +907,10 @@ void llm_load_hparams( ml.get_key(LLM_KV_DFLASH_BACKBONE_ROTARY_BASE, hparams.dflash_backbone_rotary_base, false); load_dflash_target_layer_ids(ml, LLM_KV(model.arch)(LLM_KV_DFLASH_TARGET_LAYER_IDS), hparams, false); ml.get_key(LLM_KV_ATTENTION_VALUE_SCALE, hparams.f_attn_v_scale, false); + ml.get_key(LLM_KV_DFLASH_LAGUNA, hparams.dflash_laguna, false); + if (hparams.dflash_laguna) { + ml.get_key(LLM_KV_ATTENTION_CAUSAL, hparams.causal_attn); + } // DFlash drafts may be trained with sliding-window attention (for long-context). // Read the window + per-layer pattern so the SWA mask path activates; absent keys // leave n_swa=0 / swa_layers all-zero (dense behavior, unchanged). diff --git a/src/llama-hparams.h b/src/llama-hparams.h index f21272dd..ec8daf2d 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -171,6 +171,7 @@ struct llama_hparams { uint32_t dflash_n_target_layers = 0; uint32_t dflash_target_layer_ids[8] = {}; float dflash_backbone_rotary_base = 0.0f; + bool dflash_laguna = false; // needed by encoder-decoder models (e.g. T5, FLAN-T5) // ref: https://github.com/ggerganov/llama.cpp/pull/8141 @@ -195,6 +196,7 @@ struct llama_hparams { if (this->dflash_mask_token_id != other.dflash_mask_token_id) return true; if (this->dflash_n_target_features != other.dflash_n_target_features) return true; if (this->dflash_n_target_layers != other.dflash_n_target_layers) return true; + if (this->dflash_laguna != other.dflash_laguna) return true; if (this->n_layer != other.n_layer) return true; if (this->n_rot != other.n_rot) return true; if (this->n_swa != other.n_swa) return true; diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index 3ece0f96..a75642b0 100644 --- a/src/llama-load-tensors.cpp +++ b/src/llama-load-tensors.cpp @@ -2302,6 +2302,20 @@ bool create_tensors_helper::create_dflash_tensors(const LLM_TN & tn) { model.output_mtp = model.output; model.dflash_fc = create_tensor(ctx_output, tn(LLM_TENSOR_DFLASH_FC, "weight"), {(int64_t) hparams.dflash_n_target_features, n_embd}, 0); model.dflash_hidden_norm = create_tensor(ctx_output, tn(LLM_TENSOR_DFLASH_HIDDEN_NORM, "weight"), {n_embd}, 0); + model.dflash_aux_hidden_norms.clear(); + if (hparams.dflash_laguna) { + GGML_ASSERT(hparams.dflash_n_target_layers > 0); + GGML_ASSERT(hparams.dflash_n_target_features % hparams.dflash_n_target_layers == 0); + const int64_t aux_width = hparams.dflash_n_target_features / hparams.dflash_n_target_layers; + model.dflash_aux_hidden_norms.reserve(hparams.dflash_n_target_layers); + for (uint32_t i = 0; i < hparams.dflash_n_target_layers; ++i) { + model.dflash_aux_hidden_norms.push_back(create_tensor( + ctx_output, + tn(LLM_TENSOR_DFLASH_AUX_HIDDEN_NORM, "weight", i), + {aux_width}, + 0)); + } + } for (int i = 0; i < n_layer; ++i) { ggml_context * ctx_split = use_split_ctx ? ctx_for_layer_split(i) : ctx_for_layer(i); @@ -2314,6 +2328,9 @@ bool create_tensors_helper::create_dflash_tensors(const LLM_TN & tn) { layer.wk = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_gqa}, 0); layer.wv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_gqa}, 0); layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_v * n_head, n_embd}, 0); + if (hparams.dflash_laguna) { + layer.wqkv_gate = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_head}, 0); + } layer.attn_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0); layer.attn_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index eab0e9b9..21d7b265 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -857,6 +857,7 @@ static const std::map> LLM_TENSOR_NA { LLM_TENSOR_ATTN_K, "blk.%d.attn_k" }, { LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" }, { LLM_TENSOR_ATTN_V, "blk.%d.attn_v" }, + { LLM_TENSOR_ATTN_GATE, "blk.%d.attn_gate" }, { LLM_TENSOR_ATTN_SINKS, "blk.%d.attn_sinks" }, { LLM_TENSOR_ATTN_OUT, "blk.%d.attn_output" }, { LLM_TENSOR_ATTN_POST_NORM, "blk.%d.post_attention_norm" }, @@ -865,6 +866,7 @@ static const std::map> LLM_TENSOR_NA { LLM_TENSOR_FFN_UP, "blk.%d.ffn_up" }, { LLM_TENSOR_DFLASH_FC, "dflash_fc" }, { LLM_TENSOR_DFLASH_HIDDEN_NORM, "dflash_hidden_norm" }, + { LLM_TENSOR_DFLASH_AUX_HIDDEN_NORM, "dflash_aux_hidden_norm.%d" }, }, }, { diff --git a/src/llama-model.h b/src/llama-model.h index f3f05bbf..ae1f06e9 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -473,6 +473,7 @@ struct llama_model { struct ggml_tensor * mtp_centroids = nullptr; struct ggml_tensor * dflash_fc = nullptr; struct ggml_tensor * dflash_hidden_norm = nullptr; + std::vector dflash_aux_hidden_norms; struct ggml_tensor * output_norm; struct ggml_tensor * output_norm_b; diff --git a/src/llama-spec-features-dflash.cpp b/src/llama-spec-features-dflash.cpp index 8424b0b5..01ad482c 100644 --- a/src/llama-spec-features-dflash.cpp +++ b/src/llama-spec-features-dflash.cpp @@ -403,15 +403,41 @@ static int llama_dflash_capture_eval_callback(struct ggml_tensor * tensor, bool if (capture.layer_seen_batch_id.size() != capture.layer_ids.size()) { capture.layer_seen_batch_id.assign(capture.layer_ids.size(), 0); } + if (capture.layer_rows_written.size() != capture.layer_ids.size() || capture.expected_rows <= 0) { + capture.invalid = true; + LLAMA_LOG_WARN("%s: DFlash capture state is not initialized (expected_rows=%d layers=%zu written=%zu)\n", + __func__, capture.expected_rows, capture.layer_ids.size(), capture.layer_rows_written.size()); + return 2; + } + if (capture.row_width != 0 && capture.row_width != row_width) { + capture.invalid = true; + LLAMA_LOG_WARN("%s: DFlash capture width mismatch for layer %d: got=%d expected=%d\n", + __func__, layer_id, row_width, capture.row_width); + return 2; + } auto & rows = capture.layer_rows[(size_t) layer_idx]; - rows.resize((size_t) row_count * (size_t) row_width); + auto & rows_written = capture.layer_rows_written[(size_t) layer_idx]; + if (rows_written + row_count > capture.expected_rows) { + capture.invalid = true; + LLAMA_LOG_WARN("%s: DFlash capture rows overflow for layer %d: written=%d add=%d expected=%d\n", + __func__, layer_id, rows_written, row_count, capture.expected_rows); + return 2; + } + const size_t expected_floats = (size_t) capture.expected_rows * (size_t) row_width; + if (rows.size() != expected_floats) { + rows.resize(expected_floats); + } auto backend = ggml_backend_sched_get_tensor_backend(ctx->sched, tensor); GGML_ASSERT(backend); - ggml_backend_tensor_get_async(backend, tensor, rows.data(), 0, ggml_nbytes(tensor)); + ggml_backend_tensor_get_async(backend, tensor, + rows.data() + (size_t) rows_written * (size_t) row_width, + 0, (size_t) row_count * (size_t) row_width * sizeof(float)); + rows_written += row_count; capture.row_width = row_width; - capture.row_count = row_count; + capture.row_count = std::max(capture.row_count, rows_written); capture.layer_seen_batch_id[(size_t) layer_idx] = capture.capture_batch_id; + capture.readback_pending = true; return 2; } @@ -423,9 +449,14 @@ bool llama_set_dflash_capture_layers( return false; } + if (ctx->dflash.capture && ctx->dflash.capture->readback_pending) { + llama_synchronize(ctx); + } + auto capture = std::make_unique(); capture->layer_ids.assign(layer_ids, layer_ids + n_layers); capture->layer_rows.resize((size_t) n_layers); + capture->layer_rows_written.assign((size_t) n_layers, 0); capture->layer_seen_batch_id.assign((size_t) n_layers, 0); capture->prev_cb_eval = ctx->cparams.cb_eval; capture->prev_cb_eval_user_data = ctx->cparams.cb_eval_user_data; @@ -449,6 +480,9 @@ void llama_clear_dflash_capture(struct llama_context * ctx) { ggml_backend_sched_eval_callback prev_cb_eval = nullptr; void * prev_cb_eval_user_data = nullptr; if (ctx->dflash.capture) { + if (ctx->dflash.capture->readback_pending) { + llama_synchronize(ctx); + } prev_cb_eval = ctx->dflash.capture->prev_cb_eval; prev_cb_eval_user_data = ctx->dflash.capture->prev_cb_eval_user_data; } @@ -465,15 +499,22 @@ void llama_clear_dflash_capture(struct llama_context * ctx) { } } -void llama_begin_dflash_capture_batch(struct llama_context * ctx) { +void llama_begin_dflash_capture_batch(struct llama_context * ctx, int32_t expected_rows) { if (ctx == nullptr || !ctx->dflash.capture) { return; } auto & capture = *ctx->dflash.capture; + if (capture.readback_pending) { + llama_synchronize(ctx); + capture.readback_pending = false; + } capture.capture_batch_id++; capture.row_count = 0; capture.row_width = 0; + capture.expected_rows = expected_rows; + capture.invalid = expected_rows <= 0; + std::fill(capture.layer_rows_written.begin(), capture.layer_rows_written.end(), 0); std::fill(capture.layer_seen_batch_id.begin(), capture.layer_seen_batch_id.end(), 0); } @@ -504,10 +545,18 @@ static bool llama_spec_prepare_dflash_capture( llama_synchronize(ctx); auto & capture = *ctx->dflash.capture; + capture.readback_pending = false; row_count = capture.row_count; row_width = capture.row_width; n_layers = (int32_t) capture.layer_ids.size(); - if (row_count <= 0 || row_width <= 0 || n_layers <= 0 || capture.layer_rows.size() != (size_t) n_layers) { + if (capture.invalid || row_count <= 0 || row_width <= 0 || n_layers <= 0 || + capture.expected_rows <= 0 || capture.layer_rows.size() != (size_t) n_layers || + capture.layer_rows_written.size() != (size_t) n_layers) { + return false; + } + if (row_count != capture.expected_rows) { + LLAMA_LOG_WARN("%s: DFlash capture rows incomplete: got=%d expected=%d\n", + __func__, row_count, capture.expected_rows); return false; } @@ -533,9 +582,11 @@ static bool llama_spec_prepare_dflash_capture( } const auto & rows = capture.layer_rows[(size_t) layer_idx]; - if (rows.size() != (size_t) row_count * (size_t) row_width) { - LLAMA_LOG_WARN("%s: DFlash capture rows mismatch for layer %d: got=%zu expected=%zu (rows=%d width=%d)\n", - __func__, capture.layer_ids[(size_t) layer_idx], rows.size(), + if (capture.layer_rows_written[(size_t) layer_idx] != row_count || + rows.size() != (size_t) row_count * (size_t) row_width) { + LLAMA_LOG_WARN("%s: DFlash capture rows mismatch for layer %d: got=%d/%zu expected=%d/%zu (rows=%d width=%d)\n", + __func__, capture.layer_ids[(size_t) layer_idx], + capture.layer_rows_written[(size_t) layer_idx], rows.size(), row_count, (size_t) row_count * (size_t) row_width, row_count, row_width); return false; } diff --git a/src/llama-spec-features-dflash.h b/src/llama-spec-features-dflash.h index cec99c0b..f0a1fde6 100644 --- a/src/llama-spec-features-dflash.h +++ b/src/llama-spec-features-dflash.h @@ -115,7 +115,7 @@ bool llama_set_dflash_target_features_view( bool llama_set_dflash_capture_layers(struct llama_context * ctx, const int32_t * layer_ids, int32_t n_layers); void llama_clear_dflash_capture(struct llama_context * ctx); -void llama_begin_dflash_capture_batch(struct llama_context * ctx); +void llama_begin_dflash_capture_batch(struct llama_context * ctx, int32_t expected_rows); void llama_finish_dflash_capture_batch(struct llama_context * ctx, bool is_prompt_warmup); bool llama_spec_get_dflash_feature_view( diff --git a/src/llama.cpp b/src/llama.cpp index 657dfdea..23b57cac 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -3733,7 +3733,8 @@ static std::pair, double> get_layer_sizes(const llama_model_ name == "rope_freqs.weight") { continue; } - if (name == "dflash_fc.weight" || name == "dflash_hidden_norm.weight") { + if (name == "dflash_fc.weight" || name == "dflash_hidden_norm.weight" || + name.rfind("dflash_aux_hidden_norm.", 0) == 0) { output_misc_size += size; continue; } @@ -5851,7 +5852,7 @@ static int llama_decode_internal( const auto & hparams = model.hparams; const auto & cparams = lctx.cparams; - llama_begin_dflash_capture_batch(&lctx); + llama_begin_dflash_capture_batch(&lctx, (int32_t) n_tokens_all); GGML_ASSERT((!batch_all.token && batch_all.embd) || (batch_all.token && !batch_all.embd)); // NOLINT