Tối ưu hóa suy luận trên thiết bị cho Apple Silicon
Một công cụ cục bộ tùy chỉnh giúp cải thiện thông lượng điền trước và giải mã.
Tính toán hỗn hợp trên Apple silicon (Hybrid Compute on Apple silicon) điều phối một tác vụ giữa trí tuệ biên trên đám mây và một mô hình cục bộ trên máy Mac. Các mô hình trên đám mây xử lý việc nghiên cứu và suy luận, trong khi mô hình cục bộ làm việc với các tệp và ứng dụng riêng tư trên máy Mac.
Để sự phân chia lao động này mang lại cảm giác mượt mà, suy luận cục bộ phải theo kịp phần còn lại của tác vụ. Điều đó đòi hỏi một công cụ có thể xử lý các câu lệnh nhắc nhanh chóng và duy trì tốc độ tạo token cao.
Lily, công cụ suy luận cục bộ nhẹ của chúng tôi, được xây dựng dành riêng cho Apple silicon và Qwen3.6-35B-A3B, với các tối ưu hóa riêng biệt cho giai đoạn điền trước và giải mã. Công cụ này sẽ sớm được mã nguồn mở.
Giới thiệu
Một cách phổ biến để chạy các mô hình ngôn ngữ lớn (LLM) trên máy Mac là sử dụng MLX, khung học máy mã nguồn mở của Apple dành cho Apple silicon. Thư viện đi kèm của nó, MLX-LM, bổ sung các thành phần cần thiết để tải và tạo văn bản với nhiều loại mô hình ngôn ngữ khác nhau. Cùng nhau, MLX và MLX-LM cung cấp một ngăn xếp sẵn sàng, đa mục đích cho suy luận LLM cục bộ.
Qwen3.6-35B-A3B là một mô hình hỗn hợp thưa (sparse, hybrid): mô hình sử dụng định tuyến hỗn hợp chuyên gia (MoE) và kết hợp các trạng thái hồi quy kích thước cố định với cơ chế chú ý đầy đủ. Những lựa chọn kiến trúc này làm giảm lượng tính toán cần thiết, nhưng chúng cũng tạo ra các khối lượng công việc không đều. Các token được định tuyến đến các trọng số chuyên gia khác nhau, và các trạng thái hồi quy có bản chất tuần tự.
MLX-LM đã chọn sẵn các nhân được tối ưu hóa cho các giai đoạn suy luận và các hình dạng khối lượng công việc phổ biến, nhưng các thao tác có thể tái sử dụng của nó phải hỗ trợ nhiều kiến trúc mô hình. Một công cụ chuyên dụng cho Qwen có thể chuyên biệt hóa ở cấp độ mô hình và thời gian chạy, điều phối các nhân, việc di chuyển dữ liệu và lập lịch xung quanh cấu trúc cố định của mô hình.
Lily thực hiện sự chuyên biệt hóa này một cách toàn diện trong một tiến trình duy nhất. Thời gian chạy bằng Rust tải điểm kiểm tra mô hình và quản lý trạng thái phiên cũng như vòng lặp sinh, API hoàn tất trò chuyện tương thích với OpenAI chấp nhận các yêu cầu và truyền phát token, và các nhân Metal tùy chỉnh thực thi các thao tác dành riêng cho Qwen. Cả PyTorch lẫn MLX đều không nằm trong đường dẫn thực thi.
Chúng tôi đo lường hiệu năng điền trước và giải mã một cách riêng biệt. Thông lượng điền trước nắm bắt tốc độ công cụ xử lý câu lệnh nhắc; thông lượng giải mã nắm bắt tốc độ công cụ tạo ra các token đầu ra.
Chúng tôi đánh giá hiệu năng Qwen3.6-35B-A3B trên một chiếc MacBook Pro duy nhất trang bị chip M5 Max với GPU 40 nhân và 128 GB bộ nhớ thống nhất. Qua mười độ dài câu lệnh nhắc cho giai đoạn điền trước và mười độ dài ngữ cảnh cho giai đoạn giải mã, từ 256 đến 128K token (K = 1.024), công cụ này đạt tốc độ trung bình gấp 1,23 lần thông lượng điền trước của MLX-LM và gấp 1,35 lần thông lượng giải mã của MLX-LM. Tại câu lệnh nhắc 4K token và ngữ cảnh giải mã 4K token, công cụ tùy chỉnh đạt 5.749,9 token điền trước mỗi giây và 186,6 token giải mã mỗi giây, so với mức 4.737,5 và 140,9 của MLX-LM. Trong một phiên nhiều lượt, các khoản tiết kiệm thời gian này tích lũy theo từng lệnh gọi mô hình bổ sung.
Tiếp theo, chúng tôi giải thích cách kiến trúc của Qwen tạo ra các cơ hội tối ưu hóa dành riêng cho mô hình trên Apple silicon. Sau đó, chúng tôi điểm qua các thay đổi về điền trước và giải mã kết quả. Chúng tôi cũng đề cập đến việc tối ưu hóa bổ sung sẽ ngừng mang lại hiệu quả ở điểm nào, trước khi kết thúc bằng một phép so sánh toàn diện (end-to-end) với MLX-LM.
Các cơ hội tối ưu hóa dành riêng cho Qwen trên Apple silicon
Qwen tạo ra ba hình dạng khối lượng công việc khác nhau
Qwen3.6-35B-A3B chứa 35 tỷ tham số nhưng chỉ kích hoạt khoảng 3 tỷ tham số cho mỗi token. Một bộ định tuyến chấm điểm 256 mạng con chuyên gia và chọn ra 8 chuyên gia, bên cạnh một chuyên gia được chia sẻ xử lý mọi token. Thiết kế MoE thưa này làm giảm tính toán nhưng tạo ra công việc không đều: các chuyên gia nhận được số lượng token khác nhau và mỗi token đòi hỏi trọng số từ một tổ hợp chuyên gia khác nhau.
Qwen cũng kết hợp 10 lớp chú ý đầy đủ với 30 lớp Gated DeltaNet. Hai loại lớp này giữ lại thông tin trước đó theo các cách khác nhau.
Các lớp chú ý sử dụng cơ chế chú ý theo truy vấn nhóm (GQA). Qwen có 16 đầu truy vấn và hai đầu khóa–giá trị (KV), với tám đầu truy vấn chia sẻ mỗi đầu KV. Việc chia sẻ làm cho bộ nhớ đệm KV nhỏ hơn và cho phép dữ liệu được lưu trong bộ nhớ đệm tái sử dụng trên các đầu truy vấn. Bộ nhớ đệm vẫn lưu trữ các khóa và giá trị mới cho mọi token, do đó mỗi bước giải mã đọc nhiều dữ liệu hơn khi ngữ cảnh phát triển.
Thay vào đó, Gated DeltaNet nén thông tin trước đó thành một trạng thái hồi quy có kích thước cố định. Một cổng đã học kiểm soát lượng trạng thái hiện có cần giữ lại, trong khi bản cập nhật delta kết hợp thông tin từ token hiện tại. Mô hình định nghĩa các bản cập nhật này theo kiểu hồi quy, do đó mỗi token phụ thuộc vào trạng thái được tạo ra bởi token trước đó. Tuy nhiên, trong quá trình điền trước, một công cụ có thể đánh giá cùng một phép tính theo hai cách. Nó có thể quét trực tiếp qua các token trong khi chuyển trạng thái về phía trước, hoặc tổ chức lại các bản cập nhật thành các khối làm lộ ra nhiều phép toán ma trận và tính song song cấp độ token hơn. Phương pháp nào nhanh hơn phụ thuộc vào kích thước mô hình, khối lượng công việc và phần cứng.
Cùng nhau, các cấu trúc này tạo ra ba mẫu tính toán: các nhóm chuyên gia không đồng đều, cơ chế chú ý trên bộ nhớ đệm đang tăng trưởng, và sự hồi quy kích thước cố định có thể được đánh giá trực tiếp hoặc theo khối.
Apple silicon cung cấp các đường dẫn khác nhau cho các khối lượng công việc khác nhau
Giai đoạn điền trước xử lý nhiều hàng kích hoạt token của câu lệnh nhắc cùng một lúc. Khối lượng công việc cục bộ được xem xét ở đây thường giải mã từng yêu cầu một (batch 1) và xử lý một hàng mới cho mỗi bước. Sự khác biệt này làm thay đổi cách sử dụng các trọng số mô hình tương tự nhau. Điền trước có thể tái sử dụng mỗi khối trọng số trên hàng trăm hoặc hàng nghìn hàng. Giải mã phần lớn không thể làm như vậy, vì mỗi token mới đòi hỏi một lượt đi qua các trọng số khác nhau.
Apple silicon đặt CPU và GPU đằng sau bộ nhớ thống nhất, một nhóm bộ nhớ vật lý duy nhất có thể truy cập bởi cả hai. Điều này cho phép mô hình duy trì trạng thái thường trú mà không cần duy trì một bản sao GPU riêng biệt, nhưng nó không làm cho việc di chuyển dữ liệu trở nên miễn phí. Việc đọc trọng số và giá trị trung gian vẫn tiêu thụ băng thông bộ nhớ, trong khi các thanh ghi và vùng lưu trữ trên chip khác nhanh hơn nhưng nhỏ hơn rất nhiều.
GPU M5 cũng cung cấp các đường dẫn tính toán khác nhau. Các lớp tuyến tính của giai đoạn điền trước sử dụng phép nhân ma trận–ma trận chung (GEMM), áp dụng ma trận trọng số cho nhiều hàng cùng một lúc. Các GEMM tương thích có thể sử dụng Bộ tăng tốc thần kinh (Neural Accelerator) trong mỗi lõi GPU thông qua các thao tác ten-xơ Metal 4. Ngược lại, giai đoạn giải mã batch-1 sử dụng phép nhân ma trận–vectơ chung (GEMV), áp dụng cùng trọng số cho một hàng. Với ít khả năng tái sử dụng trọng số, GEMV bị giới hạn chủ yếu bởi băng thông bộ nhớ và thích hợp hơn với các đơn vị logic số học (ALU) vectơ của GPU thay vì các Bộ tăng tốc thần kinh được thiết kế cho các phép toán ma trận có khả năng tái sử dụng dữ liệu lớn hơn.
Các đường dẫn thực thi này không phải là độc quyền của Lily. MLX hoạt động trên cùng một bộ nhớ thống nhất và chọn các nhân ma trận và vectơ được tối ưu hóa theo hình dạng khối lượng công việc. Bản cài đặt Qwen của MLX-LM đã nhóm công việc chuyên gia, đánh giá Gated DeltaNet bằng một nhân Metal hồi quy hợp nhất, và sử dụng cơ chế chú ý nhận thức GQA. Các khả năng này là điểm khởi đầu chung cho suy luận Qwen hiệu quả trên Apple silicon.
Chiến lược tối ưu hóa
Phạm vi hẹp hơn của Lily cho phép nó điều phối các đường dẫn thực thi được chia sẻ này xung quanh kiến trúc và kích thước chính xác của Qwen. Nó sử dụng các đường dẫn GPU chuyên biệt theo giai đoạn, ánh xạ khối lượng công việc chuyên gia, hồi quy và chú ý của Qwen để giảm thiểu việc di chuyển dữ liệu, đồng thời chọn các nhân và bố cục từ hình dạng khối lượng công việc đã đo đạc. Chiến lược này có ba phần:
- Ghép đôi đường dẫn GPU với giai đoạn suy luận. Sử dụng việc thực thi định hướng ma trận khi giai đoạn điền trước có thể tái sử dụng trọng số trên nhiều hàng, và việc thực thi định hướng vectơ khi giai đoạn giải mã batch-1 xử lý từng hàng một.
- Ánh xạ cấu trúc của Qwen lên GPU đồng thời giảm thiểu việc di chuyển dữ liệu. Giữ các trọng số ở trạng thái nén cho đến khi chúng được sử dụng, tổ chức công việc chuyên gia được định tuyến mà không cần quay lại CPU, giữ lại trạng thái Gated DeltaNet trên chip thông qua quá trình quét hồi quy của nó, và tái sử dụng dữ liệu KV được chia sẻ bởi cơ chế chú ý theo truy vấn nhóm.
- Thích ứng các nhân theo hình dạng khối lượng công việc. Trong mỗi giai đoạn, chọn kích thước tile, bố cục thực thi và đường dẫn chú ý từ số lượng hàng có sẵn, sự phân phối các hàng qua các chuyên gia, kích thước của thao tác, và độ dài ngữ cảnh hiện tại.
Các phần sau giải thích những lựa chọn này. Đối với các tối ưu hóa được đánh giá trong các thử nghiệm bóc tách khớp nhau trên M5 Max, chúng tôi ước tính hiệu ứng của chúng bằng cách so sánh các cấu hình công cụ giống hệt nhau ngoại trừ điểm khác biệt duy nhất là phép tối ưu hóa đang được nghiên cứu. Bởi vì các thí nghiệm này so sánh các phiên bản của công cụ chúng tôi với chính nó, chúng giải thích các cơ chế thay vì phân tách các kết quả cuối cùng so với MLX-LM.
Điền trước: tái sử dụng trọng số và giữ định tuyến trên GPU
Điền trước làm lộ ra nhiều hàng token cùng một lúc, nhưng Qwen định tuyến các hàng đó một cách không đều qua các chuyên gia và cập nhật trạng thái hồi quy qua chuỗi. Các tối ưu hóa của nó được chia thành ba nhóm: tổ chức công việc chuyên gia thưa xoay quanh các hàng được định tuyến, giữ quá trình quét Gated DeltaNet trên chip, và chia các câu lệnh nhắc dài thành các đoạn có giới hạn.
Tối ưu hóa tính toán chuyên gia thưa
Giải lượng tử hóa trọng số trong quá trình nhân ma trận
Điểm kiểm tra Qwen3.6-35B-A3B sử dụng phương pháp lượng tử hóa 4-bit affine theo nhóm. Mỗi trọng số được lưu trữ dưới dạng mã số nguyên 4-bit, trong khi mỗi nhóm 64 trọng số chia sẻ một tỷ lệ (scale) và độ lệch (bias) bfloat16 được sử dụng để khôi phục các giá trị của nó. Điều này làm giảm mô hình 35 tỷ tham số từ khoảng 70 GB trọng số bfloat16 xuống điểm kiểm tra 19,4 GB, giúp việc giữ mô hình thường trú trên máy Mac trở nên khả thi.
Thao tác ten-xơ Metal 4 được sử dụng để nhân ma trận tiêu thụ các toán hạng bfloat16 thay vì biểu diễn 4-bit đã đóng gói. Trước khi nhân, GPU phải tái cấu trúc các trọng số về định dạng bfloat16. GEMM nhóm được tối ưu hóa trong Lily thực hiện việc chuyển đổi này trên từng tile trọng số nhỏ tại một thời điểm và giữ kết quả trong bộ nhớ nhóm luồng trên chip đủ lâu để nhân nó với các hàng kích hoạt được định tuyến. Quá trình tích lũy sử dụng dấu phẩy động 32-bit và đầu ra được ghi dưới định dạng bfloat16. Mảng trọng số mở rộng hoàn chỉnh không bao giờ được tạo ra trong bộ nhớ thống nhất.
Trong thử nghiệm bóc tách, việc giải lượng tử hóa chạy như một thao tác riêng biệt: nó mở rộng các trọng số 4-bit thành một mảng bfloat16 trong bộ nhớ thống nhất, sau đó nhân ma trận đọc lại mảng đó. Tại câu lệnh nhắc 512 token, việc chuyển quá trình giải lượng tử hóa vào GEMM nhóm đã làm tăng thông lượng điền trước toàn diện thêm 77.4% bằng cách loại bỏ việc ghi và đọc trung gian này.
Giữ định tuyến chuyên gia trên GPU
GEMM nhóm yêu cầu các hàng kích hoạt được gán cho mỗi chuyên gia phải được lưu trữ cùng nhau. Sau khi chọn 8 chuyên gia cho mỗi token, một biểu đồ tần suất đếm số lượng phép gán được chuyển đến mỗi chuyên gia. Quá trình quét tiền tố biến các số đếm đó thành các phần bù bắt đầu (starting offsets), một bước phân tán (scatter) đặt các hàng vào các nhóm chuyên gia của chúng, và một sơ đồ khối liệt kê các khối ma trận kích thước cố định mà GEMM nhóm phải xử lý.
Đường dẫn được tối ưu hóa giữ toàn bộ trình tự này trong một vùng đệm lệnh duy nhất (command buffer), một tập hợp các thao tác GPU có thứ tự, cho mỗi phân đoạn câu lệnh nhắc. Thay vào đó, một thử nghiệm bóc tách sẽ tạm dừng để CPU có thể kiểm tra các giá trị trung gian định tuyến và gửi thao tác tiếp theo. Việc giữ biểu đồ tần suất (histogram) và quét tiền tố trên GPU bổ sung thêm hai nhân nhưng loại bỏ quá trình đồng bộ hóa CPU–GPU bên trong mỗi lớp MoE.
Tại câu lệnh nhắc 512 token, việc bật định tuyến thường trú trên GPU đã làm tăng hiệu năng điền trước toàn diện thêm 89%. Điều này cũng cho thấy lý do tại sao chỉ riêng số lượng nhân có thể gây hiểu lầm: tuyến đường nhanh hơn khởi chạy nhiều nhân hơn nhưng không bao giờ phải chờ CPU bên trong lớp đó.
Ghép đôi kích thước tile với tải của chuyên gia
Tại câu lệnh nhắc 2K token, việc định tuyến mỗi token đến 8 trong số 256 chuyên gia tạo ra 16.384 phép gán token–chuyên gia, hoặc trung bình 64 hàng kích hoạt cho mỗi chuyên gia. Phân phối thực tế không đồng đều: một số chuyên gia nhận được nhiều hàng, trong khi những chuyên gia khác nhận được ít hàng.
GEMM nhóm chia đầu ra của mỗi chuyên gia thành các tile, là các khối hình chữ nhật nhỏ của đầu ra phép nhân ma trận. Mỗi tile được gán cho một nhóm luồng GPU. Trên các GPU của Apple silicon, một nhóm luồng chứa một hoặc nhiều nhóm simd, mỗi nhóm bao gồm 32 chuỗi thực thi các lệnh theo nhịp khóa (lockstep).
Các tile lớn hơn phân tán chi phí thiết lập qua nhiều hàng hơn và mở ra nhiều công việc song song hơn, nhưng một phần của tile lớn vẫn ở trạng thái nhàn rỗi khi một chuyên gia chỉ nhận được một vài hàng. Do đó, kích thước tile và số lượng nhóm simd có mối liên hệ mật thiết với nhau.
Một thử nghiệm bóc tách cố định tile ở mức 16 hàng. So với đối chứng đó, việc bật tile 32 hàng với bốn nhóm simd đã cải thiện hiệu năng điền trước toàn diện thêm 13,2% ở mức 2K token.
Giữ trạng thái hồi quy trên chip
Trong quá trình điền trước, mỗi lớp Gated DeltaNet quét câu lệnh nhắc theo thứ tự trong khi chuyển trạng thái hồi quy về phía trước. Khi vô hiệu hóa tính năng cư trú trong thanh ghi, thử nghiệm bóc tách sử dụng quá trình quét theo khối. Tại câu lệnh nhắc 2K token, đường dẫn đó di chuyển 256 MiB (mebibyte) trạng thái cho mỗi lớp và liên tục dừng các chuỗi đang phối hợp tại các rào cản (barriers), là các điểm đồng bộ hóa nơi tất cả các chuỗi tham gia phải chờ đợi lẫn nhau.
Trạng thái hồi quy là một ma trận. Nhân được tối ưu hóa gán mỗi cột cho một nhóm simd. Nhóm simd chia cột cho các chuỗi của nó, tải cột vào thanh ghi của chúng một lần, và mang trạng thái qua toàn bộ quá trình quét. Các chuỗi trao đổi kết quả trung gian thông qua các thao tác nhóm simd thay vì bộ nhớ nhóm luồng, một dạng lưu trữ trên chip được chia sẻ giữa các nhóm luồng. Trạng thái hoàn thành chỉ được ghi ngược lại sau khi quá trình quét kết thúc.
Trạng thái và cổng của nó sử dụng định dạng dấu phẩy động 32-bit bởi vì các lỗi làm tròn nhỏ sẽ tích tụ qua các bản cập nhật tuần tự. Các kích hoạt truy vấn và khóa vẫn ở định dạng bfloat16.
Tại câu lệnh nhắc 2K token, việc bật tính năng quét cư trú trong thanh ghi đã cải thiện hiệu năng điền trước toàn diện thêm 5,6%. Các phép toán chuyên gia GEMM chiếm khoảng 90% thời gian điền trước. Quá trình quét tuần tự không tạo ra đủ công việc ma trận có thể tái sử dụng để hưởng lợi từ các bộ tăng tốc thần kinh (Neural Accelerators).
Giới hạn bộ nhớ tạm thời bằng việc phân đoạn câu lệnh nhắc
Thời gian chạy xử lý một câu lệnh nhắc dài dưới dạng một chuỗi các đoạn bị giới hạn thay vì giữ dữ liệu tạm thời cho mọi token câu lệnh nhắc trong bộ nhớ cùng một lúc. Trọng số mô hình vẫn thường trú trong bộ nhớ thống nhất, trong khi trạng thái hồi quy và bộ nhớ đệm KV mang ngữ cảnh từ đoạn này sang đoạn tiếp theo. Không có ngữ cảnh trước đó nào bị loại bỏ.
Nếu không có phân đoạn, các mảng kích hoạt tạm thời sẽ phát triển cùng với toàn bộ câu lệnh nhắc và cạnh tranh với trọng số mô hình, trạng thái hồi quy và bộ nhớ đệm KV để giành bộ nhớ thống nhất. Việc phân đoạn chỉ giữ lại các giá trị tạm thời của một phân đoạn hoạt động tại một thời điểm, sau đó giải phóng hoặc tái sử dụng vùng lưu trữ đó trước khi xử lý phân đoạn tiếp theo. Điều này giới hạn đỉnh bộ nhớ làm việc và cho phép công cụ xử lý các câu lệnh nhắc dài hơn mà không làm thay đổi đầu ra của mô hình.
Điền trước theo phân đoạn rất phổ biến trong nhiều công cụ và đóng vai trò quan trọng trong việc phục vụ các quỹ đạo nhiều lượt dài trong các môi trường bị giới hạn bộ nhớ này. Tổng thời gian điền trước cho các lớp chú ý vẫn mang tính bậc hai theo độ dài câu lệnh nhắc, với một số chi phí bổ sung từ việc tải lặp lại bộ nhớ đệm KV của các đoạn trước đó.
Giải mã: thiểu hóa số byte được di chuyển cho mỗi token
Giải mã batch-1 xử lý từng hàng mới một lúc. Với khả năng tái sử dụng trọng số ít, thông lượng của nó phụ thuộc chủ yếu vào số lượng byte mà công cụ di chuyển cho mỗi token. Các thay đổi về giải mã được chia thành bốn nhóm: tối ưu hóa đường dẫn trọng số một hàng, giữ từng bước trên GPU, giảm lưu lượng trung gian và trạng thái, và đọc bộ nhớ đệm chú ý một cách hiệu quả.
Tối ưu hóa đường dẫn trọng số một hàng
MLX đã điều phối công việc một hàng đến các nhân ma trận–vectơ chuyên biệt. Do Lily không sử dụng MLX, thời gian chạy tùy chỉnh phải cung cấp cùng một chiến lược cơ bản. Phép GEMV song song hàng của chúng tôi được thiết kế cho một hàng kích hoạt. Một nhóm simd hợp tác trên đầu ra trong khi đọc các phần khác nhau của ma trận trọng số một cách song song.
Giữ từng bước giải mã trên GPU
Giữ bàn giao token trên GPU
Mỗi bước giải mã kết thúc bằng việc chọn token tiếp theo; bước tiếp theo bắt đầu với token đó làm đầu vào. Việc gửi lựa chọn đến CPU và sau đó quay lại GPU thêm một điểm đồng bộ hóa vào mọi token. Thay vào đó, thời gian chạy của chúng tôi luân phiên giữa hai vùng đệm lệnh và hai khe token thường trú trên GPU. GPU chọn token có điểm số cao nhất và ghi ID token của nó trực tiếp vào khe đầu vào cho bước giải mã tiếp theo, trong khi CPU chuẩn bị công việc tiếp theo.
Chạy chồng chéo công việc GPU độc lập
Trong một bước giải mã batch-1 được ghi nhận, việc tạo một token đã khởi chạy 795 nhân GPU. Các sự phụ thuộc của chúng tạo thành 555 giai đoạn tuần tự, để lại một số nhân tự do chạy đồng thời. Tuy nhiên, chế độ thực thi nối tiếp của Metal vẫn chạy từng nhân theo thứ tự.
Đường dẫn giải mã được tối ưu hóa ghi lại các sự phụ thuộc dữ liệu thực tế trong một lượt truyền Metal đồng thời. Các lần khởi chạy nhân độc lập có thể chạy cùng lúc khi tài nguyên GPU cho phép. Một rào cản chỉ được chèn vào khi công việc sau đó đòi hỏi kết quả trước đó.
Giảm lưu lượng trạng thái và trung gian
Các nhân riêng biệt thường tạo ra các giá trị trung gian: một nhân ghi kết quả tạm thời vào bộ nhớ, và nhân tiếp theo đọc lại kết quả đó. Đường dẫn giải mã được tối ưu hóa hợp nhất bốn chuỗi: hai phép chiếu đầu vào chuyên gia với kích hoạt có cổng của chúng; phép chiếu đầu ra chuyên gia với điểm số định tuyến của nó và kết quả chuyên gia được chia sẻ; chuẩn bị truy vấn và khóa trước khi chú ý; và bản cập nhật hồi quy với quá trình chuẩn hóa của nó. Mỗi nhân hợp nhất giữ các giá trị tạm thời trong thanh ghi thay vì gửi chúng qua bộ nhớ.
Sự hợp nhất cũng làm ngắn biểu đồ phụ thuộc: khi một lượt ghi trung gian biến mất, rào cản bảo vệ trình tiêu thụ của nó cũng biến mất.
Đọc bộ nhớ đệm chú ý một cách hiệu quả
Gộp các thao tác đọc bộ nhớ đệm chú ý
Cơ chế chú ý đọc các khóa và giá trị từ bộ nhớ đệm KV trong mỗi bước giải mã. Trong thử nghiệm bóc tách, các chuỗi GPU lân cận không phải lúc nào cũng yêu cầu các byte lân cận, buộc hệ thống bộ nhớ phải phục vụ nhiều giao dịch riêng biệt hơn. Việc bật các lần tải được gộp (coalesced loads) khiến các chuỗi lân cận yêu cầu các byte lân cận để phần cứng có thể kết hợp các lần đọc của chúng.
Trên cấu hình bfloat16, việc gộp (coalescing) đã làm tăng băng thông khóa (key bandwidth) từ 33,8 lên 47,9 GB/s, tăng băng thông giá trị (value bandwidth) từ 42,0 lên 61,8 GB/s, và cải thiện quá trình giải mã toàn diện thêm 2,1% ở ngữ cảnh 3.840 token.
Đóng gói các đầu truy vấn để tái sử dụng các hàng KV
Cơ chế chú ý theo truy vấn nhóm (grouped-query attention) cho phép tám đầu truy vấn chia sẻ một đầu KV. Trong quá trình thử nghiệm bóc tách (ablation), mỗi đầu truy vấn chạy trong một nhóm simd (simdgroup) riêng biệt, do đó cả tám đầu đều độc lập yêu cầu cùng một hàng KV được lưu trong bộ nhớ đệm. Nhân được tối ưu hóa sẽ gom bốn đầu truy vấn vào một nhóm luồng (threadgroup), tải từng hàng KV một lần và tái sử dụng nó trên bốn phép tính chú ý. Nhóm luồng thứ hai xử lý bốn đầu còn lại.
Kỹ thuật này, thường được gọi là đóng gói GQA (GQA packing), thực hiện cùng một phép tính số học và tạo ra các byte đầu ra giống hệt nhau trong khi giảm tám yêu cầu KV độc lập xuống còn hai lần tải được chia sẻ. So với thử nghiệm bóc tách chưa đóng gói, nó cải thiện thông lượng giải mã toàn diện thêm 23,8% ở ngữ cảnh 32K token.
Chuyển đổi bố cục chú ý ở các ngữ cảnh dài
Mỗi bước giải mã trong một lớp chú ý đầy đủ sẽ quét bộ nhớ đệm KV hiện có. Bố cục khối cố định chia bộ nhớ đệm đó thành các phần bằng nhau mà GPU có thể xử lý song song. Việc lên lịch bổ sung của bố cục này không đáng kể khi bộ nhớ đệm nhỏ, nhưng bố cục khối cố định cân bằng công việc đều đặn hơn khi ngữ cảnh phát triển.
Đối với mô hình này, thời gian chạy giữ đường dẫn chú ý chung dưới 32K token và sử dụng đường dẫn khối cố định ở mức 32K token trở lên. Việc chuyển đổi áp dụng khi mỗi đầu có 256 giá trị và tám đầu truy vấn chia sẻ một đầu KV; các hình dạng khác vẫn ở trên đường dẫn chung. Một thử nghiệm bóc tách vô hiệu hóa công tắc này và luôn sử dụng đường dẫn chung. Việc bật tuyến đường khối cố định đã cải thiện hiệu năng giải mã toàn diện thêm 7,7% ở mức 32K, 27,4% ở mức 64K, và 40,2% ở mức 128K.
Giới hạn của việc tối ưu hóa thêm
Một số thay đổi đã cải thiện một thao tác đơn lẻ nhưng không cải thiện hiệu năng suy luận toàn diện.
Giải mã phỏng đoán (Speculative decoding), sử dụng một mô hình nhỏ hơn để đề xuất các token cho mô hình đầy đủ xác minh, làm cho quá trình giải mã batch-1 chậm hơn 18%. Quá trình xác minh xử lý các nhóm từ hai đến năm hàng, một hình dạng không hiệu quả cho phần cứng này, và các hàng thường chọn các chuyên gia khác nhau, làm tăng lượng dữ liệu trọng số chuyên gia được đọc. Việc giảm từ vựng đầu ra của mô hình nháp đã cải thiện thông lượng của nó thêm 4,7–5,1%, nhưng không làm cho vòng lặp phỏng đoán hoàn chỉnh nhanh hơn. Kết quả này mang tính đặc thù của khối lượng công việc: quá trình triển khai Qwen được gom nhóm trên Blackwell của chúng tôi sử dụng giải mã phỏng đoán trong các điều kiện khác nhau.
Các thí nghiệm khác bao gồm giảm số lần khởi chạy GPU, chạy chồng chéo toàn bộ các giai đoạn, sử dụng các tile điền trước lớn hơn, áp dụng sự hợp nhất rộng hơn, tăng tốc bộ định tuyến, và kết hợp phép chiếu đầu ra với lựa chọn token. Không có cải tiến nào làm cho vòng lặp suy luận hoàn chỉnh nhanh hơn.
Các phép đo về giới hạn phần cứng cũng cho thấy rất ít không gian dư (headroom) còn lại trong các thao tác điền trước và giải mã chính. Các phép toán MoE GEMM và GEMV đạt 97,9% và 90,3% tốc độ đọc trọng số bền vững nhanh nhất cho các mẫu truy cập của chúng. Việc loại bỏ các phép toán số học khỏi GEMV thưa chỉ làm thay đổi thông lượng 0,2%, xác nhận rằng việc đọc trọng số chứ không phải tính toán mới là tài nguyên giới hạn. Phép nhân ma trận của giai đoạn điền trước cũng đạt 93% giới hạn ma trận lý thuyết ở trạng thái cô lập và 80–86% bên trong các mô hình được kiểm tra.
Hiệu năng toàn diện (End-to-end)
Phép so sánh toàn diện tải các byte điểm kiểm tra 4-bit giống hệt nhau trong cả hai công cụ và chạy từng yêu cầu một trên một chip M5 Max 40 nhân, 128 GB. Trong mỗi vòng, hai công cụ chạy theo thứ tự luân phiên để giảm độ lệch từ tải nền và sự thay đổi nhiệt độ chip. Chúng tôi so sánh với đường dẫn sinh trực tiếp nhanh nhất của MLX-LM, chứ không phải máy chủ của nó, do đó phép đo tập trung vào việc thực thi mô hình hơn là chi phí phục vụ.
Quá trình quét bao gồm mười độ dài câu lệnh nhắc cho điền trước và mười độ dài ngữ cảnh cho giải mã, từ 256 đến 128K token. Thông lượng điền trước ban đầu tăng lên khi công cụ phân tán chi phí thiết lập cố định cho nhiều token hơn. Thông lượng điền trước đạt đỉnh ở khoảng câu lệnh nhắc 4K token, sau đó giảm xuống vì mười lớp chú ý đầy đủ thực hiện nhiều công việc hơn khi câu lệnh nhắc dài ra. Quá trình giải mã duy trì ở mức gần như phẳng tại các ngữ cảnh ngắn và giảm dần một khi việc đọc bộ nhớ đệm KV đang phát triển trở nên đáng kể. Công cụ tùy chỉnh nhanh hơn ở mọi độ dài được ghi nhận.
Bởi vì việc thực thi chuyên biệt có thể làm thay đổi thứ tự các phép toán dấu phẩy động, chúng tôi cũng kiểm tra tính nhất quán về mặt số học so với MLX-LM. Trong một phép so sánh có sự hỗ trợ của giáo viên (teacher-forced), cả hai công cụ đều dự đoán token tiếp theo từ cùng một tiền tố tham chiếu tại mỗi vị trí trong số 192 vị trí, ngăn không cho các khác biệt trước đó ảnh hưởng đến các đầu vào sau này. Độ bối rối (perplexity) của Lily chỉ cao hơn 0,04% và nó đã chọn cùng một token xếp hạng cao nhất ở 96,35% các vị trí được kiểm tra.
Được xây dựng cho nền tảng cục bộ
Apple silicon không phải là một GPU trung tâm dữ liệu thu nhỏ. Đây là một nền tảng suy luận cục bộ hoàn chỉnh với các đặc tính phần cứng và phần mềm riêng. Bộ nhớ thống nhất mang lại cho một nút đơn giới hạn rất cao về lượng mô hình và trạng thái mà nó có thể chứa. Các bộ tăng tốc thần kinh M5 tiếp nhận công việc ma trận dày đặc trong giai đoạn điền trước. Các ALU vectơ xử lý phần còn lại bị giới hạn băng thông và có khả năng tái sử dụng thấp trong giai đoạn giải mã.
Qwen mang lại các cơ hội chuyên biệt hóa bổ sung: giữ định tuyến chuyên gia và trạng thái hồi quy trên GPU, loại bỏ các giá trị trung gian không cần thiết, chạy chồng chéo công việc độc lập, tái sử dụng dữ liệu KV được chia sẻ, và thích ứng các nhân theo hình dạng khối lượng công việc.
Với việc tối ưu hóa theo mô hình và nền tảng cụ thể, một máy Mac có thể chạy một mô hình thưa lớn một cách hiệu quả. Công việc trong tương lai sẽ mở rộng phạm vi phủ sóng trên các mô hình, chip và khối lượng công việc phục vụ, đồng thời biến các cơ chế được xác thực ở đây trên một cấu hình thành chính sách thời gian chạy tổng quát hơn.
Nguyên tắc rộng hơn là ghép đôi công cụ với cả kiến trúc của mô hình lẫn các đường dẫn tính toán và bộ nhớ cụ thể của phần cứng. Khi các mô hình trọng số mở tiên tiến và phần cứng phát triển, suy luận cục bộ hiệu năng cao sẽ ngày càng phụ thuộc vào các công cụ được thiết kế riêng cho cả hai thay vì những công cụ che giấu sự khác biệt giữa chúng.