DFlash: add Laguna XS 2.1 support

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 <matching target model directory>")
@@ -2426,7 +2429,7 @@ def set_vocab(self):
     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 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter
         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

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 @@ class MODEL_TENSOR(IntEnum):
     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 @@ class MODEL_TENSOR(IntEnum):
         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 @@ class MODEL_TENSOR(IntEnum):
         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,

src/graphs/build_dflash.cpp

@@ -30,15 +30,61 @@ 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);
 
     for (int il = 0; il < n_layer; ++il) {
         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);
 

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;
     }
 

src/llama-arch.cpp

@@ -168,6 +168,7 @@ static const std::map<llm_kv, const char *> 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"          },

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

src/llama-context.h

@@ -362,10 +362,14 @@ struct llama_context {
         struct capture_state {
             std::vector<int32_t> layer_ids;
             std::vector<std::vector<float>> layer_rows;
+            std::vector<int32_t> 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<uint64_t> 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;
         };