# Learning Reasoning World Models for Parallel Code

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Reasoning for agents, planning and code, one of 17 directions in LLM reasoning.
- Builds on 21 of its 43 references that are in the map.
- First posted on arXiv 22 Apr 2026.

## Placed in

- LLM reasoning / [Reasoning for agents, planning and code](https://mlsanity.com/lr/reasoning-for-agents-planning-and-code/)

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

- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Code Llama: Open Foundation Models for Code](https://mlsanity.com/paper/2308.12950/) (2023)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [OpenHands: An Open Platform for AI Software Developers as Generalist Agents](https://mlsanity.com/paper/2407.16741/) (2024)
- [Teaching Large Language Models to Self-Debug](https://mlsanity.com/paper/2304.05128/) (2023)
- [OpenThoughts: Data Recipes for Reasoning Models](https://mlsanity.com/paper/2506.04178/) (2025)
- [Abductive Commonsense Reasoning](https://mlsanity.com/paper/1908.05739/) (2019)
- [Quiet-STaR: Language Models Can Teach Themselves to Think Before Speaking](https://mlsanity.com/paper/2403.09629/) (2024)
- [RLEF: Grounding Code LLMs in Execution Feedback with Reinforcement Learning](https://mlsanity.com/paper/2410.02089/) (2024)
- [CWM: An Open-Weights LLM for Research on Code Generation with World Models](https://mlsanity.com/paper/2510.02387/) (2025)
- [A Survey on LLM-based Code Generation for Low-Resource and Domain-Specific Programming Languages](https://mlsanity.com/paper/2410.03981/) (2024)
- [Comparing Llama-2 and GPT-3 LLMs for HPC kernels generation](https://mlsanity.com/paper/2309.07103/) (2023)
- [ProGraML: Graph-based Deep Learning for Program Optimization and Analysis](https://mlsanity.com/paper/2003.10536/) (2020)
- [The Landscape and Challenges of HPC Research and LLMs](https://mlsanity.com/paper/2402.02018/) (2024)
- [Performance-Aligned LLMs for Generating Fast Code](https://mlsanity.com/paper/2404.18864/) (2024)
- [What I cannot execute, I do not understand: Training and Evaluating LLMs on Program Execution Traces](https://mlsanity.com/paper/2503.05703/) (2025)
- [LLM & HPC:Benchmarking DeepSeek's Performance in High-Performance Computing Tasks](https://mlsanity.com/paper/2504.03665/) (2025)
- [Beyond Code Pairs: Dialogue-Based Data Generation for LLM Code Translation](https://mlsanity.com/paper/2512.03086/) (2025)
- [Leveraging AI for Productive and Trustworthy HPC Software: Challenges and Research Directions](https://mlsanity.com/paper/2505.08135/) (2025)
- [From Evidence to Trajectory: Abductive Reasoning Path Synthesis for Retrieval-Augmented Generation Agents Development](https://mlsanity.com/paper/2509.23071/) (2025)

## Abstract (arXiv)

Large language models have shown remarkable ability in serial code generation, but they still struggle with parallel code for which training data is comparatively scarce. A common remedy is to use coding agents that interact with external tools, but tool calls can be costly and sometimes impractical, e.g., for partially written code. We propose Parallel-Code World Models (PCWMs), reasoning LLMs that aim to predict tool outcomes directly from parallel source code. To train PCWMs, we design a novel exploration and data generation pipeline that samples diverse parallel-coding problems and candidate implementations across multiple domains, then executes them via tools to record data races and performance profiles. From these, we synthesize reasoning traces that causally connect source code to observed tool outcomes. Fine-tuning on the resulting data yields noticeable gains, with a 7B-parameter world model improving from 64.3% to 72.8% accuracy in race-outcome prediction, while an 8B-parameter model improves in a performance profiling task from 49.3% to 58.6% accuracy. Furthermore, when open-weight models were tasked with fixing data races, world-model feedback improved their race-fixing rates relative to self-feedback by 2.7%-9.1% using our 7B-parameter world model and by 6.1%-11.1% using our 14B-parameter world model. Our results suggest that reasoning world models may have the potential to serve alongside external tool calls in parallel-coding agents.