Skip to content
Open
Changes from 5 commits
Commits
Show all changes
27 commits
Select commit Hold shift + click to select a range
75baabc
docs(readme): update feature and platform support
frank-2077 Sep 1, 2026
9ca6008
docs(readme): highlight release results and accelerator support
Flink-ddd Sep 5, 2026
b790d81
docs(readme): restore global architecture diagram
Flink-ddd Sep 5, 2026
2f13b0c
docs(readme): clarify current scope and roadmap
Flink-ddd Sep 5, 2026
76bf455
docs(readme): streamline project overview
Flink-ddd Sep 5, 2026
4c0aaf6
docs(readme): align introduction with project positioning
Flink-ddd Sep 5, 2026
11ee862
docs(readme): remove opening benchmark summary
Flink-ddd Sep 5, 2026
fb25ffe
docs(readme): split quick start by accelerator
Flink-ddd Sep 5, 2026
f5b14c5
build: defer C++ standard selection to PyTorch
Flink-ddd Sep 5, 2026
d4588da
docs(readme): specify supported accelerator targets
Flink-ddd Sep 5, 2026
e535a66
docs(readme): record validated H100 build command
Flink-ddd Sep 5, 2026
fe594b7
docs(readme): record MI300X validation
Flink-ddd Sep 5, 2026
edfa45f
docs(readme): use validated MI300X build command
Flink-ddd Sep 5, 2026
a740c89
docs(readme): simplify hardware support language
Flink-ddd Sep 5, 2026
267d856
docs(readme): remove inline code styling
Flink-ddd Sep 5, 2026
b69c5d5
docs(readme): rename benchmark comparison
Flink-ddd Sep 5, 2026
7dc1a9b
docs(readme): link Qwen3-8B integration runbook
Flink-ddd Sep 5, 2026
9af6f99
docs(readme): fix benchmark chart label overlap
Flink-ddd Sep 5, 2026
3583246
docs(readme): simplify hardware support copy
Flink-ddd Sep 5, 2026
ad00e55
docs(readme): fix VIME repository link
Flink-ddd Sep 5, 2026
9f8e6db
docs(readme): add X profile badge
Flink-ddd Sep 5, 2026
77d8a14
Merge branch 'main' into docs/readme-platform-support
Flink-ddd Sep 5, 2026
590a087
docs(readme): use official vime casing
Flink-ddd Sep 5, 2026
09b34b9
docs(readme): show only validated vime path
Flink-ddd Sep 6, 2026
edf90ed
docs(readme): separate framework and hardware support
Flink-ddd Sep 6, 2026
4430b8d
docs(readme): remove MUSA from vime diagram
Flink-ddd Sep 6, 2026
c36aab7
docs(readme): add consistency curves
Flink-ddd Sep 6, 2026
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
243 changes: 138 additions & 105 deletions README.md
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
<p align="center">
<img src="docs/assets/logo.png" width="220" alt="RL-Kernel Logo">
<img src="docs/assets/logo.png" width="220" alt="RL-Kernel logo">
</p>

<h1 align="center">RL-Kernel</h1>

<p align="center">
<strong>Extreme Infrastructure for GRPO & Large-Scale Reinforcement Learning.</strong>
<strong>Bitwise train–inference consistency. Faster RL post-training.</strong>
</p>

<p align="center">
Expand All @@ -14,145 +14,178 @@
<a href="./docs/community/wechat.md"><img src="https://img.shields.io/badge/WeChat-Join%20Group-07C160?logo=wechat&logoColor=white" alt="WeChat"></a>
<a href="./docs/assets/whatsapp-group.png"><img src="https://img.shields.io/badge/WhatsApp-Join%20Group-25D366?logo=whatsapp&logoColor=white" alt="WhatsApp"></a>
<a href="https://deepwiki.com/RL-Align/RL-Kernel"><img src="https://img.shields.io/badge/Ask-DeepWiki-7B3FE4" alt="Ask DeepWiki"></a>
<a href="https://github.com/RL-Align/RL-Kernel"><img src="https://img.shields.io/badge/Hardware-NVIDIA%20CUDA%20%7C%20AMD%20ROCm-orange" alt="Hardware"></a>
<a href="https://opensource.org/licenses/Apache-2.0"><img src="https://img.shields.io/badge/License-Apache%202.0-blue.svg" alt="License"></a>
<a href="#hardware-support"><img src="https://img.shields.io/badge/Supported-CUDA%20%7C%20ROCm-2ea44f" alt="CUDA and ROCm supported"></a>
<a href="https://opensource.org/licenses/Apache-2.0"><img src="https://img.shields.io/badge/License-Apache%202.0-blue.svg" alt="Apache 2.0 license"></a>
</p>


**RL-Kernel** is a high-performance, memory-efficient infrastructure for Reinforcement Learning post-training. It eliminates the memory and latency bottlenecks in Large Language Model alignment, This project targets AI infrastructure engineers, algorithm researchers, and enterprise-level large model alignment scenarios, providing specialized kernels for algorithms like **GRPO**, **PPO**, and **DPO**.


---

# Our Core Philosophy

**1. Operator-Level Train-Inference Consistency**
The biggest hidden barrier in large-scale RL is the subtle numerical divergence between rollout engines (e.g., vLLM) and training engines (e.g., Megatron/DeepSpeed). RL-Kernel provides mathematically rigorous, fused operators that lock down the computational graph. By guaranteeing absolute numerical consistency and deterministic reduction orders across the entire RL loop, we prevent reward hacking and distribution drift at the operator level.

**2. Extreme Memory & Compute Efficiency**
We replace naive PyTorch paths—which suffer from $O(G \cdot L \cdot V)$ memory explosion—with specialized industrial-grade kernels (like `prefix_shared_attention` and `fused_logp`). This reduces VRAM consumption by up to 10x, unlocking massive batch sizes for GRPO workloads without triggering Out-Of-Memory (OOM) errors.

---

# Global Architecture

RL-Kernel sits strictly at the operator layer, acting as a non-intrusive bridge between high-level alignment orchestration (e.g., vime, slime) and foundational execution engines. We ensure maximum throughput and rigorous numerical parity without modifying upstream framework source code.

<p align="center">
<img src="docs/assets/RL-Kernel underlying operator library technical architecture.png" alt="RL-Kernel Global Architecture" width="800">
<a href="#architecture">Architecture</a> ·
<a href="#current-scope-and-roadmap">Current scope</a> ·
<a href="#benchmark-highlights">Results</a> ·
<a href="#hardware-support">Hardware support</a> ·
<a href="#quick-start">Quick start</a> ·
<a href="https://rl-align.github.io/RL-Kernel/">Documentation</a>
</p>

*Note: RL-Kernel integrates natively into Rollout Engines (vLLM, sglang, LMDeploy) and Training Engines (Megatron, DeepSpeed) via non-intrusive custom operator hooks, powered by underlying CUDA, Triton, and ROCm backends.*

---

# Performance Benchmarks: Breaking the Memory Wall
**RL-Kernel** is a high-performance operator library for **GRPO and PPO-style RL
post-training**. It aligns numerical computation between rollout and training engines
with deterministic attention, dense FFN, log-probability, and collective operations.
Hardware-specific kernels optimize execution and memory use beneath the RL framework.

With **VIME + vLLM + Megatron-LM**, the published Qwen3-8B experiment achieves **zero
train–rollout LogP mismatches at every one of 200 steps**, **68.7% higher rollout
throughput**, and **8.2% lower end-to-end step time** than VIME's native production
operator path. [See the configuration and results below](#benchmark-highlights).

## Why RL-Kernel?

Rollout and training engines can produce different log probabilities for the same tokens
and model weights because their kernels, batching, and reduction orders differ. Those
differences enter the policy ratios and KL terms used by RL algorithms.

- **Exact train–inference consistency:** deterministic operator contracts align the strict
dense-model path. The published experiment verifies exact runtime LogP agreement across
all 200 training steps.
- **RL-native operator stack:** deterministic attention and dense FFN, fused and
batch-invariant LogP, GRPO/PPO objectives, and deterministic collectives cover the
numerical boundaries that matter to RL post-training.
- **Efficient RL execution:** fused log-probability computation, optimized attention and FFN
paths, and deterministic collectives target rollout, memory, and synchronization costs.
- **Current integration:** VIME is the supported RL orchestration layer, with vLLM for
rollout and Megatron-LM for training. RL-Kernel supplies the operators beneath them.
- **Multiple accelerators:** dense-model train–inference consistency is supported on CUDA
and ROCm, with Ascend and Moore Threads adaptation in progress.

## Architecture

RL-Kernel sits between framework execution engines and accelerator backends. Runtime
adapters select operators through a hardware-aware registry; strict routes enforce the
required numerical contract and expose execution provenance.

The complete architecture below is a layer map, not a current support matrix. It shows
the broader external scheduling and engine ecosystem, operator-library layers, and
hardware abstraction boundary.

RL-Kernel is designed to solve the $O(G \cdot L \cdot V)$ memory explosion in DeepSeek-style **GRPO** training. A typical scenario is as follows:

### 1. Logprob Computation (Training Stability)
By implementing **Pre-allocated Chunking**, RL-Kernel maintains constant additional VRAM overhead regardless of the group size ($G$).

**Testbed**: NVIDIA A100 80GB | **Model**: Llama-3-8B | **Vocab**: 128,256 | **SeqLen**: 512

| Group Size ($G$) | TRL (Standard) | PyTorch Native | **RL-Kernel (Ours)** | Status |
| :--- | :--- | :--- | :--- | :--- |
| **G = 64** | OOM | 15.66 GB | **16.15 GB** | Success |
| **G = 128** | OOM | 31.31 GB | **31.80 GB** | Success |
| **G = 256** | **FAILED (OOM)** | 62.63 GB | **63.12 GB** | **Optimized** |
<p align="center">
<img src="docs/assets/RL-Kernel underlying operator library technical architecture.png" alt="RL-Kernel global architecture" width="800">
</p>

*Note: RL-Kernel is the only solution that successfully scales G=256 on a single A100 by keeping extra VRAM usage to a constant ~0.5GB.*
The following diagram is a concise view of the validated runtime path, planned
orchestration integrations, and accelerator coverage. Solid arrows represent the current
VIME path; dashed arrows from Miles and AReaL represent roadmap work.

```mermaid
flowchart TB
VIME["VIME · integrated"] --> ORCH["RL orchestration integration"]
MILES["Miles · roadmap"] -.-> ORCH
AREAL["AReaL · roadmap"] -.-> ORCH
ORCH --> VLLM["vLLM · rollout"]
ORCH --> MEGATRON["Megatron-LM · training"]
VLLM --> RLK["RL-Kernel · deterministic and optimized operators"]
MEGATRON --> RLK
RLK --> CUDA["CUDA · supported"]
RLK --> ROCM["ROCm · supported"]
RLK -.-> ASCEND["Ascend · partial adaptation"]
RLK -.-> MUSA["Moore Threads / MUSA · in progress"]
```

The published benchmark validates **VIME + vLLM + Megatron-LM on CUDA**. Framework and
backend coverage are documented independently. See [runtime dispatch](./docs/design/runtime-dispatch.md)
for operator selection and strict execution contracts.

<p align="center">
<img src="docs/assets/1. VRAM Saving.png" alt="Logprob VRAM Benchmark">
</p>
## Current Scope and Roadmap

### 2. Sampling Latency (Rollout Speed)
Integrating **FlashInfer** fused kernels to accelerate the bottleneck of RL training: the sampling phase.
Current end-to-end support is deliberately scoped to **Qwen3-8B Dense** and **VIME**.
Items in development or on the roadmap are not yet part of the supported path.

| Batch Size ($G$) | Native PyTorch | **RL-Kernel (Fused)** | **Speedup** |
| :--- | :--- | :--- | :--- |
| 32 | 176.79 ms | **1.08 ms** | **163x** |
| 64 | 10.54 ms | **1.31 ms** | **8x** |
| 128 | 18.89 ms | **1.86 ms** | **10x** |
| 256 | 36.23 ms | **2.94 ms** | **12x** |
| Dimension | Available today | In development / roadmap |
| :--- | :--- | :--- |
| **Model architecture** | Qwen3-8B Dense | [DeepSeek-V4-Flash-0731 MoE](./docs/blog/2026-08-09-dsv4-flash-moe-consistency-roadmap.md) — active development |
| **RL orchestration** | VIME | Miles and AReaL |
| **Execution engines** | vLLM rollout + Megatron-LM training | Additional engine integrations will follow validated operator coverage |

## Benchmark Highlights

<p align="center">
<img src="docs/assets/2. sampling latency.png" alt="Sampling Latency Benchmark">
</p>
### VIME native operators vs. VIME + RL-Kernel

### 3. Real Model Validation (Qwen3-30B-A3B MoE)
The completed experiment in [PR #377](https://github.com/RL-Align/RL-Kernel/pull/377)
compares **G10**, VIME's native production operator path, with **optimized G11**, the
strict RL-Kernel path. Both use VIME with vLLM rollout and Megatron-LM training, and both
enable rollout-logp reuse. G11 applies RL-Kernel attention, FFN, and LogP on both paths.

**Testbed**: NVIDIA A100 80GB | **Model**: Qwen3-30B-A3B | **Vocab**: 151,936 | **dtype**: fp16
**Setup:** Qwen3-8B BF16 · GRPO · 1 node with 8×H100 80GB · actor TP4/CP2/PP1 ·
two TP4 rollout engines · 8 prompts × 16 samples (batch 128) · 200 steps · seed 1234 ·
maximum response length 7,168 · KL-loss coefficient 0.001.

Model weights consume 56.9 GB — only 23 GB headroom remaining for training computation.
| Metric | VIME native (G10) | VIME + RL-Kernel (G11) | G11 result |
| :--- | ---: | ---: | :--- |
| Steps with nonzero train–rollout LogP mismatch | 200 / 200 | **0 / 200** | **Exact agreement at every step** |
| Maximum absolute Δlogp across the run | 1.591547 | **0** | **Zero measured difference** |
| Mean rollout time | 130.22 s/step | **82.75 s/step** | **36.5% lower** |
| Mean rollout throughput | 672.39 tok/GPU/s | **1,134.00 tok/GPU/s** | **68.7% higher** |
| Mean reference LogP time | 20.90 s/step | 20.92 s/step | Approximately equal |
| Mean actor training time | **80.51 s/step** | 107.18 s/step | 33.1% higher |
| Mean end-to-end step time | 251.99 s/step | **231.27 s/step** | **8.2% lower** |

G11 saves **47.47 seconds per rollout step**, offsetting the additional actor training
cost for a net saving of **20.72 seconds per end-to-end step**.

<p align="center">
<img src="docs/assets/3. moe .png" alt="Real Model MoE Benchmark">
</p>
![Qwen3-8B performance comparison: stage times, throughput, and relative changes for VIME native G10 and optimized RL-Kernel G11](https://raw.githubusercontent.com/RL-Align/RL-Kernel/40db4d31982cd4a7ba28fbc96982b2af1f62921d/examples/vime_qwen3_8b_tp4_cp2_200/results/scale_reference_s1234_g10_g11_optimized/performance-summary.png)

# Key Features
## Hardware Support

- **Zero-Growth Memory Pool**: Uses pre-allocated buffers and micro-chunking to prevent VRAM spikes during advantage calculation.
- **Fused Sampling Pipeline**: Direct integration with **FlashInfer** and **vLLM** backends for sub-2ms sampling latency.
- **Universal Backend Abstraction**: Unified API supporting both **NVIDIA (CUDA/FlashInfer)** and **AMD (ROCm/AITER)**.
- **Post-Training Ready**: Drop-in replacement for standard sampling and logprob operators in TRL or DeepSpeed-Chat.
The following matrix tracks accelerator coverage for the current dense-model path. It
does not extend the end-to-end model claim beyond Qwen3-8B Dense.

---
| Vendor | Accelerator | Software stack | Dense train–inference consistency | Coverage / progress |
| :--- | :--- | :--- | :---: | :--- |
| **NVIDIA** | GPU | CUDA | ✅ **Supported** | Dense-model strict path; Qwen3-8B H100 end-to-end results above |
| **AMD** | GPU | ROCm | ✅ **Supported** | Dense-model strict path; backend-specific setup and validation |
| **Huawei** | Ascend NPU | CANN / Ascend C | 🟡 **Partially adapted** | Selected operators adapted; broader dense-model integration in progress |
| **Moore Threads** | GPU | MUSA | 🚧 **In progress** | Accelerator adaptation and dense-model integration underway |

# Architecture
✅ **Supported** · 🟡 **Partial adaptation** · 🚧 **Active development**

RL-Kernel sits between high-level alignment libraries and low-level GPU kernels, ensuring maximum throughput without sacrificing flexibility.
Support applies to the implemented dense-model paths; model, dtype, operator, and
parallelism coverage varies by backend. The performance numbers above are CUDA/H100
measurements. See the [installation guide](./docs/getting_started/installation.md) and
[operator catalog](./docs/operators/README.md) for backend requirements and contracts.

---
## Quick Start

# Quick Start
Install a PyTorch build matching your accelerator runtime, then install RL-Kernel from
source. Python 3.10 or newer is required.

### Installation
```bash
# Clone the repository
git clone https://github.com/RL-Align/RL-Kernel.git
cd RL-Kernel

# CPU-only / pure-Python fallback
python -m pip install -e .

# Native CUDA or ROCm extension (install a matching PyTorch build first)
# Native CUDA or ROCm extension
RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e .
python -c "import rl_engine._C as _C; assert hasattr(_C, 'fused_logp'); print(_C.__file__)"
```

### Contributions
Inspired by the kernel designs of vLLM and DeepSpeed. As an active contributor to the AI Infrastructure ecosystem, RL-Kernel aims to push the boundaries of RL efficiency.

Target: Building the most efficient RLHF toolchain for the open-source community.

# Support
Don’t hesitate to ask!

Contact the developers and community in [Slack](https://rl-align.slack.com/join/shared_invite/zt-46bxj7uyt-gEK3xzwSJr_lppJsZolR~g#/shared-invite/email) if you need any help.

[Open an issue](https://github.com/RL-Align/RL-Kernel/issues) if you find a bug in **RL-Kernel**.

# Documentation

The documentation of **RL-Kernel** is located on the website: [https://rl-align.github.io/RL-Kernel](https://rl-align.github.io/RL-Kernel)
or in the [docs](./docs) directory of the source code.

Featured docs:
For CPU-only or pure-Python development, use `python -m pip install -e .`.
Strict train–inference consistency requires the corresponding operators and runtime
configuration in both engines. Follow the [installation guide](./docs/getting_started/installation.md)
and [quick-start guide](./docs/getting_started/quickstart.md); the full benchmark
configuration and companion VIME revision are linked in [PR #377](https://github.com/RL-Align/RL-Kernel/pull/377).

- [Announcing RL-Kernel for vime: Faster and Leaner `linear_logp` for Full RL Rollouts](./docs/blog/2026-07-08-announcing-rl-kernel-linear-logp-for-vime.md)
- [中文版:发布 vime + RL-Kernel](./docs/blog/2026-07-08-announcing-rl-kernel-linear-logp-for-vime-zh.md)
## Community and Contributions

# Acknowledgments
Join us on [Slack](https://rl-align.slack.com/join/shared_invite/zt-46bxj7uyt-gEK3xzwSJr_lppJsZolR~g#/shared-invite/email)
or [WeChat](./docs/community/wechat.md), and
[open an issue](https://github.com/RL-Align/RL-Kernel/issues) for bugs and feature requests.
Contributions to kernels, framework integrations, hardware adaptation, and benchmarks
are welcome. See the [contributing guide](./docs/contributing/README.md).

RL-Kernel builds on the shoulders of excellent open-source projects:
## Acknowledgments

- **[FlashInfer](https://github.com/flashinfer-ai/flashinfer)** — We integrate FlashInfer's fused sampling kernels as the NVIDIA backend for our sampling pipeline. The sub-2ms sampling latency results are enabled by FlashInfer's highly optimized CUDA operators.
- **[vLLM](https://github.com/vllm-project/vllm)** — Inspired by vLLM's kernel design philosophy and hardware-aware scheduling approach.
- **[DeepSpeed](https://github.com/microsoft/DeepSpeed)** — Inspired by DeepSpeed's approach to memory-efficient training infrastructure.
RL-Kernel builds on the work of the open-source AI infrastructure community, including
[VIME](https://github.com/RL-Align/vime), [vLLM](https://github.com/vllm-project/vllm),
[Megatron-LM](https://github.com/NVIDIA/Megatron-LM),
[DeepSpeed](https://github.com/deepspeedai/DeepSpeed), and
[FlashInfer](https://github.com/flashinfer-ai/flashinfer).
We thank their contributors and everyone helping bring RL-Kernel to new accelerators.

We are grateful to these teams for their contributions to the open-source AI infrastructure ecosystem.
Licensed under the [Apache License 2.0](./LICENSE).
Loading