mlxcel-core: MLX C++ FFI 브리지¶

목차¶

• 직접 C++ 바인딩이 필요한 이유
• cxx 브리지 아키텍처
• 모듈 개요
• 핵심 타입
• 핵심 모듈
• 빌드 시스템
• 새 연산 추가하기

직접 C++ 바인딩이 필요한 이유¶

MLX는 C++로 작성된 Apple Silicon 네이티브 ML 프레임워크입니다. 안정적인 ABI를 제공하는 C 래퍼 레이어(mlx-c)가 이미 있지만 여기에는 오버헤드가 따릅니다:

1. 할당: C 래퍼는 모든 배열 참조에 대해 힙에 불투명 핸들 객체를 할당합니다.
2. 복사: 일부 연산은 C ABI 경계를 넘기 위해 중간 복사가 필요합니다.
3. 누락된 연산: C 래퍼가 모든 MLX 연산을 노출하지 않아 우회책이 필요합니다.

mlxcel은 C 래퍼를 완전히 우회하고 cxx 크레이트를 통해 MLX C++ API에 직접 바인딩합니다. 그 결과:

• 배열 전달에 오버헤드가 없습니다 (C++ UniquePtr<MlxArray>가 기본 mlx::core::array에 직접 매핑).
• 전체 MLX 연산 세트에 접근할 수 있습니다.
• Rust에서 커스텀 MLX 커널을 작성하고 등록할 수 있습니다.

mlxcel-core 크레이트는 src/lib/mlxcel-core/에 있습니다.

cxx 브리지 아키텍처¶

브리지는 cxx 크레이트의 #[cxx::bridge] 매크로를 사용하여 src/lib/mlxcel-core/src/lib.rs에 정의되어 있습니다:

#[cxx::bridge(namespace = "mlx_cxx")]
mod ffi {
    unsafe extern "C++" {
        include!("mlx_cxx_bridge.h");

        /// Opaque wrapper for mlx::core::array
        type MlxArray;

        /// Opaque wrapper for mlx::core::Stream
        type MlxStream;

        // Operations are declared here and implemented in C++
        fn matmul(a: &MlxArray, b: &MlxArray) -> UniquePtr<MlxArray>;
        // ...
    }
}

mlx_cxx_bridge.h 헤더는 mlx::core 연산을 래핑하는 C++ 심 함수를 선언합니다. 이 심 함수들은 UniquePtr/참조 시맨틱을 내부에서 쓰이는 mlx::core::array 값 시맨틱으로 변환합니다.

graph LR
    subgraph "Rust (mlxcel-core)"
        R[Model code<br/>e.g., Llama3Model] -->|ffi::matmul| B[cxx bridge<br/>lib.rs]
    end

    subgraph "C++ Shim Layer"
        B -->|mlx_cxx::matmul| SH[mlx_cxx_bridge.cpp]
    end

    subgraph "MLX C++ Core"
        SH -->|mlx::core::matmul| MLX[mlx::core]
        MLX -->|Metal GPU kernel| GPU[Apple GPU]
    end

모듈 개요¶

핵심 타입¶

MlxArray¶

MlxArray는 mlx::core::array에 대한 불투명 핸들입니다. Rust에서는 항상 UniquePtr<MlxArray>로 소유하며 이 포인터는 배열 값을 보유하는 힙 할당 래퍼입니다.

// Creating arrays
let zeros = ffi::zeros(&[batch, seq_len, hidden], ffi::FLOAT16);
let ones  = ffi::ones(&[heads, head_dim], ffi::FLOAT32);

// Operations return UniquePtr<MlxArray>
let result: UniquePtr<MlxArray> = ffi::matmul(&a, &b);

// Extract scalar
let value: f32 = ffi::item_f32(&logits);

// Get shape
let shape: Vec<i32> = ffi::array_shape(&arr);

MlxStream¶

MlxStream은 MLX 연산 스트림을 나타냅니다. 동일한 스트림에 제출된 연산은 순서대로 실행되며 다른 스트림의 연산은 병렬로 실행될 수 있습니다.

mlxcel은 전용 생성 스트림을 사용하여 연산과 메모리 전송을 겹칩니다 (룩어헤드 파이프라이닝):

let gen_stream = ffi::new_stream_on_device(true); // GPU stream
// Operations on gen_stream run asynchronously
let output = ffi::matmul_stream(&a, &b, &gen_stream);
ffi::synchronize_stream(&gen_stream); // wait for completion

KVCache¶

KVCache 구조체는 하나의 트랜스포머 레이어에 대한 키 및 값 행렬을 보유합니다:

pub struct KVCache {
    pub keys: Option<UniquePtr<MlxArray>>,   // [batch, heads, seq_len, head_dim]
    pub values: Option<UniquePtr<MlxArray>>,  // [batch, heads, seq_len, head_dim]
    pub offset: usize,                         // current sequence position
}

핵심 모듈¶

engine/과 generate.rs¶

generate.rs는 모든 모델이 구현하는 LanguageModel 트레잇을 정의하고, 디코드 루프는 engine/에 있습니다. Engine은 모델과 함께 열린 모든 시퀀스의 KV 상태를 담는 CachePool을 소유하며, 모델 순전파와 샘플러 추출, 종료 단계가 모두 이 엔진을 거칩니다. 서버의 BatchScheduler는 Engine을 통해서만 모델을 구동하고, DirectEngine(engine/direct.rs)은 mlxcel generate, mlxcel-bench-decode, 추론 세션이 쓰는 단일 시퀀스 raw-completion 클라이언트입니다. 단일 시퀀스는 크기 1인 스텝 배치로 처리하므로 별도의 단일 시퀀스 디코드 루프는 없습니다. 오류는 시퀀스 단위로 처리되어, 한 행이 실패하면 그 행만 오류로 끝나고 나머지 배치는 계속 디코딩합니다.

flowchart TD
    A[Input tokens] --> B[Embed tokens]
    B --> C{First step?}
    C -->|Yes| D[Prefill: full prompt forward pass]
    C -->|No| E[Decode: single token forward pass]
    D --> F[Sample token from logits]
    E --> F
    F --> G{Stop condition?}
    G -->|EOS or max_tokens| H[Return generated text]
    G -->|No| I[Append token to input]
    I --> E

핵심 최적화: 룩어헤드 파이프라이닝. CPU가 토큰 n의 샘플링 처리를 마치기 전에 토큰 n+1의 순전파가 GPU 스트림에 제출되므로, CPU 측 샘플링 시간이 GPU 실행 시간과 겹칩니다. 엔진은 이를 분할 스텝(Engine::submit 후 Engine::finish_rows)으로 제공하고 DirectEngine과 BatchScheduler가 모두 이 스텝을 사용하며, MLXCEL_FORCE_SYNC를 설정하면 꺼집니다.

layers.rs¶

모든 모델 구현이 사용하는 공유 레이어 연산:

• rms_norm(x, weight, eps): RMSNorm
• layer_norm(x, weight, bias, eps): LayerNorm
• linear(x, weight, bias): 선형 프로젝션
• rope(x, cos, sin, traditional): Rotary Position Embeddings
• attention(q, k, v, mask, scale): 스케일 내적 어텐션
• scaled_dot_product_attention(q, k, v, scale, mask): 플래시 어텐션 변형
• silu(x) / gelu(x) / relu(x): 활성화 함수
• take(x, indices, axis): 임베딩 조회
• 4비트 / 8비트 가중치를 위한 다양한 양자화 행렬 곱셈 연산

sampling.rs¶

샘플링 스택을 다음 순서로 구현합니다:

1. 온도 스케일링 적용: logits / temperature
2. 로짓 바이어스 적용 (제공된 경우)
3. top-k 필터링 적용: 상위 k개의 로짓만 유지
4. top-p (핵) 필터링 적용
5. min-p 필터링 적용
6. XTC (Exclude Top Choices) 샘플링 적용
7. DRY (Don't Repeat Yourself) 페널티 적용
8. 반복 페널티 적용
9. 빈도/존재 페널티 적용
10. 결과 분포에서 샘플링 (또는 그리디의 경우 argmax)

weights.rs¶

SafeTensors 파일 로딩을 처리합니다:

// Load all .safetensors shards from a directory into a WeightMap
let weights = load_weights_from_dir(&model_path)?;

speculative/¶

투기적 디코딩의 드래프터를 구현합니다. 자체 draft-model 디코드 루프는 없습니다. 토큰 전용 드래프터(PromptLookupDrafter, FullModelDrafter)는 엔진의 Engine::verify_round가 구동하고, MTP와 DFlash 드래프터는 각자의 라운드 루프(speculative/mtp/, drafter/)를 가집니다. 한 라운드는 다음과 같습니다.

1. 드래프터가 최대 k개의 후보 토큰을 제안합니다. FullModelDrafter는 작은 체크포인트를 자체 KV 캐시로 실행해 argmax를 제안하고, PromptLookupDrafter는 시퀀스 끝부분이 앞서 나온 위치의 뒤를 이어 복사합니다.
2. 타깃 모델이 단일 순전파로 현재 토큰과 제안 토큰을 검증합니다.
3. 제안 토큰은 해당 위치의 타깃 샘플러 출력(temperature 0에서는 argmax)과 같을 때 수용됩니다(sampler-match). 경로별 규칙은 docs/speculative-acceptance.md에 있습니다.
4. 처음 거부된 위치에서는 타깃 자신의 샘플을 내보내고, 거부된 제안은 타깃과 드래프터 캐시에서 trim합니다.
5. 대부분의 라운드에서 토큰이 둘 이상 수용되면 실효 처리량이 올라갑니다.

빌드 시스템¶

빌드 시스템은 src/lib/mlxcel-core/build.rs에 있습니다. src/lib/mlxcelverse/의 CMake 프로젝트로 MLX를 빌드한 뒤, 그 헤더를 대상으로 C++ 심 레이어를 컴파일합니다.

주요 단계:

1. MLX 빌드: src/lib/mlxcelverse/의 CMake 프로젝트를 실행합니다. CMakeLists.txt가 고정된 커밋(FetchContent_Declare(mlx ...) 블록의 GIT_TAG)의 upstream MLX를 가져오고, overlays/{cpu,metal,cuda,rocm}/의 백엔드별 소스 오버레이를 가져온 트리 위에 복사한 다음 libmlx를 컴파일합니다.
2. 백엔드 선택: macOS에서는 Metal과 Accelerate가 기본으로 켜지고(MLXCEL_BUILD_METAL=OFF로 Metal을 끌 수 있음), cuda와 rocm Cargo feature가 CUDA와 ROCm 백엔드를 선택합니다.
3. 컴파일: cpp/mlx_cxx_bridge.cpp를 비롯한 브리지 소스와 src/lib/mlxcelverse/kernels/의 mlxcel 자체 커널 런처를 -std=c++20으로 컴파일합니다.
4. 링킹: 정적 libmlx.a에 링크하고, macOS에서는 Metal과 Accelerate 프레임워크에도 링크합니다.

src/lib/mlxcel-core/
  build.rs              # Build script
  src/
    lib.rs              # cxx bridge declaration
    *.rs                # Rust modules
  cpp/
    mlx_cxx_bridge.h    # C++ header (declares shim functions)
    mlx_cxx_bridge.cpp  # C++ implementation (calls mlx::core)

lib.rs 수정 후 cxx 바인딩을 재생성하려면 다음을 실행합니다:

cargo build -p mlxcel-core

cxx 크레이트는 프로시저럴 매크로를 통해 빌드 시 필요한 글루 코드를 생성하므로 별도의 코드 생성 단계는 필요하지 않습니다.

새 연산 추가하기¶

Rust에 새 MLX 연산을 노출하려면:

1. 심 함수를 선언합니다. cpp/mlx_cxx_bridge.h에:

   namespace mlx_cxx {
       std::unique_ptr<MlxArray> my_new_op(
           const MlxArray& a,
           const MlxArray& b,
           float scale
       );
   }
   

2. 심을 구현합니다. cpp/mlx_cxx_bridge.cpp에:

   std::unique_ptr<MlxArray> mlx_cxx::my_new_op(
       const MlxArray& a,
       const MlxArray& b,
       float scale
   ) {
       auto result = mlx::core::my_new_op(a.arr, b.arr, scale);
       return std::make_unique<MlxArray>(std::move(result));
   }
   

3. 바인딩을 선언합니다. src/lib.rs에:

   #[cxx::bridge(namespace = "mlx_cxx")]
   mod ffi {
       unsafe extern "C++" {
           // ...existing declarations...
           fn my_new_op(a: &MlxArray, b: &MlxArray, scale: f32) -> UniquePtr<MlxArray>;
       }
   }
   

4. Rust 래퍼를 추가합니다. 적절한 모듈(예: src/lib/mlxcel-core/src/ops.rs)에:

   pub fn my_new_op(a: &MlxArray, b: &MlxArray, scale: f32) -> UniquePtr<MlxArray> {
       crate::ffi::my_new_op(a, b, scale)
   }
   

5. 테스트를 추가합니다. src/lib/mlxcel-core/src/ffi_tests.rs에:

   #[test]
   fn test_my_new_op() {
       let a = ffi::zeros(&[4, 4], ffi::FLOAT32);
       let b = ffi::ones(&[4, 4], ffi::FLOAT32);
       let result = crate::my_new_op(&a, &b, 1.0);
       assert_eq!(ffi::array_shape(&result), vec![4, 4]);
   }
   

여러 모델 계열이 사용하는 연산을 추가할 때는 src/lib/mlxcel-core/src/layers.rs의 사용 주석을 갱신해 모든 호출 위치를 나열하십시오. 이 주석은 해당 연산을 나중에 변경할 때 재테스트 목록으로 사용됩니다.