DiffusionGemma 26B M2 Max 상륙: MLX 처리량 실측 및 Context 한계 도전

로컬 환경에서도 Agent에게 무한한 토큰 자유를 줄 수 있는 방법을 찾기 위해, 원래 가지고 있던 M4 24GB Mac에서 DiffusionGemma 26B를 실행해 보려 했으나, 1,000 tokens의 Context조차 버티지 못하고 바로 OOM(Out of Memory, 메모리 부족) 비극을 맞이했습니다.

M2 Max 96GB로 옮긴 후, 드디어 제 성능을 발휘할 수 있을까요? 저는 MLX(mlx-vlm 0.6.3)를 사용했습니다. 과정 중에 MXFP4 양자화(Quantization) 버그를 겪어 수동으로 패치(Patch)를 처리해야 했지만, 결과적으로 4-bit 형식으로 전체 벤치마크(Benchmark)를 성공적으로 완료했습니다.

본문은 지난 며칠간 Apple Silicon에서 DiffusionGemma 26B의 처리량(Throughput) 한계, 프롬프트(Prompt) 로드 비용, 그리고 Context 길이와 메모리 비용 사이의 관계를 기록합니다. 동시에, 이 실측 데이터를 향후 GH200 및 GB10 크로스 플랫폼 성능 비교를 위한 베이스라인(Baseline) 시리즈의 첫 번째 문서로 활용할 예정입니다.

두 가지 양자화 버전

첫 번째 배포(Deploy) 시 MXFP4의 역양자화(Dequantize) 버그를 겪었으며, 이후 4-bit로 교체하여 전체 벤치마크를 안정적으로 실행했습니다:

항목	MXFP4 (초기 버전)	standard 4-bit (최종)
하드웨어	Apple M2 Max, 96 GB 통합 메모리 (38-core GPU)
...

mlx-vlm은 MLX 커뮤니티에서 VLM(Vision Language Model) 전용으로 만든 추론(Inference) 프레임워크이며, DiffusionGemma의 블록 확산 디코더(Block diffusion decoder) 또한 지원 범위 내에 있습니다. 고수들의 글을 참고하여 이것으로 결정했습니다 XD

두 가지 양자화 버전의 선택

처음 시도한 버전은 MXFP4(mlx-community/diffusiongemma-26B-A4B-it-mxfp4)였습니다. 로딩은 성공한 것 같았으나 첫 번째 생성(Generation) 단계에서 즉시 오류가 발생했습니다:

ValueError: [dequantize] Biases must be provided for affine quantization

mlx-vlm의 _diffusion_soft_embedding_weight는 dequantize embed_tokens 시 기본적으로 affine 모드를 사용하지만, DiffusionGemma의 MXFP4 형식에는 bias 파라미터가 없습니다. 현재 찾아낸 해결책은 biases is None을 감지했을 때 mode="mxfp4"로 변경하는 것입니다.

MXFP4를 패치한 후 실행이 가능해졌습니다. 처음에는 짧은 Context(~8K)에서 피크(Peak) 메모리가 19 GB밖에 되지 않아 어떻게 이렇게 메모리를 아끼나 싶었습니다 XD. 하지만 나중에 확인해 보니 Context가 8K를 넘어서면 속도가 선형적으로 떨어지며, 16K에서는 거의 움직이지 못했습니다.

그래서 이후에는 standard 4-bit(mlx-community/diffusiongemma-26B-A4B-it-4bit)로 테스트를 변경했습니다. 메모리는 19 GB에서 45.7 GB로 뛰었지만, 짧은 Context에서의 속도가 두 배 이상 빨랐고 안정성도 훨씬 좋았습니다.

항목	MXFP4	standard 4-bit
모델 크기	14.8 GB	16.18 GB
...

MXFP4는 메모리를 절약하고 긴 Context에서 비교적 안정적이지만, standard 4-bit는 생성 속도가 두 배 빠르지만 메모리를 많이 사용합니다. 결국 다른 플랫폼과의 베이스라인 비교를 위해 저는 standard 4-bit를 기준으로 속도를 비교했습니다. 만약 정말로 더 긴 Context가 필요하다면 MXFP4로 돌아가는 것을 고려할 수 있습니다.

Prompt Encoding

프롬프트 인코딩(Prompt encoding)의 속도 곡선은 매우 흥미롭습니다:

프롬프트 길이	인코딩 속도
14 tokens	198 tok/s
...

짧은 프롬프트의 인코딩은 매우 느리지만(198 tok/s), 500 tokens를 넘어서면 650-700 tok/s 정도로 안정화됩니다. 이는 MLX가 짧은 시퀀스(Sequence)일 때 Metal GPU의 병렬 메커니즘을 충분히 활용하지 못해 오버헤드(Overhead)가 상대적으로 크게 나타나기 때문으로 보입니다. 초기 1K tokens의 콜드 스타트(Cold start) 비용은 실제 사용 시 큰 영향이 없습니다. 어차피 인코딩 단계(Encoding phase)는 생성 단계(Generation phase)보다 두 자릿수(Two orders of magnitude) 정도 빠르기 때문입니다.

Generation Throughput

Standard 4-bit 버전의 생성 속도와 MXFP4 버전의 차이는 매우 뚜렷합니다.

출력 길이	생성 속도	지연 시간
32 tokens	7.1 tok/s	4.5s
...
피크는 256 tokens(31.6 tok/s)에서 나타나며, 이는 하나의 diffusion canvas에 딱 들어맞는 수준입니다. MXFP4 버전의 14.7 tok/s보다 115% 더 빠릅니다. 512 tokens는 canvas를 넘어가야 하므로 29.1 tok/s로 속도가 떨어집니다.

더 높은 처리량(throughput)을 원한다면 max_denoising_steps=16(기본값 48)을 시도해 보세요. 품질은 낮아지지만 속도는 두 배로 빨라집니다.

Context Length

Standard 4-bit는 장점이 있지만, 메모리 소모가 45.7 GB로 급증하여 KV cache 공간이 MXFP4 버전보다 훨씬 적어지는 비극을 초래했습니다.

Context 길이	생성 속도	지연 시간
~1.8K tokens	0.61 tok/s	52.1s
...
데이터상으로는 시간이 꽤 많이 걸리는 것처럼 보이지만, 이 수치들은 생성이 느린 것이 아니라 주로 prompt encoding 단계에서 대부분의 시간이 소모된 것입니다. 이 약 45 GB에 달하는 모델 점유로 인해 KV cache는 남은 50 GB 내에서 짜내야 하는데, mlx_vlm.server의 메모리 관리(memory management)가 이러한 대규모 모델에 최적화되어 있지 않은 것으로 보입니다(곧 업데이트될까요?). 이로 인해 긴 prompt의 encoding phase가 거의 선형적으로 급증하는 현상이 발생합니다.

병렬 처리: standard 4-bit의 스케일링 (scaling)

Standard 4-bit 버전은 병렬 처리(concurrency) 테스트에서 MXFP4보다 성능이 다소 좋았습니다. 이 부분은 제가 Mac의 저수준(low-level) 구조에 익숙하지 않아 정확한 원인은 모르겠지만, 관찰된 결과입니다.

병렬 수	총 처리량	Wall time
Sequential	—	—
...
Concurrent 2(병렬 2)의 총 처리량은 단일 요청의 피크치와 비슷합니다(31.2 vs 31.6 tok/s). 이는 스케줄링 메커니즘의 오버헤드(overhead)가 크지 않음을 의미합니다. 반면 Concurrent 4는 약 26.9 tok/s로 떨어지며, 스케일링(scale) 효율은 약 85% 정도입니다.

또한, Concurrent 4의 wall time은 16.4s에서 38.1s로 급증하며, 마지막 요청(Request)은 처리를 시작하기까지 거의 22s를 기다려야 했습니다. 이는 DiffusionGemma의 문제가 아니라 MLX server의 설계 한계(design limitation)입니다. Metal backend는 CUDA의 병렬 커널 실행(concurrent kernel execution) 방식과는 달라 보이며, 모든 요청이 순서대로 대기해야 합니다. 따라서 Mac을 프로덕션 엔드포인트(Production endpoint)로 사용하는 것은 아직 권장하지 않습니다.

M4 24GB와의 비교

앞서 언급했듯이 M4 24GB에서 동일한 모델을 테스트한 결과는 그야말로 비극이었습니다.

항목	M4 24GB	M2 Max 96GB (standard 4-bit)
모델 footprint	16.18 GB	16.18 GB
...
가장 큰 병목(bottleneck)은 여전히 메모리입니다. M4 24GB는 모델조차 겨우 올릴 수 있는 수준이라 KV cache를 위한 공간이 전혀 없었지만, M2 Max 96GB는 Standard 4-bit가 약 45.7 GB를 점유하더라도 최소한 추론(inference)을 수행할 공간은 확보됩니다.

M2 Max 96GB는 DiffusionGemma 26B를 로컬에서 원활하게 실행할 수 있는 것처럼 보이지만(Standard 4-bit 피크 31.6 tok/s), 메모리와 백엔드 스케줄링 메커니즘이 긴 Context 및 병렬 처리 성능을 여전히 제한하고 있습니다.

실제로 CLI나 개발 환경에서 접하게 될 경우, 현재 온라인 서비스가 제공하는 사용자 경험과는 여전히 차이가 큽니다. 이어지는 두 번째 글에서는 GH200을 통해 vLLM으로 1180 tok/s라는 극한의 속도를 뽑아내는 과정을 다룰 예정이며, 세 번째 글에서는 GB10에서 32K Context의 한계에 도전하겠습니다.

대규모 모델이 다양한 하드웨어 아키텍처에서 보여주는 한계에 관심이 있다면, 이어지는 크로스 플랫폼 종합 리뷰를 계속 지켜봐 주세요!

최종 채택 모델: mlx-community/diffusiongemma-26B-A4B-it-4bit (Standard 4-bit)

초기 테스트 모델: mlx-community/diffusiongemma-26B-A4B-it-mxfp4 (MXFP4, dequantize 버그 수동 수정 필요)

테스트 환경: mlx 0.31.2 + mlx-vlm 0.6.3