AMD MI355X перевищує NVIDIA B200 у розгортанні Kimi K3

Це може бути момент, якого AMD чекала довго.

Останнім часом Wafer AI розгорнув на AMD MI355X Kimi K3. В результаті модель, яка раніше вимагала 16 NVIDIA B200, розподілених між двома серверами, тепер може бути розгорнута на одному сервері AMD з 8 MI355X.

Ще важливіше, це не просто вбудовування моделі.

У тесті з введенням 1024 токенів та виведенням 400 токенів MI355X показав загальну пропускну здатність 952 токени/с, а швидкість генерації для одного користувача досягла 118 токенів/с.

За розрахунком на один вузол, його пропускна здатність приблизно в 3,8 раза вища, ніж у схемі з 16 картами B200, а також перевищує співвідношення ціна/якість B200 і B300.

А найбільш несподіваним було те, що ROCm цього разу не вдавався в особливі труднощі.

Kimi K3 має 2,8 трильйона параметрів, і лише ваги моделі вимагають понад 1,5 ТБ відеопам’яті, не враховуючи KV Cache, необхідний для мільйонів токенів контексту.

Сервер з 8 GPU B200, кожен з яких має 192 ГБ відеопам’яті, загальний об’єм становить близько 1,5 ТБ. Це означає, що важко розмістити навіть ваги моделі повністю, не кажучи вже про виділення простору для KV Cache. Тому B200 повинен використовувати дві сервери з 16 GPU.

Кожна карта B300 має 288 ГБ відеопам’яті, що дозволяє розмістити модель в одному вузлі. Цікаво, що AMD MI355X також має 288 ГБ відеопам’яті; вісім карт MI355X разом дають приблизно 2,3 ТБ — цього достатньо для одного сервера.

Це не просто менше використання одного пристрою. Після запуску моделі через кілька вузлів, кожен згенерований токен може вимагати синхронізації даних через мережу. Навіть при використанні мережі RoCE v2 зі швидкістю приблизно 195 Гб/с, обмін даними між вузлами все ще сповільнює декодування.

MI355X завдяки більшому об’єму відеопам’яті зберігає всю модель в одному вузлі.

Згідно з кінцевими результатами, загальна пікова пропускна здатність 8 MI355X досягла 952 токенів/с, а швидкість генерації на одному каналі — 118 токенів/с.

Для порівняння, загальна пропускна здатність двохвузлової конфігурації з 16 B200 становить 498 токенів/с, що дорівнює приблизно 249 токенів/с на вузол.

Тобто, однокластерна пропускна здатність MI355X становить приблизно 3,8 рази більше, ніж середня однокластерна пропускна здатність розгортання B200 з двома вузлами. Щодо швидкості генерації для одного користувача, 118 токенів/с у MI355X також вищі, ніж 90 токенів/с у B200.

B300 залишається найпотужнішим рішенням за продуктивністю. Загальна пропускна здатність вузла з 8 B300 досягає 1568 Token/s, швидкість генерації на один шлях — 172 Token/s, загальна пропускна здатність приблизно в 1,65 раза вища, ніж у MI355X.

Але зміна ціни змінила висновок. Wafer розраховується за ціною 2,5 долара за карту за годину для MI355X, 4,25 долара для B200 та 6 доларів для B300.

За цією ціновою гіпотезою, MI355X забезпечує приблизно 48 Token/s пікової пропускної здатності за долар; B200 — приблизно 7 Token/s; B300 — приблизно 33 Token/s.

B300 швидший, але MI355X має вищу ефективність витрат на одиницю. Для центрів обробки даних, які потребують масштабного запуску відкритих моделей, це може бути важливішим, ніж просто боротьба за титул лідера за продуктивністю.

Протягом тривалого часу найбільшою проблемою GPU AMD для центрів обробки даних часто було не обладнання, а програмне забезпечення.

Та сама модель може працювати без змін на CUDA, але на ROCm може знадобитися змінити фреймворк, додати оператори або навіть переписати нижчі ядра.

але Kimi Ситуація з K3 інша.

AMD надала підтримку, близьку до синхронізації при запуску. Wafer зазначив, що моделі можна запускати безпосередньо на MI355X, а подальша робота зосереджена на невеликій кількості проблем сумісності та оптимізації продуктивності.

Одна з проблем виникла під час етапу припущення декодування. Kimi K3 сам по собі не надає параметрів чернеткової моделі, необхідних для MTP або EAGLE, тому Wafer використав зовнішню чернеткову модель на основі блоків.

Ця схема працює безпосередньо на CUDA, але в середовищі ROCm перший реальний запит призводить до помилки планувальника. Причина полягає в тому, що в гілці ROCm не визначено функцію під назвою top_k_renorm_prob.

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

Wafer використав звичайну функцію PyTorch, щоб заповнити цю логіку, не потребуючи написання GPU-ядер вручну чи переробки системи прогнозування декодування.

Після виправлення оцінена декодування збільшило продуктивність однієї лінії приблизно в 2,2 рази, продуктивність однієї потоку при середньому рівні паралелізму — приблизно в 1,7 рази, а пикову загальну пропускну здатність — приблизно на 18%.

Ще важливіше, система може досягати пікової пропускної здатності при більш високій паралельності.

Звичайно, пропускна здатність — це не все для сервісу виведення. Для реальних користувачів іншим показником, що безпосередньо впливає на досвід, є TTFT — час очікування від надсилання запиту до отримання першого токена.

На цьому завданні MI355X спочатку показав поганий результат. Для виконання завдання з попереднім заповненням у 172 000 токенів MI355X витрачає приблизно 51 секунду, тоді як B300 — лише приблизно 23 секунди.

У моделях, що підтримують контекст з мільйонами токенів, завдання попереднього заповнення може бути дуже об’ємним. Якщо під час роботи з довгим контекстом користувачам доводиться чекати по десятки секунд або навіть більше, будь-яка висока швидкість декодування важко зможе компенсувати проблеми з досвідом.

Wafer виявив, що різниця в продуктивності майже повністю походить від одного атенційного ядра. Kimi K3 у конфігурації з 8-канальним тензорним паралелізмом розподіляє 12 увагових голов на кожну GPU. А швидший префіллінговий ядро MLA в AMD AITER підтримує лише форми, кратні 4, 8 або 16.

12 голов не збіглися, тому система повернулася до повільнішої загальної реалізації Triton.

Рішення дуже просте: додайте нулі до 12 голов уваги, щоб отримати 16, використовуйте наявний високопродуктивний ядро, а після обчислення візьміть назад лише справжні 12 голов. Не було змінено структуру моделі і не написано нових асемблерних ядер — просто додано чотири нулі.

Після оптимізації стабільна швидкість попереднього заповнення ядра AITER MLA досягла приблизно 13 000 токенів/с, тоді як початковий шлях Triton становив лише 4 000–7 000 токенів/с, що скоротило час холодного попереднього заповнення приблизно у два-три рази.

Ця оптимізація не змінює кінцеву пропускну здатність декодування, але значно скорочує час очікування користувача на з’явлення першого символу.

Це також свідчить, що на перший погляд велика різниця в програмному забезпеченні між AMD та NVIDIA іноді не пов’язана з недостатністю базових можливостей, а лише з тим, що існуючі високопродуктивні ядра ще не підтримують певні нові форми моделей.

Одне тестування, звичайно, не може довести, що AMD повністю наздогнав NVIDIA.

B200 через недостаток відеопам’яті змушений працювати між вузлами; абсолютна продуктивність B300 залишається лідером; інструментарій, підтримка фреймворків та екосистема розробників ROCm все ще поступаються CUDA.

Але відкриті моделі швидко входять у еру трильйонів параметрів. Коли модель стає настільки великою, що її не можна розмістити на одному сервері, обсяг відеопам’яті перестає бути просто цифрою в таблиці параметрів — він безпосередньо впливає на вартість комунікації, складність розгортання та кінцеву пропускну здатність.

Стратегія AMD щодо надання більшої кількості HBM для однієї карти перетворюється на реальний системний перевагу.

Якщо AMD зможе продовжувати підвищувати стабільність ROCm, розширювати підтримку форматів високопродуктивних ядер і забезпечувати швидшу сумісність з новими моделями в день випуску, то центрами обробки даних слід серйозно розглянути ці GPU. Вони дешевші, мають більше відеопам’яті, достатню продуктивність і не вимагають місяців підлаштування програмного забезпечення.

Як ви думаєте про це?

Посилання для довідки:

https://x.com/wafer_ai/status/2083628389903315406

https://x.com/ChiragAsarpota/status/2083864019870634151

Цей матеріал зі сторінки WeChat «Machine Heart» (ID: almosthuman2014), автор: той, хто стежить за LLM