mizorewww/laya-mlx

Native MLX runtime for Laya typed decision models — 7–14 ms short decisions on M3 Max. No text generation, PyTorch, or cloud API.

Laya‑MLX – Apple Silicon에서 빠른 타입 기반 결정 추론

무엇인가요 – macOS Apple Silicon (M-시리즈) GPU에서 Laya 계열의 의사결정 언어 모델을 완전히 로컬로 실행할 수 있는 Python 패키지입니다. 원본 Convai Innovations 체크포인트를 MLX(Apple의 Metal 가속 텐서 라이브러리)로 이식하여, PyTorch, 🤗 Transformers, 또는 클라우드 호출 없이도 짧은 질문 하나에 대해 15ms 미만의 지연 시간을 제공합니다.


핵심 아이디어

개념 Laya‑MLX의 구현 방식
타입 기반 결정 자유형 텍스트 생성이 아니라, 단일 포워드 패스에서 구조화된 답변(선택, 점수, 이진 "noul")을 반환합니다. 이는 토큰 단위 디코딩을 피하고 결정론적, 저지연 출력을 가능하게 합니다.
양방향 인코더 입력(상태 + 질문)은 ModernBERT-large 또는 mmBERT-base 백본으로 인코딩된 후, 전용 헤드가 요청된 타입에 대한 확률을 생성합니다.
로컬, 런타임 의존성 없음 모든 추론은 MLX 내에서 실행됩니다. 토큰화는 Hugging Face Rust 토크나이저가 wheel에 컴파일되어 있습니다. PyTorch/Transformers 바이너리는 필요하지 않습니다.
Apple‑silicon 최적화 선택적 compile=True, 프리픽스 캐싱 및 패딩 기법으로 M3 Max에서 약 6%의 속도 향상이 가능합니다. 또한 GPU용 사전 변환된 FP16 체크포인트도 함께 제공됩니다.

빠른 시작 (30초)

pip install laya-mlx            # 핵심 라이브러리
pip install 'laya-mlx[demo]'    # 선택적 데모 유틸리티
import laya_mlx as laya
agent = laya.load("aac6fef/laya-mlx")   # 처음 사용 시 FP16 체크포인트 다운로드
result = agent.predict(
    "I was billed twice. Please refund the duplicate.",
    {
        "department": {
            "type": "choice",
            "instructions": "Who should handle this?",
            "criteria": ["billing", "technical", "sales"]
        }
    },
)
print(result["answers"]["department"])   # → "billing"

macOS 14+, Python 3.11+, M1–M3 모든 Apple‑silicon GPU에서 작동. 첫 번째 호출 시 모델 다운로드, 이후 호출은 완전히 오프라인으로 작동.


사용 가능한 체크포인트

모델 ID (로드) 인코더 파라미터 수 컨텍스트 언어
aac6fef/laya-mlx ModernBERT‑large 421 M 512 영어
aac6fef/laya-multilingual-mlx mmBERT‑base 322 M 1 024 다국어
aac6fef/laya-typed-decisions-mlx ModernBERT‑large 421 M 1 024 영어 (타입 기반 결정 워크플로우)

세 모델 모두 상류 Convai Innovations 가중치의 정확한 FP16 변환이며, Hugging Face에 호스팅되어 있습니다. 또한 laya.load에 원본 허브 ID(convaiinnovations/laya 등)를 지정할 수 있습니다. 라이브러리가 자동으로 다운로드하고 변환합니다.


성능 (M3 Max, FP16)

메트릭 영어 (Laya 421M) 다국어 (Laya-multilingual 322M)
단일 짧은 질문의 중앙 지연 13.4 ms 7.4 ms
95번째 백분위 지연 13.9 ms 7.8 ms
50질문 배치(배치 크기 64)의 처리량 146 q/s 395 q/s
요청당 최대 GPU 메모리 944 MiB 688 MiB
최적화된 laya‑snake --optimize --max‑speed 75.4 move/s (eager보다 약 6.5% 빠름)

수치에는 토큰화, 텐서 준비, 추론, 캘리브레이션, 결과 포맷팅이 포함되며, 모델 로딩은 제외됩니다.


주요 API 표면

agent = laya.load(
    checkpoint,          # HF 리포지토리 ID 또는 로컬 경로
    dtype="float16",   # 또는 "float32", "bfloat16"
    batch_size=16,       # 한 번의 포워드 패스당 최대 질문 수
    device="gpu",       # "cpu"도 가능 (매우 느림)
    compile=False,       # 속도 향상을 위해 MLX 컴파일 활성화
    cache_prompts=False, # 재사용을 위해 토큰화된 프롬프트 유지
)

# 예측 – `system_one`은 별칭
answers = agent.predict(state, questions)

*state는 단순한 문자열, JSON 사전, 또는 이전 메시지 리스트가 될 수 있습니다. questions는 각 항목이 원하는 답변 타입(choice, score, noul)을 설명하는 사전입니다. 반환 값에는 다음이 포함됩니다:

  • answers (구조화된 결과)
  • action.act_probability (원시 헤드 확률)
  • 토큰 사용 통계
  • 상류 모델과 동일하게 소수점 넷째 자리까지 반올림됩니다.*

라우터 헬퍼

언어별 체크포인트를 자동 선택해야 하는 애플리케이션용:

router = laya.Router(dtype="float16", max_loaded=2)
out = router.predict(state, triage_questions())
print(out["routing"])   # 예: "multilingual"

라우터는 최대 max_loaded개의 모델을 메모리에 유지할 수 있으며, Router(preload=True)로 사전 로드할 수 있습니다.


명령줄 유틸리티

명령어 목적
laya-mlx predict … JSON 파일 또는 인라인 문자열에서 단일 추론 실행.
laya-mlx convert … Hugging Face 체크포인트를 MLX 호환 디렉터리(safetensors + config)로 변환.
laya-snake 모델이 각 움직임에 호출되는 인터랙티브 터미널 데모(클래식한 뱀 게임). --optimize로 컴파일된 빠른 경로 사용.

모든 CLI는 laya.load에서 사용되는 것과 동일한 --model 인수를 수락합니다.


개발 및 테스트

  • 의존성: 프로젝트는 재현 가능한 환경을 위해 uv를 사용합니다. uv sync --extra dev로 테스트, 벤치마크, 참조 추가를 가져옵니다.
  • 테스트: 유닛 테스트는 작은 랜덤 모델에서 MLX 구현과 원본 Transformers 헤드를 비교합니다. 전체 체크포인트 검증은 토큰화, 캘리브레이션된 확률, 결정론적 반복, 메모리 성장 여부를 확인합니다.
  • 벤치마크: benchmarks/run은 지연 시간/처리량을 측정합니다. 결과는 benchmarks/results에 저장되며, BENCHMARKS.md에서 요약됩니다.
  • 연구: docs/ 폴더에는 성능 병목 현상에 대한 심층 보고서와 10배 속도 향상 아이디어(수학적, 공학적, 구현 수준)가 포함되어 있습니다. 스크립트는 experiments/에 있습니다.

라이선스 및 출처

  • 코드 – Apache‑2.0 (see LICENSE).
  • 가중치 – 원본 Laya 가중치는 © Convai Innovations 소유이며, 위에 링크된 Hugging Face 리포지토리를 통해 동일한 라이선스 하에 재배포됩니다.
  • 이식 – MLX 재구현 및 주변 유틸리티는 mizorewww가 작성했으며, 상류 NandhaKishorM/laya 리포지토리의 일부를 수정했습니다 (NOTICE에 MIT 스타일 출처 기재).

누가 사용할 수 있나요?

  • 제품 팀 – 결정론적이고 저지연의 라우팅 또는 분류를 디바이스 내에서 필요로 하는 경우 (예: 티켓 트라이아지, 긴급도 스코어링, 이진 정책 검사).
  • 개발자 – 클라우드로 데이터를 전송하는 것을 원치 않는 macOS 전용 AI 어시스턴트 또는 엣지 서비스를 개발하는 경우.
  • 연구자 – Apple GPU에서 MLX를 PyTorch/Transformers와 비교하거나, 더 빠른 속도 향상 기술을 탐색하고자 하는 경우.

TL;DR – Laya‑MLX는 Apple Silicon에서 Laya 의사결정 모델을 즉시 실행 가능한 고성능 추론 라이브러리를 제공하며, 깔끔한 Python API, 터미널 Snake 데모, 변환, 벤치마크, 모델 카드 게시를 위한 완전한 도구 스택을 갖추고 있습니다.

이 프로젝트를 다룬 글

관련

  • 프로젝트
  • 프로젝트
  • 프로젝트
  • 프로젝트