Быстрые бэкенды

Переносимость TileLang

Все пять ядер в laya/tl_kernels.py могут выполняться на CPU-цели (c) TileLang через compile_cpu. Это явная скалярная специализация fp32, а не CPU-бэкенд ускорения для Agent.accelerate. Она требует TileLang и локальный компилятор C++. Она не вводит зависимости от службы времени выполнения.

import torch
from laya import tl_kernels as K

A = torch.randn(17, 67)
W = torch.randn(70, 67)
b = torch.randn(70)
C = torch.empty(17, 70)
kernel = K.compile_cpu(K.gemm_kernel, 70, 67, bias=True, act="gelu")
kernel(A, W, b, C)

compile_cpu(factory, *args, **kwargs) принимает те же опции формы и операции, что и пять GPU-фабрик. Он выбирает cpu=True, dtype="float32", цель c и отключает векторизацию. Все вещественные входы и выходы должны быть непрерывными CPU-тензорами float32; длины внимания остаются int32. Преобразуйте 16-битные активации с помощью .cpu().float().contiguous() перед вызовом. LayerNorm по-прежнему обновляет свой остаточный тензор на месте, а RoPE по-прежнему обновляет столбцы Q/K на месте. Сохраняйте скомпилированное ядро, чтобы переиспользовать его динамические размеры.

CPU-специализация заменяет выделения fragment и shared локальными буферами, использует последовательные циклы T.grid и опускает GPU-аннотацию swizzle для внимания. GEMM используют скалярную CPU-реализацию TileLang, а редукции используют локальные буферы. Исходные блочные алгоритмы, предикаты padding, маски внимания и online softmax остаются общими с GPU-реализацией. Промежуточные результаты CPU остаются fp32, включая вероятности внимания; промежуточные результаты GPU сохраняют исходный dtype. Никакого ускорения на CPU или CPU-бэкенда для всей модели здесь не заявляется.

Повторная проверка на TileLang 0.1.14

Наблюдалось на Linux, Python 3.12.13, torch 2.11.0+cu130, TileLang 0.1.14 и RTX 4070 Ti SUPER. Прежнее опасение о переносимости относится к компиляции неизменённой GPU-специализации, а не к возможности понижения для CPU. На базовом коммите fa9a2a7 этой страницы документации не было в checkout.

Тест компилирует factory.get_tir(...) с target="c" и конфигурацией прохода FAST для GPU. Это точные тексты диагностики (расположения в исходниках и трассировки стека опущены). Для каждого ядра проверялись и bf16, и fp16.

Первые четыре сбоя fragment/редукции — это tvm.error.InternalError. Сбой генерации кода bf16 тоже InternalError. RoPE во FP16 возбуждает RuntimeError: Compilation Failed!, за которым следуют вызов компилятора и исходный код; приведённая выше диагностика выводится в stderr. Его сгенерированные векторные приведения также не могут преобразовать векторы float обратно в векторы half.

Чтобы отделить поддержку dtype от поддержки fragment, тесты компилируют каждое ядро снова с cpu=True (локальные буферы и последовательные циклы), сохраняя bf16 и отключая векторизацию. Тогда все пять падают в точности с:

Cannot convert type bfloat16 to C type

Таким образом, фрагменты — первый блокиратор для GEMM, GEGLU и внимания, редукции фрагментов — для LayerNorm, а bf16 независимо блокирует все пять. У RoPE нет ни выделения фрагмента, ни редукции; у GEMM и GEGLU нет явной операции T.reduce_*. Редукции внимания поначалу замаскированы его сбоем раскладки. Отдельные проверки только fp32-фрагментов воспроизводят ошибку заполнения и ошибку редукции как для reduce_sum, так и для reduce_max, без всякого GEMM или bf16. Замена только scope также дала эту семантическую ошибку для буфера GEMM (другие затронутые буферы — Ci, x и s):

[Tilelang Semantic Check] Local buffer `C_l` is indexed by T.Parallel loop variable `i`. Local buffers are thread-private and do not participate in parallel layout inference. Use T.serial/T.vectorized/T.unroll for per-thread local indexing, or T.alloc_fragment when the indexed dimension should be distributed across threads.

Последовательная специализация fp32 устраняет эти блокираторы. Диагностические утверждения привязаны к версии 0.1.14 и пропускаются на другой версии, где ошибки следует проверить заново. Числовые CPU-тесты продолжают выполняться на других версиях.

Числовые проверки

Запустите python -m pytest tests/test_fast_cpu.py -q -s. В указанном выше окружении: 58 прошли. Проверки только на CPU не требуют CUDA; без неё пропускаются лишь сравнения с GPU. Покрытие включает неравномерные тайлы GEMM M/N/K, все эпилоги GEMM, GEGLU, все комбинации residual/bias, большие значения residual, зацикливание позиции RoPE и нетронутые столбцы V, статические/динамические формы внимания, скользящие окна, частичные тайлы, неравные длины, пустые последовательности и конечные выходы padding.

Зерно — 1234. Сравнения с GPU используют идентичные входные значения, округлённые до bf16 или fp16, а затем повышенные до fp32 для выполнения на CPU. Это абсолютные допуски для ограниченных фикстур, а не гарантия для произвольных величин или глубины модели. Сравнения внимания используют допустимые строки query, как в tests/test_fast.py.

Допуск CPU/эталон — 2e-5 (2e-6 для RoPE). Обновления остаточного потока в точности равны, включая случай без residual.

Свидетельства сохранения GPU

Ключевое слово cpu по умолчанию равно false. Существующие значения по умолчанию bf16/fp16, scope выделения GPU, параллельные циклы, swizzle и опции FAST не изменены. Вызовы GPU в fp32 по умолчанию по-прежнему возбуждают ValueError.

Прогоны «до/после» использовали соответственно исходный модуль, извлечённый через git show fa9a2a7:laya/tl_kernels.py, и изменённый модуль. Исходный загружался как laya.tl_kernels через importlib для базовых прогонов; весь остальной код Laya и окружение Python остались прежними.

• python -m pytest tests/test_fast.py -q, при этом full-forward тесты настроены загружать указанный ниже кэшированный английский чекпойнт: 13 прошли до; 13 прошли после. Изначально без чекпойнта было 11 прошли / 2 пропущены. Оба полных прогона выдали существующее предупреждение чекпойнта об ограничении температуры (choice:11+=0.10058280825614929 -> 0.5).
• Захват с зерном существующих тестов ядер: 24 выходных тензора бит-в-бит идентичны, включая обновления residual; максимальная разница до/после 0.
• Сгенерированный исходный код CUDA для всех пяти проверочных форм выше, в обоих dtype: 10/10 побайтово идентичны. Это также покрывает RoPE, отсутствующий в исходном fast-наборе.
• python benchmarks/parity_fast.py --model "$MODEL" --dtype bf16 --json ... и эквивалентная команда для fp16: 288 вопросов по 60 состояниям на каждый dtype. Все записи JSON до/после (вероятности fp32, stock и fast) сравниваются в точности равными; максимальная разница вероятностей до/после 0.

MODEL — это кэшированный английский снимок convaiinnovations/laya 55cf4c4ebb4ebe31b2550e8bdf3bd21b99753851; прогоны использовали HF_HUB_OFFLINE=1. Сравнение с многоязычным чекпойнтом или чекпойнтом типизированных решений здесь не заявляется.

Проверки репозитория

Все команды использовали /home/ckl/projects/S/laya/.venv/bin/python; ruff и zensical брались из каталога bin этого virtualenv.

Новый набор тестов зарегистрирован как требующий необязательного extra TileLang и компилятора C++ в существующих исключениях теста упаковки. Тест контракта API фиксирует аддитивный CPU-селектор только по ключевому слову, неизменный dtype GPU по умолчанию и сигнатуру compile_cpu, не импортируя TileLang в базовое окружение CI.

AOTInductor

DecisionModel можно экспортировать, скомпилировать в пакет .pt2 и загрузить с помощью torch._inductor.aoti_load_package. Пакет возвращает как логиты решения, так и логиты действия. Токенизация, padding, температурная калибровка и форматирование ответов остаются задачей вызывающего; это не добавляет бэкенд Agent и не меняет его путь выполнения по умолчанию.

Закрытый PR #472 задокументировал прежний блокиратор dtype. Объединённое состояние и признаки уверенности головы действий вычисляются в fp32, даже когда веса модели явно bf16. Без autocast первый линейный слой действия поэтому получал вход fp32 и веса bf16. Теперь прямой проход приводит конкатенированный вход к dtype весов головы только вне autocast. Softmax, энтропия и возвращаемые логиты решения сохраняют вычисления в fp32. Существующий eager-инференс с параметрами fp32, включая fp16/bf16 AMP, сохраняет свою численную схему; смешанные dtype весов/autocast тоже сохраняют исходное преобразование autocast.

Экспортируйте отдельную оценочную копию, преобразованную в bf16, вне autocast. Объявите ось token как 16 * Dim("tokens16", ...), чтобы удовлетворить защиты выравнивания внимания. Строки и количества маркеров могут быть динамическими независимо. Не предполагайте, что экспорт, снятый под autocast, можно упаковать вне этого контекста: для этого рецепта используйте явные dtype весов.

Воспроизведение офлайн

Проверка на основе assert по умолчанию использует крошечный случайно инициализированный ModernBERT, без доступа к сети. Передайте локальный каталог чекпойнта для измерения с реальной моделью. Отсутствие CUDA, API AOTInductor или компилятора C++ даёт явный SKIP; сбои компиляции и паритета на поддерживаемой установке приводят к сбою проверки.

python scripts/check_aoti.py --output-dir /tmp/laya-aoti-smoke
HF_HUB_OFFLINE=1 TORCHINDUCTOR_CACHE_DIR=/tmp/laya-aoti/cache \
  python scripts/check_aoti.py --model /path/to/local/multilingual \
  --output-dir /tmp/laya-aoti

Выходной каталог содержит decision.pt2, results.json, inputs.pt и eager.pt. JSON включает полные снимки nvidia-smi, время экспорта/упаковки/загрузки, байты пакета, задержку и максимальные абсолютные дельты логитов/вероятностей для обеих голов. Бинарные артефакты и кэши компилятора намеренно не хранятся в репозитории.

Измерено на RTX 4070 Ti SUPER

2026-10-03, Linux x86-64, Python 3.12.13, torch 2.11.0+cu130, transformers 5.17.0, драйвер NVIDIA 615.71.09, 16 376 MiB VRAM. Чекпойнт: convaiinnovations/laya, подкаталог multilingual на ревизии 1c5edc17a7acd8701df6fc341c0d179f1c62c982. Базовая линия — main fa9a2a7. Полные снимки машины и неокруглённые измерения — в aoti_multilingual_rtx4070.json.

Воспроизведённый сбой происходит на упаковке, после успешного экспорта: mat1 and mat2 must have the same dtype, but got Float and BFloat16. Каждый прогон использовал отдельный изначально пустой кэш Inductor; компиляция AMP выполнялась до упаковки, поэтому время упаковки — не измерение холодного старта свежего интерпретатора. Проверка в свежем процессе загружала только пакет и сохранённые входы (без чекпойнта), при этом точки входа torch.compile и упаковки были заблокированы. Она воспроизвела ответы обеих голов. PyTorch всё же скомпилировал небольшую проверку возможностей CPU AVX в изначально пустой кэш; ни граф модели, ни ядро CUDA не компилировались. Поэтому обобщённое утверждение «при загрузке нет активности компилятора» было бы на этой среде выполнения неточным.

Те же сохранённые входы содержат семь строк решений в трёх батчах, с реальным токенизированным текстом о выставлении счетов/возвратах, вопросами choice/score/noul, padding и неравными количествами допустимых маркеров. Каждая задержка — медиана пяти групп по 30 прямых проходов после десяти прогревов, с синхронизированным измерением по настенным часам. torch.compile(dynamic=True) использует режим по умолчанию. В эти числа задержки не входят явные графы CUDA, токенизация, экспорт или компиляция.

Здесь AMP означает параметры fp32 с bf16 autocast. Явный bf16 означает, что веса и остаточный поток — bf16, без autocast; используйте эти столбцы для наиболее близкого к пакету сравнения выполнения. Явный bf16 eager/compiled и AOTI были недоступны до этого исправления.

И argmax решения, и argmax действия совпадают на 7/7 строках для обоих сравнений AOTI. Вероятности — это сырые выходы softmax, без калибровки. Распределение действия насыщено на этих входах, поэтому его нулевая дельта вероятности не означает идентичных лежащих в основе логитов против AMP. Явный bf16 compiled прямой проход тоже отличается от явного bf16 eager (максимальная дельта логита решения 0.5, дельта вероятности 0.00659859). Компиляция с пониженной точностью не бит-в-бит.

GPU использовался совместно с настольными приложениями и другими задачами Python; компиляция на CPU тоже была совместной, а частоты не фиксировались. Это наблюдаемые времена, а не заявление об изолированном ускорении. Снимки nvidia-smi, обрамляющие прогоны:

Оставшиеся ограничения

• Артефакт специфичен для этого чекпойнта, точности, стека PyTorch/среды выполнения и цели GPU; переносимость на другое оборудование или версии PyTorch не проверялась.
• Эта проверка экспортирует строки 1–8, токены 32–512 с шагом 16 и слоты маркеров 2–8. Она задействует три формы, включая формы, отличные от примера экспорта, а не каждую точку в этих диапазонах. Один допустимый вариант может использовать дополненные слоты маркеров; настоящий однослотовый тензор идёт по отдельной ветви одного варианта и требует отдельного экспорта. Произвольные длины токенов и производственный слой бакетирования/диспетчеризации — вне этого изменения.
• Семь строк проверяют регрессию упаковки, а не широкую точность чекпойнта или паритет калиброванной уверенности. Приведение экспортной копии к bf16 отличается от сохранения весов fp32 под autocast; сам существующий eager-путь остаётся неизменным.
• Скрипт требует сборки PyTorch с поддержкой CUDA и локального набора инструментов компилятора, чтобы создать артефакт. .pt2 встраивает модель и ядра CUDA; никакой размещённый сервис не задействован.