Tạo các kernel tùy chỉnh cho AMD MI300

Kernel AMD

Hợp tác giữa Hugging Face và AMD để cung cấp hiệu suất cao nhất trên các nền tảng AMD, mang lại lợi ích cho cộng đồng mã nguồn mở. Bài viết này trình bày cách tối ưu hóa hiệu suất cho MI300X và cách mỗi kernel được tinh chỉnh riêng lẻ.

Giới thiệu

ChatGPT xử lý hơn một tỷ yêu cầu mỗi ngày, đòi hỏi phải tối ưu hóa mô hình ở mọi cấp độ. Các kernel, thuật toán được thực thi khi thực hiện bất kỳ thao tác nào trong mạng, đóng vai trò quan trọng trong việc suy luận mô hình. Tuy nhiên, các kernel mới thường được tối ưu hóa cho phần cứng Nvidia mới nhất, bỏ qua các thiết bị khác, đặc biệt là GPU AMD. Hugging Face hợp tác với AMD để cung cấp hiệu suất cao nhất trên các nền tảng AMD và mang lại lợi ích cho cộng đồng mã nguồn mở. Sự hợp tác này tập trung vào việc cung cấp các kernel mã nguồn mở được tối ưu hóa để cải thiện hiệu suất của việc phân phối Llama 3.1 405B trong FP8 trên một nút gồm 8 MI300X bằng VLLM.

Bằng cách kết hợp ba kernel được tối ưu hóa:

• Kernel kết nối lại dư, chuẩn hóa RMS và chuyển đổi FP8 hợp nhất
• Kernel kích hoạt SwiGLU hợp nhất và chuyển đổi FP8
• Kernel GEMM Skinny

Đã đạt được tốc độ đáng kể khi chạy VLLM trên một nút được cung cấp bởi GPU MI300X.

Các phép đo được thực hiện với kích thước đầu vào là 1 và kích thước đầu ra là 128 để bắt chước chế độ giải mã. Đo độ trễ giải mã bằng cách sử dụng giá trị trung bình trên 30 lần lặp.

Các cải thiện về hiệu suất này được đo lường trong VLLM, nhưng bạn cũng có thể sử dụng các kernel riêng biệt.

Cách sử dụng các kernel này

Các kernel này có sẵn trên kho lưu trữ hf-rocm-kernels tại đây. Bạn sẽ tìm thấy hướng dẫn về cách cài đặt gói, mã nguồn cho từng kernel, các liên kết Python tương ứng của chúng, các tập lệnh đo điểm chuẩn khác nhau và bộ thử nghiệm.

Các kernel sẽ sớm được tích hợp vào nhánh AMD của dự án VLLM.

Quy trình tối ưu hóa

1. Tìm hiểu về MI300X: Xem lại kiến trúc của thiết bị đang làm việc.
2. Phân tích hiệu suất ngày 0: Đánh giá hiệu suất suy luận mô hình ban đầu để xác định các nút thắt.
3. Kernel tùy chỉnh: Tạo các kernel tùy chỉnh để giải quyết các nút thắt và cải thiện hiệu suất tổng thể.

Giới thiệu nhanh về MI300X

Để tối ưu hóa mã GPU, cần phải hiểu cách hoạt động của GPU.

• Luồng (Thread): Đơn vị công việc nhỏ nhất. Mỗi luồng có bộ nhớ riêng, được gọi là thanh ghi (VGPR), chỉ nó mới có thể truy cập.
• Warp: Bao gồm 64 luồng. Tất cả các luồng trong một warp phải thực thi cùng một lệnh cùng một lúc.
• Đơn vị tính toán (CU): Các warp được gói lại với nhau thành các khối luồng (thread block), là các trừu tượng phần mềm, nhưng chạy trên một thành phần phần cứng được gọi là đơn vị tính toán (CU). Mỗi đơn vị tính toán có bộ nhớ cache L1 chuyên dụng và bộ nhớ chia sẻ.
• XCD: Các đơn vị tính toán sau đó được nhóm thành khuôn mạch phức hợp tăng tốc (XCD), mỗi XCD chứa 38 đơn vị tính toán.
• MI300X: Bằng cách lắp ráp 8 XCD (tổng cộng 304 CU), thêm bộ nhớ cache cấp cuối (bộ nhớ cache vô cực, với 256MB) và một lượng lớn ram video (192GB), chúng ta có được MI300X.

Khi tối ưu hóa một kernel, luôn có một sự cân bằng giữa thực hiện nhiều thao tác hoặc tải nhiều dữ liệu, nhưng nói chung, bạn muốn truy cập VRAM (thường được gọi là bộ nhớ toàn cục) càng ít càng tốt.

Phân tích hiệu suất ngày 0

Việc lập hồ sơ công việc là rất quan trọng trước khi viết bất kỳ mã nào. Trong trường hợp này, hãy lập hồ sơ suy luận mô hình trong VLLM để hiểu thời gian mà mỗi thao tác sử dụng.

Hầu hết độ trễ đến từ GEMM và giao tiếp. Các lệnh gọi đến kernel “chuẩn hóa RMS” khá tốn kém. Cùng với kernel SwiGLU, chúng chiếm 15% tổng độ trễ.

Kernel chuẩn hóa RMS

Mỗi khối giải mã có một khối chú ý và một khối MLP. Cả hai đều bắt đầu bằng một kết nối dư giữa trạng thái ẩn hiện tại (x) và phần dư (r). Sau khi chúng được cộng lại với nhau, áp dụng chuẩn hóa Root Mean Square (RMS) theo từng hàng cho x và lượng tử hóa x thành FP8 bằng cách sử dụng tỷ lệ s. Các hoạt động cần thực hiện như sau:

x ← x + r
r ← x
V = Σ(i=1 đến d) xi^2
x ← x / √(V + ϵ)
xq = Q_fp8(s * x * w)

Việc hợp nhất các hoạt động này thành một kernel duy nhất có thể mang lại hiệu suất tốt.

Khi so sánh với Pytorch, phiên bản barebones của kernel này mang lại tốc độ nhanh hơn 10 lần.

Tối ưu hóa: liên quan đến bộ nhớ

Các nguyên tắc để giảm chi phí tải dữ liệu:

• Số lượng: Tận dụng tối đa lượng dữ liệu mà một luồng có thể tải trong một lệnh.
• Kết hợp: Đảm bảo truy cập bộ nhớ được kết hợp. Truy xuất dữ liệu liền kề sẽ tối ưu hóa việc tải dữ liệu.
• Giảm số lượng cửa hàng: Giảm số lượng thao tác lưu trữ bằng cách lưu trữ các phiên bản đã sửa đổi trong bộ nhớ chia sẻ.

Kết quả

Việc áp dụng các tối ưu hóa bộ nhớ sẽ giúp đạt được những kết quả sau:

Số lượng hàng	Torch (μs)	VLLM (μs)	Của chúng tôi (μs)
1	38.8998	5.5145	4.18138
2	43.2469	5.65645	4.36976
4	41.1304	5.6893	4.37628
8	43.8883	5.72275	4.39081
16	46.8876	5.85667	4.48165
32	55.2276	6.08502	4.72017
64	75.6086	6.4629	5.54214
128	98.1122	7.49166	6.27341
256	119.727	11.8812	10.739
512	195.782	23.1595	18.5549
1024	355.42	44.8143	34.7204
2048	671.513	81.2089	73.35

với một tensor đầu vào FP16 có hình dạng [X, 16384].

Kernel SwiGLU

Sau kernel chuẩn hóa RMS, một phép chiếu gọi là phép chiếu “Gate / Up”. Các phép chiếu “Gate” và “Up” được kết hợp và sau đó là hàm kích hoạt SwiGLU, được theo sau bởi phép chiếu “Down” ở FP8.

y = σ(xG) ⋅ xU
yQ = Q_FP8(s * y)

Trong đó σ là hàm sigmoid.

Tối ưu hóa: liên quan đến tính toán

Hai phương pháp chính để tăng tốc các kernel:

• Tăng khối lượng công việc trên mỗi hướng dẫn: Sử dụng các lệnh được đóng gói để áp dụng cùng một toán tử cho nhiều phần tử.
• Sử dụng hướng dẫn nhanh hơn: Thay thế các hướng dẫn tốn kém bằng các hướng dẫn hiệu quả hơn. Ví dụ: thay vì sử dụng hướng dẫn exp, hãy sử dụng hướng dẫn exp2 nhanh hơn.

Kết quả

Số lượng hàng	1	2	4	8	16	32	64	128	256	512	1024	2048
Torch (μs)	40.2731	29.923	35.305	23.5763	22.4738	25.3445	31.5829	40.3194	53.5369	79.8037	124.873	243.202
VLLM (μs)	3.84116	3.86192	3.92937	3.94151	4.01047	4.02421	4.08943	4.20317	4.48755	7.48465	13.7389	25.4306
Của chúng tôi (μs)	1.92981	1.93904	1.93524	1.99316	2.00415	1.91563	2.04498	2.61763	3.57726	5.47608	10.0482	19.8957
Tăng tốc (VLLM / Của chúng tôi)	1.990434291	1.991665979	2.030430334	1.977518112	2.001082753	2.100724044	1.999740829	1.605715857	1.254465708	1.366789747	1.367299616	1.278195791

Với các tối ưu hóa bộ nhớ và tính toán được điều chỉnh cho MI300X, thu được một kernel nhanh hơn Torch trung bình hơn 14 lần và nhanh hơn từ 27% đến 100% so với kernel VLLM.

Kernel GEMM Skinny

Khoảng 60% độ trễ suy luận của mô hình đến từ các phép chiếu, dựa trên kernel GEMM. Viết một kernel tùy chỉnh hoạt động tốt hơn trong mọi trường hợp là khá khó. Nhưng nếu chúng ta tập trung vào một số trường hợp biên mà có liên quan đến chúng ta và viết một kernel GEMM cho những trường hợp cụ thể đó, thì có một cơ hội là kernel tùy chỉnh của chúng ta có thể nhanh hơn các kernel trong các thư viện chuyên dụng.

Tối ưu hóa: chia tách K

Vì vấn đề chính của GEMM skinny là chúng ta sử dụng quá ít đơn vị tính toán, điều đầu tiên chúng ta có thể làm là tìm ra một cách để sử dụng nhiều hơn. Để làm điều này, chúng ta có thể khai thác công thức sau:

cij = Σ(k=1 đến K) aik bkj = (Σ(k=1 đến K/2) aik bkj) + (Σ(k=1+K/2 đến K) aik bkj)

Nhờ tính kết hợp của tổng, chúng ta có thể chia tách GEMM chính dọc theo trục chia sẻ (thường được gọi là trục K) và thay thế một GEMM bằng một số GEMM con được thực hiện đồng thời.

Tối ưu hóa: loại bỏ phần đệm

Để tận dụng tốt nhất tensor core trên MI300X, chúng ta cần điều chỉnh kích thước ma trận đầu vào để phù hợp với yêu cầu của tensor core. Các định dạng tensor core là mfma_MxNxK... trong đó mfma là viết tắt của matrix fused multiply-add, M là số hàng của ma trận bên trái, N là số cột của ma trận bên phải và K là chiều chia sẻ của cả hai.

Để tận dụng hướng dẫn thưa, chúng ta có thể sử dụng kỹ thuật sau:

1. Thêm đệm: Đảm bảo rằng ma trận đầu vào phù hợp với các yêu cầu về kích thước của hướng dẫn tensor core.
2. Áp dụng thủ thuật độ thưa: Tận dụng lợi thế về cấu trúc độ thưa 4:2 trong các hướng dẫn tensor core thưa.

Sử dụng thủ thuật này, chúng ta có thể tăng tốc đáng kể GEMM của mình cho bất kỳ đầu vào nào có 8 hàng trở xuống, dẫn đến việc giảm độ trễ trên mỗi mã thông báo cho bất kỳ lô giải mã nào có ít hơn 8 yêu cầu.

Tối ưu hóa: chuyên môn hóa warp và thực thi không đồng bộ

Chúng ta có thể sử dụng một kỹ thuật gọi là chuyên môn hóa warp. Thay vì tất cả các warp trong khối luồng thực hiện cùng một hướng dẫn, chúng ta sẽ dành riêng một số warp cho chỉ tải dữ liệu và một số cho chỉ tính toán kết quả. Các warp chịu trách nhiệm tải dữ liệu được gọi là người sản xuất (producer) và những warp tính toán kết quả được gọi là người tiêu dùng (consumer). Để làm cho toàn bộ quá trình hoạt động, một hàng đợi sẽ được lưu trữ trong bộ nhớ dùng chung. Khi một người sản xuất hoàn tất việc lưu trữ dữ liệu trong bộ đệm bộ nhớ dùng chung, nó sẽ thay đổi trạng thái của biến hàng đợi thứ i để báo hiệu dữ liệu có sẵn ở đó. Người tiêu dùng đang xem xét điều này và bắt đầu tải dữ liệu sau đó.

Thông qua chuyên môn hóa warp và công việc không đồng bộ, chúng ta có thể điều chỉnh kernel của mình cho khối lượng công việc cường độ số học thấp, nhưng điều đó có đủ để vượt qua các thư viện như hipBLASLT không?

Kết quả

M (hàng)	N (cột)	K (chiều sâu)	Thời gian Torch (μs)	Thời gian SkG (μs)	Tăng tốc
1	2304	16384	14.938 ± 0.292	11.685 ± 0.299	127.84 %
8	2304	16384	16.300 ± 0.282	12.342 ± 0.375	132.07 %
16	2304	16384	16.693 ± 0.233	13.909 ± 0.295	120.02 %
32	2304	16384	16.817 ± 0.124	17.021 ± 0.133	98.80 %
1	13312	16384	77.636 ± 0.364	54.717 ± 0.628	141.88 %
8	13312	16384	80.031 ± 0.449	58.355 ± 0.612	137.15 %
16	13312	16384	75.236 ± 0.378	59.973 ± 1.922	125.45 %
32	13312	16384	82.198 ± 0.590	69.483 ± 1.672	118.30 %
1	16384	6656	31.066 ± 0.193	27.613 ± 0.218	112.51 %
8	16384	6656	31.559 ± 0.200	28.134 ± 0.209	112.17 %
16	16384	6656	31.671 ± 0.250	30.233 ± 0.267	104.76 %
32	16384	6656	35.561 ± 0.335	35.052 ± 1.365	101.45 %

Đối với những chiều mà chúng ta có thể sử dụng thủ thuật độ thưa (M = 1, 8), chúng ta thấy tốc độ tăng tốc đáng kể so với Torch.

Kết luận

Bài viết này chỉ khám phá một số kỹ thuật tối ưu hóa kernel. Repository hf-rocm-kernels có thể được sử dụng để thử nghiệm các kỹ thuật.

Link bài viết gốc

• Tags:
• Ai
• July 9, 2025
• Huggingface.co