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

• URL: http://arxiv.org/abs/2602.11937v1
• Date: Thu, 12 Feb 2026 13:36:19 GMT
• Title: Extending Puzzle for Mixture-of-Experts Reasoning Models with Application to GPT-OSS Acceleration
• Abstract summary: Reasoning-focused LLMs improve answer quality by generating longer reasoning traces, but the additional tokens dramatically increase serving cost.<n>We extend and apply Puzzle, a post-training neural architecture search (NAS) framework, to produce gpt-oss-puzzle-88B.<n>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.
• Score: 17.929617807312635
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• A.X K2 Technical Report [16.892621349469206]
  We introduce A.X K2, a language model trained from scratch as a high-performance foundation for emphagentic applications.<n>A.X K2 improves over A.X K1 across the board, by over 30 percentage points on some benchmarks.<n>A.X K2 performs competitively against strong open-weight baselines, matching or exceeding them on math and Korean-language benchmarks.
  arXiv  Detail & Related papers  (2026-08-31T03:04:30Z)
• LaTER: Efficient Test-Time Reasoning via Latent Exploration and Explicit Verification [35.08680804423239]
  Chain-of-thought (CoT) reasoning improves large language models (LLMs) on difficult tasks, but it also makes inference expensive.<n>We propose Latent-Then-Explicit Reasoning (LaTER), a two-stage paradigm that first performs bounded exploration in a continuous latent space.<n>LaTER projects final-layer hidden states back to the input embedding space, preserves the latent KV cache, and uses entropy and model-native stop-token probes to decide when to switch.
  arXiv  Detail & Related papers  (2026-05-08T06:23:58Z)
• Self-Adversarial One Step Generation via Condition Shifting [11.426065945265647]
  APEX is a framework for adversarial correction signals extracted endogenously from a flow model.<n>Our 0.6B model surpasses FLUX-Schnell 12B (20$times$ more parameters) in one step quality.<n>With LoRA tuning on Qwen-Image 20B, APEX reaches a GenEval score of 0.89 at NFE=1 in 6 hours, surpassing the original 50-step teacher (0.87) and providing a 15.33$times$ inference speedup.
  arXiv  Detail & Related papers  (2026-04-14T05:54:33Z)
• $\ abla$-Reasoner: LLM Reasoning via Test-Time Gradient Descent in Latent Space [71.23672814629448]
  $nabla$-Reasoner is an iterative generation framework that integrates differentiable optimization over token logits into the decoding loop.<n>$nabla$-Reasoner achieves over 20% accuracy improvement on a challenging mathematical reasoning benchmark.
  arXiv  Detail & Related papers  (2026-03-05T08:42:54Z)
• LLaDA2.1: Speeding Up Text Diffusion via Token Editing [72.92893150459909]
  We unveil LLaDA2.1, a paradigm shift designed to transcend the trade-off between decoding speed and generation quality.<n>By seamlessly weaving Token-to-Token (T2T) editing into the conventional Mask-to-Token (M2T) scheme, we introduce a joint, threshold-decoding scheme.<n>This structural innovation gives rise to two distinct personas: the Speedy Mode (S Mode), which audaciously lowers the M2T threshold to bypass traditional constraints while relying on T2T to refine the output; and the Quality Mode (Q Mode), which leans into conservative thresholds to secure superior benchmark
  arXiv  Detail & Related papers  (2026-02-09T14:00:07Z)
• LACONIC: Dense-Level Effectiveness for Scalable Sparse Retrieval via a Two-Phase Training Curriculum [73.82125917416067]
  LACONIC is a family of learned sparse retrievers based on the Llama-3 architecture.<n>The 8B variant achieves a state-of-the-art 60.2 nDCG on the MTEB Retrieval benchmark, ranking 15th on the leaderboard.
  arXiv  Detail & Related papers  (2026-01-04T22:42:20Z)
• DTS: Enhancing Large Reasoning Models via Decoding Tree Sketching [54.98126916293868]
  Large Reasoning Models (LRMs) produce excessively long chain-of-thought traces that degrade accuracy.<n>We propose a model-agnostic decoding framework that sketches the reasoning space by branching at high-entropy tokens and applies early stopping to select the shortest completed reasoning path.<n>This approach approximates the optimal solution that enhances both efficiency and accuracy, without requiring additional training or supervision.
  arXiv  Detail & Related papers  (2025-11-01T17:41:28Z)
• Enabling Vibration-Based Gesture Recognition on Everyday Furniture via Energy-Efficient FPGA Implementation of 1D Convolutional Networks [11.481972015296812]
  This study proposes an energy-efficient solution deploying compact NNs on low-power Field-Programmable Gate Arrays (FPGAs)<n>We replace complex spectral preprocessing with raw waveform input, eliminating complex on-board preprocessing while reducing input size by 21x without sacrificing accuracy.<n>We design two lightweight architectures (1D-CNN and 1D-SepCNN) tailored for embedded FPGAs, reducing parameters from 369 million to as few as 216 while maintaining comparable accuracy.
  arXiv  Detail & Related papers  (2025-10-27T09:30:36Z)
• DLER: Doing Length pEnalty Right - Incentivizing More Intelligence per Token via Reinforcement Learning [134.03095505580276]
  Doing Length pEnalty Right (DLER) is a training recipe combining batch-wise reward normalization, higher clipping, dynamic sampling, and a simple truncation length penalty.<n>DLER achieves state-of-the-art accuracy--efficiency trade-offs, cutting output length by over 70 percent while surpassing all previous baseline accuracy.
  arXiv  Detail & Related papers  (2025-10-16T20:05:57Z)
• Intra-request branch orchestration for efficient LLM reasoning [52.68946975865865]
  Large Language Models (LLMs) increasingly rely on inference-time reasoning algorithms to improve accuracy on complex tasks.<n>Prior work has largely focused on reducing token usage, often at the expense of accuracy, while overlooking other latency factors.<n>We present DUCHESS, an LLM serving system that reduces cost and latency without sacrificing accuracy through intra-request branch orchestration guided by predictions.
  arXiv  Detail & Related papers  (2025-09-29T15:52:08Z)
• LAWCAT: Efficient Distillation from Quadratic to Linear Attention with Convolution across Tokens for Long Context Modeling [27.045621004239067]
  We propose LAWCAT, a novel linearization framework designed to efficiently transfer the capabilities of pre-trained transformers into a performant linear attention architecture.<n> LAWCAT integrates causal Conv1D layers to enhance local dependency modeling and employs normalized gated linear attention to improve generalization across varying context lengths.<n>Our evaluations demonstrate that, distilling Mistral-7B with only 1K-length sequences yields over 90% passkey retrieval accuracy up to 22K tokens.
  arXiv  Detail & Related papers  (2025-09-22T22:43:44Z)
• R-Stitch: Dynamic Trajectory Stitching for Efficient Reasoning [80.104336426172]
  Chain-of-thought (CoT) enhances problem-solving ability of large language models.<n>CoT incurs substantial inference cost due to long autoregressive trajectories.<n>We introduce R-Stitch, a training-free hybrid decoding framework.
  arXiv  Detail & Related papers  (2025-07-23T08:14:36Z)
• Scaling Attention to Very Long Sequences in Linear Time with Wavelet-Enhanced Random Spectral Attention (WERSA) [1.7622426179653563]
  Transformer models are computationally costly on long sequences since regular attention has quadratic $O(n2)$ time complexity.<n>We introduce Wavelet-Enhanced Random Spectral Attention (WERSA), a novel mechanism of linear $O(n)$ time complexity.<n>By significantly reducing computational loads without compromising accuracy, WERSA makes possible more practical, more affordable, long-context models.
  arXiv  Detail & Related papers  (2025-07-11T14:40:40Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.