Skip to content
Open
Show file tree
Hide file tree
Changes from all 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
238 changes: 148 additions & 90 deletions README.md
Original file line number Diff line number Diff line change
@@ -1,158 +1,216 @@
<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>Building cross-hardware and multi-model RL post-training infrastructure for kernel-level train–inference consistency.</strong>
</p>

<p align="center">
<a href="https://rl-align.github.io/RL-Kernel/"><img src="https://img.shields.io/badge/Documentation-Docs-2ea44f" alt="Documentation"></a>
<a href="https://rl-align.slack.com/join/shared_invite/zt-46bxj7uyt-gEK3xzwSJr_lppJsZolR~g#/shared-invite/email"><img src="https://img.shields.io/badge/Slack-Join%20Us-4A154B" alt="Slack"></a>
<a href="https://www.linkedin.com/company/rl-align"><img src="https://img.shields.io/badge/LinkedIn-Follow-0A66C2?logo=linkedin&logoColor=white" alt="Follow RL-Align on LinkedIn"></a>
<a href="https://x.com/RLKernel"><img src="https://img.shields.io/badge/X-Follow-000000?logo=x&logoColor=white" alt="Follow RL-Kernel on X"></a>
<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>

<p align="center">
<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>

**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**.

**RL-Kernel** is high-performance infrastructure for RL post-training. It provides
deterministic operators for consistent numerical computation across rollout and training
engines, together with hardware-specific kernels for faster execution and lower memory
use in GRPO, PPO, and related workloads.

---
Today, the end-to-end path covers **Qwen3-8B Dense with vime, vLLM, and Megatron-LM**.
Work on DeepSeek-V4 Flash MoE, Miles, and AReaL is ongoing.

# Our Core Philosophy
## Why RL-Kernel?

**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.
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.

**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.
- **Exact train–inference consistency:** deterministic operators keep rollout and training
computations aligned. The published experiment records exact runtime LogP agreement
across all 200 training steps.
- **RL operators:** deterministic attention, dense FFN, LogP, GRPO and PPO objectives,
and collectives cover the numerical boundaries in RL post-training.
- **Performance:** fused computation and hardware-specific kernels reduce rollout time,
memory use, and synchronization costs.
- **vime integration:** vime orchestrates vLLM rollout and Megatron-LM training, with
RL-Kernel supplying the operators used by both engines.
- **Hardware:** NVIDIA SM90 and AMD gfx942 are supported. Ascend dav_c220 has partial
operator coverage. Support for other hardware is in progress.

---
## Architecture

# Global Architecture
RL-Kernel sits between execution engines and accelerator backends. Its runtime adapters
select the operator implementation for each backend while keeping the same numerical
contract across rollout and training.

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.
The architecture below shows how orchestration frameworks, execution engines, RL-Kernel
operators, and hardware backends fit together.

<p align="center">
<img src="docs/assets/RL-Kernel underlying operator library technical architecture.png" alt="RL-Kernel Global Architecture" width="800">
<img src="docs/assets/RL-Kernel underlying operator library technical architecture.png" alt="RL-Kernel global architecture" width="800">
</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.*
The smaller diagram shows the current vime integration.

```mermaid
flowchart TB
vime["vime · RL orchestration"] --> VLLM["vLLM · rollout"]
vime --> MEGATRON["Megatron-LM · training"]
VLLM --> RLK["RL-Kernel · deterministic and optimized operators"]
MEGATRON --> RLK
RLK --> CUDA["CUDA · SM90<br/>H100, H200, GH200"]
RLK --> ROCM["ROCm · gfx942<br/>MI300A, MI300X, MI325X"]
RLK -.-> ASCEND["Ascend · dav_c220<br/>partial adaptation"]
```

---
The benchmark below uses **vime, vLLM, Megatron-LM, and CUDA**. See
[runtime dispatch](./docs/design/runtime-dispatch.md) for operator selection.

# Performance Benchmarks: Breaking the Memory Wall
## Current Scope and Roadmap

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:
The current end-to-end path uses Qwen3-8B Dense with vime.

### 1. Logprob Computation (Training Stability)
By implementing **Pre-allocated Chunking**, RL-Kernel maintains constant additional VRAM overhead regardless of the group size ($G$).
| Area | Current | Next |
| :--- | :--- | :--- |
| **Model** | Qwen3-8B Dense | [DeepSeek-V4-Flash-0731 MoE](./docs/blog/2026-08-09-dsv4-flash-moe-consistency-roadmap.md) |
| **Orchestration** | vime | Miles and AReaL |
| **Engines** | vLLM rollout and Megatron-LM training | More rollout and training engines |

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

| 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** |
### vime native vs. RL-Kernel + vime

*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.*
[PR #377](https://github.com/RL-Align/RL-Kernel/pull/377) compares vime native G10 with
RL-Kernel G11. Both runs use vime, vLLM rollout, Megatron-LM training, and rollout LogP
reuse. G11 uses RL-Kernel attention, FFN, and LogP in rollout and training.

**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.

<p align="center">
<img src="docs/assets/1. VRAM Saving.png" alt="Logprob VRAM Benchmark">
</p>
| Metric | vime native (G10) | vime + RL-Kernel (G11) | G11 result |
| :--- | ---: | ---: | :--- |
| Steps with nonzero train–rollout LogP mismatch | 200 of 200 | **0 of 200** | **Exact agreement at every step** |
| Maximum absolute Δlogp across the run | 1.591547 | **0** | **Zero measured difference** |
| Mean rollout time | 130.22 seconds per step | **82.75 seconds per step** | **36.5% lower** |
| Mean rollout throughput | 672.39 tokens per GPU per second | **1,134.00 tokens per GPU per second** | **68.7% higher** |
| Mean reference LogP time | 20.90 seconds per step | 20.92 seconds per step | Approximately equal |
| Mean actor training time | **80.51 seconds per step** | 107.18 seconds per step | 33.1% higher |
| Mean end-to-end step time | 251.99 seconds per step | **231.27 seconds per step** | **8.2% lower** |

### 2. Sampling Latency (Rollout Speed)
Integrating **FlashInfer** fused kernels to accelerate the bottleneck of RL training: the sampling phase.
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**.

| 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** |
![Qwen3-8B performance comparison: stage times, throughput, and relative changes for vime native G10 and optimized RL-Kernel G11](./docs/assets/qwen3-8b-performance-summary.png)

![Qwen3-8B training, reward, and train–rollout consistency curves for G10 and optimized G11](./examples/vime_qwen3_8b_tp4_cp2_200/results/scale_reference_s1234_g10_g11_optimized/consistency-reward.png)

<p align="center">
<img src="docs/assets/2. sampling latency.png" alt="Sampling Latency Benchmark">
</p>
## Hardware Support

### 3. Real Model Validation (Qwen3-30B-A3B MoE)
RL-Kernel currently supports the following hardware targets.

**Testbed**: NVIDIA A100 80GB | **Model**: Qwen3-30B-A3B | **Vocab**: 151,936 | **dtype**: fp16
| Hardware | Architecture | Software | Status |
| :--- | :--- | :--- | :--- |
| NVIDIA H100, H200, GH200 | SM90 | CUDA | **Supported** |
| AMD Instinct MI300A, MI300X, MI325X | gfx942 | ROCm | **Supported** |
| Huawei Ascend dav_c220 | dav-2201 | CANN 9.1.0 and Ascend C | **Partial** |
| Moore Threads | In development | MUSA | **In progress** |

Model weights consume 56.9 GB — only 23 GB headroom remaining for training computation.
The published end-to-end benchmark was run on H100. The ROCm extension and backend
checks were verified on MI300X. Ascend support is limited to dav_c220. Support for other
hardware models is in progress.

## Quick Start

<p align="center">
<img src="docs/assets/3. moe .png" alt="Real Model MoE Benchmark">
</p>
Install Python 3.10 or newer, a PyTorch build matching your accelerator runtime, and the
corresponding CUDA or ROCm compiler toolchain. Then clone RL-Kernel:

# Key Features
```bash
git clone https://github.com/RL-Align/RL-Kernel.git
cd RL-Kernel
```

- **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.
For the Qwen3-8B train–rollout command and setup for vime with RL-Kernel, see the
[8×H100 integration runbook](https://github.com/RL-Align/RL-Kernel/issues/342).

---
### NVIDIA CUDA

# Architecture
Build against a visible NVIDIA GPU. Set TORCH_CUDA_ARCH_LIST when you want to pin the
target architecture instead of relying on device detection.

RL-Kernel sits between high-level alignment libraries and low-level GPU kernels, ensuring maximum throughput without sacrificing flexibility.
```bash
# NVIDIA SM90: H100, H200, GH200
MAX_JOBS=8 \
RL_KERNEL_REQUIRE_EXT=1 \
TORCH_CUDA_ARCH_LIST="9.0+PTX" \
python3 -m pip install --no-build-isolation --no-deps -e .
```

---
The CUDA build targets SM90 and has been tested on an NVIDIA H100 80GB HBM3. H100, H200,
and GH200 use SM90. Support for other CUDA architectures is in progress.

# Quick Start
Verify the loaded extension, GPU, SM capability, and required native symbol:

### 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)
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__)"
python3 -c "import torch, rl_engine._C as C; print('GPU:', torch.cuda.get_device_name(0)); print('Capability:', torch.cuda.get_device_capability(0)); print('Extension:', C.__file__); print('fused_logp:', hasattr(C, 'fused_logp')); assert hasattr(C, 'fused_logp'); print('H100 build: PASS')"
```

### 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.
### AMD ROCm

# Support
Don’t hesitate to ask!
The gfx942 build targets AMD Instinct MI300A, MI300X, and MI325X:

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.
```bash
PYTORCH_ROCM_ARCH=gfx942 python3 setup.py develop
```

[Open an issue](https://github.com/RL-Align/RL-Kernel/issues) if you find a bug in **RL-Kernel**.
Verify the ROCm environment and required native symbol:

# Documentation
```bash
python3 scripts/check_rocm_env.py
python3 -c "import torch, rl_engine._C as C; print('GPU:', torch.cuda.get_device_name(0)); print('HIP:', torch.version.hip); print('Extension:', C.__file__); print('fused_logp:', hasattr(C, 'fused_logp')); assert hasattr(C, 'fused_logp'); print('MI300X build: PASS')"
```

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.
The extension and environment checks have been tested on AMD Instinct MI300X. Support for
other ROCm architectures is in progress.

Featured docs:
For CPU-only or pure-Python development, use an editable pip installation. Ascend has
partial operator support on dav_c220 with the dav-2201 target. Moore Threads support is
in progress. See the [installation guide](./docs/getting_started/installation.md) for
backend dependencies and troubleshooting.

- [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/vllm-project/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).
Binary file added docs/assets/qwen3-8b-performance-summary.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
4 changes: 3 additions & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -211,7 +211,9 @@ def get_extensions():
nvcc_flags.append("-allow-unsupported-compiler")
nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH")

cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"]
# Let PyTorch's BuildExtension choose the C++ standard required by the
# installed PyTorch release (for example, C++20 in PyTorch 2.14).
cxx_flags = ["-O3", "-DKERNEL_ALIGN_WITH_CUDA"]
extra_link_args = list(torch_rpath)
if os.name != "nt":
# CUDA IPC metadata queries use the driver API (cuPointerGetAttribute).
Expand Down
Loading