Skip to content

Commit f10be44

Browse files
linuxid10tclaude
andcommitted
laguna: generalize converter + model to Laguna-M.1 (all full attention)
Laguna-M.1 is the all-full-attention sibling of Laguna-XS.2: 70 layers, no SWA layers, per-element attention gate, full rotary (~226B). The XS.2-only converter/model could not load it (empty layer_types arrays, per-head-only gate with a shape mismatch, hardcoded half-rotary). - conversion/laguna.py: default layer_types to all full_attention when absent; write data-driven rope.dimension_count[_swa] from partial_rotary_factor; write the new attention.gate_per_head KV. - laguna.cpp: drop the hardcoded n_rot_full /= 2; support per-element gating (shape + element-wise application); set swa_type from is_swa_any() and branch the attention input path (build_attn_inp_kv when there are no SWA layers, required by create_memory's swa_type <=>> is_swa_any() assert). - New laguna.attention.gate_per_head KV across gguf-py + llama-arch.h/.cpp + llama-hparams.h (default true = per-head, keeps old XS.2 GGUFs working). Verified: test-llama-archs laguna still OK (XS.2 mixed-SWA path); a 420 GiB f16 Laguna-M.1 GGUF converts cleanly, has correct metadata (gate_per_head= false, dimension_count=128, gate tensor 4096x8192), and loads + runs. Pre-existing XS.2 GGUFs now need reconversion (they predate the required rope.dimension_count). Co-Authored-By: Claude <noreply@anthropic.com>
1 parent 9404cf2 commit f10be44

8 files changed

Lines changed: 97 additions & 22 deletions

File tree

CLAUDE.md

Lines changed: 26 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ IMPORTANT: Ensure you've thoroughly reviewed the [AGENTS.md](AGENTS.md) file bef
44

55
# Task: Add Laguna (Poolside) Support to llama.cpp
66

7-
**Goal:** Implement `LLM_ARCH_LAGUNA` targeting `poolside/Laguna-XS.2` — a 33B-A3B MoE model with 40 layers, 256 routed + 1 shared expert, head-wise sigmoid attention gate, per-layer head counts, dual RoPE (YaRN for global, default for SWA), and per-layer-type partial rotary.
7+
**Goal:** Implement `LLM_ARCH_LAGUNA` covering the Laguna family. Initially targeted `poolside/Laguna-XS.2` (33B-A3B: 40 layers, 256 routed + 1 shared expert, head-wise softplus attention gate, per-layer head counts, dual RoPE (YaRN for global, default for SWA), per-layer-type partial rotary). Later generalized to also support **`poolside/Laguna-M.1`** (~226B: 70 layers, all full attention, per-element attention gate, full rotary) — see "Laguna-M.1 generalization" below.
88

99
**Base commit:** `b4c0549` (2026-05-27)
1010

@@ -23,6 +23,7 @@ IMPORTANT: Ensure you've thoroughly reviewed the [AGENTS.md](AGENTS.md) file bef
2323
- **Step 1 (bonus) — Model saver**: Added laguna to unsupported arch list in `llama-model-saver.cpp` (same as step35) since roundtrip serialization loses SWA pattern keys.
2424
- **Step 6 — SWA + long context**: 10K prompt (SWA) clean; 35K needle-in-haystack ("SWORDFISH99") correctly retrieved, confirming YaRN global layers work.
2525
- **Chat template + streaming parser**: `models/templates/poolside-Laguna.jinja` added; name registration and auto-detection via `laguna_glm_thinking_v5` marker in `src/llama-chat.cpp`/`.h`. Autoparser workaround in `common/chat-diff-analyzer.cpp` fixes two streaming bugs (see learnings 9–10).
26+
- **Laguna-M.1 generalization**: Extended converter + `laguna.cpp` from XS.2-only to the all-full-attention sibling `poolside/Laguna-M.1` (~226B). (1) `conversion/laguna.py` defaults `layer_types` to all-`full_attention` when absent and writes data-driven `rope.dimension_count`/`rope.dimension_count_swa` plus a new `attention.gate_per_head` KV. (2) `laguna.cpp` removed the hardcoded `n_rot_full /= 2`, supports per-element gating, sets `swa_type` from `is_swa_any()`, and branches the attention input path. A 420 GiB f16 GGUF was produced and verified to load + run; XS.2 regression (`test-llama-archs laguna`) still `OK`.
2627

2728
### Remaining
2829

@@ -40,7 +41,10 @@ IMPORTANT: Ensure you've thoroughly reviewed the [AGENTS.md](AGENTS.md) file bef
4041
7. **EOT token from `eos_token_id` list**: `config.json` has `eos_token_id: [2, 24]`. SpecialVocab only registers index 0 as EOS; token 24 (`</assistant>`) must be written as EOT via `add_eot_token_id` so it gets added to `special_eog_ids`. Without this the model loops past turn boundaries.
4142
8. **ABI mismatch after cparams change**: Adding fields to `llama_cparams` (Step 2) requires a full rebuild of all binaries that link against `libllama-common`. A partial rebuild will crash with unexpected behavior.
4243
9. **Prefilled thinking token breaks autoparser**: Laguna prefills `<think>` as the generation prompt prefix (thinking=on). Auto-detection finds `<think>` as `reasoning.start` from template diffs, but it never appears in the generated stream — the parser never enters reasoning mode. Fix: set `reasoning.start = ""` (delimiter-style) in the workaround. The `analyze_reasoning::build_parser` was extended to return `p.eps()` when `start.empty() && !ctx.inputs.enable_thinking`, preventing the streaming parser from misclassifying all plain content as reasoning when thinking is disabled (`--reasoning off`).
43-
10. **EOT token as regular vocab token**: `</assistant>` (token 24) is a regular vocabulary token, not a special token. `preserved_tokens` only controls special-token decode and doesn't suppress it. Use `additional_stops` instead — this feeds into `task->params.antiprompt` which the `process_token` stop-word erase logic checks, stripping the text before it's sent. Added `additional_stops` field to `autoparser` struct; apply in `cli.cpp` as `task.params.antiprompt` additions.
44+
10. **EOT token as regular vocab token**: `</assistant>` (token 24) is a regular vocabulary token, not a special token. `preserved_tokens` only controls special-token decode and doesn't suppress it. Use `additional_stops` instead — this feeds into `task->params.antiprompt` which the `process_token` erase logic checks, stripping the text before it's sent. Added `additional_stops` field to `autoparser` struct; apply in `cli.cpp` as `task.params.antiprompt` additions.
45+
11. **`swa_type` must match `is_swa_any()` (`create_memory` assert)**: `llama_model::create_memory` asserts `swa_type != NONE iff is_swa_any()` — the iswa KV cache is only allocated when there are SWA layers. Laguna-M.1 has zero SWA layers, so `swa_type` must be `NONE` and the graph builder must use `build_attn_inp_kv()` (+ its `build_attn` overload), not `build_attn_inp_kv_iswa()`. `laguna.cpp` sets `swa_type = is_swa_any() ? STANDARD : NONE` and branches the attention input on `has_swa`. (Leaving `swa_type = STANDARD` with zero SWA layers was the first M.1 smoke-test crash.)
46+
12. **Per-element vs per-head attention gate**: Laguna-M.1 uses `gating: "per-element"` (g_proj output = `num_heads*head_dim`, element-wise `attn *= softplus(g_proj(x))`); XS.2 uses per-head (g_proj output = `num_heads`, broadcast across head_dim). Signal via a new `laguna.attention.gate_per_head` bool KV (default `true` = per-head, so old XS.2 GGUFs keep working). Declare the gate tensor shape from it (`{n_embd, n_head_l}` vs `{n_embd, n_head_l*n_embd_head_v}`) and branch the application — per-element is a plain `ggml_mul`.
47+
13. **`layer_types` defaults to all full_attention (Laguna-M); rotary is data-driven**: `configuration_laguna.py` defaults `layer_types` to `["full_attention"]*n_layers` when absent (Laguna-M); Laguna-XS ships an explicit mix. The converter must mirror this or the per-layer `head_count`/`sliding_window_pattern` arrays come out empty. Partial rotary is now data-driven via `rope.dimension_count`/`rope.dimension_count_swa` (written from each rope config's `partial_rotary_factor`: XS.2 0.5 global / 1.0 SWA; M.1 1.0 everywhere), replacing the old hardcoded `n_rot_full /= 2`.
4448

4549
---
4650

@@ -63,18 +67,24 @@ IMPORTANT: Ensure you've thoroughly reviewed the [AGENTS.md](AGENTS.md) file bef
6367
| Context length | 131,072 |
6468
| Checkpoint layout | split-per-expert, singular `shared_expert`, `e_score_correction_bias` under `experts` not `gate` |
6569

70+
> **Laguna-M.1** differs from the XS.2 values above: 70 layers (3 dense + 67 sparse), uniform 64 heads / 8 KV heads (no per-layer head counts), `sliding_window = 0` with **no SWA layers** (`layer_types` all `full_attention`), **per-element** attention gate on all layers, **full rotary** (partial_rotary 1.0), top-16 routing, 3 dense-lead layers, YaRN factor 64 / context 262144. Same tensor layout and MoE structure as XS.2.
71+
6672
---
6773

6874
## Key Design Decisions
6975

7076
**YaRN per-layer approach:** Use **Path A** — add per-layer-type YaRN siblings to hparams/cparams (`yarn_ext_factor_swa`, `yarn_attn_factor_swa`, `yarn_beta_fast_swa`, `yarn_beta_slow_swa`). Initialize them to non-SWA values so other architectures are unaffected. In the graph builder, select via `hparams.is_swa(il)`.
7177

72-
**Partial rotary:** Keep step35's `n_rot_full / 2` mechanism — it exactly matches Laguna's pattern (global layers half-rotated, SWA layers full-rotated).
78+
**Partial rotary:** Data-driven via `rope.dimension_count` (global layers) and `rope.dimension_count.swa` (SWA layers), written by the converter from each rope config's `partial_rotary_factor` (XS.2: 0.5 global / 1.0 SWA; M.1: 1.0 everywhere). The earlier hardcoded `n_rot_full /= 2` was removed — it only worked for XS.2.
7379

7480
**Dense lead:** Layer 0 is dense FFN; layers 1–39 are MoE. Branch on `static_cast<uint32_t>(i) < hparams.n_layer_dense_lead`.
7581

7682
**Attention gate:** `TENSOR_NOT_REQUIRED` on `wqkv_gate`; guard with `if (model.layers[il].wqkv_gate)` in the graph builder.
7783

84+
**Attention gate mode (M.1):** Per-element (`gate_per_head=false`) applies a plain element-wise `ggml_mul(attn_out, gate)`; per-head (XS.2) reshapes the gate to `[1, n_head, n_tokens]` and broadcasts. Mode comes from the `laguna.attention.gate_per_head` KV.
85+
86+
**SWA path selection (M.1):** `swa_type` and the attention input builder are chosen from `is_swa_any()` — iswa path for XS.2 (mixed SWA), plain `build_attn_inp_kv()` for M.1 (no SWA). Required by `create_memory`'s `swa_type ⇔ is_swa_any()` assert.
87+
7888
---
7989

8090
## Files Added (2)
@@ -113,10 +123,21 @@ Subclass `TextModel`; registered as `"LagunaForCausalLM"`. Overrides `set_gguf_p
113123
16. **`common/chat-diff-analyzer.cpp`** ✅ — Laguna workaround: `reasoning.start=""`, `reasoning.end="</think>"`, `mode=TAG_BASED`, `additional_stops=["</assistant>"]`.
114124
17. **`tools/cli/cli.cpp`** ✅ — Apply `chat_params.additional_stops` to `task.params.antiprompt` so the CLI stop-word erase logic strips them from streaming output.
115125

126+
### Laguna-M.1 generalization (additional)
127+
128+
18. **`conversion/laguna.py`** ✅ — Default `layer_types` to all-`full_attention` when absent; write `rope.dimension_count`/`dimension_count_swa` from `partial_rotary_factor`; write `attention.gate_per_head`.
129+
19. **`src/models/laguna.cpp`** ✅ — Removed hardcoded `n_rot_full /= 2`; per-element vs per-head gate (shape + application); `swa_type` from `is_swa_any()` with branched attention input path (`build_attn_inp_kv` vs `_iswa`).
130+
20. **`gguf-py/gguf/constants.py` + `gguf-py/gguf/gguf_writer.py`** ✅ — New `attention.gate_per_head` key (`Keys.Attention.GATE_PER_HEAD`, `KEY_ATTENTION_GATE_PER_HEAD`) + `add_attention_gate_per_head()`.
131+
21. **`src/llama-arch.h` + `src/llama-arch.cpp`** ✅ — `LLM_KV_ATTENTION_GATE_PER_HEAD` enum value + `"attention.gate_per_head"` string mapping.
132+
22. **`src/llama-hparams.h`** ✅ — `bool attn_gate_per_head = true;` field.
133+
116134
---
117135

118136
## Reference Output
119-
Greedy decoding on `"The capital of France is"``" Paris.\nThe capital of Germany is"` (from sgl-project empirical reproducer).
137+
Greedy decoding on `"The capital of France is"``" Paris.\nThe capital of Germany is"` (from sgl-project empirical reproducer). *(XS.2 only; M.1 reference output not captured — the 226B MoE is I/O-bound on 60 GB RAM, ~50 s/token via mmap.)*
120138

121139
## Open Items
122-
None. `moe_router_logit_softcapping` is absent from `poolside/Laguna-XS.2` config.json — no plumbing needed.
140+
- `moe_router_logit_softcapping` is absent from both Laguna-XS.2 and Laguna-M.1 `config.json` — no plumbing needed.
141+
- **Laguna-M.1 chat template**: ships `laguna_glm_thinking_v4`, but auto-detection in `src/llama-chat.cpp` keys on the `laguna_glm_thinking_v5` marker, so the built-in matcher rejects it (`custom template not supported`). Workaround: pass `--jinja`. Follow-up: add the v4 marker to auto-detection.
142+
- **Step 5 / Step 7** still deferred — numerical validation needs CUDA; the 226B M.1 is also impractical for the quant top-1 check on current hardware.
143+
- **Pre-existing XS.2 GGUFs need reconversion** — they predate the now-required `rope.dimension_count` and would otherwise get full rotary instead of half.

conversion/laguna.py

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ def set_gguf_parameters(self):
3535
n_head_base = hparams["num_attention_heads"]
3636
n_kv_base = hparams.get("num_key_value_heads", n_head_base)
3737

38-
layer_types = hparams.get("layer_types", [])
38+
layer_types = hparams.get("layer_types", ["full_attention"] * n_layers)
3939

4040
# Per-layer head counts: use num_attention_heads_per_layer if available
4141
heads_per_layer = hparams.get("num_attention_heads_per_layer", None)
@@ -66,6 +66,23 @@ def set_gguf_parameters(self):
6666
head_dim = hparams.get("head_dim", hparams["hidden_size"] // n_head_base)
6767
self.gguf_writer.add_value_length(head_dim)
6868

69+
# Partial rotary is data-driven: write rope.dimension_count for the global
70+
# (full-attention) layers and rope.dimension_count_swa for sliding layers.
71+
# Laguna-M uses full rotary (1.0) on every layer; Laguna-XS uses half rotary
72+
# (0.5) on global layers and full rotary on SWA layers.
73+
full_rope = self.rope_parameters.get("full_attention", {})
74+
partial_full = full_rope.get("partial_rotary_factor", 1.0)
75+
self.gguf_writer.add_rope_dimension_count(int(head_dim * partial_full))
76+
swa_rope = self.rope_parameters.get("sliding_attention")
77+
if swa_rope is not None:
78+
partial_swa = swa_rope.get("partial_rotary_factor", 1.0)
79+
self.gguf_writer.add_rope_dimension_count_swa(int(head_dim * partial_swa))
80+
81+
# Attention output gate mode: per-head (broadcasts across head_dim, Laguna-XS)
82+
# vs per-element (one gate per (head, head_dim) channel, Laguna-M). Default is
83+
# per-element to match configuration_laguna.py (gating defaults to True/per-element).
84+
self.gguf_writer.add_attention_gate_per_head(hparams.get("gating", "per-element") == "per-head")
85+
6986
# MoE params
7087
n_experts = hparams["num_experts"]
7188
n_experts_used = hparams["num_experts_per_tok"]

gguf-py/gguf/constants.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -195,6 +195,7 @@ class Attention:
195195
SHARED_KV_LAYERS = "{arch}.attention.shared_kv_layers"
196196
SLIDING_WINDOW_PATTERN = "{arch}.attention.sliding_window_pattern"
197197
TEMPERATURE_SCALE = "{arch}.attention.temperature_scale"
198+
GATE_PER_HEAD = "{arch}.attention.gate_per_head"
198199

199200
class Indexer:
200201
HEAD_COUNT = "{arch}.attention.indexer.head_count"
@@ -4774,6 +4775,7 @@ class VisionProjectorType:
47744775
KEY_ATTENTION_CLAMP_KQV = Keys.Attention.CLAMP_KQV
47754776
KEY_ATTENTION_LAYERNORM_EPS = Keys.Attention.LAYERNORM_EPS
47764777
KEY_ATTENTION_LAYERNORM_RMS_EPS = Keys.Attention.LAYERNORM_RMS_EPS
4778+
KEY_ATTENTION_GATE_PER_HEAD = Keys.Attention.GATE_PER_HEAD
47774779

47784780
# RoPE
47794781
KEY_ROPE_DIMENSION_COUNT = Keys.Rope.DIMENSION_COUNT

gguf-py/gguf/gguf_writer.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -772,6 +772,11 @@ def add_key_length(self, length: int) -> None:
772772
def add_value_length(self, length: int) -> None:
773773
self.add_uint32(Keys.Attention.VALUE_LENGTH.format(arch=self.arch), length)
774774

775+
def add_attention_gate_per_head(self, per_head: bool) -> None:
776+
# Laguna attention output gate mode: True = per-head (gate broadcasts across
777+
# head_dim), False = per-element (one gate per (head, head_dim) channel).
778+
self.add_bool(Keys.Attention.GATE_PER_HEAD.format(arch=self.arch), per_head)
779+
775780
def add_key_length_mla(self, length: int) -> None:
776781
self.add_uint32(Keys.Attention.KEY_LENGTH_MLA.format(arch=self.arch), length)
777782

src/llama-arch.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -240,6 +240,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
240240
{ LLM_KV_ATTENTION_RELATIVE_BUCKETS_COUNT, "%s.attention.relative_buckets_count" },
241241
{ LLM_KV_ATTENTION_SLIDING_WINDOW, "%s.attention.sliding_window" },
242242
{ LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, "%s.attention.sliding_window_pattern" },
243+
{ LLM_KV_ATTENTION_GATE_PER_HEAD, "%s.attention.gate_per_head" },
243244
{ LLM_KV_ATTENTION_SCALE, "%s.attention.scale" },
244245
{ LLM_KV_ATTENTION_OUTPUT_SCALE, "%s.attention.output_scale" },
245246
{ LLM_KV_ATTENTION_VALUE_SCALE, "%s.attention.value_scale" },

src/llama-arch.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -245,6 +245,7 @@ enum llm_kv {
245245
LLM_KV_ATTENTION_RELATIVE_BUCKETS_COUNT,
246246
LLM_KV_ATTENTION_SLIDING_WINDOW,
247247
LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN,
248+
LLM_KV_ATTENTION_GATE_PER_HEAD,
248249
LLM_KV_ATTENTION_SCALE,
249250
LLM_KV_ATTENTION_OUTPUT_SCALE,
250251
LLM_KV_ATTENTION_VALUE_SCALE,

src/llama-hparams.h

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -191,6 +191,11 @@ struct llama_hparams {
191191
bool attn_soft_cap = false;
192192
bool use_kq_norm = false;
193193

194+
// Laguna attention output gate: true = per-head (broadcast across head_dim),
195+
// false = per-element (one gate per (head, head_dim) channel). Default true
196+
// preserves Laguna-XS GGUFs that predate the gate_per_head key.
197+
bool attn_gate_per_head = true;
198+
194199
// for Classifiers
195200
uint32_t n_cls_out = 1;
196201

0 commit comments

Comments
 (0)