Skip to content

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

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

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>
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.

1 participant