Skip to content

model: add Poolside Laguna XS.2 support#1911

Merged
ikawrakow merged 9 commits into
ikawrakow:mainfrom
joelfarthing:laguna-xs2-support
Jun 8, 2026
Merged

model: add Poolside Laguna XS.2 support#1911
ikawrakow merged 9 commits into
ikawrakow:mainfrom
joelfarthing:laguna-xs2-support

Conversation

@joelfarthing

@joelfarthing joelfarthing commented Jun 1, 2026

Copy link
Copy Markdown
Contributor

Summary

This adds support for Laguna-architecture GGUFs matching the public Laguna XS.2 schema. It has been validated on Q6_K and Q4_K_M GGUFs generated from official Laguna XS.2 weights. Other Laguna XS.2-derived GGUFs may work if they preserve the same metadata and tensor layout, but are not separately validated here. This compatibility PR includes:

  • laguna architecture registration and model metadata loading
  • Laguna tensor mapping/loading
  • graph support for Laguna's mixed global/sliding-window attention pattern
  • per-layer query head counts
  • Laguna's intended global/SWA RoPE and YaRN behavior
  • softplus attention gating
  • routed/shared MoE handling
  • tokenizer/GGUF metadata support, including tokenizer.ggml.pre = laguna
  • converter support for official poolside/Laguna-XS.2 weights
  • resolved Poolside chat-template embedding during conversion
  • partial-offload --cpu-moe handling where routed experts are CPU-hosted while shared experts remain on their normal split device

This PR does not include DFlash, MTP, speculation, performance tuning, custom quantization work, or broad chat/server UX changes.

Architecture Notes

Laguna-XS.2 is close to existing Step-style hybrid MoE/SWA models, but it is not a pure alias.

The correctness-sensitive parts are:

  • global attention layers use YaRN-scaled RoPE with base 500000
  • sliding-window layers use the SWA RoPE base 10000
  • global and SWA layers have different rotary dimensions
  • layer query-head counts follow the Laguna [48, 64, 64, 64, ...] pattern
  • the SWA window is 512
  • attention output gating uses softplus
  • the model uses one leading dense block plus routed and shared experts

The converter preserves the source checkpoint's Laguna metadata. The runtime then applies Laguna-specific global-vs-SWA RoPE/YaRN semantics in the graph path instead of requiring users to disable or rewrite the model's intended scaling metadata.

Validation

Official Poolside source revision:

1162d99d501157f6da4a59ab9fd49eb312c3afdb

Reviewer/community convenience GGUFs generated from official Poolside BF16 weights:

https://huggingface.co/ji-farthing/Laguna-XS.2-ik-llama-validation-GGUF

Files:

laguna-xs2-ikllama-Q6_K.gguf
laguna-xs2-ikllama-Q4_K_M.gguf

These GGUFs require an ik_llama.cpp build with this Laguna-XS.2 support, or another runtime with equivalent Laguna architecture handling. They are not expected to load or decode correctly with current ikawrakow/main without this PR.

Validation highlights:

  • official BF16 tensor payload was compared against the Lucebox BF16 reference across 678/678 tensors
  • fresh official BF16 conversion completed from the Poolside checkpoint
  • fresh converter output embeds the resolved Poolside chat template instead of an unresolved {% include %} stub
  • BF16 metadata preserves native Laguna YaRN/RoPE metadata, including yarn_ext_factor=64
  • Q6_K and Q4_K_M were quantized from the corrected BF16 GGUF with current llama-quantize
  • Q6_K and Q4_K_M metadata was inspected after quantization
  • Q6_K and Q4_K_M both load in llama-server
  • server metadata dumps show Laguna YaRN/SWA metadata and the expected per-layer head-count pattern
  • Q6_K passes an OpenAI-compatible server chat sanity check
  • Q6_K passes a topological-sort coding gate:
    • first Python block extracted from the response
    • python3 -m py_compile passes
    • DAG test returns a valid topological order
    • cycle test raises a clear ValueError
  • Q4_K_M passes an OpenAI-compatible server chat/load sanity check
  • Q4_K_M passes a local Pi Agent word-statistics coding smoke at 8K context
  • CPU and CUDA build targets pass locally
  • converter/GGUF Python files pass py_compile
  • git diff --check passes locally

Representative local server shapes used for the fresh artifacts:

./build-cuda/bin/llama-server \
  -m /models/laguna-xs2-ikllama-Q6_K.gguf \
  --host 127.0.0.1 \
  --port 18092 \
  -c 8192 \
  -ngl 8 \
  -b 64 \
  -ub 64 \
  --jinja \
  --no-warmup

Partial-offload work after addressing #1911 comments:

./build-cuda-laguna/bin/llama-server \
  -m /models/laguna-xs2-ikllama-Q6_K.gguf \
  --host 127.0.0.1 \
  --port 18092 \
  -c 4096 \
  -ngl 8 \
  -b 64 \
  -ub 64 \
  --jinja \
  --no-warmup \
  --cpu-moe

The mixed placement and the shared-experts-on-host control both repeat exactly within placement, but their top-token probabilities differ by 0.0357007384 on the same prompt. Treat --cpu-moe partial offload as stable-but-non-exact.

Further bisection narrowed the delta to the shared expert up/down_shexp matmuls for the eight GPU-offloaded layers under -ngl 8. Disabling fused MoE/multi-add/up-gate paths did not restore parity, gate_shexp placement was irrelevant, and the -ngl 0 control was exact. Current evidence suggests a CUDA-vs-host shared-expert up/down execution delta rather than a routed/shared add-reduce bug.

  • Forced cuBLAS and forced MMQ CUDA builds produced the same mixed-vs-shared-host delta as the default build on the short deterministic /completion harness: max_abs_common_top_prob_delta=0.0013219118118286133, same selected tokens ( Paris, .).
  • Within each placement, default-vs-forced-cuBLAS and default-vs-forced-MMQ were exact (max_abs_common_top_prob_delta=0), so the short-prompt delta does not follow that matrix-matrix selector.
  • Forced DMMV did not change the top-prob vectors in the same harness. Source inspection shows GGML_CUDA_FORCE_DMMV does not bypass the earlier MMVQ/MMQ fast return in ggml_cuda_mul_mat, so treat that as a weak/inconclusive DMMV scout rather than proof that DMMV was fully exercised.
  • A fresh Q6_K placement recheck on the current build-cuda/bin/llama-server reproduced the same short-prompt delta exactly. Hosting only blk\..*\.ffn_gate_shexp.weight was identical to mixed placement (max_abs_common_top_prob_delta=0), while hosting only blk\..*\.ffn_(up|down)_shexp.weight was identical to hosting all shared expert tensors. This confirms that gate_shexp and the routed/shared add are not the source of the observed drift.
  • The local laguna-xs2-official-Q6_0.gguf control also showed nonzero mixed-vs-up/down_shexp-host probability drift on the same short harness (max_abs_common_top_prob_delta=0.00404695887118578, same selected tokens). This makes a Laguna-specific Q6_K-only issue less likely and supports treating the remaining delta as quantized CPU-vs-CUDA matmul drift for the shared expert up/down projections.
  • A llama-perplexity --kl-divergence run over 18 chunks / 270 evaluated positions compared mixed placement against a shared-host base:
    • Mean PPL(Q)/PPL(base)=1.004195 ± 0.001132
    • Mean KLD=0.000056 ± 0.000003
    • Maximum KLD=0.000269
    • RMS Δp=0.089 ± 0.006%
    • Maximum Δp=0.228%
    • Same top p=100.000 ± 0.000%

The partial-offload path is stable and keeps shared experts on GPU as requested, but it is not exact probability parity against a shared-experts-on-host control. The remaining evidence looks like small deterministic CUDA-vs-host numeric drift localized to offloaded shared expert up/down_shexp, rather than Laguna graph-composition bug.

Notes / Caveats

  • Q6_K is the primary reviewer-facing validation artifact.

  • Q4_K_M is provided as a smaller secondary artifact.

  • No performance claims are made.

  • The local validation profile intentionally uses explicit conservative placement (-ngl 8) rather than --fit; --fit tuning is left out of scope for this PR.

  • The --cpu-moe partial-offload path is stable but not exact parity: current mixed routed-host/shared-GPU placement loads and decodes, while deterministic probability comparisons against a shared-experts-on-host control show small nonzero drift.

  • The implementation focuses on Laguna architecture correctness rather than speculative decoding, DFlash, or server UX changes.

  • GitHub may show a hidden/bidirectional Unicode warning on src/llama-vocab.cpp. Audit result: this PR adds no Unicode format or direction-control characters; the warning comes from pre-existing tokenizer regex Unicode in that file, while the Laguna-added lines are ASCII-only and needed for GGUFs with tokenizer.ggml.pre = laguna.

  • I have read the contributing guidelines

  • Self-reported review complexity:

    • Low
    • Medium
    • High

@joelfarthing joelfarthing changed the title llama: add Poolside Laguna XS.2 support model: add Poolside Laguna XS.2 support Jun 3, 2026
@joelfarthing joelfarthing marked this pull request as ready for review June 4, 2026 12:38
Comment thread src/llama-load-tensors.cpp Outdated
// Laguna's MoE path adds a shared expert beside the routed experts.
// Keep both on CPU for --cpu-moe; mixing CUDA-host routed experts with
// GPU shared experts currently produces non-parity logits.
pattern = "blk\\." + std::to_string(i) + "\\.ffn_((up|down|gate|gate_up)_exps|(up|down|gate)_shexp)\\.(weight|scale)";

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No, we don't want to have this. For partial offload it needs to work with shared experts on the GPU.

Comment thread src/llama-build-context.cpp Outdated
Qcur = ggml_rope_multi(ctx0, Qcur, inp_pos, rope_factors,
n_rot_l, sections, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
ext_factor, attn_factor, beta_fast, beta_slow);
ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l);

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So, either you make the Laguna implementation work with build_std_attention, then perhaps you can make changes to it, or you don't use build_std_attention (as the PR currently stands), and in that case you don't modify build_std_attention.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Gah! Apologies, I'll clean that up, and investigate where non-parity in GPU shared experts is coming from. It looks like CUDA-vs-host execution for up/down_shexp on the offloaded layers.

@joelfarthing

joelfarthing commented Jun 5, 2026

Copy link
Copy Markdown
Contributor Author

Hi, @ikawrakow!

I cleaned up the build_std_attention edits, and reverted shared expert tensor placement to the normal split device, so shared experts stay on GPU for GPU-offloaded layers. As noted in the edited description above, the partial-offload path is stable and keeps shared experts on GPU, but is not exact probability parity against a shared-experts-on-host control. I narrowed the discrepancy down to likely numerical drift in offloaded shared expert up/down_shexp, but I didn’t want to touch broader CPU-vs-CUDA quantized matmul behavior in this compatibility PR.

Changed code:

  • src/llama-load-tensors.cpp (line 240): removed the Laguna-only --cpu-moe shared-expert override. Routed expert tensors still use the normal CPU-host override pattern; shared experts now stay on their normal split device.
  • src/llama-build-context.cpp (line 2540): removed Laguna-specific RoPE/gate/precision edits from build_std_attention. Laguna still dispatches to build_laguna() for its dedicated graph behavior.

Comment thread src/llama.cpp Outdated
return scale <= 1.0f ? 1.0f : (0.1f * mscale * logf(scale) + 1.0f);
};

if (cparams.yarn_ext_factor != 0.0f) {

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This code here will run for all models. On a scale of 1 to 10, how confident are you that this will not break other models? If it was copy-pasted from mainline llama.cpp, things are different there.

Comment thread src/llama-hparams.h Outdated
float yarn_attn_factor = 1.0f;
float yarn_beta_fast = 32.0f;
float yarn_beta_slow = 1.0f;
float yarn_ext_factor_swa = -1.0f;

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here we set it to 1.0f. But in the Laguna section in llama-hparams.cpp we set it to 0.0f. Why? The *_swa values have been added to llama_hparams specifically for the Laguna architecture and are not used for anything else, so we might as well set them to their final value and discard the corresponding code in llama-hparams.cpp.

@joelfarthing

joelfarthing commented Jun 5, 2026

Copy link
Copy Markdown
Contributor Author

@ikawrakow

Understood! New commit: b63a5c174 model: localize Laguna SWA YaRN defaults.

With this cleanup, I’m much more confident, at least 9/10, that this won’t affect other models because the SWA-specific YaRN values are no longer being globally inferred in llama.cpp; they're just carried from hparams, and currently only Laguna consumes them.

The regular YaRN mscale logic was adapted from the existing llama.cpp context initialization path. My mistake was extending that pattern to the new SWA-specific fields, even though those fields are Laguna-specific and can just carry final hparams values. Your point is correct: the SWA-specific part shouldn't live in the global initialization path.

@ikawrakow ikawrakow left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As far as I can tell, you can delete everything related to YaRN and SWA, and you get the exact same result.

Comment thread src/llama.cpp Outdated
if (cparams.yarn_ext_factor < 0.0f) { // negative indicates 'not set'
cparams.yarn_ext_factor = rope_scaling_type == LLAMA_ROPE_SCALING_TYPE_YARN ? 1.0f : 0.0f;
}
static auto get_mscale = [](float scale, float mscale) {

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remove lines 6819-6845 from here and do this in build_laguna(). We really do not want the addition of a new model to affect other models. And we are not a copy of llama.cpp, so there is no need to copy/paste code from there.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed, thanks. New commit: 2f7e43968 model: localize Laguna SWA RoPE constants.

I removed the copied global mscale block from llama.cpp and removed the SWA-only hparams/cparams fields. build_laguna() now uses the SWA RoPE constants directly for SWA layers, so this no longer adds global initialization behavior for other models.

I also tested this against the previous PR commit b63a5c174 with the Q4 Laguna GGUF using a deterministic short completion/probability probe. The generated text and full completion_probabilities JSON were identical.

One note: I kept Laguna’s model-level YaRN key reads. I tested removing those too, but that changed output probabilities materially, so the narrower cleanup appears to be the parity-preserving version.

@ikawrakow ikawrakow merged commit bbe1a51 into ikawrakow:main Jun 8, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants