Insights
AI가 자동으로 큐레이션·번역·정리하는 기술 동향 피드입니다.
KAIO는 Rust 기반의 DSL을 사용하여 CUDA C++나 복잡한 툴킷 없이 커스텀 GPU 커널을 작성할 수 있게 해주는 도구입니다. 이 도구는 빌드 시점에 코드를 PTX로 컴파일하고, 실행 시점(JIT)에 GPU에서 로드하여 높은 성능을 보여줍니다. NVIDIA 드라이버만 있으면 Windows와 Linux 모두에서 작동하는 것이 특징입니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 GitHub ML Hardware의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기Crates.io
Documentation
Build Status
Coverage
License
Rust
Rust로 커스텀 GPU 커널 작성 — CUDA C++, Python, 툴킷 불필요. KAIO는 Rust 기반의 DSL(Domain Specific Language)을 빌드 시점에 PTX로 컴파일하고 실행 시점(JIT)에 GPU에서 로드합니다. NVIDIA 디스플레이 드라이버만 설치되어 있으면 Windows와 Linux 모두에서 작동합니다.
KAIO(그리스어 καιω — 태우다, 점화하다)는 오늘날 커스텀 GPU 커널이 필요한 Rust 엔지니어를 위한 도구입니다. 예를 들어, 퓨즈드 어텐션 변형(fused attention variants), 양자화 연산(quantization ops), 새로운 활성화 함수(novel activations) 등이 필요하지만 현재 사용 중인 프레임워크가 이를 지원하지 않아 기존에는 CUDA C++를 작성해야 했던 경우에 유용합니다.
Tensor-core matmul이 RTX 4090에서 4096³ 크기 cuBLAS sgemm 대비 115% 도달 (연속 10회 실행 중 최악의 관측치; 전체 분포는
docs/performance.md참조).
rust-cuda, Triton 대비 갖는 가장 큰 차별점입니다.cargo build만으로 어디서든 빌드가 가능합니다.#[gpu_kernel] 매크로는 컴파일 시점에 Rust 코드를 PTX IR 모듈로 낮추고(lowers); 실행 시점에는 해당 모듈을 현재 GPU의 SM 타겟과 검증하고, PTX 텍스트로 출력한 뒤 CUDA 드라이버에 JIT 컴파일을 위해 전달합니다.타입 안전한 커널(Type-safe kernel)
시그니처가 컴파일 시점에 데이터 타입 불일치를 포착하여 런타임에 발생하는 조용한 GPU 손상(silent GPU corruption)을 방지합니다.
kaio-candle를 통해 Candle의 텐서 그래프에 통합됩니다 —CustomOp) 바인딩(matmul_tc f16 + bf16, sync + async, matmul_int8, matmul_int4, 융합된 QKV 변형, attention, FlashAttention)과 네 가지 matmul TC 변형 및 FlashAttention(일반 + 인과적)에 대한 모든 역방향 바인딩을 제공하며, CUDA-Graph와 호환되는 이벤트 기반 스트림 동기화가 가능합니다.Rust ML 생태계는 Python을 따라잡지 못하고 있습니다. 새로운 모델 아키텍처가 커스텀 연산(custom operation)과 함께 출시될 때마다 — 새로운 attention 변형, 융합된 활성화 함수, 커스텀 양자화 커널 등 — candle이나 burn 같은 프레임워크는 누군가 GPU 함수를 작성할 때까지 이를 지원할 수 없습니다. 오늘날 이는 CUDA C++을 작성하고, FFI 바인딩과 씨름하며, Windows에서는 포기해야 함을 의미합니다.
한편, Python 개발자들은 오후 만에 Triton 커널을 작성하고 다음 작업으로 넘어갑니다. Triton은 Windows를 지원하지 않고, Python이 필요하며, 런타임에 JIT 컴파일되지만 — 작동합니다. 그리고 Rust에는 이에 상응하는 것이 없습니다.
KAIO가 바로 그 상응하는 솔루션입니다.
저장소(repo)를 클론하고, 한 개의 명령을 실행하면 실제 ML 커널이 GPU에서 빌드되고 실행되는 것을 확인할 수 있습니다:
git clone https://github.com/dmriding/kaio.git
cd kaio
cargo xtask showcase
활성화 함수(fused_silu_gate, gelu_comparison), 정규화(normalization, rms_norm, layer_norm), 축소(reduction, softmax), W8A8 / W4A16 양자화 파이프라인(int8_dequant, int8_matmul, int4_matmul), 그리고 엔드투엔드 quantized_attention을 포함하는 9개의 커널 각각은 컴파일되고, CPU 참조 값과 검증되며, 중앙값 지연 시간(median latency)을 보고합니다.
성능 시연을 원하시나요? cargo xtask bench는 다섯 가지 크기에 걸쳐 cuBLAS sgemm 대비 텐서 코어 matmul 벤치마크를 실행합니다. 또는 둘 다 보려면 cargo xtask all을 사용하세요. 전체 도구 체계에 대한 도움말은 cargo xtask --help에서 확인할 수 있습니다.
설치된 디스플레이 드라이버가 있는 NVIDIA GPU가 필요합니다 (NVIDIA 525 이상 — 표준 Game Ready 또는 Studio 드라이버면 모두 작동). CUDA 툴킷 설치는 필요하지 않습니다.
디스플레이 드라이버가 설치된 NVIDIA GPU(SM 7.0+, Volta 이상)가 필요합니다 — NVIDIA 525 이상 (표준 Game Ready 또는 Studio 드라이버면 모두 작동; Tesla/TCC 전용 드라이버는 필요하지 않습니다).
cargo add kaio
use kaio::prelude::*;
#[gpu_kernel(block_size = 256)]
...
$ cargo run
result: [4.5, 4.5, 4.5, 4.5, 4.5, 4.5, 4.5, 4.5]
커널 버퍼 매개변수는 장치 포인터 의미론(device-pointer semantics)을 반영하기 위해 *const [T] / *mut [T]를 사용합니다 (수천 개의 스레드, 비할당 계약 없음); &[T] / &mut [T]는 편리한 문법적 설탕(ergonomic sugar)으로 받아들여지며 동일하게 작동합니다 — RFC-0001 참조.
SAXPY는 DSL을 학습하기 위한 것입니다. 실제 가치는 다음과 같은 형태를 띱니다:
use kaio::prelude::*;
// Gated SiLU — 모든 LLaMA / Mistral / 에서의 피드포워드 활성화 함수
...
디렉터리 변경 없이 리포지토리 루트에서 실행합니다:
$ cargo xtask showcase silu
=== fused_silu_gate ===
...
또는 `cargo xtask showcase`로 아홉 가지 데모를 순차적으로 실행할 수 있습니다:
[fused SiLU-gate](examples/fused_silu_gate/) ,
[exact vs fast GELU](examples/gelu_comparison/) ,
[single-block RMSNorm](examples/rms_norm/) ,
[single-block LayerNorm](examples/layer_norm/) ,
[single-block softmax](examples/softmax/) ,
[INT8 dequantization](examples/int8_dequant/) ,
[INT8 matmul](examples/int8_matmul/) ,
[INT4 matmul](examples/int4_matmul/) ,
[quantized attention](examples/quantized_attention/). 각 예제는 독립적인 완벽한 프로젝트이며 정확성 테스트와 타이밍 테스트를 포함하고 있습니다 (자체 `Cargo.toml`을 가지고 있어 레포지토리에서 디렉토리를 복사하여 자체 커널의 참고 자료로 사용할 수 있습니다). 특히 GELU 비교 예제의 README에는 ML 워크로드에서 산술 최적화보다 커널 퓨전(kernel fusion)이 왜 더 중요한지에 대한 설명(대역폭 제한에 의한 학습 포인트)이 담겨 있습니다.
## KAIO를 사용해야 할 때
다음과 같은 상황일 때 KAIO를 활용하세요:
- **프레임워크가 커스텀 연산자(custom op)** (새로운 attention, 퓨즈된 활성화 함수, 양자화)를 지원하지 못하지만, 하나의 커널을 위해 CUDA C++로 전환하고 싶지 않을 때.
- WSL2나 Triton의 리눅스 전용 런타임 없이 **Windows에서 GPU 추론**이 필요할 때.
- CI가 GPU 또는 CUDA 툴킷 접근 권한이 없는 표준 GitHub 러너에서 실행될 때. KAIO는 GPU 없이 호스트 테스트를 통과하며, `#[ignore]`로 게이팅된 통합 테스트만 GPU가 필요합니다. 나중에 행렬을 뒤집으면 됩니다.
- 고수준 프레임워크가 추상화하는 **결정론적 VRAM 사용량, 명시적인 버퍼 재사용 또는 제로 카피 전송**이 필요할 때.
- 이미 알고 있는 언어(Rust)로 GPU 코드를 프로토타이핑하면서 CUDA C++를 학습하고 싶지 않을 때.
KAIO는 [Candle](https://github.com/huggingface/candle)나 [Burn](https://github.com/tracel-ai/burn)을 대체하는 것이 아닙니다. 이들은 제공하는 것보다 더 많은 제어력이 필요할 때 사용하는 레이어입니다.
<table>
<thead>
<tr>
<th></th>
<th>KAIO</th>
<th>cudarc</th>
<th>Candle / Burn</th>
<th>Triton (Python)</th>
<th>Raw CUDA</th>
</tr>
</thead>
<tbody>
<tr>
<td>Rust로 커널 작성</td>
<td>예</td>
<td>아니요 (PTX 로드)</td>
<td>아니요</td>
<td>아니요 (Python)</td>
<td>아니요</td>
</tr>
</tbody>
</table>
## 성능(Performance)
성능은 대규모 ML 워크로드(트랜스포머 규모의 텐서 모양)에 최적화되어 있습니다. 작은 크기는 실행 오버헤드에 지배되며 cuBLAS보다 뒤처질 수 있습니다 — 아래에서 공정한 비교를 참고하십시오.
아래 숫자는 RTX 4090 sm_89, 릴리스 빌드, 워밍업된 GPU 환경에서 연속 10회 `cargo xtask bench` 실행을 통해 관찰된 **최악의 값**입니다 (실행당 20개 타이밍 반복의 중앙값, 이후 실행 간 최소값). 이 수치들은 README 주장 시 사용되는 보수적인 로컬 수치입니다. 열적 스로틀링이 발생했거나 정상 상태가 아닌 카드는 더 낮은 값을 보일 수 있으며, 방법론은 [docs/benchmarks.md](docs/benchmarks.md)에 나와 있습니다.
**텐서 코어 행렬 곱셈 (f16 × f16 → f32):**
<table>
<thead>
<tr>
<th>크기</th>
<th>TC 동기식 최악</th>
<th>TC 비동기식 최악</th>
<th>cuBLAS sgemm 최악</th>
<th>동기식 vs cuBLAS</th>
<th>비동기식 vs cuBLAS</th>
</tr>
</thead>
<tbody>
<tr>
<td>256³</td>
<td>0.09 TF</td>
<td>0.09 TF</td>
<td>1.71 TF</td>
<td>5.3%</td>
<td>5.3%</td>
</tr>
</tbody>
</table>
**4096³에서의 양자화된 행렬 곱셈 (동일 벤치마크, 10회 중 최악):**
<table>
<thead>
<tr>
<th>커널</th>
<th>최악 TOPS</th>
<th>중앙값 TOPS</th>
<th>최고 TOPS</th>
</tr>
</thead>
<tbody>
<tr>
<td>`matmul_int8` (W8A8)</td>
<td>84.07</td>
<td>92.58</td>
<td>93.38</td>
</tr>
<tr>
<td>`matmul_int4` (W4A16)</td>
<td>52.02</td>
<td>57.52</td>
<td>58.04</td>
</tr>
</tbody>
</table>
**사과와 사과 비교를 위한 면책 조항:** KAIO 텐서 코어 행렬 곱셈(matmul)은 fp16 입력을 사용하고 fp32로 누적합니다. cuBLAS sgemm은 f32 입력 / f32 출력이입니다. INT8 및 INT4 열은 `cublasGemmEx`의 INT8이 `cudarc` 0.19에서 깔끔하게 노출되지 않기 때문에 cuBLAS sgemm과 비교됩니다. 가중치 대역폭만으로도 (INT4의 경우 sgemm 대비 0.5 B/가중치 vs 4 B) 이 비율을 지배합니다. 이러한 수치들을 회귀 감지(regression detection)를 위해 스프린트별로 비교하십시오. "vs cuBLAS sgemm" 열은 정밀도 동일성 주장(precision-identity claim)이 아닌 프로젝트 로컬 기준선입니다. 전체 분포(모든 크기에 대한 최소값 / 중앙값 / 최대값)는 [docs/performance.md](docs/performance.md)를 참조하십시오.
## 이것이 아닌 것들 (What this is not)
- **컴파일된 Rust가 아닙니다.** `#[gpu_kernel]` 본문은 Rust 구문을 사용하지만, KAIO 자체 IR로 파싱되어 PTX로 직접 낮춰집니다(lowered). rustc의 백엔드(LLVM, MIR, borrow checker)는 커널 본문을 결코 보지 못합니다. 따라서 커널 내부에서 외부에서 선언된 Rust 함수를 호출할 수 없습니다.
- **CUDA 바인딩이 아닙니다.** KAIO가 자체적으로 PTX를 생성합니다. cuDNN, cuBLAS, CUTLASS 또는 어떤 CUDA C++ 라이브러리도 래핑(wrap)하지 않습니다. 이 README에서 cuBLAS sgemm과 비교하는 것은 종속성(dependency)이 아니라 *측정 기준*입니다.
- **완전한 ML 프레임워크가 아닙니다.** 자동 미분(autograd, `kaio-candle`의 소수의 역방향 바인딩을 제외하고는), 모델 동물원(model zoo), 훈련 루프가 없습니다. KAIO는 [Candle](https://github.com/huggingface/candle) 또는 [Burn](https://github.com/tracel-ai/burn)이 필요한 연산자(op)를 가지고 있지 않을 때 사용하는 레이어입니다.
## 패턴 (Patterns)
다음 골격 코드를 복사하여 로직을 채우십시오.
**경계 검사 요소별 처리:**
```rust
#[gpu_kernel(block_size = 256)]
fn my_kernel(input: *const [f32], output: *mut [f32], n: u32) {
let idx = thread_idx_x() + block_idx_x() * block_dim_x();
...
공유 메모리 타일링:
#[gpu_kernel(block_size = 256)]
fn tiled(data: *const [f32], out: *mut [f32], n: u32) {
let tid = thread_idx_x();
...
블록 축소:
#[gpu_kernel(block_size = 256)]
fn reduce(input: *const [f32], out: *mut [f32], n: u32) {
let idx = thread_idx_x() + block_idx_x() * block_dim_x();
...
| 기능 (Feature) | 설명 (Notes) |
|---|---|
#[gpu_kernel] 절차 매크로 (proc macro) | 컴파일 시점: Rust → PTX IR; PTX가 방출되고 실행 시작 시점에 JIT 로드됨. 타입 안전한 런처 래퍼 자동 생성. |
| ... | |
| See docs.rs/kaio에서 전체 API 표면 및 내부 IR 타입을 확인하세요. |
KAIO는 1.0 미만의 소프트웨어입니다. 현재 엔지니어링 제약 사항: