Від cuda до mlx: Експертиза K-Search на службі Apple Silicon

Нова ера в обчисленнях

Світ обчислень швидко змінюється. Інновації в апаратному забезпеченні виходять за межі швидших GPU, вводячи різноманітні чіпи від різних виробників, кожен з яких має свою архітектуру, часто розроблену для специфічних навантажень, пов'язаних з штучним інтелектом (ШІ). Паралельно, розробка програмного забезпечення також прогресує швидко, з інструментами кодування в ШІ, здатними генерувати те, що раніше займало місяці, за кілька хвилин.

У середовищі, де ШІ займає центральне місце, ядра GPU стають необхідними. Ці програми низького рівня, які виконуються всередині GPU, вимагають значної експертизи для ефективної оптимізації. Міграція ядра з одного апаратного забезпечення на інше є ще більшим викликом, часто вимагаючи повторного відкриття оптимізацій, специфічних для кожної архітектури.

Експертиза ядер CUDA

Екосистема CUDA, підтримувана десятиліттями досліджень та оптимізацій, дозволила розробити високоефективні реалізації критичних функцій, таких як увага або моделі простору стану. Ці досягнення представляють тисячі годин інженерії. Однак нові апаратні екосистеми, такі як Apple Silicon та інші спеціалізовані прискорювачі ШІ, швидко розвиваються, але не мають цієї глибини експертизи.

Тому стає критично важливим запитати, чи можна цю експертизу автоматично перенести.

K-Search: Ривок в оптимізації ядер

K-Search, еволюційна рамка оптимізації ядер, розроблена в UC Berkeley, є відповіддю на цей виклик. Цей метод використовує штучний інтелект для оптимізації ядер GPU. Розширивши K-Search з бекендом для MLX, машинного навчання Apple, ми створили шар перекладу з CUDA на MLX. Цей підхід дозволяє K-Search використовувати існуючі ядра CUDA для генерації нових для Apple Silicon, зменшуючи таким чином потребу в реконструкції ядер з нуля.

Результати говорять самі за себе: наш метод досягає продуктивності, близької до експертної, на Apple Silicon, пропонуючи прискорення 0,97x у порівнянні з рідним ядром уваги MLX і до 20x в швидкості попереднього заповнення в порівнянні з реалізацією спільноти mlx-lm на ядрі Mamba SSM.

Чому обрати MLX?

З кінця 2023 року рамка MLX від Apple зазнала значного прийняття. З мільйонами MacBook та Mac Studios, оснащених Apple Silicon, MLX дозволяє локальну інференцію ШІ без витрат на хмару. Його архітектура єдиної пам'яті робить його особливо підходящим для моделей середнього розміру (від 7 до 70 мільярдів параметрів на чіпах M).

Однак залишається значний розрив: багато критичних для продуктивності ядер, які зазвичай доступні в екосистемі NVIDIA, виявляються або відсутніми, або недостатньо оптимізованими на MLX без специфічного налаштування під апаратне забезпечення. Хоча MLX правильно виконує моделі, він часто залишає значний потенціал продуктивності на столі.

Як працює K-Search

K-Search вирізняється своєю здатністю ітерувати на наївному ядрі на основі специфікацій апаратного забезпечення. Модель ШІ міркує про оптимізації, які потрібно застосувати, генерує кандидатні ядра, які потім компілюються та тестуються на реальному апаратному забезпеченні. Отримані вимірювання живлять процес пошуку, дозволяючи постійне вдосконалення, поки продуктивність не конвергує.

Цей процес керується документом, специфічним для домену, який кодує правила апаратного забезпечення та математичні обмеження, запобігаючи генерації недійсного коду та забезпечуючи, щоб кандидати компілювалися та виконувалися ефективно. На практиці єдина модель (наприклад, Gemini 3.5 Pro Preview) виконує як роль міркування, так і генерації коду, аналізуючи ядро перед тим, як запропонувати оптимізації.

Висновок

Підсумовуючи, поєднання накопиченої експертизи в екосистемі CUDA та інновацій, принесених K-Search, відкриває нові перспективи для оптимізації продуктивності на Apple Silicon. Застосування цього методу не обмежується MLX, але поширюється на будь-яку екосистему, де експертиза CUDA може бути корисно перенесена.

Якщо ви хочете дізнатися більше про оптимізацію продуктивності вашої інфраструктури ШІ або про впровадження подібних стратегій, не соромтеся Зв'язатися зі мною.