قم ببناء ومشاركة حبات ROCm بسهولة مع عناق الوجه

التعريف

تُعد النوى المخصصة العمود الفقري للتعلم العميق عالي الأداء، مما يتيح عمليات GPU المصممة خصيصًا لتناسب عبء العمل الخاص بك؛ سواء كان ذلك معالجة الصور، أو تحويلات الموتر، أو غيرها من المهام الحاسوبية الثقيلة. لكن تجميع هذه النوى للبنيات الصحيحة، وتوصيل جميع إشارات البناء، ودمجها بشكل نظيف في امتدادات PyTorch يمكن أن يصبح سريعًا فوضى من CMake/Nix، وأخطاء المترجم، ومشكلات ABI، وهو أمر غير ممتع. احتضان الوجه حبات المكتبة تجعل من السهل البناء (باستخدام منشئ النواة) ومشاركة هذه النوى مع نواة المجتمع، مع دعم العديد من واجهات GPU والمسرعات الخلفية، بما في ذلك CUDA وROCm وMetal وXPU. وهذا يضمن أن تكون حباتك سريعة ومحمولة ومتكاملة بسلاسة مع PyTorch.

في هذا الدليل، نركز حصريًا على النوى المتوافقة مع ROCm ونعرض كيفية بنائها واختبارها ومشاركتها باستخدام النوى. ستتعلم كيفية إنشاء نوى تعمل بكفاءة على وحدات معالجة الرسومات AMD، إلى جانب أفضل الممارسات المتعلقة بإعادة الإنتاج والتعبئة والنشر.

تعد هذه الإرشادات التفصيلية الخاصة بـ ROCm نسخة مبسطة من دليل منشئ kernel الأصلي. إذا كنت تبحث عن الإصدار الأوسع الذي يركز على CUDA، فيمكنك العثور عليه هنا: دليل لبناء وتوسيع نطاق نواة CUDA الجاهزة للإنتاج.

خطوات البناء

سوف نستخدم نواة GEMM من RadeonFlow_Kernels كمثال. إذا كنت تريد الذهاب مباشرة إلى الدليل، انقر هنا.

حول النواة

هذا القسم كتبه راديون فلو جيمم مؤلفو النواة لتقديم النواة.
المؤلفون: ColorsWind، وZesen Liu، وAndy

ال راديون فلو جيمم kernel عبارة عن تطبيق مضاعفة مصفوفة عالي الأداء من طراز FP8 مُحسّن لوحدة معالجة الرسومات AMD Instinct MI300X. GEMM (ضرب المصفوفة العامة) هو لبنة البناء الأساسية وراء معظم أحمال عمل التعلم العميق: بالنظر إلى المصفوفتين A وB، يمكنك حساب منتجهما C = A × B. هنا يتم تنفيذه في FP8، وهو تنسيق فاصلة عائمة منخفض الدقة يتاجر ببعض الدقة من أجل إنتاجية أعلى بكثير ونطاق ترددي أقل للذاكرة. تم تطوير هذه النواة لتحدي AMD Developer Challenge 2025، وحصلت على جائزة 🏆 الجائزة الكبرى في يونيو 2025، تقديرًا لتميزها في الأداء والابتكار على أجهزة AMD.

تعمل النواة على المدخلات الكمية باستخدام e4m3fnuz تنسيق الفاصلة العائمة ويطبق مقياسًا لكل كتلة للحفاظ على الدقة أثناء العمليات الحسابية منخفضة الدقة. ال e4m3fnuz التنسيق هو متغير FP8 مع 4 بتات أسية و3 بتات الجزء العشري، وهو مصمم ليكون فعالاً في التعامل مع أحمال عمل الشبكة العصبية. نظرًا لأن FP8 يحتوي على نطاق ديناميكي أصغر بكثير من FP16/FP32، فإننا نطبق عوامل قياس لكل كتلة (a_scale وb_scale) بحيث يتم إعادة قياس كل كتلة من القيم إلى نطاق “مريح” عدديًا قبل الحساب وبعده، مما يساعد في الحفاظ على الدقة على الرغم من الدقة المنخفضة. يستغرق الحجج التالية:

(a, b, a_scale, b_scale, c)

أين a و b هي مصفوفات الإدخال، a_scale و b_scale هي عوامل التحجيم ل a و b على التوالي، و c هي مصفوفة الإخراج:

• a هو K × M في e4m3fnuz
• b هو K × N في e4m3fnuz
• a_scale هو (K // 128) × M في fp32
• b_scale هو (K // 128) × (N // 128) في fp32
• c هو M × N في bf16

تم تجميع النواة مسبقًا لأشكال مصفوفة محددة وتفترض تخطيط ذاكرة منقولًا (كما هو مطلوب في المنافسة). لدعم أشكال إضافية أو تخطيطات ذاكرة بديلة، يجب عليك تعديل مشغل kernel.

والآن بعد أن أصبح لدينا نواة ROCm عالية الأداء، فإن السؤال الطبيعي هو: كيف يمكننا دمجها في سير عمل PyTorch الحقيقي ومشاركتها مع الآخرين؟ هذا هو بالضبط ما سنغطيه بعد ذلك، باستخدام kernels لهيكلة وبناء ونشر نواة ROCm.

هذا دليل تقني إلى حد ما، ولكن لا يزال بإمكانك اتباعه خطوة بخطوة دون فهم كل التفاصيل وكل شيء سوف يعمل بشكل جيد. إذا كنت فضوليًا، يمكنك دائمًا العودة لاحقًا للتعمق أكثر في المفاهيم.

الخطوة 1: هيكل المشروع

يتوقع Hugging Face Kernel Builder أن يتم تنظيم ملفاتك على النحو التالي:

gemm/
├── build.toml
├── gemm
│   └── gemm_kernel.h
├── flake.nix
└── torch-ext
    ├── torch_binding.cpp
    ├── torch_binding.h
    └── gemm
        └── __init__.py

• build.toml: بيان المشروع إنه عقل عملية البناء.
• جيمم/: كود مصدر CUDA الخام الخاص بك حيث يحدث سحر GPU.
• flake.nix: المفتاح لبيئة بناء قابلة للتكرار تمامًا.
• الشعلة تحويلة/gemm/: غلاف Python لمشغلي PyTorch الخام

في بعض الأحيان قد يعتمد مشروعك على ملفات أخرى، مثل الاختبارات أو البرامج النصية المساعدة، ويمكنك إضافتها دون أي مشاكل. في حالتنا، سيتم تنظيم مشروعنا على النحو التالي:

gemm/
├── build.toml
├── gemm
│   ├── gemm_kernel.h
│   ├── gemm_kernel_legacy.h
│   ├── transpose_kernel.h
│   └── gemm_launcher.hip
├── include
│   ├── clangd_workaround.h
│   ├── gpu_libs.h
│   ├── gpu_types.h
│   └── timer.h
├── src/utils
│   ├── arithmetic.h
│   └── timer.hip
├── tests/checker
│   ├── checker.cpp
│   ├── metrics.h
│   └── checker.h
├── flake.nix
└── torch-ext
    ├── torch_binding.cpp
    ├── torch_binding.h
    └── gemm
        └── __init__.py

إذا نظرت إلى الملفات الأصلية لنواة Gemm في RadeonFlow Kernels، فهي ملفات مصدر HIP ذات .cpp ملحقات. كخطوة أولى، تحتاج إلى تغيير هذه الامتدادات إلى .h أو .hip حسب محتواها واستخدامها:

• يستخدم .h لملفات الرأس التي تحتوي على إعلانات kernel أو الوظائف المضمنة أو رمز القالب الذي سيتم تضمينه في ملفات أخرى
• يستخدم .hip لملفات التنفيذ التي تحتوي على رمز HIP/GPU الذي يجب تجميعه بشكل منفصل (على سبيل المثال، مشغلات kernel، ووظائف الجهاز مع عمليات التنفيذ المعقدة)

في مثالنا، gemm_kernel.h, gemm_kernel_legacy.h، و transpose_kernel.h هي ملفات رأس، في حين gemm_launcher.hip هو ملف تنفيذ HIP. يساعد اصطلاح التسمية هذا منشئ النواة (kernels/builder) تحديد وتجميع كل نوع ملف بشكل صحيح.

الخطوة 2: إعداد ملفات التكوين

ال build.toml يظهر

ينسق هذا الملف البناء بأكمله. إنه يخبر منشئ النواة بما يجب تجميعه وكيف يرتبط كل شيء.

[general]
name = "gemm"
universal = false

[torch]
src = [
  "torch-ext/torch_binding.cpp",
  "torch-ext/torch_binding.h",
]

[kernel.gemm]
backend = "rocm"
rocm-archs = [
    "gfx942",
]

depends = ["torch"]

src = [
  "include/clangd_workaround.h",
  "include/gpu_libs.h",
  "include/gpu_types.h",
  "include/timer.h",
  "gemm/gemm_kernel.h",
  "gemm/gemm_kernel_legacy.h",
  "gemm/gemm_launcher.hip",
  "gemm/transpose_kernel.h",
  "src/utils/arithmetic.h",
  "src/utils/timer.hip",
  "tests/checker/metrics.h",
]

include = ["include"]

عام

يحتوي هذا القسم على إعدادات تكوين المشروع العامة.

• اسم (مطلوب): اسم مشروعك. يجب أن يتطابق هذا مع اسم النواة الخاص بك وسيتم استخدامه لحزمة Python.
• عالمي (اختياري): النواة هي نواة عالمية عند ضبطها على true. النواة العالمية هي حزمة بايثون خالصة (لا توجد ملفات مجمعة). لا تستخدم النواة العالمية الأقسام الأخرى الموضحة أدناه. من الأمثلة الجيدة على النواة الشاملة نواة تريتون. تقصير: false

الشعلة

يصف هذا القسم تكوين ملحق Torch. فهو يحدد روابط Python التي ستعرض النواة الخاصة بك لـ PyTorch.

• src (مطلوب): قائمة بالملفات المصدرية والرؤوس الخاصة بامتداد PyTorch. في حالتنا، يتضمن ذلك ملفات ربط C++ التي تنشئ واجهة Python.

kernel.gemm

مواصفات النواة المسماة “gemm”. يمكنك تحديد أقسام نواة متعددة في نفس ملف build.toml إذا كان لديك نواة متعددة.

• الخلفية (مطلوب): الواجهة الخلفية للحوسبة للنواة. نحن نستخدم “rocm” لدعم AMD GPU.
• أقواس rocm (مطلوب لـ ROCm): قائمة بنيات ROCm التي يجب تجميع النواة لها. يستهدف “gfx942” وحدات معالجة الرسومات من سلسلة MI300.
• يعتمد على (مطلوب): قائمة التبعيات. نحن نعتمد على “الشعلة” لاستخدام عمليات موتر PyTorch.
• يشمل (اختياري): قم بتضمين الدلائل المتعلقة بجذر المشروع. وهذا يساعد المترجم في العثور على ملفات الرأس.

ال flake.nix ملف الاستنساخ

للتأكد من أن أي شخص يمكنه إنشاء النواة الخاصة بك على أي جهاز، فإننا نستخدم ملف flake.nix. يقوم بتأمين الإصدار الدقيق لمنشئ النواة وتبعياته. (يمكنك فقط نسخ هذا المثال ولصقه وتغيير الوصف)

{
  description = "Flake for GEMM kernel";

  inputs = {
    kernel-builder.url = "github:huggingface/kernels";
  };

  outputs =
    {
      self,
      kernel-builder,
    }:

    kernel-builder.lib.genFlakeOutputs {
      inherit self;
      path = ./.;
    };
}

كتابة النواة

الآن للحصول على رمز GPU. داخل gemm/gemm_launcher.hipنحدد كيفية إطلاق نواة GEMM. اعتمادا على التكوين، فإننا نسمي إما الجديد الأمثل gemm/gemm_kernel أو الرجوع إلى التنفيذ القديم (gemm/gemm_kernel_legacy).

extern "C" void run(
    void *a, void *b, void *as, void *bs, void *c,
    int m, int n, int k,
    PerfMetrics *metrics, hipStream_t job_stream0
) {
    const __FP8_TYPE *a_ptr = static_cast<const __FP8_TYPE *>(a);
    const __FP8_TYPE *b_ptr = static_cast<const __FP8_TYPE *>(b);
    __BF16_TYPE *c_ptr = static_cast<__BF16_TYPE *>(c);
    const float *as_ptr = static_cast<const float *>(as);
    const float *bs_ptr = static_cast<const float *>(bs);

    KernelTimerScoped timer(timers, 2LL * m * n * k,
        metrics ? &metrics->entries[0].time : nullptr,
        metrics ? &metrics->entries[0].gflops : nullptr, job_stream0);

    
    switch (pack_shape(m, n, k)) {
        DISPATCH_GEMM(1024, 1536, 7168, 256, 128, 128, 4, 2, 512, 4, 16);
        DISPATCH_GEMM(6144, 7168, 2304, 256, 128, 128, 4, 2, 512, 1, 16);
        default: {
            printf("Error: Unsupported shape M=%d, K=%d, N=%d\n", m, k, n);
            abort();
        }
    }
}

تسجيل مشغل PyTorch الأصلي

هذه الخطوة هي المفتاح. نحن لا نجعل الوظيفة متاحة في بايثون فحسب؛ نحن نحوله إلى مشغل PyTorch الأصلي. وهذا يعني أنه يصبح جزءًا من الدرجة الأولى من PyTorch نفسها، ويمكن الوصول إليه من خلاله torch.ops.

الملف torch-ext/torch_binding.cpp يعالج هذا التسجيل.

#include <torch/all.h>
#include <torch/library.h>
#include <hip/hip_runtime.h>

#include "registration.h"
#include "torch_binding.h"


extern "C" {
    struct PerfMetrics;
    void run(void *a, void *b, void *as, void *bs, void *c, int m, int n, int k, PerfMetrics *metrics, hipStream_t job_stream0);
}

void gemm(torch::Tensor &out, torch::Tensor const &a, torch::Tensor const &b, 
          torch::Tensor const &as, torch::Tensor const &bs) {
    
    
    TORCH_CHECK(a.device().is_cuda(), "Input tensor a must be on GPU device");
    TORCH_CHECK(b.device().is_cuda(), "Input tensor b must be on GPU device");
    TORCH_CHECK(as.device().is_cuda(), "Scale tensor as must be on GPU device");
    TORCH_CHECK(bs.device().is_cuda(), "Scale tensor bs must be on GPU device");
    TORCH_CHECK(out.device().is_cuda(), "Output tensor out must be on GPU device");
    
    TORCH_CHECK(a.is_contiguous(), "Input tensor a must be contiguous");
    TORCH_CHECK(b.is_contiguous(), "Input tensor b must be contiguous");
    TORCH_CHECK(as.is_contiguous(), "Scale tensor as must be contiguous");
    TORCH_CHECK(bs.is_contiguous(), "Scale tensor bs must be contiguous");
    TORCH_CHECK(out.is_contiguous(), "Output tensor out must be contiguous");
    
    
    
    int M = a.size(0);
    int K = a.size(1);
    int N = b.size(1);
    
    TORCH_CHECK(b.size(0) == K, "Matrix dimensions mismatch: a.size(1) != b.size(0)");
    TORCH_CHECK(out.size(0) == M, "Output tensor dimension mismatch: out.size(0) != M");
    TORCH_CHECK(out.size(1) == N, "Output tensor dimension mismatch: out.size(1) != N");
    
    
    const hipStream_t stream = 0;
    
    
    run(a.data_ptr(), b.data_ptr(), as.data_ptr(), bs.data_ptr(), out.data_ptr(),
        M, N, K, nullptr, stream);
}

TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def("gemm(Tensor! out, Tensor a, Tensor b, Tensor a_scale, Tensor b_scale) -> ()");
  ops.impl("gemm", torch::kCUDA, &gemm);
}

REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

ال torch_binding.h يحتوي الملف على إعلانات الوظائف. على سبيل المثال، gemm يحتوي kernel على الإعلان التالي في torch_binding.h:

#pragma once

#include <torch/torch.h>

void gemm(torch::Tensor &out, torch::Tensor const &a, torch::Tensor const &b, 
          torch::Tensor const &as, torch::Tensor const &bs);

إعداد __init__.py إزار

في torch-ext/gemm/ نحن بحاجة إلى __init__.py ملف لجعل هذا الدليل حزمة Python ولكشف عامل التشغيل المخصص لدينا بطريقة سهلة الاستخدام.

from typing import Optional
import torch
from ._ops import ops

def gemm(a: torch.Tensor, b: torch.Tensor, as_: torch.Tensor, bs: torch.Tensor, 
         out: Optional[torch.Tensor] = None) -> torch.Tensor:
         
    if out is None:
        
        M, K = a.shape
        K_b, N = b.shape
        assert K == K_b, f"Matrix dimension mismatch: A has {K} cols, B has {K_b} rows"
        
        
        out = torch.empty((M, N), dtype=torch.bfloat16, device=a.device)
    
    ops.gemm(out, a, b, as_, bs)
    return out

الخطوة 3: بناء النواة

يستخدم منشئ النواة Nix لبناء النواة. يمكنك إنشاء النوى أو تشغيلها مباشرةً إذا كان Nix مثبتًا على نظامك. نوصي بتثبيت Nix بالطريقة التالية:

الشروع في العمل مع نيكس

أولاً، قم بتشغيل هذا:

nix flake update

هذا يولد أ flake.lock الملف الذي يثبت منشئ النواة وجميع تبعياته المتعدية. ارتكاب كليهما flake.nix و flake.lock إلى مستودعك للتأكد من أن إصدارات kernel قابلة للتكرار.

نظرًا لأن منشئ kernel يعتمد على العديد من الحزم (على سبيل المثال، كل إصدار مدعوم من PyTorch)، فمن المستحسن تمكين ذاكرة التخزين المؤقت Hugging Face لتجنب عمليات إعادة البناء باهظة الثمن:

cachix use huggingface

أو قم بتشغيله مرة واحدة دون تثبيت cachix بشكل دائم:

nix run nixpkgs

بناء النواة مع Nix

النواة التي تحتوي على flake.nix يمكن إنشاء الملف باستخدام أمر البناء والنسخ:

cd Build_RadeonFlow_Kernels/gemm
nix build . -L

ستكون النواة المترجمة بعد ذلك باللغة المحلية build/ دليل.

شركة شل للتنمية للتنمية المحلية

يوفر منشئ النواة الأصداف اللازمة لتطوير النواة. في مثل هذه الصدفة، تتوفر جميع التبعيات المطلوبة أيضًا build2cmake لإنشاء ملفات المشروع:

$ nix develop
$ build2cmake generate-torch build.toml
$ cmake -B build-ext
$ cmake --build build-ext

إذا كنت تريد اختبار النواة كحزمة بايثون، فيمكنك القيام بذلك. nix develop سيتم تلقائيًا إنشاء بيئة افتراضية في .venv وتفعيله:

$ nix develop
$ build2cmake generate-torch build.toml
$ pip install --no-build-isolation -e .

تتوفر قذائف التطوير لكل تكوين بناء. على سبيل المثال، يمكنك الحصول على غلاف تطوير Torch 2.7 مع ROCm 6.3 باستخدام:

$ rm -rf .venv  
$ nix develop .

الخطوة 4: تحميل النواة إلى المركز

الآن بعد أن قمنا ببناء النواة الخاصة بنا، يمكننا اختبارها وتحميلها إلى المركز.

بناء النواة لجميع إصدارات PyTorch وROCm

أحد الأشياء الصغيرة التي نرغب في القيام بها قبل أن نشارك هو تنظيف جميع عناصر التطوير التي تم إنشاؤها أثناء عملية الإنشاء لتجنب تحميل الملفات غير الضرورية.

build2cmake clean build.toml 

لبناء النواة لجميع الإصدارات المدعومة من PyTorch وROCm، تقوم أداة منشئ النواة بأتمتة العملية:

nix build . -L

ملحوظة:
قد تستغرق هذه العملية بعض الوقت، لأنها ستبني النواة لجميع الإصدارات المدعومة من PyTorch وROCm.
سيكون الإخراج في result دليل.

الخطوة الأخيرة هي نقل النتائج إلى دليل البناء المتوقع (هذا هو المكان الذي ستبحث عنه مكتبة kernels).

mkdir -p build
rsync -av --delete --chmod=Du+w,Fu+w result/ build/

الدفع إلى محور الوجه المعانق

إن دفع عناصر البناء إلى Hub سيجعل من السهل على المطورين الآخرين استخدام النواة الخاصة بك. يمكننا استخدام kernels upload أمر لهذا:

kernels upload <path_to_kernel> --repo_id hub-username/img2gray

hf repo create gemm

git init
git remote add origin https://huggingface.co/<your-username>/gemm


git pull origin main
git xet install
git checkout -b main


git xet track "*.so"



git add \
  build/ gemm/ include/ src/utils tests/checker \
  torch-ext/torch_binding.cpp torch-ext/torch_binding.h torch-ext/gemm \
  flake.nix flake.lock build.toml

git commit -m "feat: Created a compliant gemm kernel"
git push -u origin main

رائع! النواة الخاصة بك موجودة الآن على Hugging Face Hub، وهي جاهزة للاستخدام من قبل الآخرين ومتوافقة تمامًا مع مكتبة kernels.

الخطوة 5: دعونا نستخدمها 🙂

مع حبات المكتبة، فأنت لا “تثبت” النواة بالمعنى التقليدي. يمكنك تحميله مباشرة من مستودع Hub الخاص به، والذي يقوم تلقائيًا بتسجيل المشغل الجديد.

import torch
from kernels import get_kernel


gemm = get_kernel("kernels-community/gemm")


M, N, K = 1024, 1536, 7168
QUANT_SIZE = 128


device = torch.device("cuda")


A_fp32 = torch.randn(M, K, device=device)
B_fp32 = torch.randn(K, N, device=device)


A_fp8 = A_fp32.to(torch.float8_e4m3fnuz)
B_fp8 = B_fp32.to(torch.float8_e4m3fnuz)


A_scale = torch.ones(K // QUANT_SIZE, M, device=device, dtype=torch.float32)
B_scale = torch.ones(K // QUANT_SIZE, N // QUANT_SIZE, device=device, dtype=torch.float32)

C = torch.zeros(M, N, device=device, dtype=torch.bfloat16)


result = gemm.gemm(A_fp8, B_fp8, A_scale, B_scale, C)

هذا كل شيء! نواة ROCm الخاصة بك جاهزة الآن للاستخدام من Hugging Face Hub.

خاتمة

أصبح الآن إنشاء حبات ROCm ومشاركتها مع Hugging Face أسهل من أي وقت مضى. من خلال سير عمل نظيف وقابل للتكرار مدعوم من Nix والتكامل السلس في PyTorch، يمكن للمطورين التركيز على تحسين الأداء بدلاً من الإعداد. بمجرد إنشائها، يمكن مشاركة النواة المخصصة الخاصة بك على Hugging Face Hub؛ مما يجعلها في متناول المجتمع على الفور ويمكن استخدامها عبر المشاريع باستخدام بضعة أسطر فقط من التعليمات البرمجية. 🚀

المكتبات والمركز ذو الصلة