Vì sao quay lại mổ M1 Neural Engine
Tác giả từng dừng viết driver reverse-engineered cho Apple Neural Engine (ANE) vì nhận ra phần cứng này quá chuyên biệt để trở thành accelerator đa dụng. M5 sau đó đưa neural accelerator vào GPU và quảng bá hiệu năng LLM, khiến kiến trúc ANE độc lập trông giống sản phẩm của một giai đoạn khác. Việc quay lại M1 không nhằm biến ANE thành GPU, mà để đọc các lựa chọn silicon năm 2017: chúng giả định workload ML nào, và transformer đã phá giả định nào.
Kết luận ngắn là phép nhân không phải vấn đề. Cả CNN lẫn attention đều dùng dot product. Khác biệt nằm ở dataflow và memory movement quanh các MAC. CNN tái sử dụng kernel theo mẫu ổn định; autoregressive decode phải stream khối lượng weight lớn cho từng token.
Khối tính toán
M1 ANE có 16 core. Mỗi core chứa 128 lane MAC cho FP16 hoặc 256 lane cho INT8, tương đương 2.048 phép tích lũy FP16 chạy song song trên toàn chip. Mỗi lane nhân hai toán hạng rồi cộng kết quả vào accumulator 32-bit qua nhiều chu kỳ. Việc ánh xạ tensor và lịch chạy quyết định đó là convolution, matrix multiplication hay dạng reduction khác; MAC tự nó không hiểu khái niệm layer.
Các phép thử cho thấy accumulator dùng fixed-point Q16.16 rồi mới đưa kết quả về FP16. Khi tổng vượt miền số signed 32-bit tương ứng, phần cứng saturate trước bước xuất FP16.
Sau MAC là khối activation, nhờ vậy bias, scale và nonlinear function có thể được fusion mà không cần ghi kết quả trung gian ra memory. Với tanh, firmware chứa 33 mẫu FP16 chia thành 32 khoảng và nội suy tuyến tính giữa hai điểm kề. ReLU dùng mode riêng không cần bảng. Compiler cũng có thể gộp scale và offset hằng vào weight cùng bias của convolution trước khi chạy.
Scheduler không nhận opcode neural
Driver ANE không gửi lệnh kiểu CONV, MATMUL hay RELU. Model đã được compiler biến thành một command stream gồm các task descriptor cố định. Driver đặt địa chỉ và kích thước stream, chọn queue rồi ghi thanh ghi “doorbell”; phần cứng tự fetch descriptor, cấu hình datapath, chạy và phát interrupt khi hoàn tất.
Có tám task queue. Mỗi queue giữ priority, trạng thái, 32 base-address register và hai slot descriptor theo kiểu ping-pong để một task chạy trong khi software chuẩn bị task kế tiếp. Command stream chỉ tham chiếu offset tương đối; bảng địa chỉ cung cấp base IOVA đã quyết định khi compile.
Mỗi task descriptor thực chất là bản serialize của các thanh ghi điều khiển cho từng khối: DMA nạp kernel, DMA đọc tile input, L2, processing engine, MAC cùng post-processing, rồi DMA ghi tile output. Một lượt chạy gần đúng theo thứ tự:
- Control DMA chép descriptor vào register.
- KernelDMA đưa weight từ DRAM vào kernel memory.
- TileDMA đưa input từ DRAM vào L2.
- Các core MAC giảm từng hàng theo weight.
- Post-processing chạy và giữ kết quả trong L2.
- TileDMADst ghi output từ L2 trở lại DRAM.
Đây là dataflow engine cố định, không phải processor chạy instruction tùy ý. Thiết kế đó tiết kiệm diện tích, năng lượng và latency khi model khớp workload dự kiến. Nó không cấm transformer về mặt opcode; giới hạn chính xuất hiện ở cách dữ liệu phải di chuyển.
Memory hierarchy được tối ưu cho CNN
ANE dùng chung DRAM với CPU và GPU, nhưng phải copy dữ liệu vào SRAM cục bộ trước khi MAC tiêu thụ. Kiến trúc M1 có:
- 2 MiB L2 dùng chung cho 16 core.
- 16 bank kernel memory, mỗi bank 64 KiB, tổng 1 MiB.
- Vùng L1 riêng cho input của MAC trên từng core.
Ba đường DMA phục vụ datapath. Mười sáu lane logic của KernelDMASrc đưa weight vào bank riêng của từng core. TileDMASrc chép activation vào L2; TileDMADst ghi kết quả ra DRAM.
Việc tách weight khỏi tile là lựa chọn hợp lý cho convolution: một kernel nhỏ được giữ gần MAC và tái sử dụng trên nhiều cửa sổ input. Kernel path chỉ đọc, memory lại riêng từng core; intermediate tile có thể quay vòng trong L2. Nhưng weight mới phải đi thẳng từ DRAM vào kernel memory, không có đường L2 → KMem. Điều đó trở thành bất lợi khi dữ liệu resident cần liên tục đóng vai trò weight, như KV và weight streaming của transformer.
Từ peak compute và bandwidth DRAM, tác giả tính workload phải đạt khoảng 162 phép toán trên mỗi byte lấy từ DRAM thì MAC mới không bị memory giới hạn. CNN có thể đạt mức reuse cao này; single-token decode thường không thể vì mỗi token cần quét gần như toàn bộ model weight.
Số đo bandwidth giải thích chênh lệch với GPU
Trên M3, tác giả tạo workload đọc buffer lớn hơn cache và dùng dữ liệu giả ngẫu nhiên để đo slope giữa kích thước đọc và thời gian thực thi. Kết quả sustained bandwidth:
- ANE KernelDMA: khoảng 37,99 GB/s.
- ANE TileDMA: khoảng 59,08 GB/s.
- GPU qua Metal: khoảng 77,70 GB/s.
Nếu kernel và tile DMA chạy song song, tổng hai đường có thể gần chạm trần DRAM 100 GB/s của M3. Thực nghiệm lại cho thấy thời gian khi chạy cả hai gần bằng tổng thời gian chạy riêng, nghĩa là request được gửi tuần tự. Với transformer decode vốn đã chạm memory roofline, ANE chịu hai bất lợi: từng đường đọc đều chậm hơn GPU, và thời gian kernel cùng tile cộng dồn.
Một con số TOPS lớn vì thế không dự đoán được token/giây. Khi compute rảnh chờ weight, thêm MAC không cải thiện throughput. GPU thắng không nhất thiết vì có nhiều phép toán hơn mà vì duy trì được luồng đọc DRAM cao hơn.
Ý nghĩa với dev
Reverse-engineering này không nói ANE vô dụng. Với model có shape biết trước, tensor reuse cao và pipeline Core ML phù hợp, fixed dataflow vẫn đổi lấy hiệu quả năng lượng tốt. Nó cũng cho thấy compiler có thể giả lập transformer bằng tensor 4D và convolution 1×1; giới hạn nằm ở hiệu quả, không phải khả năng biểu diễn.
Khi chọn phần cứng local AI, dev nên hỏi bốn câu thay vì chỉ nhìn TOPS:
- Sustained bandwidth cho đúng operand path là bao nhiêu?
- Các DMA có thực sự overlap không?
- Compiler có giữ tensor trong SRAM hay buộc quay lại DRAM?
- Framework hiện tại ánh xạ kiến trúc model vào accelerator tốt đến đâu?
Bài viết kết thúc bằng một hướng mở: một lỗi hoặc đường DMA khác có thể lấy thêm khoảng 50 GB/s từ ANE. Điều đó càng củng cố kết luận chính—hiệu năng transformer trên accelerator chuyên dụng được quyết định bởi memory path nhiều hơn bảng thông số MAC.