[ollama] Ollama Gemma4 모델의 CUDA Prefill 성능 최적화: MLX SDPA 도입기

PR 링크: ollama/ollama#18809 상태: Merged | 변경: +7 / -43

들어가며

최근 Ollama 레포지토리에 Gemma4 모델의 추론 성능을 비약적으로 향상시키는 PR이 병합되었습니다. 기존 Gemma4 모델은 CUDA 환경에서 head dimension이 128을 초과할 경우, cuDNN의 SDPA(Scaled Dot Product Attention) 제약으로 인해 수동으로 구현된 Matmul 기반의 Attention 연산을 수행해야 했습니다. 이 방식은 연산 효율이 떨어져 특히 Prefill 단계에서 병목 현상을 유발했습니다. 이번 개선은 MLX의 최적화된 SDPA를 활용하여 이 문제를 해결했습니다.

코드 분석

mlxrunner/model/gemma4/gemma4.go

핵심 변경 사항은 Forward 함수 내에서 조건부로 실행되던 수동 Attention 로직을 제거하고, 일관되게 nn.ScaledDotProductAttention을 사용하도록 리팩토링한 것입니다.

Before:

if headDim > 128 && L > 1 && !mlx.MetalIsAvailable() {
    // Manual attention for CUDA prefill with head_dim > 128.
    // ... (수십 줄에 걸친 Matmul, Softmax, Transpose 수동 구현)
    scores := mlx.Matmul(q, kT)
    scores = mlx.Add(scores, maskArr)
    scores = mlx.SoftmaxAxis(scores, -1, true)
    out = mlx.Matmul(scores, v)
} else {
    out = nn.ScaledDotProductAttention(b, q, scale, opt, nn.WithMask(mask))
}

After:

var opt nn.SDPAOption
mask := baseMask
if kv.history != nil {
    opt = nn.WithKVHistory(kv.history)
} else {
    opt = nn.WithKV(kv.k, kv.v, b.SeqQueryLens)
    mask = mask.Intersect(kv.mask)
}
out = nn.ScaledDotProductAttention(b, q, scale, opt, nn.WithMask(mask))

기존에는 headDim > 128인 경우 CUDA 환경에서 수동 연산을 강제했으나, MLX의 SDPA가 이제 더 넓은 head dimension을 지원함에 따라 복잡한 수동 로직을 완전히 삭제하고 표준화된 API를 호출하도록 변경되었습니다.

왜 이게 좋은가

이번 최적화의 핵심은 '수동 구현된 연산의 제거'와 '하드웨어 가속 라이브러리의 활용'입니다.

1. 성능 향상: 수동으로 구현된 Matmul 기반 Attention은 메모리 접근 패턴이 최적화되지 않아 성능 저하가 컸습니다. MLX의 SDPA를 사용함으로써 Gemma4 모델의 Prefill 속도가 12B 모델 기준 2~4배, e2b 모델 기준 약 12배 향상되었습니다.
2. 유지보수성: 수십 줄에 달하는 복잡한 텐서 조작(Reshape, Transpose 등) 코드가 제거되어 코드 베이스가 훨씬 간결해졌습니다. 이는 향후 모델 업데이트 시 발생할 수 있는 버그 가능성을 크게 낮춥니다.
3. 일반적 교훈: 하드웨어 가속기(GPU)를 다룰 때는 가급적 프레임워크에서 제공하는 최적화된 커널(SDPA 등)을 우선적으로 사용해야 합니다. 수동 구현은 특정 조건에서만 동작하거나 최신 하드웨어 아키텍처의 이점을 충분히 활용하지 못할 가능성이 높기 때문입니다.

이번 PR은 복잡한 로직을 제거하는 것만으로도 성능과 코드 품질을 동시에 잡을 수 있다는 좋은 사례를 보여줍니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html

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