KO | EN

GitLite

Minimalist Open Source Repository Search Engine

gitlite — search

#dsh #dsh-plugin #codex #typescript #developer-tools #ai-agents #deepseek-harness #claude-code #ai-agent #cli #react #mcp

knowledge_flow

★ 35 Open GitHub ↗

Official Implementation of Knowledge Flow Prompting

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# Knowledge Flow: Scaling Reasoning Beyond the Context Limit

Yufan Zhuang1,2, Liyuan Liu1, Dinghuai Zhang1, Chandan Singh1, Yelong Shen1, Jingbo Shang2, Jianfeng Gao1

1Microsoft Research 2UC San Diego

We explore how **Knowledge Flow Prompting**—iteratively updating a knowledge list between rollouts—overcomes the context limit of LLMs in test-time scaling. This iterative refinement process mirrors human deliberation: the model progressively distills insights from attempts, refining a knowledge list, thereby empowering subsequent rollouts.
Knowledge Flow enables both gpt-oss-120b and Qwen3-235B-A22B-Thinking to achieve **100% accuracy on AIME25 without any training/tools/external feedback**.

For detailed insights into the methodology and results, please refer to our [blog post](https://yufanzhuang.notion.site/knowledge-flow).

---

## 🚀 Quick Start

### Environment Setup

#### 1. Prerequisites
- CUDA >= 12.8
- VLLM == 0.10.2

#### 2. Installation

**Option A: Minimal Installation** (recommended for quickstarts)
```bash
pip install -r requirements.txt
```
This installs the core dependencies.

**Option B: Full Installation** (includes all dependency versions)
```bash
pip install -r requirements_full.txt
```

This includes additional packages for extended functionality including CUDA libraries, evaluation tools, and more.

---

## 📂 Project Structure

### Main Implementation

The core Knowledge Flow implementation is available in two variants:

#### **For GPT-OSS Models:**
- **Main Script:** [`scripts/vllm_kflow_oss.py`](scripts/vllm_kflow_oss.py)
- **Bash Script:** [`bash_scripts/gpt_oss.sh`](bash_scripts/gpt_oss.sh)
- **Model:** `openai/gpt-oss-120b`

#### **For Qwen Models:**
- **Main Script:** [`scripts/vllm_kflow_qwen.py`](scripts/vllm_kflow_qwen.py)
- **Bash Script:** [`bash_scripts/qwen3.sh`](bash_scripts/qwen3.sh)
- **Model:** `Qwen/Qwen3-235B-A22B-Thinking-2507`

### Key Parameters

Both main scripts support the following arguments:

```bash
--model_name        # Hugging Face model identifier
--max_new_tokens    # Maximum generation length (e.g., 131072 for GPT-OSS, 262144 for Qwen)
--temperature       # Sampling temperature (default: 0.6)
--reflex_size       # Number of iterations (default: 64)
--split             # Dataset split: train/test/validation
--output_postfix    # Custom postfix for output files
```

---

## 🔬 Ablation Studies

We provide several ablation variants to study the impact of different components:

### 1. **Knowledge Regeneration Variants**
These scripts test the effect of regenerating knowledge at each step:

- [`scripts/vllm_kflow_oss_knowledgeRegen.py`](scripts/vllm_kflow_oss_knowledgeRegen.py)
- [`scripts/vllm_kflow_qwen_knowledgeRegen.py`](scripts/vllm_kflow_qwen_knowledgeRegen.py)

**Key Difference:** Uses NLTK sentence tokenization to segment and regenerate knowledge descriptions more systematically.

### 2. **Positive Knowledge Variants**
These scripts explore using positive reinforcement (correct solutions) instead of mistake tracking:

- [`scripts/vllm_kflow_oss_positiveKnowledge.py`](scripts/vllm_kflow_oss_positiveKnowledge.py)
- [`scripts/vllm_kflow_qwen_positiveKnowledge.py`](scripts/vllm_kflow_qwen_positiveKnowledge.py)

**Key Difference:** Builds a knowledge base from successful solutions rather than learning from mistakes.

### 3. **Markovian Thinking Baseline**
These scripts implement a baseline that doesn't maintain long-term knowledge:

- [`scripts/vllm_oss_markovian_thinking.py`](scripts/vllm_oss_markovian_thinking.py)
- [`scripts/vllm_qwen_markovian_thinking.py`](scripts/vllm_qwen_markovian_thinking.py)

**Key Difference:** Tests the Markovian Thinking where each iteration only depends on the immediate previous step, without accumulated knowledge.

---

## 💻 Usage Examples

### Running Knowledge Flow on GPT-OSS

```bash
cd bash_scripts
bash gpt_oss.sh
```

Or run directly with Python:
```bash
model_name=openai/gpt-oss-120b
BACKEND=TRITON_ATTN_VLLM_V1
export HF_HUB_ENABLE_HF_TRANSFER=1
export VLLM_FLASHINFER_ALLREDUCE_FUSION_THRESHOLDS_MB='{"2":32,"4":32,"8":8}'
export VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8=1

VLLM_ATTENTION_BACKEND=$BACKEND TOKENIZERS_PARALLELISM=false PYTHONPATH="../":"$PYTHONPATH" python3 ../scripts/vllm_kflow_oss.py \
   --model_name "$model_name" \
   --max_new_tokens 131072 \
   --temperature 0.6 \
   --reflex_size 64 
```

---

## 📊 Output Format

Results are saved in `./results/vllm_{max_new_tokens}_{output_postfix}/{model_name}/`:
- `predictions_{step}.json` - Predictions at each step
- Includes: questions, generated solutions, extracted answers, gold answers, and correctness

---

## 🔧 Technical Details

### vLLM Backend Configuration

We used this setting on B200 in our experiments, adjust accordingly to your compute platforms.

**For GPT-OSS:**
```bash
export VLLM_ATTENTION_BACKEND=TRITON_ATTN_VLLM_V1
export HF_HUB_ENABLE_HF_TRANSFER=1
export VLLM_FLASHINFER_ALLREDUCE_FUSION_THRESHOLDS_MB='{"2":32,"4":32,"8":8}'
export VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8=1
```

**For Qwen:**
```bash
export VLLM_ATTENTION_BACKEND=FLASH_ATTN
export HF_HUB_ENABLE_HF_TRANSFER=1
```

---

## 📚 Citation

If you find this work helpful, please cite us:

```bibtex
@misc{zhuang2025knowledgeflow,
  title = {Knowledge Flow: Scaling Reasoning Beyond the Context Limit},
  url = {https://yufanzhuang.notion.site/knowledge-flow},
  author = {Zhuang, Yufan and Liu, Liyuan and Zhang, Dinghuai and Singh, Chandan and Shen, Yelong and Shang, Jingbo and Gao, Jianfeng},
  journal = {Notion Blog},
  year = {2025},
  month = Oct,
}
```

---

## 📄 License

We release our code under Apache 2.0, see [LICENSE](LICENSE) file for details.

knowledge_flow

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README