ハクソク

世界を動かす技術を、日本語で。

AppleのMLXがCUDAサポートを追加

2025年7月15日原文(github.com)

概要

CUDAバックエンドの最適化によって、トレーニング速度が大幅に向上。 cudaEventやprefetchコール削減、cudaLaunchHostFuncの呼び出し最適化を実施。 prefetchの削除により、カーネル実行の遅延が解消し、さらなる高速化を実現。 今後はバッファキャッシュやカーネル最適化などが課題。 現状でPyTorchを超える性能を達成。

CUDAバックエンド最適化の経緯

  • Event実装 をcuda::std::atomicから cudaEvent へ切り替え、トレーニング速度が 500 it/sから900 it/s に向上
  • prefetchコール の削減で、速度が 900 it/sから1100 it/s へ
  • Metalバックエンドでは、各op評価後に オペランドや一時変数 をカーネル終了まで保持
    • コマンドバッファの必要時にcommitし、完了ハンドラで解放・エラーチェック
  • CUDAではcudaLaunchHostFunc を用いて類似実装
    • ただし、 20µs以上のストリーム遅延 が発生し、各カーネルが待たされる課題
    • 対策として、 mx::eval_impl()完了時 までオペランド・一時変数を保持
      • これにより cudaLaunchHostFuncの呼び出し回数を大幅削減
      • 2カーネル間の遅延が 1µs未満 となり、PyTorch以上のパフォーマンス
      • デメリットは 配列の解放遅延によるメモリ使用量増加
      • MAX_ACTIVE_TASKSを超えるタスク時の待機も廃止
  • この最適化で速度が 1100 it/sから1600 it/s へ

カーネル遅延の原因と解決

  • 一部カーネルの 遅延共通点 は、実行前に cudaMemPrefetch API を呼んでいた点
  • CUDAバックエンドは Unified Memory API を利用し、自動でデータ転送
    • 明示的な cudaMemPrefetch でデバイス側に転送指示をしていた
  • プロファイリングで prefetchがカーネル実行を大幅遅延 させていることを発見
  • prefetch削除 により、速度が 1600 it/sから2100 it/s に向上
    • プロファイラのタイムラインも大幅に改善

今後の最適化課題

  • バッファキャッシュ の実装
    • 主要opの実装とサードパーティライブラリ統合後に着手予定
  • 残る課題は主に カーネル最適化 と CPUコードの高速化
    • すべてのop実装完了後に本格対応予定

まとめ

  • CUDAバックエンドの最適化により、 PyTorchを超える性能 を実現
  • 今後は バッファキャッシュ・カーネル最適化・CPU高速化 が焦点
  • 継続的なプロファイリングとボトルネック解消が重要

Hackerたちの意見

「待って、AppleのプラットフォームにはファーストパーティのCUDAサポートがない！」って思うかもしれないけど、このパッチセットは「CUDA 12とSM 7.0（Volta）以上のLinuxプラットフォーム」もサポートしてるからね。

MLXの重要な機能の一つが統一メモリアーキテクチャを使うことなんだけど、これってどうなるの？（リポジトリのREADMEに書いてある箇条書きを見てみて: https://github.com/ml-explore/mlx）全てのUMA APU（どのベンダーでも）にそれを持ってくるのは面白いと思うけど、ディスクリートGPUは別のアプローチが必要だよね？ 編集: PRのコメントを読んでみたら、CUDAはUMA APIを直接サポートしていて、必要に応じて自動的にコピーするみたいだね。

└

うん、そうだけど、私の経験から言うと、プリフェッチがないとメモリの待機時間がかなり長くなって、コピーを待つことになるんだ。もしデータセット全体がVRAMに収まるなら「手動プリフェッチ」をすれば適してるかもしれないけど、私のアプリケーション（MLトレーニング）ではパフォーマンスがめちゃくちゃ悪くなって、ストリーミングロードに移行する時間ができたくらいだよ。

これはワクワクするね。これってCUDAの統一メモリを使ってるの？どれくらい上手く機能するんだろう。CUDAの統一メモリの動作はAppleシリコンのと同じなのかな？Appleシリコンの場合、GPUとCPUの間でメモリが共有されてるって理解してるけど、CUDAではそうじゃないよね。だから、CPUにテンソルがあるとき、どうやってGPUに送るの？何らかの形でコピーが必要だよね。それとも、CUDAが全部隠してくれてるのかな？

└

これは私の推測だけど、彼らが販売している高性能なハードウェア、例えばAI用のサーバーラックのやつには、統一メモリがあるのかな？標準のGPUにはないって知ってるけど。パッチによると、その理由の一つはMacでの開発を簡単にして、スパコンで動かせるようにするためだったみたい。だから、統一メモリを持ってるハードウェアはそのクラスに入るかもしれないね。

└

ハードウェアの統一メモリがない場合、CUDAはページフォルトが発生したときに自動的にCPUとGPUの間でデータをコピーするよ。

これはElectronを作ったzcbenzから来てるよ、他にも色々やってるね。https://zcbenz.com/ いいね。

これってStrix Haloの影響がどれくらいあるんだろう。仕事用のコンピュータのために結構標準的な手当をもらってたんだけど、しばらく使わなかったから最近EVO-X2に使ったんだ。マジで、あれは中堅のznver5 EPYCマシンと簡単に競争できるよ。普通のEC2やGCEインスタンスなんて、r1337.xxxxlarge.metal.metalとかがちょっと優位かもしれないけど、z1d.metalやc6.2xlargeみたいなやつ（速いコア、いいNIC、基本的な条件）はそれを圧倒してる。あれは月に3-10Kかかるけど、こっちは本物のNVMEで1800ドルだった。ローカル推論はあまりやってないけど、いろんなYouTuberがDGX SparkはStrix Haloに比べて過剰/高すぎるって言い始めてる。もちろん、ROCｍはまだそこにないけど（今は真剣に取り組んでるみたいだから、時間の問題だね）。Appleの機材で完璧なCUDAが使えたら、Strixが安くて良いのとは違った魅力があるよね。

└

PRの説明では、クラウドクラスターのトレーニングを明確にターゲットにしてるね。

└

もちろん、問題はROCmがまだ整ってないってことだね（でも今は真剣に取り組んでるみたいだから、時間の問題だと思う）。競争力のあるAMDのGPUニューラルコンピュートは、もう10年以上前から「いつでも出る」って言われてる。

└

初めての人に説明すると、Strix HaloはAMD Ryzen AI Max+ 395と同じもので、Framework Desktopに搭載される予定で、最近はミニPCにも出てきてるよ。そのメモリ帯域幅は200GB/sなんだ。ほとんどの消費者向けx86プラットフォームと比べるとすごく良いけど、NvidiaのGPU（5090は1792GB/sだし、プロレベルのカードはどうかは知らないけど）やAppleの最高モデル（M3 Ultraは800GB/s）には遠く及ばないね。確かにお得感はあるけど、LLMのようなメモリ帯域幅を重視するアプリケーションには、やっと「十分なレベル」に入ってきたって感じ。

└

M4 Mac Miniと比べてどうなの？

Hacker Newsで議論の続きを見る