diff --git a/README.md b/README.md index e2788c6a..0686d494 100644 --- a/README.md +++ b/README.md @@ -1,158 +1,184 @@

- RL-Kernel Logo + RL-Kernel logo

RL-Kernel

- Extreme Infrastructure for GRPO & Large-Scale Reinforcement Learning. + Building cross-hardware and multi-model RL post-training infrastructure for kernel-level train–inference consistency.

Documentation Slack + Follow RL-Align on LinkedIn + Follow RL-Kernel on X WeChat WhatsApp Ask DeepWiki - Hardware - License + CUDA and ROCm supported + Apache 2.0 license

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

- RL-Kernel Global Architecture + Architecture · + Current scope · + Results · + Hardware support · + Quick start · + Documentation

-*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.* +**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. -# Performance Benchmarks: Breaking the Memory Wall +## Why RL-Kernel? -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: +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. -### 1. Logprob Computation (Training Stability) -By implementing **Pre-allocated Chunking**, RL-Kernel maintains constant additional VRAM overhead regardless of the group size ($G$). +- **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. -**Testbed**: NVIDIA A100 80GB | **Model**: Llama-3-8B | **Vocab**: 128,256 | **SeqLen**: 512 +## Architecture -| 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** | - -*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.* +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. +The architecture below shows how orchestration frameworks, execution engines, RL-Kernel +operators, and hardware backends fit together.

- Logprob VRAM Benchmark + RL-Kernel global architecture

-### 2. Sampling Latency (Rollout Speed) -Integrating **FlashInfer** fused kernels to accelerate the bottleneck of RL training: the sampling phase. +## Current Scope and Roadmap -| 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** | +The current end-to-end path uses Qwen3-8B Dense with vime. +| 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 | -

- Sampling Latency Benchmark -

+## Benchmark Highlights -### 3. Real Model Validation (Qwen3-30B-A3B MoE) +### CUDA H100 -**Testbed**: NVIDIA A100 80GB | **Model**: Qwen3-30B-A3B | **Vocab**: 151,936 | **dtype**: fp16 +[![Qwen3-8B CUDA training, reward, and train–rollout consistency curves](./examples/vime_qwen3_8b_tp4_cp2_200/results/scale_reference_s1234_g10_g11_optimized/consistency-reward.png)](https://github.com/RL-Align/RL-Kernel/pull/377) -Model weights consume 56.9 GB — only 23 GB headroom remaining for training computation. +[![Qwen3-8B CUDA mean absolute train–rollout LogP difference on NVIDIA H100](./examples/vime_qwen3_8b_tp4_cp2_200/results/scale_reference_s1234_g10_g11_optimized/mean-logp-diff.png)](https://github.com/RL-Align/RL-Kernel/pull/377) +### ROCm MI300X -

- Real Model MoE Benchmark -

+[![Qwen3-8B ROCm training, reward, and train–rollout consistency curves](./examples/vime_qwen3_8b_tp4_cp2_200/results/pr396_rocm_s1234_g10_g11_200/consistency-reward.png)](https://github.com/RL-Align/RL-Kernel/pull/400) -# Key Features +[![Qwen3-8B ROCm mean absolute train–rollout LogP difference on MI300X](./examples/vime_qwen3_8b_tp4_cp2_200/results/pr396_rocm_s1234_g10_g11_200/mean-logp-diff.png)](https://github.com/RL-Align/RL-Kernel/pull/400) -- **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. +## Hardware Support ---- +RL-Kernel currently supports the following hardware targets. -# Architecture +| 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** | -RL-Kernel sits between high-level alignment libraries and low-level GPU kernels, ensuring maximum throughput without sacrificing flexibility. +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 -# Quick Start +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: -### Installation ```bash -# Clone the repository git clone https://github.com/RL-Align/RL-Kernel.git cd RL-Kernel +``` + +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/blob/main/examples/vime_qwen3_8b_tp4_cp2_200/REPRODUCTION.md). -# CPU-only / pure-Python fallback -python -m pip install -e . +### NVIDIA CUDA -# 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__)" +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. + +```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 . ``` -### 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. +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. + +Verify the loaded extension, GPU, SM capability, and required native symbol: -Target: Building the most efficient RLHF toolchain for the open-source community. +```bash +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')" +``` -# Support -Don’t hesitate to ask! +### AMD ROCm -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. +The gfx942 build targets AMD Instinct MI300A, MI300X, and MI325X: -[Open an issue](https://github.com/RL-Align/RL-Kernel/issues) if you find a bug in **RL-Kernel**. +```bash +PYTORCH_ROCM_ARCH=gfx942 python3 setup.py develop +``` -# Documentation +Verify the ROCm environment and required native symbol: + +```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), 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). diff --git a/examples/vime_qwen3_8b_tp4_cp2_200/results/pr396_rocm_s1234_g10_g11_200/mean-logp-diff.png b/examples/vime_qwen3_8b_tp4_cp2_200/results/pr396_rocm_s1234_g10_g11_200/mean-logp-diff.png index c188a526..22e52c32 100644 Binary files a/examples/vime_qwen3_8b_tp4_cp2_200/results/pr396_rocm_s1234_g10_g11_200/mean-logp-diff.png and b/examples/vime_qwen3_8b_tp4_cp2_200/results/pr396_rocm_s1234_g10_g11_200/mean-logp-diff.png differ diff --git a/examples/vime_qwen3_8b_tp4_cp2_200/results/scale_reference_s1234_g10_g11_optimized/mean-logp-diff.png b/examples/vime_qwen3_8b_tp4_cp2_200/results/scale_reference_s1234_g10_g11_optimized/mean-logp-diff.png new file mode 100644 index 00000000..b4400f34 Binary files /dev/null and b/examples/vime_qwen3_8b_tp4_cp2_200/results/scale_reference_s1234_g10_g11_optimized/mean-logp-diff.png differ diff --git a/setup.py b/setup.py index ee560db0..542b4077 100644 --- a/setup.py +++ b/setup.py @@ -1,444 +1,444 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -import importlib.util -import os -import sysconfig -import warnings -from distutils.errors import CompileError -from distutils.spawn import find_executable -from pathlib import Path - -from setuptools import Extension, find_packages, setup - - -def _load_envs_module(): - envs_path = Path(__file__).with_name("envs.py") - spec = importlib.util.spec_from_file_location("_rl_kernel_envs", envs_path) - if spec is None or spec.loader is None: - raise RuntimeError(f"failed to load environment helpers from {envs_path}") - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -envs = _load_envs_module() - - -def _load_torch_extension_tools(): - try: - import torch - except ModuleNotFoundError as exc: - if exc.name != "torch": - raise - return None, None, None - - from torch.utils.cpp_extension import BuildExtension, CUDAExtension - - # CUDAExtension is also the supported extension entry point for ROCm - # PyTorch builds. BuildExtension dispatches .cu/.hip sources to hipcc when - # torch.version.hip is set. - return torch, BuildExtension, CUDAExtension - - -def _native_extension_required() -> bool: - """Whether the caller explicitly requested a native extension build.""" - return ( - envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) - or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) - or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) - or envs.env_flag("FORCE_CUDA") - ) - - -def _cuda_define_from_env(name: str, macro: str) -> list[str]: - value = os.environ.get(name) - if value is None: - return [] - parsed = int(value) - if parsed <= 0: - raise ValueError(f"{name} must be positive, got {value!r}") - return [f"-D{macro}={parsed}"] - - -_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES = ( - "-Xfatbin", - "-compress-all", - "-gencode", - "--generate-code", - "--expt-", - "-lineinfo", - "-allow-unsupported-compiler", - "-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH", -) -_ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE = { - "-Xfatbin", - "-gencode", - "--generate-code", -} - - -def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: - """Remove CUDA-only device compiler flags before BuildExtension calls hipcc.""" - filtered_flags = [] - skip_next = False - for flag in flags: - if skip_next: - skip_next = False - continue - if flag in _ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE: - skip_next = True - continue - if flag.startswith(_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES): - continue - filtered_flags.append(flag) - return filtered_flags - - -def get_extensions(): - torch, _, CUDAExtension = _load_torch_extension_tools() - if torch is None: - message = ( - "PyTorch is unavailable, so rl_engine._C cannot be built. Install a matching " - "CUDA/ROCm PyTorch build first, then run " - "`RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e .`." - ) - if _native_extension_required(): - raise RuntimeError(message) - warnings.warn( - f"{message} Continuing with the pure-Python fallback because no native extension " - "was explicitly requested.", - RuntimeWarning, - stacklevel=2, - ) - return [] - - extensions = [] - torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib") - torch_rpath = ["-Wl,-rpath,$ORIGIN/../torch/lib"] - if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": - torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") - is_rocm = getattr(torch.version, "hip", None) is not None - - # CUDAExtension is intentionally used for both CUDA and ROCm. On ROCm, - # PyTorch's BuildExtension hipifies CUDA sources and invokes hipcc; it also - # consumes PYTORCH_ROCM_ARCH (one or more ';'-separated gfx targets) to add - # --offload-arch. Do not require a visible GPU when a ROCm target was - # explicitly selected. - no_rocm_arch = not os.environ.get("PYTORCH_ROCM_ARCH", "").strip() - if is_rocm and no_rocm_arch and torch.cuda.device_count() == 0: - raise RuntimeError( - "ROCm builds without a visible GPU require PYTORCH_ROCM_ARCH. " - "Set one or more ';'-separated targets, for example " - "PYTORCH_ROCM_ARCH='gfx942;gfx950'." - ) - - if is_rocm or torch.cuda.is_available(): - cuda_sources = [ - "csrc/ops.cpp", - "csrc/fused_logp_kernel.cu", - "csrc/deterministic_logp_kernel.cu", - "csrc/cuda/gemm/det_gemm_kernel.cu", - "csrc/cuda/rmsnorm.cu", - "csrc/cuda/activation.cu", - "csrc/cuda/attention/deterministic_attention.cu", - ] - if is_rocm: - # ROCm-tuned WS2 vocab-parallel logprob kernels; the shared - # deterministic_logp_kernel.cu keeps the SM90-tuned CUDA path. - cuda_sources.extend( - [ - "csrc/hip/hip_deterministic_logp_kernel.hip", - "csrc/rocm/distributed/deterministic_collective.hip", - ] - ) - else: - # CUDA IPC and the fixed-tree collective implementation are not - # part of the ROCm extension. - cuda_sources.append("csrc/cuda/distributed/deterministic_collective.cu") - # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). - # The ROCm dispatcher falls back to PyTorch SDPA for this operator. - cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") - - nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] - if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): - nvcc_flags.append("--use_fast_math") - if not is_rocm: - cc_major, cc_minor = torch.cuda.get_device_capability() - enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" - if not enable_sm90: - # SM90 build emits 90a below; mixing plain compute_90 breaks TMA ptxas. - nvcc_flags.append( - f"-gencode=arch=compute_{cc_major}{cc_minor},code=sm_{cc_major}{cc_minor}" - ) - nvcc_flags.append("--expt-relaxed-constexpr") - nvcc_flags.append("--expt-extended-lambda") - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_TWOPASS_BLOCK_SIZE", - "FUSED_LOGP_TWOPASS_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_BLOCK_SIZE", - "FUSED_LOGP_ONLINE_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", - "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", - "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", - "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", - ) - ) - if is_rocm: - for tile_knob in ( - "DETERMINISTIC_LOGP_TILE_BLOCK_SIZE", - "DETERMINISTIC_LOGP_TILE_VECTOR_ELEMENTS", - "DETERMINISTIC_LOGP_BACKWARD_BLOCK_SIZE", - ): - nvcc_flags.extend(_cuda_define_from_env(tile_knob, tile_knob)) - else: - # Same idea for the CUDA tile-stats kernel: the defaults are tuned for - # sm_90, and a different architecture or vocabulary split may prefer - # another block size or vector width. - for tile_knob in ( - "DETERMINISTIC_LOGP_TILE_BLOCK_SIZE_NARROW", - "DETERMINISTIC_LOGP_TILE_BLOCK_SIZE_WIDE", - "DETERMINISTIC_LOGP_TILE_VECTOR_BYTES", - ): - nvcc_flags.extend(_cuda_define_from_env(tile_knob, tile_knob)) - if not is_rocm and envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): - nvcc_flags.append("-lineinfo") - if ( - not is_rocm - and os.name == "nt" - and envs.env_flag(envs.KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC) - ): - nvcc_flags.append("-allow-unsupported-compiler") - nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") - - platform_define = "-DKERNEL_ALIGN_WITH_ROCM" if is_rocm else "-DKERNEL_ALIGN_WITH_CUDA" - cxx_flags = ["-O3", "-std=c++17", platform_define] - extra_link_args = list(torch_rpath) - if os.name != "nt" and not is_rocm: - # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). - extra_link_args.append("-lcuda") - - if not is_rocm: - sm90_srcs = [ - "csrc/cuda/fused_logp_sm90.cu", - "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob - "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp - "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build - # Single-card batch-invariant embedding/lm-head. - "csrc/cuda/embedding_lm_head_sm90.cu", - ] - enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) - present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] - if enable_sm90 and present_sm90: - tma_arch = f"{cc_major}{cc_minor}a" # WGMMA/TMA require the arch-native 'a' variant - cuda_sources.extend(present_sm90) - nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") - cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") - if "-lcuda" not in extra_link_args: - extra_link_args.append("-lcuda") - - # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp - # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in - # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. - enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" - if enable_det_gemm_sm90: - tma_arch = f"{cc_major}{cc_minor}a" - arch_flag = f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}" - if arch_flag not in nvcc_flags: - nvcc_flags.append(arch_flag) - if "-lcuda" not in extra_link_args: - extra_link_args.append("-lcuda") - nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") - cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") - - if is_rocm: - nvcc_flags = _filter_rocm_incompatible_nvcc_flags(nvcc_flags) - - extensions.append( - CUDAExtension( - name="rl_engine._C", - sources=cuda_sources, - include_dirs=[], - extra_compile_args={ - "cxx": cxx_flags, - "nvcc": nvcc_flags, - }, - extra_link_args=extra_link_args, - ) - ) - - extensions.extend(_ascend_extensions()) - - if _native_extension_required() and not extensions: - raise RuntimeError( - "rl_engine._C was requested but no CUDA/ROCm build environment is available. " - "Use a matching GPU-enabled PyTorch build; for a GPU-less ROCm build, set " - "PYTORCH_ROCM_ARCH to the target architecture." - ) - - return extensions - - -def _ascend_extensions(): - """Ascend C (CANN) kernels, built with bisheng. Gated on KERNEL_ALIGN_FORCE_ASCEND=1. - - Follows the official torch_npu cpp_extension_asc pattern: .asc sources - (kernel + host + pybind) are compiled by the CANN bisheng compiler into a - single rl_engine._C_npu extension module. Requires CANN toolkit (bisheng on - PATH or ASCEND_HOME_PATH set) and torch_npu. - """ - if not envs.env_flag(envs.KERNEL_ALIGN_FORCE_ASCEND): - return [] - try: - import torch # noqa: F401 - import torch_npu # noqa: F401 - except ImportError as e: - raise RuntimeError( - "KERNEL_ALIGN_FORCE_ASCEND=1 requires torch and torch_npu to be installed" - ) from e - - asc_srcs = sorted(str(p) for p in Path("csrc/ascend").glob("*.asc")) - if not asc_srcs: - raise RuntimeError("KERNEL_ALIGN_FORCE_ASCEND=1 but no .asc sources under csrc/ascend/") - return [Extension(name="rl_engine._C_npu", sources=asc_srcs, language="asc")] - - -def _bisheng_compile_cmd(ext, ext_fullpath): - """Single-command bisheng build for an Ascend C extension (see op-plugin example).""" - import torch - import torch.utils.cpp_extension as cpp_extension - import torch_npu - - if find_executable("bisheng") is None: - raise RuntimeError( - "bisheng compiler not found on PATH; source the CANN toolkit environment first" - ) - - soc = os.environ.get(envs.KERNEL_ALIGN_ASCEND_ARCH, "dav-2201") # A2/A3; A5: dav-3510 - abi_value = "1" if torch._C._GLIBCXX_USE_CXX11_ABI else "0" - module_name = ext.name.rsplit(".", 1)[-1] - - torch_npu_dir = os.path.dirname(os.path.realpath(torch_npu.__file__)) - ascend_home = os.environ.get("ASCEND_HOME_PATH", "/usr/local/Ascend/ascend-toolkit/latest") - - include_dirs = [ - *cpp_extension.include_paths(), - sysconfig.get_config_var("INCLUDEPY"), - os.path.join(torch_npu_dir, "include"), - os.path.join(torch_npu_dir, "include", "third_party", "acl", "inc"), - os.path.join(ascend_home, "include"), - ] - lib_dirs = [ - sysconfig.get_config_var("LIBDIR"), - os.path.join(os.path.dirname(torch.__file__), "lib"), - os.path.join(torch_npu_dir, "lib"), - os.path.join(ascend_home, "lib64"), - ] - - cmd = [ - "bisheng", - "-x", - "asc", - f"--npu-arch={soc}", - "-shared", - "-fPIC", - "-std=c++17", - "-O2", - f"-D_GLIBCXX_USE_CXX11_ABI={abi_value}", - f"-DTORCH_EXTENSION_NAME={module_name}", - "-lascendcl", - "-ltorch_npu", - "-ltorch", - "-ltorch_cpu", - "-ltorch_python", - "-lc10", - *ext.sources, - "-o", - ext_fullpath, - ] - cmd += [f"-I{d}" for d in include_dirs if d] - cmd += [f"-L{d}" for d in lib_dirs if d] - return cmd - - -def get_cmdclass(): - _, BuildExtension, _ = _load_torch_extension_tools() - if BuildExtension is None: - return {} - - class AscendBuildExtension(BuildExtension): - """torch BuildExtension + bisheng path for language="asc" extensions.""" - - def build_extension(self, ext): - if getattr(ext, "language", None) != "asc": - super().build_extension(ext) - return - ext_fullpath = self.get_ext_fullpath(ext.name) - os.makedirs(os.path.dirname(ext_fullpath), exist_ok=True) - try: - self.spawn(_bisheng_compile_cmd(ext, ext_fullpath)) - except Exception as e: - raise CompileError(str(e)) from e - - return {"build_ext": AscendBuildExtension} - - -setup( - name="rl-engine", - version="0.1.0", - packages=find_packages(include=["rl_engine", "rl_engine.*"]), - install_requires=[ - "torch>=2.4.1", - "tabulate", - "numpy", - "accelerate", - "transformers==5.13.1", - ], - ext_modules=get_extensions(), - cmdclass=get_cmdclass(), - extras_require={ - "cuda": ["flashinfer"], - "rocm": ["aiter"], - "vllm": ["vllm>=0.6.0"], - "drift-viewer": ["Pillow>=10", "PySide6>=6.6"], - }, - entry_points={ - "console_scripts": [ - "rlk-drift-view=rl_engine.alignment.cross_config.drift_viewer:main", - ], - }, - python_requires=">=3.10", - include_package_data=True, - zip_safe=False, -) +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import importlib.util +import os +import sysconfig +import warnings +from distutils.errors import CompileError +from distutils.spawn import find_executable +from pathlib import Path + +from setuptools import Extension, find_packages, setup + + +def _load_envs_module(): + envs_path = Path(__file__).with_name("envs.py") + spec = importlib.util.spec_from_file_location("_rl_kernel_envs", envs_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"failed to load environment helpers from {envs_path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +envs = _load_envs_module() + + +def _load_torch_extension_tools(): + try: + import torch + except ModuleNotFoundError as exc: + if exc.name != "torch": + raise + return None, None, None + + from torch.utils.cpp_extension import BuildExtension, CUDAExtension + + # CUDAExtension is also the supported extension entry point for ROCm + # PyTorch builds. BuildExtension dispatches .cu/.hip sources to hipcc when + # torch.version.hip is set. + return torch, BuildExtension, CUDAExtension + + +def _native_extension_required() -> bool: + """Whether the caller explicitly requested a native extension build.""" + return ( + envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) + or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) + or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) + or envs.env_flag("FORCE_CUDA") + ) + + +def _cuda_define_from_env(name: str, macro: str) -> list[str]: + value = os.environ.get(name) + if value is None: + return [] + parsed = int(value) + if parsed <= 0: + raise ValueError(f"{name} must be positive, got {value!r}") + return [f"-D{macro}={parsed}"] + + +_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES = ( + "-Xfatbin", + "-compress-all", + "-gencode", + "--generate-code", + "--expt-", + "-lineinfo", + "-allow-unsupported-compiler", + "-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH", +) +_ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE = { + "-Xfatbin", + "-gencode", + "--generate-code", +} + + +def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: + """Remove CUDA-only device compiler flags before BuildExtension calls hipcc.""" + filtered_flags = [] + skip_next = False + for flag in flags: + if skip_next: + skip_next = False + continue + if flag in _ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE: + skip_next = True + continue + if flag.startswith(_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES): + continue + filtered_flags.append(flag) + return filtered_flags + + +def get_extensions(): + torch, _, CUDAExtension = _load_torch_extension_tools() + if torch is None: + message = ( + "PyTorch is unavailable, so rl_engine._C cannot be built. Install a matching " + "CUDA/ROCm PyTorch build first, then run " + "`RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e .`." + ) + if _native_extension_required(): + raise RuntimeError(message) + warnings.warn( + f"{message} Continuing with the pure-Python fallback because no native extension " + "was explicitly requested.", + RuntimeWarning, + stacklevel=2, + ) + return [] + + extensions = [] + torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib") + torch_rpath = ["-Wl,-rpath,$ORIGIN/../torch/lib"] + if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": + torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") + is_rocm = getattr(torch.version, "hip", None) is not None + + # CUDAExtension is intentionally used for both CUDA and ROCm. On ROCm, + # PyTorch's BuildExtension hipifies CUDA sources and invokes hipcc; it also + # consumes PYTORCH_ROCM_ARCH (one or more ';'-separated gfx targets) to add + # --offload-arch. Do not require a visible GPU when a ROCm target was + # explicitly selected. + no_rocm_arch = not os.environ.get("PYTORCH_ROCM_ARCH", "").strip() + if is_rocm and no_rocm_arch and torch.cuda.device_count() == 0: + raise RuntimeError( + "ROCm builds without a visible GPU require PYTORCH_ROCM_ARCH. " + "Set one or more ';'-separated targets, for example " + "PYTORCH_ROCM_ARCH='gfx942;gfx950'." + ) + + if is_rocm or torch.cuda.is_available(): + cuda_sources = [ + "csrc/ops.cpp", + "csrc/fused_logp_kernel.cu", + "csrc/deterministic_logp_kernel.cu", + "csrc/cuda/gemm/det_gemm_kernel.cu", + "csrc/cuda/rmsnorm.cu", + "csrc/cuda/activation.cu", + "csrc/cuda/attention/deterministic_attention.cu", + ] + if is_rocm: + # ROCm-tuned WS2 vocab-parallel logprob kernels; the shared + # deterministic_logp_kernel.cu keeps the SM90-tuned CUDA path. + cuda_sources.extend( + [ + "csrc/hip/hip_deterministic_logp_kernel.hip", + "csrc/rocm/distributed/deterministic_collective.hip", + ] + ) + else: + # CUDA IPC and the fixed-tree collective implementation are not + # part of the ROCm extension. + cuda_sources.append("csrc/cuda/distributed/deterministic_collective.cu") + # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). + # The ROCm dispatcher falls back to PyTorch SDPA for this operator. + cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") + + nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] + if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): + nvcc_flags.append("--use_fast_math") + if not is_rocm: + cc_major, cc_minor = torch.cuda.get_device_capability() + enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" + if not enable_sm90: + # SM90 build emits 90a below; mixing plain compute_90 breaks TMA ptxas. + nvcc_flags.append( + f"-gencode=arch=compute_{cc_major}{cc_minor},code=sm_{cc_major}{cc_minor}" + ) + nvcc_flags.append("--expt-relaxed-constexpr") + nvcc_flags.append("--expt-extended-lambda") + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_TWOPASS_BLOCK_SIZE", + "FUSED_LOGP_TWOPASS_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_BLOCK_SIZE", + "FUSED_LOGP_ONLINE_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", + "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", + "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", + "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", + ) + ) + if is_rocm: + for tile_knob in ( + "DETERMINISTIC_LOGP_TILE_BLOCK_SIZE", + "DETERMINISTIC_LOGP_TILE_VECTOR_ELEMENTS", + "DETERMINISTIC_LOGP_BACKWARD_BLOCK_SIZE", + ): + nvcc_flags.extend(_cuda_define_from_env(tile_knob, tile_knob)) + else: + # Same idea for the CUDA tile-stats kernel: the defaults are tuned for + # sm_90, and a different architecture or vocabulary split may prefer + # another block size or vector width. + for tile_knob in ( + "DETERMINISTIC_LOGP_TILE_BLOCK_SIZE_NARROW", + "DETERMINISTIC_LOGP_TILE_BLOCK_SIZE_WIDE", + "DETERMINISTIC_LOGP_TILE_VECTOR_BYTES", + ): + nvcc_flags.extend(_cuda_define_from_env(tile_knob, tile_knob)) + if not is_rocm and envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): + nvcc_flags.append("-lineinfo") + if ( + not is_rocm + and os.name == "nt" + and envs.env_flag(envs.KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC) + ): + nvcc_flags.append("-allow-unsupported-compiler") + nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") + + platform_define = "-DKERNEL_ALIGN_WITH_ROCM" if is_rocm else "-DKERNEL_ALIGN_WITH_CUDA" + cxx_flags = ["-O3", "-std=c++17", platform_define] + extra_link_args = list(torch_rpath) + if os.name != "nt" and not is_rocm: + # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). + extra_link_args.append("-lcuda") + + if not is_rocm: + sm90_srcs = [ + "csrc/cuda/fused_logp_sm90.cu", + "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob + "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp + "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build + # Single-card batch-invariant embedding/lm-head. + "csrc/cuda/embedding_lm_head_sm90.cu", + ] + enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) + present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] + if enable_sm90 and present_sm90: + tma_arch = f"{cc_major}{cc_minor}a" # WGMMA/TMA require the arch-native 'a' variant + cuda_sources.extend(present_sm90) + nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") + cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") + if "-lcuda" not in extra_link_args: + extra_link_args.append("-lcuda") + + # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp + # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in + # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. + enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" + if enable_det_gemm_sm90: + tma_arch = f"{cc_major}{cc_minor}a" + arch_flag = f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}" + if arch_flag not in nvcc_flags: + nvcc_flags.append(arch_flag) + if "-lcuda" not in extra_link_args: + extra_link_args.append("-lcuda") + nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") + cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") + + if is_rocm: + nvcc_flags = _filter_rocm_incompatible_nvcc_flags(nvcc_flags) + + extensions.append( + CUDAExtension( + name="rl_engine._C", + sources=cuda_sources, + include_dirs=[], + extra_compile_args={ + "cxx": cxx_flags, + "nvcc": nvcc_flags, + }, + extra_link_args=extra_link_args, + ) + ) + + extensions.extend(_ascend_extensions()) + + if _native_extension_required() and not extensions: + raise RuntimeError( + "rl_engine._C was requested but no CUDA/ROCm build environment is available. " + "Use a matching GPU-enabled PyTorch build; for a GPU-less ROCm build, set " + "PYTORCH_ROCM_ARCH to the target architecture." + ) + + return extensions + + +def _ascend_extensions(): + """Ascend C (CANN) kernels, built with bisheng. Gated on KERNEL_ALIGN_FORCE_ASCEND=1. + + Follows the official torch_npu cpp_extension_asc pattern: .asc sources + (kernel + host + pybind) are compiled by the CANN bisheng compiler into a + single rl_engine._C_npu extension module. Requires CANN toolkit (bisheng on + PATH or ASCEND_HOME_PATH set) and torch_npu. + """ + if not envs.env_flag(envs.KERNEL_ALIGN_FORCE_ASCEND): + return [] + try: + import torch # noqa: F401 + import torch_npu # noqa: F401 + except ImportError as e: + raise RuntimeError( + "KERNEL_ALIGN_FORCE_ASCEND=1 requires torch and torch_npu to be installed" + ) from e + + asc_srcs = sorted(str(p) for p in Path("csrc/ascend").glob("*.asc")) + if not asc_srcs: + raise RuntimeError("KERNEL_ALIGN_FORCE_ASCEND=1 but no .asc sources under csrc/ascend/") + return [Extension(name="rl_engine._C_npu", sources=asc_srcs, language="asc")] + + +def _bisheng_compile_cmd(ext, ext_fullpath): + """Single-command bisheng build for an Ascend C extension (see op-plugin example).""" + import torch + import torch.utils.cpp_extension as cpp_extension + import torch_npu + + if find_executable("bisheng") is None: + raise RuntimeError( + "bisheng compiler not found on PATH; source the CANN toolkit environment first" + ) + + soc = os.environ.get(envs.KERNEL_ALIGN_ASCEND_ARCH, "dav-2201") # A2/A3; A5: dav-3510 + abi_value = "1" if torch._C._GLIBCXX_USE_CXX11_ABI else "0" + module_name = ext.name.rsplit(".", 1)[-1] + + torch_npu_dir = os.path.dirname(os.path.realpath(torch_npu.__file__)) + ascend_home = os.environ.get("ASCEND_HOME_PATH", "/usr/local/Ascend/ascend-toolkit/latest") + + include_dirs = [ + *cpp_extension.include_paths(), + sysconfig.get_config_var("INCLUDEPY"), + os.path.join(torch_npu_dir, "include"), + os.path.join(torch_npu_dir, "include", "third_party", "acl", "inc"), + os.path.join(ascend_home, "include"), + ] + lib_dirs = [ + sysconfig.get_config_var("LIBDIR"), + os.path.join(os.path.dirname(torch.__file__), "lib"), + os.path.join(torch_npu_dir, "lib"), + os.path.join(ascend_home, "lib64"), + ] + + cmd = [ + "bisheng", + "-x", + "asc", + f"--npu-arch={soc}", + "-shared", + "-fPIC", + "-std=c++17", + "-O2", + f"-D_GLIBCXX_USE_CXX11_ABI={abi_value}", + f"-DTORCH_EXTENSION_NAME={module_name}", + "-lascendcl", + "-ltorch_npu", + "-ltorch", + "-ltorch_cpu", + "-ltorch_python", + "-lc10", + *ext.sources, + "-o", + ext_fullpath, + ] + cmd += [f"-I{d}" for d in include_dirs if d] + cmd += [f"-L{d}" for d in lib_dirs if d] + return cmd + + +def get_cmdclass(): + _, BuildExtension, _ = _load_torch_extension_tools() + if BuildExtension is None: + return {} + + class AscendBuildExtension(BuildExtension): + """torch BuildExtension + bisheng path for language="asc" extensions.""" + + def build_extension(self, ext): + if getattr(ext, "language", None) != "asc": + super().build_extension(ext) + return + ext_fullpath = self.get_ext_fullpath(ext.name) + os.makedirs(os.path.dirname(ext_fullpath), exist_ok=True) + try: + self.spawn(_bisheng_compile_cmd(ext, ext_fullpath)) + except Exception as e: + raise CompileError(str(e)) from e + + return {"build_ext": AscendBuildExtension} + + +setup( + name="rl-engine", + version="0.1.0", + packages=find_packages(include=["rl_engine", "rl_engine.*"]), + install_requires=[ + "torch>=2.4.1", + "tabulate", + "numpy", + "accelerate", + "transformers==5.13.1", + ], + ext_modules=get_extensions(), + cmdclass=get_cmdclass(), + extras_require={ + "cuda": ["flashinfer"], + "rocm": ["aiter"], + "vllm": ["vllm>=0.6.0"], + "drift-viewer": ["Pillow>=10", "PySide6>=6.6"], + }, + entry_points={ + "console_scripts": [ + "rlk-drift-view=rl_engine.alignment.cross_config.drift_viewer:main", + ], + }, + python_requires=">=3.10", + include_package_data=True, + zip_safe=False, +)