Cursor's MoK досягає зростання продуктивності на 2,37x на NVIDIA GB300 NVL72

Cursor відкрив код MoK, інтегрувавши планування токенів, між-GPU комунікацію та обчислення експертів у один GPU-ядро. Це рішення досягає на GB300 NVL72 прискорення в 2,37 рази для прямого проходу та 1,78 рази для зворотного проходу при використанні MXFP8, а також підвищує пропускну здатність навчання на 512 GPU GB300 на 41% до 1070,2 токенів/с. Затримка сигналізації зменшена з 103 мкс до 18 мкс. Аналіз показує, що сьогодні, коли пропускна здатність NVLink досягла 130 ТБ/с, час, витрачений на дані в GPU, став новим обмеженням продуктивності. Цей прорив означає, що конкуренція в галузі ШІ перейшла на етап «повної власності стеку»: компанії, чий код знаходиться ближче до пам’яті та реєстрів GPU, отримують більше контрольних прав.

Автор статті, джерело: Leiphone

21 липня NVIDIA опублікувала останні результати навчання DeepSeek-V3 на GB300 NVL72: на 256 GPU продуктивність на одиницю досягла 1 648 TFLOPS.

Менше ніж через два тижні Cursor відкрив джерела Mixture-of-Kittens (MoK). Замість подальших спроб зробити матричне множення швидшим, він повністю переписав рівень виконання MoE, об’єднавши планування токенів, між-GPU комунікацію та обчислення експертів у один GPU-ядро.

Це трохи протилежно інтуїції, оскільки GB300 NVL72 вже помістив 72 GPU в одну й ту саму NVLink-домену, загальна пропускна здатність NVLink усього стелажу досягла 130 ТБ/с. За цими параметрами передача даних між GPU повинна бути достатньо швидкою.

У великомасштабному навчанні MoE комунікація все ще може уповільнювати обчислення експертів.

Проблема в потоці даних MoE. На кожному кроці роутер повторно вирішує, якому експерту надіслати токен, а експерти розподілені між різними GPU. Токен спочатку треба відправити через карти, після обчислення — повернути назад; перед відправленням потрібно впорядкувати позиції, а після отримання — чекати, поки дані не будуть повністю доступні. Зі зменшенням часу обчислення експертів завдяки MXFP8 і Blackwell Tensor Core ці очікування, які раніше були приховані за обчисленнями, стають все більш помітними.

Cursor робить MoK, саме з цього місця варто починати. Він не прагне лише до того, щоб один Dispatch передавався якнайшвидше, а натомість перерозподіляє, як токени досягають експертів, коли починати відлік, а також як одночасно використовувати GPU для комунікації та обчислень.

Сьогодні, коли обчислювальні ресурси були доведені до максимуму за допомогою Blackwell і NVLink, розробники зрозуміли: навіть найшвидше обладнання не зможе виправити неефективне програмне забезпечення.

Відкритий код MoK — це не лише перемога ядра, а й початок ери, коли «рівень застосунків визначає операції»: щоб витиснути останні 30% обчислювальної потужності, стартапи в галузі ШІ розпочали війну за панування в нижчих шарах.

Але щоб зрозуміти, чому цей дизайн Cursor ефективний, спочатку треба зрозуміти, де саме MoE сплачує «податок на зв’язок».

01

«Комунікаційний податок» MoE

У звичайній щільній FFN ваги, які проходять токени, майже фіксовані. Після додавання Router у MoE кожен токен тимчасово вибирає кілька експертів. При використанні Expert Parallel експерти розподіляються між багатьма GPU, тому під час однієї прямого проходу відбувається щонайменше два цикли міжкартової комунікації.

Перший етап називається Dispatch: токен надсилається на GPU, де знаходиться експерт; після завершення обчислень експертом результат повертається назад до початкового місця токена за допомогою Combine. Під час навчання відбувається зворотне поширення, що вимагає ще двох етапів зв’язку у зворотному напрямку.

Якщо просто перенести великий неперервний блок даних з A в B, NVLink уже достатньо швидкий. Проблема MoE полягає в тому, що розподіл даних, що надається Router на кожному кроці, різний.

На цьому етапі експерт може отримати багато токенів, а на наступному — дуже мало. Система спочатку має підрахувати, скільки токенів має кожен експерт, потім вирішити, куди розмістити ці токени на цільовому GPU, і намагатися розташувати дані одного експерта якомога ближче один до одного. Лише так Grouped GEMM отримає упорядкований вхід і зможе ефективно подавати його на Tensor Core.

Після завершення зв’язку не можна одразу починати обчислення. Цільовий GPU має переконатися, що віддалений запис повністю завершено, і дані дійсно видимі. Навантаження експертів не є повністю рівномірним — деякі GPU можуть завершити роботу раніше, але все одно доведеться чекати, поки найзавантаженіший експерт завершить свою роботу.

Отже, на етапі «зв’язку» поєднуються кілька завдань: перенесення даних, генерація розміщення, синхронізація та нерівномірність навантаження. 130 ТБ/с описує пикову пропускну здатність, яку може забезпечити весь стенд, але не означає, що кожна динамічна та фрагментарна комунікація MoE зможе одночасно повністю завантажити ці лінії.

DeepEP вже дуже швидко виконує перенесення даних. Це високопродуктивна бібліотека комунікацій, розроблена для Expert Parallel, яка надає спеціалізовані ядра Dispatch та Combine, підтримує FP8 і дозволяє керувати кількістю SM, що використовуються для комунікації. У найновішій версії можна підтримувати високу пропускну здатність комунікації за допомогою значно меншої кількості SM.

Навіть одна шар MoE все ще повинен постійно передавати дані між Dispatch, Grouped GEMM та Combine.

Виникає досить практичний суперечка: якщо чекати, поки збереться достатня кількість токенів, перед тим як розрахувати, матриця буде великою, і Tensor Core працюватимуть ефективно, але обчислення почнуться пізно; якщо ж розраховувати відразу ж, як тільки з’являються токени, комунікація та обчислення можуть раніше перекриватися, але матриця буде занадто малою, і багато SM на GPU не матимуть достатньо роботи.

Кілька CUDA Stream можуть дозволити паралельне виконання комунікації та обчислень, але важко завжди забезпечити оптимальне виділення ресурсів GPU для обох сторін.

Дизайн після MoK у основному стосується цього питання «ритму».

02

Однією з найцікавіших змін у MoK є зміна переднього Dispatch з Push на Pull.

Традиційний Push дуже інтуїтивний: якщо джерельний GPU має токен, він активно записує його на цільовий GPU. Проблеми виникають із цільовими адресами. Один GPU одночасно отримує багато токенів від інших GPU.

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

Зі збільшенням кількості задіяних GPU цей процес планування стає все більш навантаженим. Pull запровадив інший підхід: зберігає цільовий GPU експерта, який сам отримує необхідні токени. Йому достатньо знати, на якому джерельному GPU знаходиться токен і в якій позиції джерельних даних, а розміщення на локальному рівні визначається самим GPU.

Це економить координацію між кількома відправниками щодо цільової адреси. Отримані дані також можна безпосередньо впорядкувати за місцевими експертами.

Цікаво, що Pull не передає менше даних. У мікробенчмарку Cursor для одного й того ж блоку даних BF16 розміром 256×256, Push переміщує приблизно 159,6 КБ через NVLink, тоді як Pull досягає 172,0 КБ, оскільки для читання потрібно додатково надсилати запити.

Перевага Pull у іншому: комунікація MoE дуже фрагментована, а навантаження не збалансоване. Два напрямки NVLink мають окремі канали, і Pull може одночасно використовувати обидва напрямки — запити та повернення даних. У тестах з неоднорідним навантаженням експертів Cursor зафіксував зростання використання NVLink до 29%.

Розрив синхронізації стає більш помітним. Після завершення операції Push цільовий GPU повинен чекати на сигнал завершення від інших GPU; при великому масштабі Expert Parallel один rank може бути пов’язаний з іншими 71 peer. Pull ініціюється локальним GPU, і дані можуть використовуватися безпосередньо після отримання.

У багатовузловому мікробенчмарку Cursor це затримка сигналізації зменшилася з приблизно 103 мікросекунд при Push до 18 мікросекунд при Pull.

MoK також не використовує Pull у всіх місцях. У прямому напрямку використовується Pull Dispatch та Push Combine; у зворотному — Pull Reverse-Combine та Push Reverse-Dispatch. На етапі Dispatch потрібно перегрупувати багато токенів із різних джерел у вхідні дані експертів — Pull економить координацію. Під час Combine вже зрозуміло, до якого токена має повернутися кожен результат, тому простіше просто відправити його через Push.

Після зміни напрямку зв’язку MoK продовжив поєднувати зв’язок і експертні обчислення в одному Megakernel.

Він розділяє SM GPU на дві частини. Одна частина відповідає за диспетчеризацію, об’єднання та управління станом, інша — виконує тільки Expert FFN. Після отримання пакета повних токенів комунікаційна частина повідомляє обчислювальну частину за допомогою локального лічильника GPU; після завершення обчислень вона повідомляє комунікаційну частину, щоб передати результат назад.

Це дозволяє безпосередньо визначити, скільки SM використовувати для зв’язку, а скільки — для обчислень, не покладаючись повністю на конкуренцію кількох CUDA Stream. Найважливішим параметром тут є minibatch — кількість токенів, що надаються експертам для обчислення за один раз.

Він не може бути занадто великим. Занадто великий розмір означає, що перші обчислення доведеться чекати дуже довго. Він також не може бути занадто малим. Експерт GEMM нарешті розбивається на велику кількість обчислювальних завдань, які розподіляються між SM; якщо токенів замало, кількість завдань буде недостатньою, і багато SM залишаться безробітними.

Cursor використовує хвилі для визначення цього кордону. Повну хвилю можна спростити як ситуацію, коли всі обчислення SM отримали завдання. MoK бажає, щоб міні-пакет містив принаймні дві повні хвилі, щоб Tensor Core мав достатньо завдань для безперервного виконання.

Фактичні результати говорять самі за себе. На моделі Kimi 2.5 з прихованим розміром 7168 і проміжним виміром експертів 2048, Cursor оцінює, що міні-пакету потрібно щонайменше 2368 токенів. При 512 токенах час прямого проходу MoK становить 5,981 мс; при збільшенні до 2560 токенів він знижується до 3,425 мс. Подальше збільшення не призводить до значного покращення швидкості.

Іншими словами, розбиття комунікації на менші частини не завжди призводить до прискорення. Справді ефективне перекриття вимагає раннього надання даних, одночасно не розбиваючи GEMM надто дрібно.

Але у MoE є ще одна проблема: до завершення роботи маршрутизатора неможливо знати, скільки токенів отримає кожна GPU.

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

MoK використовує Ring Token Buffer фіксованого розміру. Простір спочатку заповнюється tokenами, що надходять від Dispatch; після завершення обчислень експертами та передачі результатів Combine, цей простір негайно повторно використовується для наступної партії tokenів. Combine попереднього macrobatch може виконуватися одночасно з Dispatch наступного macrobatch.

Буфер кільцевого типу тут діє як буферний шар: коли комунікація працює швидше, дані накопичуються всередині; коли обчислення споживають швидше, потрібно чекати на наступну партію токенів. Увесь процес керується станом на GPU, не вимагаючи, щоб CPU втручався на кожному кроці для вирішення наступного кроку.

Cursor також втілив квантування MXFP8 у маршрути даних Dispatch, Grouped GEMM і SwiGLU, виключивши окремий кернел квантування та зменшивши кількість зчитувань і записів проміжних результатів у HBM.

Після об’єднання Pull, minibatch, SM-зона та Ring Buffer, MoK справді перетворився на неперервну MoE-лінійку.

03

Бенчмарк Cursor вимірює повний MoE-шар, включаючи планування, розподіл, експертний FFN, об'єднання та остаточне зважене об'єднання, у порівнянні з NCCL + PyTorch, DeepEP, TransformerEngine та HybridEP + Megatron.

На GB300 NVL72 MoK досягає максимального прискорення до 2,37 разу у прямому проході та до 1,78 разу у зворотному проході порівняно з найшвидшими відкритими базовими лініями для кожної сценарії; для BF16 максимальне прискорення становить 1,92 разу у прямому проході та 1,58 разу у зворотному проході.

Ще важливішим є кінцеве навчання. Попередній виробничий розв’язок Cursor використовував DeepEP. Після заміни на MoK на 512 GPU GB300 пропускна здатність на один чіп зросла з 760,9 токена на секунду до 1070,2, що становить зростання приблизно на 41%.

Також треба чітко розуміти межі цих результатів. Cursor не опублікував повного поетапного абляційного аналізу, тому неможливо точно визначити, яка частина з 2,37 разів припадає на Pull, яка — на Megakernel, а яка — на Ring Buffer.

Окремо можна підтвердити покращення використання NVLink та затримки сигналізації Pull, інші переваги є результатом комбінації всієї системи виконання.

MoK також сильно залежить від апаратного забезпечення. Він спрямований на швидкі NVLink домени, такі як Blackwell і NVL72. Віддалене читання, комунікація та дрібнозернисте чергування обчислень базуються на здатності GPU низькочастотно доступати до пам’яті один одного. При зміні розміру прихованого шару моделі, Top-k та розміру експертів, відповідно змінюються оптимальний розмір міні-пакету та кількість комунікаційних SM.

Це також найцікавіша частина MoK. Раніше, обговорюючи оптимізацію MoE, легко було зосереджуватися лише на двох цифрах: скільки TFLOPS у GEMM і скільки GB/s у All-to-All. У поколінні GB300 просте подальше збільшення цих двох показників вже не може пояснити всю продуктивність.

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

Ракета з пропускною здатністю NVLink 130 ТБ/с все ще вимагає переписування ядер GPU для MoE, саме тому: з’єднання вже дуже швидке, тепер потрібно економити час, який GPU витрачає на очікування даних.

Поведінка Cursor щодо переписування ядер GPU означає, що конкуренція в епоху AI 2.0 вступила в нову фазу «повної стекової суверенітету».

Раніше ми вірили у принцип «кожен своєю справою»: якщо робиш додатки — роби додатки (Cursor), якщо нижній рівень — нижній рівень (NVIDIA). Але сьогодні конкуренція в галузі ШІ перейшла до нового етапу «позбавлення посередників»: Cursor не робить ядро, бо хоче, а бо змушений.

DeepSeek відкрив нову еру інженерного витискування, а Cursor розгорнув цей вогонь на рівні застосунків. Ця тенденція «позбавлення посередників» перетворює владу над ціноутворенням штучного інтелекту: у майбутньому оцінка компаній штучного інтелекту визначатиметься не кількістю токенів, які вони володіють, а тим, наскільки їхній код близький до пам’яті відеопам’яті та реєстрів.

Ті AI-компанії, які не зможуть проникнути через нижній чорний ящик, залишаться в багнюці «податку на посередність».