엔지니어링 조사: 이 MLX 포트를 한 번 더 10× 빠르게 만들 수 있는가?

날짜: 2026-09-19. 머신: Apple M3 Max, GPU 코어 40개, 통합 메모리 128 GiB, macOS 27.2, MLX / MLX Metal 0.32.2, FP16 추론. 이 보고서에는 손으로 작성한 Metal 커널을 포함한 실제 로컬 실험이 담겨 있습니다. 프로덕션 런타임을 바꾸거나 양자화 가중치를 공개하지는 않습니다.

테스트한 엔지니어링 변경은 10×를 내지 못합니다. 인터리브 측정은 컴파일과 마지막 결정 헤드 레이어의 미사용 출력 프루닝에서 나오는, 형상에 의존하는 소폭 개선을 뒷받침합니다. 짝지은 라운드별 중앙값으로 일부 사례는 약 3–8% 개선되었습니다. 일부 큰 배치 구간에는 개선이 아예 없습니다. 맞춤형 정확-erf GELU/게이트 커널은 수치적으로는 성공했지만 MLX 컴파일 대비 일관된 추가 엔드투엔드 이득을 주지 못했습니다. 순진한 8비트와 4비트 백본 양자화는 저장량을 줄였으나 더 큰 파일럿 작업 부하를 가속하지 못했고, 예측이나 캘리브레이션된 확률을 바꿨습니다.

수학적 한계와 근사 트레이드오프는 MATH_10X_RESEARCH.md에서 따로 다룹니다. 최초 구현 검토는 PERFORMANCE_RESEARCH.md에 있으며, 릴리스 체크포인트 벤치마크는 여전히 BENCHMARKS.md입니다.

실험 통제와 한계

모든 연구 GPU 작업은 직렬로 실행했습니다. 다른 에이전트 작업은 CPU/파일시스템/네트워크만 사용했습니다. 머신은 AC 전원이었고, pmset 열/성능 경고가 기록되지 않았으며 실험 중 스왑 사용도 보고되지 않았습니다. 일반적인 데스크톱 활동은 계속되었습니다. 통제된 열 챔버나 그 외에는 유휴 상태인 전용 벤치마크 머신이 아닙니다.

첫 스크리닝 실행은 각 후보를 새 프로세스에서 실행했고, 워밍업 4–5회와 표본 12–16개를 썼습니다. 이는 실행 간 변동이 상당함을 드러냈습니다. 예를 들어 영어 단일 질문 파일럿은 1.24× 컴파일 개선을 시사했지만, 이어진 인터리브 실험은 약 1.03×만을 찾아냈습니다. 따라서 순차 파일럿 지연 시간은 스크리닝 근거이며, 주요 인과적 속도 향상 주장이 아닙니다.

확인 스크립트 paired.py:

• 각 라운드 안에서 후보 순서를 회전시키고, 그 라운드의 모든 후보에 같은 입력을 씁니다.
• 라운드 사이에 실제 상태 텍스트를 바꿉니다. 최대 16개의 상태 변형을 생성하고 텐서 형상이 같은 변형만 남깁니다. 다국어 짧은 사례에는 그런 변형이 10개이고, 보고된 다른 사례에는 16개입니다.
• 서로 다른 자연어 질문을 쓰며, 가장 큰 짧은 작업 부하에는 50개의 서로 다른 지시문이 들어갑니다. 입력 해시를 검사하고, 답을 캐시하거나 질문을 중복 제거하거나 문맥적 인코더 상태를 재사용하지 않습니다.
• 각 타이머를 멈추기 전에 결과를 평가하고 GPU를 동기화합니다. 준비된 포워드 호출과 공개 예측 경로를 모두 측정하며, 토크나이즈와 출력 포매팅을 포함합니다. 모델 로딩은 제외합니다.
• 워밍업 후 영어 헤드/컴파일 실험에는 측정 라운드 32회, 다국어 및 맞춤형 Metal 실험에는 16회를 실행합니다. 각 후보는 같은 라운드 수와 입력 시퀀스를 봅니다.

이 입력들은 공개된 기준선 픽스처와 다릅니다. 아래 비교는 연구 실험 내부의 것이며, 무관한 표를 나눠 얻은 전후 비교가 아닙니다. 연구 배치 한도는 64이고, 릴리스된 API 기본값은 16입니다. 반복된 상태 변형은 의도적인 반복 측정이며, 결과 캐시는 없습니다.

analyze.py는 라운드 인덱스를 2,000회 재표본해 라운드별 eager_time / candidate_time 비율과 그 중앙값에 대한 탐색적 백분위 부트스트랩 구간을 계산합니다. 그 구간은 운영체제 잡음이나 계열 상관의 모든 원인을 반영하지 않으며, 다중 세션 반복을 대체하지 않습니다. 독립적으로 계산한 p50 값의 비율은 짝지은 중앙값 비율과 다를 수 있습니다.

원시 JSON에는 모든 타이밍, 입력 해시, 환경 메타데이터, 동등성 지표, 측정 시 기록한 소스 핑거프린트가 들어 있습니다. 실험 스크립트는 이후 포매팅되고 서로 겹치지 않는 선택적 후보들로 확장되었습니다. 이전 핑거프린트는 그 이전 스크립트 버전을 설명합니다.

컴파일과 정확한 최종 헤드 프루닝

네 경로를 비교했습니다:

1. Eager: 릴리스된 FP16 DecisionModel.
2. Compiled: 로딩·평가·고정된 모델을 일반적인 형상 특화로 mx.compile로 감쌈.
3. Selected Q + compiled: 마지막 헤드 레이어에서 전체 길이 QKV 투영과 K/V는 보존하되, CLS/선택지 마커 어텐션 쿼리만 발행합니다. 출력 투영과 FFN은 이 선택된 출력에만 실행합니다.
4. Full attention + selected outputs + compiled: 원래의 전체 길이 QKV와 SDPA 호출을 보존한 뒤, 출력 투영과 FFN 전에 CLS/선택지 출력을 모읍니다. 이는 원래 어텐션 커널 형상을 유지하면서 미사용 최종 헤드 dense 작업 대부분을 제거합니다.

두 프루닝 프로토타입 모두 모델의 수학적 의존성을 보존합니다. 여전히 모든 QKV 투영을 계산하며, 수학적 상한 계산에 있는 추가적인 Q-only 투영 절감은 실현하지 않습니다. GEMM과 SDPA 형상을 바꾸면 부동소수점 반올림이 달라질 수 있습니다. 어느 프로토타입도 디코더 캐시, 얼리 엑시트, 이전 transformer 레이어를 버리는 근사가 아닙니다.

엔드투엔드 p50 지연 시간, 밀리초:

출처: 영어 짝지은 데이터와 다국어 짝지은 데이터.

full-attention/selected-output 경로의 짝지은 중앙값 속도 향상과 탐색적 95% 구간은 다음과 같습니다:

영어 질문 50개와 다국어 긴 배치 구간은 1을 포함합니다. 반복 가능한 개선을 확립하지 못합니다. selected-Q 경로는 다국어 짧은-16과 긴-1 사례에서 다소 낫지만, 어떤 프루닝 경로도 모든 형상을 지배하지는 않습니다. 모든 후보 구간, 포워드 측정, 라운드별 원시 비율은 paired_analysis.json에 있습니다.

컴파일은 이 헤드/컴파일 실험에서 두 모델 계열에 걸친 변경 입력 질문 비교 1,530건에서 eager 로짓, 액션 로짓, 캘리브레이션된 확률과 정확히 일치했습니다. 두 프루닝 경로는 모두 1,530개 argmax 결정에서 일치했고, 캘리브레이션된 확률의 최대 차이는 0.0001883이었습니다. full-attention 프루닝 경로는 각 모델의 별도 63질문 픽스처 스위트도 통과했습니다: 126/126 일치, 최대 확률 차이는 영어 4.31e-5, 다국어 6.48e-6입니다. 이는 회귀 검사이며, 1,530개의 독립적으로 레이블된 예시에 대한 작업 정확도 주장이 아닙니다.

전체 모델 컴파일과 블록별 컴파일을 모두 스크리닝했습니다. 블록 실험도 영어 픽스처 출력 63개를 모두 보존했지만, 전체 모델 컴파일 대비 실질적 우위를 확립하지는 못했습니다. 형상 특화는 서비스에서 상한을 두어야 합니다. 모델이 Python의 형상 의존 reshape와 마스크를 쓰므로, shapeless=True를 무분별하게 적용하는 것은 안전하지 않습니다. 공식 컴파일 가이드는 형상 특화와 상태 캡처를 문서화합니다.

첫 영어 전체 모델 후보 호출은 2,166.7 ms가 걸렸고, 그 파일럿의 워밍 포워드 p50은 약 12.75 ms였습니다. 새 B16 형상의 첫 호출은 272.4 ms였습니다. JSON 필드 이름은 cold_forward이지만, 이는 완전히 콜드한 애플리케이션이나 새로 초기화된 Metal 드라이버가 아니라 eager 기준 추론 이후의 첫 후보 호출을 뜻합니다. 이후 후보들은 이전에 컴파일된 Metal 커널을 재사용했으므로, 그 첫 호출 시간은 콜드 스타트 비용의 통제된 순위가 아닙니다. 컴파일된 영어 짧은-1의 활성/최대 MLX 메모리는 파일럿에서 약 803.6/918.6 MiB였고, 다국어는 약 614.1/676.9 MiB였습니다. 이 할당자 측정은 호스트 측 컴파일러 할당을 모두 포함하지 않으며, 무한한 형상 변동에서의 메모리 한도를 확립하지 않습니다. 영어 컴파일 파일럿과 다국어 컴파일 파일럿을 참고하십시오.

선택적 양자화: 저장량 절감은 유용하나 속도 주장으로는 부적합

프로토타입은 dense FP16 모델을 로딩한 뒤에 nn.quantize를 호출합니다. encoder.layers.* 선형 모듈만 고르고, 아핀 그룹 크기 64를 쓰며, 그 결과 모델을 컴파일합니다. 임베딩, 정규화, 결정 헤드, 스코어러, 액션 헤드는 FP16으로 남습니다. 이는 현재의 dense 로더를 통해 패킹된 정수 가중치를 캐스팅하는 것을 피하고, 액션 헤드의 나누어떨어지지 않는 1028/772 입력 폭도 피합니다. 양자화 체크포인트 형식이나 로딩 계약은 공개하지 않습니다. 공식 MLX 양자화 레이어 구현이 이 선택 메커니즘을 제공합니다.

다국어 4비트 결과는 작은 픽스처 스위트만으로는 불충분한 이유를 잘 보여줍니다: 픽스처 argmax 63개는 그대로였지만, 별개 작업 부하 결정 26개 중 7개가 바뀌었습니다. 이는 FP16에 대한 일치 측정이며, 정답 정확도 측정이 아닙니다. 절대 확률 변화 0.8008은 80.08 퍼센트 포인트입니다.

영어 짧은-16 파일럿 입력에서 FP16 eager/compiled 엔드투엔드 p50은 91.26/87.94 ms였고, 8비트/4비트 compiled는 96.66/93.20 ms였습니다. 그 스크리닝 실행에서 짧은-1 양자화는 다소 빨라 보였지만 더 큰 형상은 그렇지 않았습니다. 다국어 큰 형상 스크리닝도 속도 이득을 보여주지 못했지만, 순차 실행의 변동이 컸습니다. 이 관찰들은 검증되지 않은 속도 향상이나 릴리스 주장을 거부할 근거가 되며, 인터리브된 양자화 반복 없이 정확한 감속 계수를 부여할 근거는 아닙니다. 추가 양자화 작업에는 활성화 인식 캘리브레이션이나 파인튜닝, 그리고 대표성 있는 레이블 품질 스위트가 필요합니다.

원시 출처: 영어 8비트, 영어 4비트, 다국어 8비트, 다국어 4비트.

손으로 작성한 Metal: 정확한 GELU/게이트 융합 구현 및 테스트

kernels.py는 실제 맞춤형 Metal 커널을 구현합니다. 이 커널은 이어 붙인 두 MLP 분기를 읽고, 같은 erf 기반 GELU를 계산하고, 게이트를 곱한 뒤, 단일 출력을 씁니다. tanh-GELU나 sigmoid 근사로 대체하지 않습니다. 커널은 MLX v0.32.2 자체의 erf와 expm1 헬퍼를 쓰고, 그 라이선스와 고지는 vendor/README.md에 보존합니다. FP16만 명시적으로 지원하고 안전한 Metal math 모드를 씁니다. 공식 맞춤형 커널 가이드는 이 API와 그 math-mode 제어를 설명합니다.

대표적인 활성화 형상 8개에 걸쳐, 무작위로 생성한 FP16 출력 원소 27,958,016개가 원래 연산과 값이 정확히 같았습니다. 마이크로벤치마크는 수치적 동등성을 비교하며, 0의 부호 비트는 비교하지 않습니다. 전체 모델 변경 입력 테스트도 정확히 일치했습니다: 두 모델 계열에 걸친 질문 비교 474/474, 그리고 63질문 픽스처 스위트 둘 모두에서 로짓, 액션-로짓, 캘리브레이션된 확률 차이가 0이었습니다.

이 정확성 결과는 MLX의 융합 컴파일 표현식 대비 일관된 속도 우위로 이어지지 않았습니다. 예를 들어 1,312토큰, 중간 폭 2,624에서 호출당 동기화된 활성화 타이밍은 eager GELU-후-게이트 0.378 ms, mx.compile 0.268 ms, 맞춤형 커널 0.280 ms였습니다. 8,192토큰, 폭 1,152에서는 각각 0.846/0.764/0.714 ms였습니다. 이 마이크로벤치마크는 디스패치와 동기화 오버헤드를 포함한 스크리닝 프로브이며, 고립된 장치 실행 시간 측정이 아닙니다. 전체 입력, 원시 타이밍, 동등성 검사는 microbench.json에 있습니다.

그다음 맞춤형 커널을 모든 인코더 MLP에 설치하고, 후보 순서를 회전시키고 입력을 바꿔가며 전체 모델에서 측정했습니다:

전체 맞춤형 커널 짝지은 실행은 동일한 가중치의 두 번째 모델 인스턴스를 써서, 수정되지 않은 구현과 맞춤형 구현이 변이 없이 또는 오래된 컴파일 캡처 없이 공존하게 합니다. 그 절대 타이밍은 이전 헤드 프루닝 실행과 비교해서는 안 됩니다. 소폭의 혼합된 결과는 맞춤형 커널을 일반적인 성능 개선으로 공개하는 것을 뒷받침하지 않습니다. 출처: 영어 Metal 짝지은 데이터와 다국어 Metal 짝지은 데이터.

맞춤형 엔지니어링이 추가 조사할 가치가 있는 곳

모델은 이미 mx.fast.scaled_dot_product_attention, mx.fast.rope, 최적화된 레이어 정규화를 호출합니다. D64 불리언 마스크 SDPA 경로는 융합되어 있으며, 10× 격차를 설명하는 빠진 Flash Attention 스위치는 없습니다. 지역 어텐션은 여전히 dense key/value 타일을 순회합니다. 진짜 양방향 윈도 커널은 양끝 포함 거리 <=64와 패딩 의미를 보존하면서 그 타일들을 건너뛸 수 있지만, 전체 모델 산술 기회는 짧은 입력에서 작고 공개된 긴 형상에서는 한정적입니다. 기존 소스 검토와 수학 보고서가 이 구분을 정량화합니다.

유용한 다음 프로젝트와 각각의 근거 요건은 다음과 같습니다:

• 장문 입력 윈도 어텐션: D64, 실제 양방향 윈도, 패딩된 배치에 맞춰 타일 경계를 특화합니다. 512/1024토큰에서 융합 dense SDPA와 비교한 뒤 전체 모델에서 비교합니다. 이 커널은 이 보고서에서 만들거나 벤치마크하지 않았습니다.
• Dense 커널 에필로그와 스케줄링: 게이트 MLP 에필로그를 GEMM에 융합하거나 짧은 M 행렬 스케줄링을 개선하는 것을 조사합니다. MLX는 이미 전용 Metal GEMM 구현을 쓰므로, 이를 대체하려면 실제 디스패치/커널 프로파일과 정확한 M/N/K 형상에서의 측정된 이득이 필요합니다. 독립 활성화 결과는 다른 요소별 커널 하나만으로는 불충분한 이유를 보여줍니다.
• 길이 인식 배칭과 공유 CPU 준비: 각 질문 시퀀스를 구성하기 전에 공유 상태 텍스트를 한 번 토크나이즈하면서 정확한 입력 ID를 보존하고, 작은 항목을 무관한 긴 항목에 패딩하지 않도록 합니다. 다국어 긴-8 파일럿은 입력 준비에 약 13.1 ms를 썼고, 엔드투엔드는 수백 밀리초였습니다. 그 준비를 완전히 없애도 이 작업 부하에서 10×가 나오지는 않습니다. 큐잉 지연 시간과 고유 추론 수는 모든 배칭 주장의 일부여야 합니다.
• 더 작은 공동 응답 학생 모델: 10×가 제품 요건이라면, 모델을 증류하거나 재설계해 dense 작업 대부분을 없애거나 하나의 문맥적 인코딩으로 많은 고정 질문에 답하게 합니다. 이는 학습된 모델을 바꾸며 대표성 있는 레이블 훈련/평가가 필요합니다. 정확한 포트 최적화가 아닙니다. 현재의 양방향 인코더에서 임의의 문맥적 상태/KV를 질문들 사이에서 재사용하는 것은 유효하지 않습니다.

독립 FP16 인코더 입력 투영 GEMM 프로브 8개는 호출당 동기화를 포함해 0.66–11.55 TFLOP/s를 달성했습니다. 큰 영어 M=4096, N=5248, K=1024 프로브는 11.55 TFLOP/s, 다국어 M=8192, N=2304, K=768 프로브는 7.96 TFLOP/s를 달성했습니다. 이는 관측된 처리량 값이며, 하드웨어 최고 사양이나 전체 그래프 처리량의 상한이 아닙니다. 작은 M 측정은 특히 제출과 동기화 비용이 지배하며, 스트리밍 그래프는 이를 다르게 상각합니다. 어떤 형상이 프로파일링할 가치가 있는지 보여줄 뿐, 더 나은 커널이 존재할 수 없다는 증명은 아닙니다. 수학 보고서의 10× 동일 작업 처리량 예산은 측정된 장치 성능이 아니라 이론적 요건으로 남습니다.

재현과 릴리스 결정

스크립트는 기존 .venv와 로컬에 고정된 체크포인트를 씁니다. GPU 명령은 순차적으로 실행하고, 공식 벤치마크와 나란히 실행하지 마십시오:

# Screening: repeat for eager, compiled, blocks, q8, q4, selected-compiled.
.venv/bin/python -m experiments.engineering.run_variants \
  --model laya --variant compiled --iterations 12 --warmup 4 --quality \
  --output experiments/engineering/reproduced-compiled.json

# Primary confirmation, including 50 genuinely different questions.
.venv/bin/python -m experiments.engineering.paired \
  --model laya --iterations 32 \
  --output experiments/engineering/reproduced-laya-paired.json
.venv/bin/python -m experiments.engineering.paired \
  --model laya-multilingual --iterations 16 --cases short1,short16,long1,long8 \
  --output experiments/engineering/reproduced-multilingual-paired.json

# Hand-written kernel microbench and complete-model comparison.
.venv/bin/python -m experiments.engineering.microbench
.venv/bin/python -m experiments.engineering.paired \
  --model laya --iterations 16 --cases short1,short16,long1 --metal \
  --output experiments/engineering/reproduced-metal-paired.json
.venv/bin/python -m experiments.engineering.run_variants \
  --model laya --variant metal-compiled --iterations 5 --warmup 3 \
  --cases short1 --quality --output experiments/engineering/reproduced-metal-quality.json

# CPU-only paired analysis.
.venv/bin/python -m experiments.engineering.analyze

모든 실험 Python 파일은 Ruff 포매팅과 린트 검사를 통과합니다. 안정적인 런타임, 원래 벤치마크 결과, 공개된 FP16 체크포인트가 릴리스 아티팩트로 남습니다. 컴파일과 정확한 최종 헤드 프루닝은 콜드 형상/캐시 정책과 더 넓은 품질 검증 이후의 신뢰할 만한 선택적 미래 최적화입니다. 측정된 이득은 기본 경로에 컴파일 지연 시간이나 맞춤형 커널을 조용히 추가하는 것을 정당화하지 못합니다. 어떤 10× 속도 향상, 프로덕션 준비 양자화 체크포인트, 측정된 지역 윈도 커널 이득도 주장하지 않습니다.