[ollama] Ollama MLX Sampler 최적화: 성능 향상과 Logprobs 지원

PR 링크: ollama/ollama#15658 상태: Merged | 변경: +None / -None

들어가며

최근 Ollama 레포지토리에 병합된 'MLX Sampler Improvements' PR은 MLX 기반 추론 엔진의 샘플링 파이프라인을 대대적으로 개선했습니다. 이번 변경은 크게 두 가지 목표를 가집니다. 첫째, 그동안 지원되지 않았던 Logprobs 기능을 MLX 러너에서도 사용할 수 있도록 외부 인터페이스를 통합했습니다. 둘째, Top-P와 Top-K 샘플링 필터가 동시에 사용될 때 발생하는 불필요한 정렬(sort) 연산을 제거하여 추론 처리량을 약 1.5% 향상시켰습니다.

코드 분석

1. x/mlxrunner/pipeline.go: 샘플링 파이프라인 구조 개선

기존 코드에서는 sample과 nextSample을 개별적으로 관리했으나, 이를 sampler.Result 구조체로 추상화하여 관리 효율성을 높였습니다. 특히 decoder 구조체를 도입하여 토큰 디코딩과 Logprobs 처리를 동기화했습니다.

// Before
mlx.Unpin(sample)
mlx.Unpin(nextSample)

// After
mlx.Unpin(sample.Arrays()...)
mlx.Unpin(nextSample.Arrays()...)

2. x/mlxrunner/sample/ (로직 최적화)

가장 핵심적인 성능 개선은 샘플링 필터링 과정에서 일어납니다. 기존에는 Top-P와 Top-K 필터를 적용할 때 각각 정렬을 수행하는 경우가 많았으나, 이를 최적화하여 중복 연산을 줄였습니다. 또한 buildLogprob 함수를 통해 샘플링된 토큰의 확률 정보를 정렬하여 API 응답에 포함하도록 변경되었습니다.

// Top-K와 Top-P 필터링 후 정렬 최적화 예시
sort.Slice(pairs, func(i, j int) bool {
    return pairs[i].Logprob > pairs[j].Logprob
})

3. x/mlxrunner/client.go: 인터페이스 통합

completionRequest 구조체를 CompletionRequest로 변경하고 api.Options를 직접 임베딩하여 다른 러너들과의 일관성을 맞췄습니다. 이를 통해 Logprobs를 활성화할 수 있는 기반을 마련했습니다.

왜 이게 좋은가

이번 최적화는 단순한 기능 추가를 넘어, 추론 엔진의 '데이터 흐름'을 최적화했다는 점에서 의미가 큽니다.

1. 성능 향상: Top-P와 Top-K 필터링 시 불필요한 정렬 연산을 제거함으로써 Gemma4 모델 기준 약 1.5%의 처리량(throughput) 향상을 달성했습니다. 이는 대규모 언어 모델 추론에서 병목이 되기 쉬운 샘플링 단계의 오버헤드를 줄인 결과입니다.
2. 일관성 확보: Logprobs 지원을 통해 다른 러너(llama.cpp 등)와 동일한 API 인터페이스를 제공하게 되어, 클라이언트 단에서의 코드 복잡도가 크게 줄었습니다.
3. 메모리 관리: mlx.Unpin과 mlx.Sweep 호출을 구조화하여 메모리 누수를 방지하고, decoder를 통해 UTF-8 시퀀스가 깨지지 않도록 안전하게 토큰을 flush하는 로직을 구현했습니다.

리뷰 과정에서 jessegross는 LogsumexpAxis를 통한 배치 안전성(batch safety) 확보를 향후 과제로 언급하며, 현재의 샘플링 로직이 향후 더 큰 배치 사이즈를 처리할 수 있는 구조로 나아가고 있음을 시사했습니다.

결론

이번 PR은 MLX 러너의 성능을 미세하게 조정하면서도, 확장 가능한 구조로 리팩토링한 좋은 사례입니다. 특히 성능과 기능(Logprobs)이라는 두 마리 토끼를 잡기 위해 기존의 복잡한 샘플링 로직을 추상화한 점이 돋보입니다.

참고 자료

• https://github.com/ollama/ollama/pull/7356
• https://github.com/ml-explore/mlx

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.