Giải mã ngược Neural Engine của Apple: Từ kỷ nguyên CNN đến kỷ nguyên Transformer

Giải mã ngược Neural Engine của Apple: Kiến trúc bên trong và lý do nó đang dần lỗi thời

Trong hơn ba năm, một kỹ sư đã âm thầm giải mã ngược driver của Apple Neural Engine (ANE) trên chip M1. Kết quả là một cái nhìn chi tiết chưa từng có về kiến trúc bên trong khối xử lý AI này — từ datapath tính toán, cơ chế lập lịch tác vụ, đến hệ thống bộ nhớ phân cấp và những hạn chế cố hữu khiến ANE không còn phù hợp với kỷ nguyên mô hình transformer.

Die shot của chip M1 với khối ANE được khoanh vùng

Từ CNN đến Transformer: Một sự chuyển dịch kiến trúc

Khi Apple ra mắt ANE lần đầu trên chip A11 Bionic (2017), mục tiêu rất rõ ràng: tăng tốc các mô hình CNN xử lý hình ảnh — vốn là những phép tính tensor dày đặc với mô hình tái sử dụng dữ liệu có thể dự đoán trước. Đến năm 2025, Apple công bố chip M5 với tính năng chủ đạo là "hiệu năng LLM", đồng thời tích hợp các lõi ANE trực tiếp vào bên trong lõi GPU.

Đây là dấu hiệu rõ ràng cho sự kết thúc của kỷ nguyên NPU độc lập. Việc Apple quyết định gộp ANE vào GPU cho thấy lõi tính toán của ANE vẫn hữu ích cho transformer, nhưng cần nằm trong một luồng dữ liệu (dataflow) khác.

Kiến trúc tính toán: 2.048 làn MAC song song

Trái tim của ANE M1 là 16 lõi tính toán song song, mỗi lõi chứa 128 làn nhân-tích lũy (MAC) hoạt động song song ở độ chính xác FP16, hoặc 256 làn ở INT8. Tổng cộng ANE M1 có 2.048 làn MAC song song.

Mỗi làn MAC thực hiện phép toán cơ bản:

s ← s + (a × b)

Điểm quan trọng là mỗi làn MAC không hề biết nó đang giảm chiều trên chiều nào của tensor. Phép tích vô hướng, nhân ma trận hay tích chập đều được tạo ra từ cách ánh xạ và lập lịch toán hạng lên lõi — chứ không phải từ phần cứng. Nói cách khác, ANE không mã hóa cứng một lớp CNN 4 kênh vào silicon.

Bộ tích lũy bên trong hoạt động ở dạng số nguyên có dấu 32-bit với 16 bit phần phân số (Q16.16). Khi thực nghiệm với vector toàn số 1, tác giả phát hiện bộ tích lũy bão hòa ở giá trị 2¹⁵ — tức 32.768. Đây chính xác là giới hạn của một giá trị fixed-point 32-bit có dấu với 16 bit phần phân số.

Hàm kích hoạt phi tuyến: Bảng tra 33 điểm

ANE triển khai hàm tanh() bằng một bảng tra (lookup table) gồm 33 phần tử FP16 liên tiếp, tương ứng với các mẫu lượng tử hóa của tanh(i/8) với i chạy từ 0 đến 32.

Đồ thị so sánh hàm tanh thực tế và bảng tra 33 điểm

Điều thú vị là ANE không dùng nội suy bậc thang đơn thuần. Qua thử nghiệm với xung impulse, tác giả xác nhận mode 2 triển khai bảng tra nội suy tuyến tính từng đoạn (piecewise-linear LUT). Hệ số R xác định khoảng cách giữa các nút, với u = 2^R × |x| ánh xạ đầu vào vào tọa độ bảng tra.

Minh họa cơ chế nội suy tuyến tính trong bảng tra tanh

Đáng chú ý hơn, phép biến đổi tuyến tính ax + b trong CoreML có thể dùng chung phần cứng nội suy này. Trình biên dịch sẽ gộp hệ số scale và offset hằng số vào trọng số tích chập — một kỹ thuật gọi là gộp trọng số (weight folding) — giúp loại bỏ các phép tính dư thừa trước khi chạy trên phần cứng.

Sơ đồ gộp trọng số và bias vào đường datapath sau MAC

Bộ lập lịch: Khi driver không hề biết gì về AI

Mã nguồn driver của ANE gây thất vọng vì quá đơn giản. Driver không bao giờ gửi cho ANE một opcode CONV, MATMUL hay RELU nào cả. Tất cả phép toán neural đã được biên dịch sẵn thành một luồng lệnh gồm các bộ mô tả tác vụ (Task Descriptor — TD).

Driver chỉ làm ba việc:

• Nạp tác vụ vào bộ nhớ
• Đặt con trỏ tới khối blob qua cặp thanh ghi (TM_ADDR, TM_SIZE)
• Gửi tác vụ bằng cách "rung chuông" qua thanh ghi TM_PUSH

Hệ thống có 8 hàng đợi tác vụ (Task Queue) giống hệt nhau, mỗi hàng đợi chứa hai bộ mô tả luồng lệnh — đây là cơ chế ping-pong cho phép một khe thực thi trong khi phần mềm sửa khe còn lại.

Điểm mấu chốt: ANE không có tập lệnh (ISA). Một TD thực chất là một bản dump thanh ghi tuần tự của datapath ANE. Mỗi "chương trình ANE" chỉ đơn thuần là cấu hình cho một lượt chạy qua datapath cố định. Ta có thể cấu hình cách datapath vận hành, nhưng không thể thay đổi nó có khả năng làm gì.

Quy trình thực thi một TD diễn ra như sau:

• ControlDMA sao chép TD vào các thanh ghi cấu hình
• KernelDMA sao chép trọng số kernel W vào bộ nhớ kernel (KMem)
• TileDMA sao chép đầu vào X từ DRAM vào L2
• Mỗi lõi MAC giảm một hàng theo trọng số, tạo ra một hàng của Y
• Hậu xử lý được áp dụng, kết quả lưu vào L2
• TileDMADst sao chép Y từ L2 trở lại DRAM

Nút thắt cổ chai: Băng thông bộ nhớ

Đây là phần quan trọng nhất để hiểu vì sao ANE gặp khó khăn với transformer.

ANE M1 đạt 11 TOP/s với băng thông DRAM hệ thống 68 GB/s. Một phép MAC thực hiện hai phép toán nhưng tiêu thụ hai toán hạng FP16 (4 byte), tức hiệu suất 0,5 OP/byte. Nếu mọi toán hạng MAC đều stream từ DRAM, để duy trì 11 TOP/s sẽ cần băng thông 22 TB/s — gấp hơn 300 lần khả năng thực tế.

Điều này đặt ra điểm gờ roofline (ridge point) ở mức 162 OP/byte. Mỗi byte đọc từ DRAM phải hỗ trợ trung bình ít nhất 162 phép toán thì băng thông DRAM mới không còn là giới hạn.

Đối với giải mã transformer từng token một — trường hợp tệ nhất cho việc tái sử dụng trọng số — ta phải stream toàn bộ trọng số của mô hình chỉ để sinh ra một token. Kết quả đo đạc thực tế cho thấy ANE chỉ đạt:

• KernelDMA: 37,99 GB/s
• TileDMA: 59,08 GB/s
• GPU: 77,70 GB/s

Tệ hơn nữa, các yêu cầu DMA kernel và tile được gửi tuần tự, không song song. Thời gian chạy kết hợp bằng đúng tổng thời gian chạy riêng lẻ — nghĩa là băng thông DRAM của ANE bị "kép" hai lần: cả hai đường DMA đều chậm hơn GPU, và chúng còn cộng dồn thay vì chồng lấp.

Hệ thống bộ nhớ phân cấp

ANE sử dụng kiến trúc bộ nhớ ba tầng:

• KMem: 16 khối 64 KiB SRAM cho kernel, tổng 1 MiB
• L1: vùng đệm đầu vào MAC cho mỗi lõi
• L2: 2 MiB SRAM chia sẻ giữa tất cả các lõi

Bố cục vật lý của ANE được tổ chức quanh hình chữ nhật L2 SRAM ở trung tâm, với 7+7 lõi dọc hai bên, 2 lõi ở cạnh trên và logic điều khiển chung ở cạnh dưới. Các lõi bên hông nhận dữ liệu L2 theo chiều ngang, nhưng hai lõi trên đỉnh cần giao diện ngang tương tự được xoay dọc — tạo ra cấu trúc "lược" đặc trưng ở phần ingress phía trên.

Tác giả cho rằng nếu Apple chịu mở rộng đường kernel đi qua L2, hiệu năng giải mã của ANE đã không tệ đến vậy. Nhưng có lẽ Apple đơn giản là không bao giờ kỳ vọng các tensor nằm trong L2 lại trở thành kernel — một giả định hoàn toàn hợp lý vào năm 2017.

Lời kết: ANE có còn tương lai?

ANE là một cỗ máy luồng dữ liệu chức năng cố định, không phải GPU thực thi lệnh tùy ý. Thiết kế này mang lại diện tích nhỏ hơn, chuyển động dữ liệu xác định, độ trễ thấp và tiêu thụ ít điện năng hơn. Đổi lại, trình biên dịch phải lập lịch tường minh mọi thứ — một sự đánh đổi hợp lý khi ta biết trước hình dạng mô hình lúc biên dịch.

Vấn đề thực sự không nằm ở khả năng tính toán, mà ở chuyển động bộ nhớ. Với việc giải mã ANE bị ghim vào roofline DRAM, một cải thiện đáng kể như từ 10 lên 25 token/giây chỉ có thể đến từ việc tăng băng thông stream bộ nhớ lên khoảng 2,5 lần.

Việc M5 gộp ANE vào GPU là lời thừa nhận rằng kỷ nguyên NPU độc lập đã khép lại. Lõi tính toán vẫn còn giá trị — nhưng phải nằm trong một luồng dữ liệu phù hợp với thời đại transformer, nơi băng thông bộ nhớ, chứ không phải năng lực tính toán, mới là thước đo thực sự của hiệu năng.