# Extending Puzzle for Mixture-of-Experts Reasoning Models with Application to GPT-OSS Acceleration

arXiv: https://arxiv.org/abs/2602.11937 · HTML: https://mlsanity.com/paper/2602.11937/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Test-time compute and efficient thinking, one of 17 directions in LLM reasoning.
- Builds on 17 of its 29 references that are in the map.
- First posted on arXiv 12 Feb 2026.

## Placed in

- LLM reasoning / [Test-time compute and efficient thinking](https://mlsanity.com/lr/test-time-compute-and-efficient-thinking/)
- Mixture-of-experts / [Efficient MoE training and inference systems](https://mlsanity.com/moe/efficient-moe-training-and-inference-systems/)

## Builds on (most cited first, 17 of 17)

- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [GPQA: A Graduate-Level Google-Proof Q&A Benchmark](https://mlsanity.com/paper/2311.12022/) (2023)
- [Gemma 2: Improving Open Language Models at a Practical Size](https://mlsanity.com/paper/2408.00118/) (2024)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [MMLU-Pro: A More Robust and Challenging Multi-Task Language Understanding Benchmark](https://mlsanity.com/paper/2406.01574/) (2024)
- [Gemma 3 Technical Report](https://mlsanity.com/paper/2503.19786/) (2025)
- [gpt-oss-120b & gpt-oss-20b Model Card](https://mlsanity.com/paper/2508.10925/) (2025)
- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://mlsanity.com/paper/2404.06654/) (2024)
- [DeepSeek-V3 Technical Report](https://mlsanity.com/paper/2412.19437/) (2024)
- [Humanity's Last Exam](https://mlsanity.com/paper/2501.14249/) (2025)
- [Generalizing Verifiable Instruction Following](https://mlsanity.com/paper/2507.02833/) (2025)
- [Llama-Nemotron: Efficient Reasoning Models](https://mlsanity.com/paper/2505.00949/) (2025)
- [Nemotron-H: A Family of Accurate and Efficient Hybrid Mamba-Transformer Models](https://mlsanity.com/paper/2504.03624/) (2025)
- [NVIDIA Nemotron 3: Efficient and Open Intelligence](https://mlsanity.com/paper/2512.20856/) (2025)
- [FFN Fusion: Rethinking Sequential Computation in Large Language Models](https://mlsanity.com/paper/2503.18908/) (2025)

## Built on by (most cited first, 1 of 1)

- [Nemotron-Labs-3-Puzzle-75B-A9B: Compressing Hybrid MoE LLMs](https://mlsanity.com/paper/2607.04371/) (2026)

## Abstract (arXiv)

Reasoning-focused LLMs improve answer quality by generating longer reasoning traces, but the additional tokens dramatically increase serving cost, motivating inference optimization. We extend and apply Puzzle, a post-training neural architecture search (NAS) framework, to gpt-oss-120B to produce gpt-oss-puzzle-88B, a deployment-optimized derivative. Our approach combines heterogeneous MoE expert pruning, selective replacement of full-context attention with window attention, FP8 KV-cache quantization with calibrated scales, and post-training reinforcement learning to recover accuracy, while maintaining low generation length. In terms of per-token speeds, on an 8XH100 node we achieve 1.63X and 1.22X throughput speedups in long-context and short-context settings, respectively. gpt-oss-puzzle-88B also delivers throughput speedups of 2.82X on a single NVIDIA H100 GPU. However, because token counts can change with reasoning effort and model variants, per-token throughput (tok/s) and latency (ms/token) do not necessarily lead to end-to-end speedups: a 2X throughput gain is erased if traces grow 2X. Conversely, throughput gains can be spent on more reasoning tokens to improve accuracy; we therefore advocate request-level efficiency metrics that normalize throughput by tokens generated and trace an accuracy--speed frontier across reasoning efforts. We show that gpt-oss-puzzle-88B improves over gpt-oss-120B along the entire frontier, delivering up to 1.29X higher request-level efficiency. Across various benchmarks, gpt-oss-puzzle-88B matches or slightly exceeds the parent on suite-average accuracy across reasoning efforts, with retention ranging from 100.8% (high) to 108.2% (low), showing that post-training architecture search can substantially reduce inference costs without sacrificing quality.