Инженерное исследование: может ли этот порт на MLX стать ещё в 10× быстрее?

Дата: 2026-09-19. Машина: Apple M3 Max, 40 ядер GPU, 128 GiB унифицированной памяти, macOS 27.2, MLX / MLX Metal 0.32.2, вывод FP16. Этот отчёт содержит реальные локальные эксперименты, включая написанное вручную ядро Metal. Он не меняет производственную среду выполнения и не публикует квантованные веса.

Протестированные инженерные изменения не дают 10×. Чередующиеся измерения поддерживают скромные, зависящие от формы улучшения от компиляции и отсечения неиспользуемых выходов последнего слоя головы решений. Отдельные случаи улучшились примерно на 3–8% по парным медианам за раунд. Некоторые интервалы для более крупных пакетов не включают улучшения. Пользовательское ядро точного GELU/вентиля на основе erf было численно успешным, но не дало устойчивой дополнительной сквозной выгоды сверх компиляции MLX. Наивное 8-битное и 4-битное квантование бэкбона уменьшило хранение, не ускорило более крупные пилотные нагрузки и изменило предсказания или калиброванные вероятности.

Математические пределы и компромиссы аппроксимации рассматриваются отдельно в MATH_10X_RESEARCH.md. Исходный обзор реализации — в PERFORMANCE_RESEARCH.md; опубликованный бенчмарк чекпойнта остаётся BENCHMARKS.md.

Экспериментальные контроли и ограничения

Вся исследовательская работа на GPU выполнялась последовательно. Другая работа агента использовала только CPU/файловую систему/сеть. Машина работала от сети питания, без записанного предупреждения pmset о тепле/производительности и без сообщённого использования swap во время эксперимента. Обычная активность рабочего стола продолжалась. Это не контролируемая термокамера и не иначе простаивающая выделенная бенчмарк-машина.

Первые отборочные запуски выполняли каждого кандидата в свежем процессе, с 4–5 прогревами и 12–16 выборками. Они выявили существенный дрейф от запуска к запуску. Например, английский пилот с одним вопросом предположил улучшение компиляции в 1.24×, тогда как последующий чередующийся эксперимент обнаружил лишь около 1.03×. Поэтому задержки последовательного пилота — отборочное свидетельство, а не основное причинное заявление об ускорении.

Скрипт подтверждения paired.py:

• Чередует порядок кандидатов внутри каждого раунда и использует одни и те же входы для каждого кандидата в этом раунде.
• Меняет фактический текст состояния между раундами. Он генерирует до 16 вариантов состояния и сохраняет варианты с той же формой тензора; многоязычные короткие случаи имеют 10 таких вариантов, тогда как остальные отчётные случаи имеют 16.
• Использует разные вопросы на естественном языке, включая 50 разных инструкций для самой крупной короткой нагрузки. Он проверяет хеши входов и не кэширует ответы, не дедуплицирует вопросы и не переиспользует контекстные состояния энкодера.
• Вычисляет результаты и синхронизирует GPU перед остановкой каждого таймера. Он измеряет как подготовленные вызовы прямого прохода, так и публичный путь предсказания, включая токенизацию и форматирование вывода. Загрузка модели исключена.
• Выполняет 32 измеренных раунда для английского эксперимента голова/компиляция и 16 для многоязычных и пользовательских Metal-экспериментов, после прогрева. Каждый кандидат видит то же число раундов и ту же последовательность входов.

Эти входы отличаются от опубликованных базовых фикстур. Сравнения ниже — внутри исследовательского эксперимента, а не сравнения до/после, полученные делением несвязанных таблиц. Лимит пакета в исследовании — 64, тогда как выпущенный API по умолчанию использует 16. Повторяющиеся варианты состояния — это намеренные повторные измерения; кэша результатов нет.

analyze.py вычисляет по раундам отношения eager_time / candidate_time и исследовательские процентильные бутстрап-интервалы для их медианы, используя 2 000 пересэмплирований индексов раундов. Эти интервалы не учитывают все источники шума операционной системы или серийную корреляцию и не являются заменой многосессионной репликации. Отношение независимо вычисленных значений p50 может отличаться от медианы парного отношения.

Сырой JSON включает все тайминги, хеши входов, метаданные окружения, метрики паритета, и отпечаток источника, записанный на момент измерения. Скрипты эксперимента впоследствии были отформатированы и расширены непересекающимися опциональными кандидатами; более ранние отпечатки описывают те более ранние версии скриптов.

Компиляция и точное отсечение финальной головы

Сравнивались четыре пути:

1. Eager: выпущенная DecisionModel в FP16.
2. Compiled: mx.compile вокруг загруженной, вычисленной, замороженной модели, с обычной специализацией по форме.
3. Selected Q + compiled: в последнем слое головы сохранить полноразмерную проекцию QKV и K/V, но выполнять только запросы внимания CLS/маркеров вариантов. Запускать проекцию вывода и FFN только на этих выбранных выходах.
4. Full attention + selected outputs + compiled: сохранить исходный полноразмерный QKV и вызов SDPA, затем собрать выходы CLS/вариантов перед проекцией вывода и FFN. Это сохраняет исходную форму ядра внимания, убирая большую часть неиспользуемой плотной работы финальной головы.

Оба прототипа отсечения сохраняют математические зависимости модели. Они всё равно вычисляют все проекции QKV; они не реализуют дополнительную экономию проекции только Q из расчёта математической верхней границы. Изменение форм GEMM и SDPA может изменить округление с плавающей точкой. Ни один прототип не является декодерным кэшем, ранним выходом или аппроксимацией, отбрасывающей более ранние слои трансформера.

Сквозная задержка p50, миллисекунды:

Источники: английские парные данные и многоязычные парные данные.

Для пути полное-внимание/выбранные-выходы парное медианное ускорение и исследовательские 95%-е интервалы включают:

Интервалы для английских 50 вопросов и многоязычного длинного пакета включают 1. Они не устанавливают воспроизводимое улучшение. Путь selected-Q несколько лучше для многоязычных случаев short-16 и long-1, но ни один путь отсечения не доминирует над всеми формами. Все интервалы кандидатов, измерения прямого прохода и сырые по-раундовые отношения — в paired_analysis.json.

Компиляция точно совпала с eager по логитам, логитам действий и калиброванным вероятностям на 1 530 сравнениях вопросов с изменённым входом по двум семействам моделей в этом эксперименте голова/компиляция. Оба пути отсечения совпали по всем 1 530 решениям argmax, с максимальным различием калиброванной вероятности 0.0001883. Путь отсечения с полным вниманием также прошёл отдельный набор из 63 вопросов-фикстур для каждой модели: 126/126 согласий, с максимальными различиями вероятностей 4.31e-5 для английского и 6.48e-6 для многоязычного. Это регрессионные проверки, а не заявление о точности на задаче на 1 530 независимо размеченных примерах.

Компиляция всей модели и по блокам обе прошли отбор. Эксперимент с блоками также сохранил все 63 выхода английских фикстур, но не установил существенного преимущества над компиляцией всей модели. Специализация по форме должна быть ограниченной в сервисе. Модель использует зависящие от формы Python-переформы и маски, поэтому применение shapeless=True без разбора небезопасно. Официальное руководство по компиляции документирует специализацию по форме и захват состояния.

Первый вызов кандидата «вся английская модель» занял 2 166.7 ms, после чего в том пилоте тёплый прямой проход p50 составил около 12.75 ms; первый вызов новой формы B16 занял 272.4 ms. Поле JSON называется cold_forward, но оно означает первый вызов кандидата после эталонного eager-вывода, а не полностью холодное приложение или только что инициализированный драйвер Metal. Последующие кандидаты переиспользовали ранее скомпилированные ядра Metal, поэтому их времена первого вызова не являются контролируемым рейтингом стоимости холодного старта. Активная/пиковая память MLX скомпилированного английского short-1 составила около 803.6/918.6 MiB в пилоте; у многоязычного — около 614.1/676.9 MiB. Эти измерения распределителя не включают каждое выделение компилятора на стороне хоста и не устанавливают пределы памяти при неограниченной смене форм. См. пилот компиляции английского и пилот компиляции многоязычного.

Выборочное квантование: полезная экономия хранения, непригодная как заявление о скорости

Прототип вызывает nn.quantize после загрузки плотной модели FP16. Он выбирает только линейные модули encoder.layers.*, с аффинным размером группы 64, затем компилирует получившуюся модель. Эмбеддинги, нормы, голова решений, scorer и голова действий остаются в FP16. Это избегает приведения упакованных целочисленных весов через текущий плотный загрузчик и избегает неделимой входной ширины 1028/772 головы действий. Никакой квантованный формат чекпойнта или контракт загрузки не поставляется. Официальная реализация квантованного слоя MLX предоставляет этот механизм выбора.

Результат многоязычного 4-битного варианта иллюстрирует, почему одного небольшого набора фикстур недостаточно: его 63 argmax на фикстурах остались теми же, но 7 из 26 решений на разных нагрузках изменились. Это измерения согласия с FP16, а не измерения точности относительно эталона истины. Абсолютное изменение вероятности 0.8008 — это 80.08 процентного пункта.

На пилотных входах английского short-16 сквозная p50 eager/compiled FP16 составила 91.26/87.94 ms; скомпилированное 8-битное/4-битное — 96.66/93.20 ms. Квантование short-1 выглядело несколько быстрее в том отборочном запуске, тогда как более крупные формы — нет. Отбор многоязычных крупных форм также не показал выигрыша по скорости, но его последовательные запуски имели существенный дрейф. Эти наблюдения оправдывают отказ от безоговорочного заявления об ускорении или релизе, а не приписывание точных коэффициентов замедления без чередующейся квантованной репликации. Дальнейшая работа по квантованию требует калибровки с учётом активаций или дообучения и представительного размеченного набора качества.

Сырые источники: английский 8 бит, английский 4 бита, многоязычный 8 бит, многоязычный 4 бита.

Написанный вручную Metal: точная fusion GELU/вентиль была реализована и протестирована

kernels.py реализует реальное пользовательское ядро Metal, которое читает две конкатенированные ветви MLP, вычисляет тот же GELU на основе erf, умножает на вентиль и записывает один выход. Оно не подставляет tanh-GELU или сигмоидную аппроксимацию. Ядро использует собственные помощники erf и expm1 MLX v0.32.2, сохраняя их лицензии и уведомления в vendor/README.md. Оно явно поддерживает только FP16 и использует безопасный математический режим Metal. Официальное руководство по пользовательским ядрам описывает этот API и его управление математическим режимом.

По восьми представительным формам активаций 27 958 016 случайно сгенерированных выходных элементов FP16 имели в точности равные значения исходной операции. Микробенчмарк сравнивает численное равенство, а не знаковый бит нуля. Тесты всей модели с изменённым входом также совпали точно: 474/474 сравнения вопросов по двум семействам моделей, плюс оба набора из 63 вопросов-фикстур, с нулевым различием логитов, логитов действий и калиброванных вероятностей.

Этот результат корректности не превратился в устойчивое преимущество по скорости над слитым скомпилированным выражением MLX. Например, при 1 312 токенах и промежуточной ширине 2 624 синхронизированное время активации на вызов составило 0.378 ms для eager GELU-затем-вентиль, 0.268 ms для mx.compile и 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, и оптимизированную нормализацию слоя. Её путь SDPA с булевой маской D64 слит; нет отсутствующего переключателя Flash Attention, который объяснял бы разрыв в 10×. Локальное внимание всё ещё проходит по плотным тайлам ключей/значений. Настоящее ядро двунаправленного окна могло бы пропускать эти тайлы, сохраняя инклюзивное расстояние <=64 и семантику дополнения, но его возможность арифметии всей модели мала на коротких входах и ограничена на опубликованных длинных формах. Существующий обзор исходников и математический отчёт количественно оценивают это различие.

Полезные следующие проекты с их требованиями к доказательствам:

• Внимание с окном для длинных входов: специализировать границы тайлов для D64, фактического двунаправленного окна и дополненных пакетов. Сравнить со слитым плотным SDPA при 512/1024 токенах, а затем в полной модели. Это ядро не было построено или измерено в этом отчёте.
• Эпилоги плотных ядер и планирование: исследовать слияние эпилога MLP с вентилем в GEMM или улучшение планирования матриц при малом M. MLX уже использует специализированные реализации Metal GEMM, поэтому их замена требует реального профиля диспетчеризации/ядер и измеренных выигрышей для точных форм M/N/K. Результат для автономной активации показывает, почему одного ещё одного поэлементного ядра недостаточно.
• Пакетирование с учётом длины и общая подготовка CPU: сохранять точные входные ID, токенизируя общий текст состояния один раз перед построением каждой последовательности вопроса, и избегать дополнения малых элементов до несвязанных длинных. Пилот многоязычного long-8 потратил около 13.1 ms на подготовку входов против сотен миллисекунд сквозного времени. Даже полное устранение этой подготовки не дало бы 10× на этой нагрузке. Задержка в очереди и число уникальных выводов должны быть частью любого заявления о пакетировании.
• Меньший студент, отвечающий совместно: если 10× — продуктовое требование, дистиллируйте или перепроектируйте модель, чтобы убрать большую часть плотной работы или отвечать на многие фиксированные вопросы одним контекстным кодированием. Это меняет обученную модель и требует представительного размеченного обучения/оценки; это не оптимизация точного порта. Переиспользование произвольного контекстного состояния/KV между вопросами в текущем двунаправленном энкодере некорректно.

Восемь автономных проб GEMM проекции входа энкодера на FP16 достигли 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×, готовый к производству квантованный чекпойнт или измеренный выигрыш локально-оконного ядра не заявляется.