Skip to content

feat(laguna): native GGUF support for poolside Laguna (S/XS) - #102

Closed
lucaspirola wants to merge 1 commit into
FlashML-org:mainfrom
lucaspirola:feat/laguna-gguf
Closed

feat(laguna): native GGUF support for poolside Laguna (S/XS)#102
lucaspirola wants to merge 1 commit into
FlashML-org:mainfrom
lucaspirola:feat/laguna-gguf

Conversation

@lucaspirola

Copy link
Copy Markdown

What

Native GGUF support for the poolside Laguna architecture (S and XS), the first
laguna implementation in FreeToken. Covers the Roadmap's "GGUF: support GGUF
checkpoints and their quantization types across model architectures" item.

The model:

  • hybrid full-attention / sliding-window layers (il % 4 == 0 full, else SWA),
    with per-layer query-head counts (48 full / 72 SWA on S)
  • QK RMSNorm before rope, a per-head softplus attention output gate,
    per-layer-type rope (YaRN partial-dim on full layers, plain rope on SWA)
  • sigmoid + score-correction-bias MoE routing, top-8/10 of 256 experts, one
    always-on shared expert; layer 0 is dense, layers 1+ sparse
  • semantics follow llama.cpp src/models/laguna.cpp

Because Poolside/Unsloth laguna checkpoints quantize per tensor (mixed ggml
types in one file), this also generalizes the GGUF plumbing:

  • adds Q3_K, Q4_K, Q5_K, IQ1_S, IQ2_S, IQ2_XXS, IQ3_XXS, IQ4_XS to the dequant
    tables and mmvq/mq/dispatch sets (the vendored CUDA kernels already had them)
  • a "gguf" expert-bank format whose quant types vary per layer: flat padded
    [E, stride] host banks + a new expert_stride_bytes argument threaded through
    the vendored moe_vec launchers (0 = previous dense behaviour; gemma4 q4_0
    unaffected)
  • moe_vec calls chunked to min(65535, 16384) rows — the kernel indexes experts
    via blockIdx.z (CUDA grid-z cap 65535) and a 16k-token × top-8 prefill chunk
    also overflowed transient VRAM (surfaced as an async "device not ready")
  • q/k/v kept as separate projections (a layer may quantize attn_v differently
    from attn_q/k — observed in the XS Q4_K_M build, so fused buffers are invalid)
  • deferred GGUF linears materialized from the file's tensor table at conversion
    time, before the engine collects the state dict

Tested on

  • Hardware: NVIDIA RTX 5080 (16 GB), 23 GB system RAM, WSL2, driver 610.62,
    CUDA 13.0, torch torch.cuda=13.0
  • Checkpoint: poolside/Laguna-XS-2.1-GGUF Laguna-XS-2.1-Q4_K_M.gguf
    (header-level) and a third-party XS build Laguna-XS-2.1-APEX-I-Mini.gguf
    (Q3_K/Q4_K/Q5_K/Q6_K/IQ2_S), run end to end
  • Command:
    ft serve --model Laguna-XS-2.1-APEX-I-Mini.gguf \
             --kv-cache-dtype fp8_e4m3 --num-tokens 262144 --kv-reserve-tokens 262144
    
    (--kv-cache-dtype is the fp8 KV cache from a separate PR; without it,
    auto/bf16 doubles the KV footprint.)

Results

  • NIAH 3/3 at 250,054 tokens — needle recovered at 10%/50%/90% depth, exact
    passcode each run, ~433 tok/s prefill.
  • decode 157–162 tok/s at 64k context (8984 expert slots); 21–23 tok/s at
    ~250k context (2441 expert slots, PCIe-bound).
  • 81 unit tests: pytest tests/models tests/kernels/test_gguf_quant_types.py

Not in this PR (documented in tasks/laguna-handover.md)

  • S-model e2e unvalidated — this host's 23 GB RAM cannot pin S's ~37 GiB
    expert banks; S geometry is identical, but no forward pass has run on S weights.
  • hybrid/cpu MoE backends for the "gguf" format (needs a CPU-side SIMD port).
  • FTW conversion is intentionally refused (a metadata-only GGUF drops the
    per-tensor quant types nullptr; serve the .gguf directly).
  • TP > 1; text-only.
Adds the laguna GGUF architecture: hybrid full/SWA attention with per-layer
query-head counts, QK RMSNorm, per-layer-type rope (YaRN partial-dim on full
layers, plain rope on SWA), a per-head softplus attention output gate, and
sigmoid + score-correction-bias MoE routing with one always-on shared expert.
Semantics follow llama.cpp src/models/laguna.cpp.

Poolside/Unsloth laguna checkpoints quantize per tensor, so this also
generalizes the GGUF plumbing:

- six more ggml types (Q3_K, Q4_K, Q5_K, IQ1_S, IQ2_S, IQ2_XXS, IQ3_XXS,
  IQ4_XS) wired into the dequant tables and mmvq/mmq/dequant dispatch
- a "gguf" expert-bank format whose per-layer quant types vary: flat padded
  [E, stride] host banks plus a new expert_stride_bytes argument threaded
  through the vendored moe_vec launchers (0 = previous dense behaviour)
- moe_vec calls chunked to min(65535, 16384) rows: the kernel indexes experts
  via blockIdx.z (CUDA grid-z cap 65535) and a 16k-token x top-8 prefill chunk
  also overflowed transient VRAM
- q/k/v kept as separate projections (a layer may quantize attn_v differently
  from attn_q/k -- observed on the XS Q4_K_M build)
- deferred GGUF linears materialized from the file's tensor table at conversion
  time, before the engine collects the state dict

Tested on: RTX 5080 (16 GB), 23 GB RAM, NVIDIA 610.62, CUDA 13,
wsl2. Checkpoints: unsloth/Laguna-S-2.1-GGUF (Laguna-S-2.1-UD-IQ1_S.gguf,
S, metadata/tensor coverage only) and poolside/Laguna-XS-2.1-GGUF
(Laguna-XS-2.1-Q4_K_M.gguf) plus a third-party APEX-I-Mini XS build
(Q3_K/Q4_K/Q5_K/Q6_K/IQ2_S) run end to end.

Validation on Laguna-XS-2.1-APEX-I-Mini.gguf:
  ft serve --model <model> --kv-cache-dtype fp8_e4m3 --num-tokens 262144 --kv-reserve-tokens 262144
- NIAH 3/3 at 250,054 tokens (needle at 10%/50%/90% depth, exact passcode),
  ~433 tok/s prefill
- decode 157-162 tok/s at 64k ctx, 21-23 tok/s at ~250k ctx (PCIe-bound)

Not done: S-model e2e (host lacked RAM for its ~37 GiB expert banks), hybrid/cpu
MoE backend for the gguf format, FTW conversion, TP>1. See tasks/laguna-handover.md.

Co-Authored-By: Claude <noreply@anthropic.com>
@pedro-moser

pedro-moser commented Aug 26, 2026

Copy link
Copy Markdown

I validated the official Laguna-S Q4_K_M checkpoint end to end on 128 GiB RAM and opened two draft follow-ups stacked on this PR:

Both are drafts because GitHub currently shows the #102 commits in their diffs against main; each body identifies the follow-up commits to review. #199 has no code dependency on #103, although the controlled benchmark branch used #103's Q8_0 KV support in both arms. I can rebase the drafts as soon as #102 lands, or split/cherry-pick further if you prefer a different integration path.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

2 participants