Extending Puzzle for Mixture-of-Experts Reasoning Models with Application to GPT-OSS Acceleration
Fuente:
arXiv
Guardado en:
| Autores principales: | , , , , , , , , , , , , , , , , , , , , , , , |
|---|---|
| Formato: | Preprint |
| Publicado: |
2026
|
| Materias: | |
| Acceso en línea: | |
| Etiquetas: |
Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
|
| _version_ | 1866918411387273216 |
|---|---|
| author | Bercovich, Akhiad Ailon, Nir Anisimov, Vladimir Asida, Tomer Assaf, Nave Dabbah, Mohammad Galil, Ido Geifman, Amnon Geifman, Yonatan Golan, Izhak Koren, Roi Levy, Itay Moshe, Zach Molchanov, Pavlo Nabwani, Najeeb Patwary, Mostofa Puny, Omri Ronen, Tomer Schen, Itamar Segal, Elad Shahaf, Ido Tropp, Oren Zilberstein, Ran El-Yaniv, Ran |
| author_facet | Bercovich, Akhiad Ailon, Nir Anisimov, Vladimir Asida, Tomer Assaf, Nave Dabbah, Mohammad Galil, Ido Geifman, Amnon Geifman, Yonatan Golan, Izhak Koren, Roi Levy, Itay Moshe, Zach Molchanov, Pavlo Nabwani, Najeeb Patwary, Mostofa Puny, Omri Ronen, Tomer Schen, Itamar Segal, Elad Shahaf, Ido Tropp, Oren Zilberstein, Ran El-Yaniv, Ran |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2602_11937 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | Extending Puzzle for Mixture-of-Experts Reasoning Models with Application to GPT-OSS Acceleration Bercovich, Akhiad Ailon, Nir Anisimov, Vladimir Asida, Tomer Assaf, Nave Dabbah, Mohammad Galil, Ido Geifman, Amnon Geifman, Yonatan Golan, Izhak Koren, Roi Levy, Itay Moshe, Zach Molchanov, Pavlo Nabwani, Najeeb Patwary, Mostofa Puny, Omri Ronen, Tomer Schen, Itamar Segal, Elad Shahaf, Ido Tropp, Oren Zilberstein, Ran El-Yaniv, Ran Machine Learning 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. |
| title | Extending Puzzle for Mixture-of-Experts Reasoning Models with Application to GPT-OSS Acceleration |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2602.11937 |