Laguna M.1/XS.2 Technical Report

TL;DR

Poolside’s technical report on Laguna M.1 (225.8B total / 23.4B active MoE) and Laguna XS.2 (33.4B total / 3B active MoE), two from-scratch foundation models built for long-horizon agentic coding. XS.2’s weights ship under Apache 2.0. The report’s main pitch is methodological: a “Model Factory” — Dagster-driven asset DAG, custom Kubernetes batch scheduler, a unified PyTorch training stack (Titan) covering pre-training and RL — that took XS.2 from start of training to release in five weeks by treating each design delta (attention layout, WSD vs cosine, expert modulation, dense-layer count) as a config flip against the M.1 baseline. On agentic benchmarks XS.2 is competitive with Devstral 2 Small / Qwen3.6-35B-A3B / Gemma 4 31B and M.1 with GLM-4.7 / DeepSeek-V4-Flash / Qwen3.5-397B-A17B / Claude Sonnet 4.6.

Key claims

Laguna M.1 has 225.8B total / 23.4B active parameters; XS.2 has 33.4B total / 3B active; both are pre-norm MoE Transformers with RMSNorm, BPE tokenizer (100,352 tokens), and were pre-trained on >30T tokens drawn from web, code, and synthetic sources [§1, §3.1, §3.2]. [Fig. 1]
M.1 on agentic benchmarks: SWE-bench Verified 79.6, SWE-bench Multilingual 73.3, SWE-Bench Pro 52.6, Terminal-Bench 2.0 56.9 — competitive with GLM-4.7, DeepSeek-V4-Flash, Qwen3.5-397B-A17B and Claude Sonnet 4.6 [Fig. 1a].
XS.2 on agentic benchmarks: SWE-bench Verified 73.4, SWE-bench Multilingual 67.2, SWE-Bench Pro 52.4, Terminal-Bench 2.0 51.5 — competitive with Devstral 2 Small, Qwen3.6-35B-A3B, Gemma 4 31B and Claude Haiku 4.5 [Fig. 1b].
XS.2 uses interleaved Sliding Window Attention (window 512, 64 Q-heads, θ=10,000) and Global Attention (48 Q-heads, θ=500,000, partial RoPE on first 50% of head dim) at a 3:1 SWA:GA ratio with softplus per-head gating; first Transformer layer is dense for stability [§3.1].
XS.2 routing: token-choice top-8 of 256 experts plus a shared expert, sigmoid router with post-top-k normalization, routed-expert output modulated by 2.5 before combining with shared expert; load balancing via the Qiu et al. global-batch auxiliary loss [§3.1].
Four targeted deltas from M.1→XS.2 — interleaved SWA+GA replacing all-GA, WSD replacing cosine, routed-expert modulation, fewer bottom dense layers (3→1) — each promoted via ablations on a small MoE proxy and then flipped on as configuration changes against the M.1 baseline [§3.1].
A WSD-specific scaling law fit across 4 MoE sizes (2B–16B total / 0.3B–2.2B active) × 6 LRs × 5 token budgets (30B–480B) at fixed 8M-token batch yields lr*(N,D) = 10^4.488 · N^-0.4639 · D^-0.2661, then a √(B/B₀) batch correction; predicts 5.5e-4 for XS.2’s 3B-active / 24M-batch run, deployed at 5e-4 for safety margin [§3.1.1, Eq. 1, Appendix A.1].
Training uses Muon (Moonlight variant) across pre-training, SFT, and RL, with a distributed implementation that assigns each parameter to one rank for Newton-Schulz, batched comm overlapped with NS compute, optional CUDA graphs — driving optimizer overhead below 1% of step time on M.1 [§3.1.2].
The MoE dispatch/combine all-to-alls are fused into the CUTLASS grouped-GEMM kernel: 8 of 132 SMs per H200 copy expert tokens over NVLink in expert order and set ready flags in HBM; the modified tile scheduler waits on per-expert flags before processing, and the modified epilogue stores GEMM outputs directly to remote HBM through shared-memory vectorized stores — overlapping compute and communication at tile granularity [§3.1.2, Fig. 2].
BF16 mixed-precision training throughout, with FP32 master weights, FP32 RMSNorm/RoPE, and per-collective configurable gradient-reduction dtypes (default FP32 where reductions are numerically sensitive); the LM-head input-gradient all-reduce is specifically called out as needing FP32 [§3.1.1, §3.1.2].
Context-extension recipe: two equal 100B-token sub-stages (4K→32K, then 32K→128K) starting from the end-of-decay checkpoint with YaRN applied only to GA layers, 24M-token batch, cosine LR decaying 2.5e-5 → 5e-6, no re-warmup; final base checkpoint is an EMA over the last 10 checkpoints and the deployed model further doubles RoPE scale to 256K with no extra training [§3.1.1].
XS.2 ships with FP8, INT4, and NVFP4 quantization recipes for the MoE layers plus FP8 KV cache, with a mixed-precision strategy to preserve quality on low-VRAM devices [§5, §1].

Method

Both models are pre-norm MoE Transformers with RMSNorm. M.1 keeps full global attention; XS.2 introduces a 3:1 SWA:GA interleaving with GQA-8 KV heads, head dim 128, and softplus per-head gating. Routing is sigmoid token-choice top-8-of-256 with a shared expert and a routed-expert modulation of 2.5; load balance is enforced by Qiu et al.’s global-batch aux loss. The first Transformer layer is dense for stability.

Training runs on PyTorch via Titan (a >2,200-change fork of TorchTitan), composing FSDP, TP, EP, PP with an (PP, DDP, FSDP, TP) mesh for non-MoE layers and (PP, DDP, FSDP, EGP, ETP) for MoE layers (Laguna uses ETP=1). Distributed Muon assigns each gradient to a single rank to gather and orthogonalize via scheduled Newton-Schulz iterations, redistributing the orthogonalized shards back; CUDA graphs hide kernel-launch overhead. The dispatch/combine collectives for EP are fused into the grouped-GEMM kernel as described above, leaving 5 SMs reserved for NCCL FSDP collectives. Pre-training uses BF16 mixed precision with selective FP32 for RMSNorm, RoPE, and sensitive reductions; the optimizer is the Moonlight Muon variant throughout pre-training, SFT, and RL.

The peak LR for XS.2 was picked from a WSD-specific scaling law fit across 120 small-MoE runs (Eq. 1), then √-scaled to the production batch — concretely 5e-4 against a predicted 5.5e-4. The schedule is linear warmup → stable phase at peak → 30% √-decay to 5% of peak. Context extension is two-stage (32K then 128K) over 200B tokens with YaRN on GA only, EMA over the last 10 checkpoints, and a final RoPE-doubling to 256K with no training. The Model Factory wraps the whole pipeline: Dagster tracks the asset DAG so any artifact (a packed token, a checkpoint, an eval result) can be traced bidirectionally; an in-house FoundationDB-backed batch scheduler replaces Volcano, giving per-job (not per-node) eviction and sticky-pod respawn for sub-minute placement at >10⁴-accelerator scale.

Results

On four agentic benchmarks (SWE-bench Verified, SWE-bench Multilingual, SWE-Bench Pro, Terminal-Bench 2.0), M.1 scores 79.6 / 73.3 / 52.6 / 56.9 — competitive with GLM-4.7 (355B-A32B), DeepSeek-V4-Flash (284B-A13B), Qwen3.5 397B-A17B, and Claude Sonnet 4.6 [Fig. 1a]. XS.2 scores 73.4 / 67.2 / 52.4 / 51.5 — competitive with Devstral 2 Small (24B dense), Qwen3.6-35B-A3B, Gemma 4 31B dense, Claude Haiku 4.5, and GPT-5.4 Nano [Fig. 1b]. Distributed-Muon overhead is reported at <1% of step time on M.1 pre-training [§3.1.2]. The custom scheduler keeps placement consistently sub-minute at full-cluster scale, which the authors call out as the change that mattered most operationally [§2.1.3].

Why it’s interesting

The architecture deltas are individually unsurprising in 2026 — 3:1 SWA hybrid attention is now the default playbook (cf. Xiaomi MiMo-V2-Flash — 309B/15B MoE with 5:1 SWA hybrid attention, 73.4 SWE-Bench Verified / Fuli Luo on MiMo-V2-Flash engineering — Hybrid SWA, MTP, MOPD post-training on MiMo-V2-Flash, Qwen3-Coder-Next Technical Report, Step 3.5 Flash); routed-expert modulation, sigmoid routing, and shared experts came from DeepSeek-V3 (Optimal Expert-Attention Allocation in Mixture-of-Experts: A Scalable Law for Dynamic Model Design); Muon-family training at trillion-token scale is the Kimi K2: Open Agentic Intelligence line. What makes Laguna worth reading is the meta-story — a fairly direct test of whether the How to Set the Learning Rate for Large-Scale Pre-training?-style WSD scaling-law fitting actually transfers when you flip it on for a production model (it did, within 10%), and the cleanest published account so far of a Model-Factory-style asset-DAG + custom-scheduler pipeline that lets a small team turn around a new pre-training run in five weeks. The dispatch/combine-fused grouped-GEMM kernel is also a concrete instance of the compute-communication-fusion thesis from Loads and Loads of Fluffy Kittens: Compute-Communication Kernels with Multi-GPU ThunderKittens, applied to standard MoE training rather than to the more exotic Kittens VM.