triton-ascend-migration: Скилл для оптимизации и миграции

--- name: triton-ascend-migration description: Преобразование GPU/CUDA операторов Triton в Triton-Ascend, или переписывание Python/PyTorch операторов для работы на NPU Ascend с использованием Triton-Ascend, с прямой выдачей оптимизированного кода, минимальных скриптов для проверки и инструкций по устранению неполадок, когда это возможно. Пользователи, упоминающие "昇腾", "Ascend", "NPU", "triton-ascend", "миграция операторов Triton", "переписывание PyTorch операторов", "coreDim", "UB overflow", "1D grid", "привязка к физическим ядрам", "block_ptr", "stride", "выравнивание доступа к памяти", "производительность маски", "уменьшение точности", "оптимизация операторов", или напрямую спрашивающие "как использовать этот навык", "как запустить это в командной строке", "как выполнить миграцию/проверку в контейнере", должны в первую очередь использовать этот навык. --- # Triton-Ascend Migration ## Quick Start При получении запроса на миграцию, выполняйте действия в следующем порядке: 1. Сначала определите способ ввода: - Путь к файлу / указанный фрагмент кода - Пользователь напрямую вставляет код 2. Затем определите источник ввода: - GPU/CUDA kernel Triton - Реализация Python/PyTorch оператора 3. Затем определите тип оператора: - `elementwise` - `broadcast / mask` - `reduce` - Содержащий `tl.dot` 4. Сначала создайте минимальную рабочую версию: - `cuda -> npu` - Добавьте `torch_npu` - Удалите логику, специфичную для GPU - Сначала используйте 1D grid - Простые примеры по умолчанию предоставляют "минимальную версию для миграции" 5. После успешного запуска выполняйте оптимизацию на стороне Ascend: - Привязка к физическим ядрам - `BLOCK_SIZE/XBLOCK` - `BLOCK_SIZE_SUB/XBLOCK_SUB` - Непрерывный/выровненный доступ к памяти - Проверка `coreDim` / UB / dtype / mask 6. Если есть четкие возможности для оптимизации, выводите оптимизированную реализацию напрямую, а не только предлагайте. ## Как использовать этот Skill Если пользователь спрашивает "как использовать этот навык", сначала не переходите сразу к подробному анализу миграции; сначала дайте краткое описание использования из 3-6 строк, а затем продолжайте выполнение на основе ввода, предоставленного пользователем. В кратком описании сохраняйте только следующие пункты: - Пользователь может предоставить код `Triton/CUDA`, реализацию `PyTorch` в качестве примера, путь к файлу или сообщения об ошибках/журналы производительности. - Пользователю лучше одновременно указать среду выполнения: командная строка локальной машины, существующий контейнер, CI, или только сгенерировать код без выполнения. - Если пользователь имеет предпочтения, он должен указать их: `минимальная миграция`, `стиль документации`, `сначала запуск, затем оптимизация`, `сразу предоставление оптимизированной версии`. - Вы будете выводить: `реализацию Triton-Ascend`, `минимальный скрипт для проверки`, `команду для выполнения`, `описание оптимизации`. Если пользователь продолжает спрашивать "как задать вопрос", "как написать команду", "как запустить в контейнере", обратитесь к `references/usage.md` и предоставьте команды для локальной машины, команды для контейнера и примеры вопросов по мере необходимости; не включайте все подробные объяснения в обычные ответы. Скопируйте этот контрольный список и отслеживайте прогресс: ```text Прогресс миграции - [ ] Определение источника ввода и типа оператора - [ ] Создание минимальной миграции или семантической переработки - [ ] Адаптация для параллелизма и grid, удобных для Ascend - [ ] Переработка block / tiling - [ ] Проверка stride / block_ptr / выравнивание - [ ] Обработка coreDim / UB / scalar деградация - [ ] Непосредственная реализация возможных оптимизаций - [ ] Создание и сохранение минимального скрипта для проверки NPU - [ ] Фактическое выполнение скрипта для проверки - [ ] Вывод результатов и описания оптимизации ``` ## Определение ввода Сначала ответьте на три вопроса: 1. Пользователь предоставляет путь к файлу или напрямую вставляет код? 2. Это полный скрипт, фрагмент или отдельный kernel? 3. Это миграция из GPU Triton или семантическая переработка Python/PyTorch? Детали способа ввода, обработка недостающей информации по умолчанию, приоритет конфликта между путем к файлу и вставкой кода, см. в `references/input-modes.md`. ### Сценарий A: GPU Triton -> Triton-Ascend Сначала проверьте: - Существует ли `device='cuda'` - Есть ли логика получения или утверждения устройства, специфичная для GPU - Сохранена ли многомерная свободная grid в стиле GPU - Используется ли `tl.dot` - Существуют ли сложные `shape/stride/block_ptr/order` ### Сценарий B: Python/PyTorch -> Triton-Ascend Сначала извлеките семантику, а затем напишите Triton: - Отношения входных и выходных тензоров - Способы индексации и broadcast - Логика mask / reduce - Требования к dtype и точности - Существует ли в исходном PyTorch естественный непрерывный доступ к памяти Если исходный оператор является только примером реализации, сначала напишите версию Triton-Ascend, семантически эквивалентную исходной, а затем продолжайте оптимизацию. ## Процесс миграции ### 1. Сбор минимально необходимой информации В первую очередь собирайте эту информацию; дополняйте недостающее: - Код ввода или минимальный воспроизводимый пример - Способ ввода: путь к файлу / указанный фрагмент кода / пользователь напрямую вставляет код - shape, dtype, stride - Есть ли mask, broadcast, reduce - Текущие ошибки или проблемы с производительностью - Требуется ли сохранение точности - Среда выполнения: командная строка локальной машины, внутри контейнера, CI или только генерация кода без выполнения Если информация неполная, дополняйте ее в следующем порядке: 1. Сначала попытайтесь вывести из существующего кода 2. Затем дополните недостающее с помощью минимально разумных предположений для скрипта проверки 3. Затем только спрашивайте пользователя о необходимой информации Если текущая информация отсутствует о "месте выполнения", предположите в следующем порядке: 1. Сначала проверьте, предоставил ли пользователь имя контейнера, `docker exec`, путь к контейнеру, информацию об образе 2. Затем проверьте, предоставил ли пользователь путь к файлу на локальной машине, текущий каталог, команду терминала 3. Если все еще невозможно определить, спросите: "Вы хотите, чтобы я написал шаги проверки для командной строки локальной машины или внутри контейнера?" ### 2. Создание минимальной миграции или семантической переработки По умолчанию сначала стремитесь к "семантической эквивалентности и работоспособности": - Для GPU Triton: замените `cuda` на `npu` - Импортируйте `torch_npu` - Удалите логику, специфичную для устройства GPU - Для простых примеров документации и учебных пособий старайтесь сохранять имена `kernel`, имена wrapper, `BLOCK_SIZE`, написание grid и основную структуру кода без изменений - В первой версии не добавляйте активно `contiguous()`, дополнительные утверждения, переименование функций, инженерную упаковку, если только пользователь явно не требует "улучшенной/производственной версии" или эти изменения необходимы для исправления определенных проблем на NPU - Для Python/PyTorch: сначала перепишите в виде самого прямого kernel Triton, соответствующего исходной семантике вычислений Не переписывайте слишком много на первом этапе. Если пользователь явно указал следующие сигнальные слова: - Стиль официальной документации - Строгая минимальная миграция - Минимальный diff - Не делайте инженерных улучшений - Только пример реализации - Только для справки Тогда этот режим "минимальной миграции" должен охватывать требования к более общим инструкциям по оптимизации: - Код должен быть изменен только в необходимых местах - "Описание оптимизации" можно ограничить 1-3 предложениями, четко указывающими, что в данном случае оптимизация не будет расширена - Не пытайтесь добавить `TRITON_ALL_BLOCKS_PARALLEL`, `multibuffer`, `care_padding=False` и другие элементы для полноты, если это не требуется - Не меняйте стиль ответа с "diff документации" на "обзор инженерной оптимизации" - Скрипт проверки также должен оставаться "минимально рабочим" и не должен по умолчанию включать в себя инженерные тестовые фреймворки Минимальная миграция в стиле документации, организация однофайловых примеров, правила именования и сохранения скриптов проверки см. в `references/output-and-validation.md`. ### 3. Переработка модели параллелизма На стороне Ascend в первую очередь следуйте этим правилам: - Сначала используйте 1D grid - Переключитесь с мышления GPU grid на мышление о физических ядрах Ascend - Для операторов, использующих `Vector-only`, сначала думайте о пути Vector Core - Для операторов, содержащих `tl.dot`, сначала думайте о пути AI Core Далее, используйте следующую группу "универсальных правил сходимости" для определения, не сохраняя все ветви реализации GPU: - Если исходная реализация содержит несколько kernel, `autotune`, ветви, зависящие от переменных среды, или автоматическое распределение по различным путям данных, сначала различайте, какие из них являются "необходимыми для семантики", а какие являются "стратегиями производительности, специфичными для GPU" - Для ветвей, которые явно не имеют преимуществ на Ascend, можно преобразовать их в один kernel или меньше; сосредоточьтесь на сохранении семантики, а не всех исторических ветвей - Если оператор по своей сути является `Vector-only`, но исходная реализация использует сложный `block_ptr`, двухмерную/трехмерную grid, дополнительное разбиение или несколько версий kernel, сначала оцените, можно ли преобразовать его в более прямую 1D grid с фиксированной конфигурацией и одним путем - Если оператор содержит `tl.dot`, не просто пытайтесь "сжать многомерную grid в 1D"; сначала определите, какие измерения grid являются просто логическими фрагментами / токенами / плитками, и лучше ли перенести их во внутренний цикл kernel, чтобы уменьшить размерность планирования - Не просто классифицируйте оператор на основе того, что "в исходном коде есть `tl.dot`"; если `tl.dot` используется для реализации таких промежуточных приемов, как prefix-sum, локальное сканирование, треугольное маскирование, его все равно следует классифицировать как `Vector-only` для сокращения/сканирования, или он действительно должен использовать путь AI Core - Не используйте механически все существующие ветви, если оператор по своей природе имеет структуру chunk, tile, window, prefix-sum, локальное сокращение и т. д.; одновременно оцените, будет ли "сначала переупорядочение макета, а затем векторизованные вычисления" более подходящим для Ascend - Если вспомогательный тензор (например, gate, mask, bias, index, state-gate) не является непрерывным в текущем направлении доступа, сначала выполните легкую `transpose/contiguous` или эквивалентную перестановку макета на стороне wrapper, а затем используйте более простой линейный указатель или более правильный `block_ptr` в kernel - Если порядок основного цикла был изменен, например, с "сначала K, затем T" на "сначала T, затем K", одновременно пересмотрите тензоры состояния, тензоры кэша и указатели на предыдущие блоки; не изменяйте только порядок выполнения, а продолжайте использовать старый вид, компенсируя это с помощью `trans` или дополнительных индексов - Если в текущей инженерной системе уже есть общие возможности, такие как `get_vectorcore_num()`, инструменты свойств устройства, распространенные помощники макета, сначала используйте общие помощники, а не создавайте встроенные альтернативные версии - Однако, если текущая выходная цель - "самостоятельный исполняемый скрипт" или "минимальный скрипт проверки", проверьте, зависят ли эти помощники от дополнительной инициализации; если они зависят от шагов инициализации инженерной системы, либо добавьте инициализацию, либо четко укажите предварительные условия в результатах - Когда вы решаете "удалить ветвь / свести реализацию", объясните причину в результатах: связано ли это с тем, что ветвь служит только для GPU autotune, только для выбора общей памяти или не имеет явных преимуществ на Ascend - Если в Triton-Ascend, работающем после миграции, в журналах выполнения появляются предупреждения, такие как `Please DO NOT tune args ['num_warps']` или `['num_stages']`, сначала проверьте, не сохраняете ли вы явно параметры запуска/настройки в стиле GPU; для минимальной рабочей реализации Ascend по умолчанию не сохраняйте эти параметры, если только вы не можете предоставить четкие требования к компиляции или измеренные преимущества - Скрипт проверки не должен использовать только один набор shape; набор тестов должен быть выведен из характеристик оператора, по крайней мере, охватывать один случай, когда размер блока не делится нацело, один случай, который с наибольшей вероятностью вызовет различия в ветвлениях, и один случай, более соответствующий реальному рабочему набору. Если вам предоставлена 2D/3D grid, сначала оцените, можно ли ее свернуть в 1D grid, а затем восстановить индекс в kernel. Подробности о `coreDim`, UB, `shape/stride/block_ptr/order`, `care_padding=False`, `TRITON_ALL_BLOCKS_PARALLEL`, `multibuffer` см. в `references/reference.md`. ## Оптимизация и устранение неполадок ### Правила прямой оптимизации Если выполнено любое из следующих условий, выведите оптимизированную реализацию напрямую: - `coreDim` явно превышает допустимое значение - UB явно слишком велик - Доступ к памяти является дискретным и может быть преобразован в непрерывный доступ - Загрузка/сохранение маски можно выполнить более оптимальным способом - dtype явно приводит к деградации векторных операций Если эти условия не выполняются, особенно для таких примеров, как простое векторное сложение, не выводите улучшенную упакованную версию по умолчанию, чтобы "выглядела более полной". Сначала предоставьте минимальную версию миграции, а затем добавьте улучшенные элементы в раздел "возможные оптимизации". ### Приоритет действий по оптимизации 1. Настройка grid и количества ядер 2. Настройка размера основного блока 3. Введение или переработка подблоков 4. Корректировка `shape/stride/block_ptr/order` 5. Оценка `care_padding=False` 6. Оценка `TRITON_ALL_BLOCKS_PARALLEL` 7. Оценка `multibuffer` и связанных с ним оптимизаций компиляции 8. Корректировка пути dtype без нарушения семантики ### Ключевые моменты, которые необходимо охватить В выводе необходимо охватить следующее: - `cuda -> npu` - `torch_npu` - 1D grid - Привязка к физическим ядрам - Различение `Vector-only` и операторов, содержащих `tl.dot` - `coreDim <= 65535` - Ограничения UB - Непрерывный/выровненный доступ к памяти - Проверка `shape/stride/block_ptr/order` - `TRITON_ALL_BLOCKS_PARALLEL` - `multibuffer` - `care_padding=False` - Деградация dtype