Bỏ qua điều hướng
chuanweb.com

Chương 5 · Operator & Runtime

Từ phép toán trong đồ thị tính toán đến chương trình chạy trên accelerator: fusion, tái sử dụng dữ liệu, song song và lập lịch giúp giảm chi phí thực thi ra sao.

~105 phút đọc · Thuật ngữ

Operator & Runtime

Cùng một mẩu dữ liệu đầu vào, vừa dùng cho một đầu ra, lại ngay lập tức được cần đến cho đầu ra kế bên: nếu nó vẫn nằm gần đơn vị tính toán, lần dùng sau có thể truy cập trực tiếp; nếu nó đã bị đuổi ra, phải nạp lại từ đầu. Một chương trình nhân ma trận phải sắp xếp thứ tự sử dụng này, đồng thời chừa chỗ phù hợp cho đầu vào và các tổng riêng phần (partial sums). Chỉ cần thay đổi số lượng đầu ra được xử lý cùng lúc, hay khối nào được tính trước, là đã làm thay đổi số lần di chuyển lặp lại cũng như số tác vụ có thể thực thi đồng thời.

Ví dụ nhân ma trận trong chương này cần khoảng 103 GFLOPs tính toán; đầu vào, trọng số và đầu ra cùng chiếm 128 MiB, vậy mà một cách triển khai chia khối (tiling) đạt lưu lượng đọc/ghi khoảng 3 GiB. Phóng to khối đầu ra khiến đầu vào được tái sử dụng nhiều lần hơn, cắt giảm một nửa khối lượng đọc/ghi, nhưng bộ nhớ cục bộ cần cho mỗi khối cũng tăng từ 24 KiB lên 80 KiB. Nắm được những con số này, ta có thể phán đoán liệu thời gian tiết kiệm được nhờ giảm di chuyển dữ liệu có bù đắp được tổn thất do có ít khối chạy đồng thời hơn hay không.

Chương 4 đã giới thiệu năng lực tính toán, năng lực lưu trữ và băng thông truyền dữ liệu của accelerator. Chương này bàn về cách chương trình sử dụng các tài nguyên đó. Với một đoạn FFN làm ví dụ xuyên suốt, trước tiên ta giải thích host khởi động tính toán thế nào và accelerator hoàn tất tác vụ ra sao, rồi phân tích vòng lặp và việc chia khối của phép nhân ma trận, cũng như cách các operator (phép toán) liền kề chia sẻ kết quả trung gian. Sau đó ta bàn cách compiler (trình biên dịch) và runtime triển khai các tối ưu hóa này, và cuối cùng tính xem các tối ưu hóa này tiết kiệm được bao nhiêu thời gian cho một request hoàn chỉnh.

Chương này chủ yếu lấy FFN của Qwen3-8B làm ví dụ xuyên suốt: chiều rộng ẩn H=4096H=4096, chiều rộng trung gian FFN F=12288F=12288. Ta xử lý M=1024M=1024 token cùng lúc, xếp vector đặc trưng 4096 chiều của mỗi token thành các hàng của ma trận đầu vào X, và tính một phép chiếu C=XWC=XW; X, W và C dùng BF16, mỗi phần tử chiếm 2 byte, và phép nhân ma trận tích lũy bằng FP32. M là số token được xử lý trong một lần gọi, mỗi token tương ứng với một hàng của ma trận đầu vào. Các token này có thể đến từ một đoạn prefill duy nhất, hoặc từ nhiều request trong cùng một batch; trong lúc decode từng token với một request riêng lẻ, M bằng 1.

Đối tượngHình dạngKích thướcVai trò trong tính toán
Đầu vào X[1024,4096][1024,4096]8 MiBVector đặc trưng của mỗi token chiếm một hàng; sau phép chiếu, ta thu được vector đặc trưng mới của token đó
Trọng số W[4096,12288][4096,12288]96 MiBĐược dùng chung cho phép tính chiếu của cả 1024 token
Phép chiếu C[1024,12288][1024,12288]24 MiBTrải qua các bước xử lý tiếp theo như activation

Chương này xoay quanh ba câu hỏi: công việc nào bị lặp lại, dữ liệu nào có thể tái sử dụng, và tính toán hay chờ đợi nào quyết định tổng thời gian. Để trả lời các câu hỏi này, đơn vị thực thi chương trình trên accelerator là kernel. Phân tích khối lượng đọc/ghi, dung lượng buffer và số lượng kernel cho biết chương trình làm bao nhiêu việc; kết hợp với thứ tự thực thi của công việc đó, ta tính được chương trình mất bao lâu.

5.1 Nộp tác vụ, di chuyển dữ liệu và đồng bộ hóa trong quá trình thực thi của accelerator

Hãy xét trước tiên một phép nhân ma trận có đầu vào và đầu ra đều nằm trong bộ nhớ CPU. CPU chuẩn bị đầu vào, gửi tới GPU, khởi động tính toán, rồi thu hồi kết quả. Quá trình này đòi hỏi hai bộ xử lý hợp tác với nhau. Hiểu được mỗi bên thực thi độc lập khi nào và chờ bên kia khi nào là nền tảng cho fusion (hợp nhất), pipelining (xếp dòng) và graph replay (phát lại đồ thị) được bàn ở phần sau.

5.1.1 Từ lời gọi framework đến kernel launch

PyTorch là một framework phần mềm để biểu diễn các phép toán tensor và tổ chức thực thi model. Khi nó gọi một phép nhân ma trận, CPU trước tiên chạy mã framework: đọc hình dạng và kiểu dữ liệu của tensor, chọn cách triển khai, chuẩn bị địa chỉ và tham số, rồi nộp tác vụ cho GPU thông qua runtime và driver. Khởi động một lần thực thi kernel trên accelerator được gọi là kernel launch. Runtime chịu trách nhiệm nộp tác vụ, quản lý bộ nhớ và đồng bộ hóa thực thi; compiler sinh mã accelerator hoặc trước thời điểm chạy, hoặc ở lần gọi đầu tiên.

Ở đây có ba cấp độ: operator mô tả công việc toán học, chẳng hạn nhân ma trận hay activation; kernel hoàn tất công việc này trên accelerator; còn launch là hành động CPU dùng để khởi động quá trình thực thi này. Một phép nhân ma trận có thể được hoàn tất bởi nhiều kernel, và nhiều operator phần tử (elementwise) cũng có thể được hoàn tất bởi một kernel duy nhất. Do đó, ta có thể tối ưu hóa hoặc tính toán trên accelerator, hoặc chi phí chuẩn bị và nộp tác vụ ở phía host.

Khi CUDA thực thi một kernel, nó khởi động nhiều thread. Một nhóm thread tạo thành thread block, và tất cả thread block hợp lại thành một grid. Phép nhân ma trận thường chia ma trận đầu ra thành nhiều khối (tile), giao cho các thread block khác nhau tính toán; các thread trong cùng một block hợp tác với nhau để đọc đầu vào và tích lũy kết quả. Accelerator lập lịch (schedule) các thread block lên những SM có sẵn tài nguyên. Khối nào hoàn tất sớm sẽ giải phóng tài nguyên, cho phép các khối tiếp theo bắt đầu; vì vậy các thread block trong một grid có thể thực thi theo từng đợt.

kernel launch được nộp một cách bất đồng bộ (asynchronous). Khi lời gọi trả về, CPU đã xong việc nộp, còn GPU thực thi theo hàng đợi tác vụ và các phụ thuộc. CPU có thể lập tức chuẩn bị tác vụ tiếp theo. Chỉ khi CPU cần dùng kết quả của accelerator, hoặc cần ghi đè dữ liệu mà accelerator vẫn đang dùng, CPU mới cần phải chờ thao tác tương ứng của accelerator hoàn tất.

Sự phân công lao động này tạo thành một đường ống đơn giản: CPU chuẩn bị và nộp tác vụ, GPU thực thi chúng. Các phép tính ma trận lớn chạy đủ lâu để CPU có thời gian chuẩn bị lần gọi kế tiếp; các operator nhỏ hoàn tất nhanh, GPU dễ hút cạn hàng đợi tác vụ và đứng chờ CPU nộp mục tiếp theo. Vì vậy ta có thể tăng tốc theo hai hướng: giảm thời gian accelerator chờ host, và rút ngắn thời gian thực thi của chính kernel.

5.1.2 Sao chép dữ liệu giữa host và accelerator

CPU có thể nộp tác vụ tiếp theo trước, nhưng accelerator chỉ tính toán được khi đầu vào đã đến nơi. Phân tích quá trình truyền dữ liệu trước tiên giúp ta thấy rõ phần việc còn lại giữa lúc nộp tác vụ và lúc tính toán bắt đầu. Ngoài H2D và D2H đã giới thiệu ở Chương 4, các GPU có bộ nhớ riêng cũng thường sao chép ngay trong bộ nhớ đó; có ba hướng sao chép phổ biến:

TênHướng sao chépĐối tượng điển hình
H2DBộ nhớ host → Bộ nhớ GPUTensor đầu vào của batch hiện tại
D2HBộ nhớ GPU → Bộ nhớ hostKết quả tính toán CPU cần dùng
D2DBộ nhớ GPU → Bộ nhớ GPUBuffer đầu vào cố định hoặc kết quả được sắp xếp lại

Quá trình thực thi của phép nhân ma trận nêu trên là:

CPU chuẩn bị đầu vào → H2D → GPU tính toán → D2H → CPU sử dụng kết quả.

Hình 5-1 vẽ vị trí của dữ liệu này trong hai bộ nhớ và các hướng sao chép.

Hình 5-1 H2D di chuyển đầu vào từ một pinned buffer của host vào bộ nhớ GPU, còn D2H đưa kết quả về host; trọng số được nạp một lần rồi lưu trú trong bộ nhớ GPU, và kernel đọc đầu vào trực tiếp từ bộ nhớ GPU và ghi đầu ra tại đó. Dữ liệu trong bộ nhớ thường phải được sao chép trước vào một pinned buffer thì copy engine mới có thể di chuyển trực tiếp được.

Khi chạy một model hoàn chỉnh, không phải lần gọi nào cũng cần truyền lại toàn bộ dữ liệu. Trọng số sau khi nạp có thể lưu trú lại trong bộ nhớ GPU; activation giữa các layer có thể được kernel kế tiếp đọc trực tiếp; và KV cache cũng có thể tiếp tục được dùng qua các bước decode sau đó. Nếu việc sampling được làm trên GPU, CPU chỉ cần nhận một số ít token ID. Vì vậy, số lần sao chép giữa host và accelerator không bằng số lần đọc bên trong accelerator: một tập trọng số có thể chỉ được truyền vào một lần nhưng lại được đọc nhiều lần trên accelerator.

Ví dụ 5-1: Truyền đầu vào từ host sang GPU mất bao lâu? Một tensor BF16 hình dạng [8192,4096][8192,4096] chiếm 8192×4096×2=648192\times4096\times2=64 MiB. RTX PRO 6000 kết nối với host qua PCIe Gen5 x16, với băng thông danh nghĩa 64 GB/s mỗi chiều, cho thời gian truyền là:

TH2D=64 MiB64 GB/s≈1.05 ms.T_{\mathrm{H2D}}=\frac{64\ \mathrm{MiB}}{64\ \mathrm{GB/s}}\approx1.05\ \mathrm{ms}.

Chia đầu vào thành hai tensor 32 MiB, mỗi cái mất khoảng 0,52 ms, tổng khối lượng truyền không đổi. Lợi ích của việc này là phần dữ liệu đầu tiên có thể đến sớm hơn: trong khi GPU bắt đầu tính toán trên nửa đầu của batch, copy engine tiếp tục truyền nửa sau. Mục 5.3 sẽ tính tiếp tổng thời gian khi việc truyền và việc tính toán chồng lấp nhau.

Loại bộ nhớ host cũng ảnh hưởng đến quá trình sao chép. Các trang bộ nhớ thường được hệ điều hành quản lý; pinned memory lưu trú tại chỗ khi đang dùng, thuận tiện cho accelerator di chuyển trực tiếp. Khi truyền dữ liệu từ bộ nhớ thường, runtime thường trước tiên sao chép dữ liệu vào một pinned buffer nội bộ: CPU sao chép một lần, rồi accelerator đọc từ đó. Đó chính là việc sao chép trong host ở phía trái của Hình 5-1. Nếu mỗi batch sao chép 64 MiB dữ liệu của mình vào một pinned buffer mới được cấp phát, sẽ có thêm một lần sao chép phía host trước H2D. Để CPU chuẩn bị đầu vào trực tiếp trong một pinned buffer dùng lại được sẽ tránh được bước trung gian này và tránh việc cấp phát lặp lại.

5.1.3 stream, event và thứ tự hoàn tất

Dữ liệu được truyền tới đâu quyết định bộ xử lý nào dùng được nó; lúc việc truyền hoàn tất quyết định khi nào tính toán tiếp theo có thể bắt đầu. Sau khi một lời gọi bất đồng bộ trả về, việc sao chép có thể vẫn đang diễn ra, nên phải dùng thứ tự thực thi để đảm bảo tính toán không đọc đầu vào chưa truyền xong. CUDA stream là một chuỗi các tác vụ accelerator thực thi theo thứ tự. Đặt việc sao chép H2D và kernel đọc đầu vào đó vào cùng một stream sẽ xếp tính toán sau việc sao chép; CPU khi đó có thể nộp cả hai phần việc liên tiếp.

Nếu việc sao chép và phép tính được đặt ở các stream khác nhau, người ta dùng event để xác định thứ tự thực thi giữa các stream. Stream sao chép ghi lại một event sau khi H2D hoàn tất, và stream tính toán chờ event đó trước khi dùng đầu vào. Việc chờ này diễn ra bên trong chuỗi tác vụ accelerator của stream tính toán; CPU vẫn có thể tiếp tục nộp các công việc khác.

Hình 5-2 vẽ hai stream và hai event trên cùng một trục thời gian.

Hình 5-2 Stream sao chép truyền tuần tự đầu vào của từng batch, còn stream tính toán đọc chúng. Tính batch 0 phải chờ event "chuyển xong batch 0"; batch 2 phải ghi đè slot A và phải chờ event "batch 0 đã được dùng". Mũi tên đứt nét biểu thị việc chờ trên event, không phải di chuyển dữ liệu. Việc luân phiên sử dụng hai slot được trình bày sâu hơn trong Mục 5.3.2.

Event cũng có thể dùng để xác định thời điểm tái sử dụng một buffer. Giả sử đầu vào hiện tại được sao chép từ một buffer của host vào slot A trong bộ nhớ GPU. Khi việc sao chép hoàn tất, CPU có thể ghi đầu vào của batch kế tiếp vào buffer host; tuy nhiên slot A chỉ bị ghi đè sau khi GPU đã dùng xong đầu vào hiện tại. Dù hai buffer chứa cùng một batch dữ liệu, chúng trở nên sẵn sàng để tái sử dụng ở các thời điểm khác nhau.

BufferNgười dùng cuốiHành động tiếp theo
Buffer nguồn H2D trên hostCopy engine đọcCPU ghi đầu vào của batch kế tiếp
Buffer đầu vào trên GPUKernel đọc đầu vào nàyCopy engine ghi đầu vào của batch kế tiếp
Buffer nguồn D2H trên acceleratorCopy engine đọcGPU ghi kết quả mới
Buffer đích D2H trên hostCopy engine ghiCPU đọc kết quả hiện tại

Do đó, nên ghi lại một event ngay sau lần đọc hoặc ghi cuối cùng, để thao tác tái sử dụng buffer này về sau có thể chờ trên event đó. Khi CPU cần kết quả D2H, nó có thể chờ event tương ứng trong khi các tác vụ accelerator không liên quan khác vẫn tiếp tục chạy. Double buffering (đệm kép) ở Mục 5.3 chính là hai chuỗi sử dụng như vậy đan xen với nhau.

Thứ tự này cũng giải thích vì sao một chương trình có thể rơi vào trạng thái kẹt. Mỗi lần chờ trỏ tới một event cụ thể; chỉ cần event đó không bao giờ được ghi lại, việc chờ không bao giờ kết thúc — điều xuất hiện là một chương trình bị kẹt, chứ không đơn thuần là chạy chậm. Cột người dùng cuối của bốn buffer trong bảng trên tương ứng với bốn event dễ bị quên ghi: nếu bỏ sót event “đã dùng xong” của buffer đầu vào GPU, copy engine không bao giờ ghi được batch kế tiếp. Điều tương tự cũng áp dụng cho các vòng phụ thuộc xuyên stream: stream A chờ một event do stream B ghi lại, trong khi tác vụ đó trong stream B lại được xếp sau công việc trong stream A chưa hoàn tất, nên không stream nào tiến được. Thực thi nhiều card bổ sung thêm một nguồn nữa: một lời gọi truyền thông tập thể (collective communication) yêu cầu mọi card tham gia phải phát ra lời gọi giống nhau; chỉ cần một card không phát ra, các card còn lại sẽ chờ vô hạn peer không bao giờ xuất hiện đó (Mục 6.4).

Khi gỡ lỗi, trước tiên hãy kiểm tra phía được chờ có còn tiến triển không: nếu nó vẫn sinh ra kết quả, thì chỉ là chậm, và bạn nên quay lại trục thời gian ở Mục 5.1.4 để đo từng đoạn một; nếu nó kẹt hoàn toàn, thì nó đang chờ một event không bao giờ đến — các node chậm trong Mục 10.4.5 thuộc trường hợp đầu. Đặt timeout cho việc chờ và in ra các tác vụ chưa hoàn tất của mỗi stream khi quá hạn sẽ chỉ rõ chính xác event nào đang thiếu.

5.1.4 Từ thời điểm nộp đến thời điểm kết quả sẵn sàng

Nối thứ tự sao chép, tính toán và sử dụng buffer chính là cách giúp ta xác định khi nào một kết quả trở nên sẵn sàng. Hình 5-3 vẽ thứ tự này thành một trục thời gian.

Ví dụ 5-2: Vì sao thời điểm nộp, thời gian accelerator và thời gian hoàn tất lời gọi lại khác nhau? Đầu vào đã được chuẩn bị sẵn trên host. CPU nộp toàn bộ tác vụ trong khoảng 0–3 μs, và công việc đầu tiên của accelerator bắt đầu ở mốc 3 μs; H2D, kernel và D2H lần lượt mất 8, 20 và 4 μs, và accelerator thực thi chúng liên tiếp theo các phụ thuộc.

Hình 5-3 CPU nộp tác vụ, sao chép đầu vào H2D, thực thi kernel và trả kết quả D2H diễn ra theo trình tự. Ở hàng trên, kernel thực thi 20 μs và kết quả sẵn sàng ở mốc 35 μs; ở hàng dưới, kernel thực thi 5 μs và kết quả sẵn sàng ở mốc 20 μs.

Việc nộp của CPU mất 3 μs, kernel thực thi 20 μs, và CPU không thể dùng kết quả cho đến mốc 35 μs. Nếu ta rút thời gian thực thi kernel từ 20 μs xuống 5 μs, kết quả sẵn sàng ở mốc 3+8+5+4=203+8+5+4=20 μs. Dù tốc độ kernel cải thiện gấp bốn lần, toàn bộ quá trình thực thi chỉ giảm từ 35 μs xuống 20 μs, tăng tốc chỉ khoảng 1,8×; 15 μs còn lại gồm việc nộp và sao chép.

Ba loại khoảng thời gian được hỏi trong Ví dụ 5-2 tương ứng chính xác với ba điểm đo thời gian. Kết quả đo phụ thuộc vào nơi bắt đầu và kết thúc đo. Đọc đồng hồ CPU trước và sau một lời gọi bất đồng bộ đo thời gian dành cho việc nộp; nếu bạn chờ kết quả trước khi dừng đồng hồ, bạn đo thời gian từ khi khởi động đến khi hoàn tất. Một CUDA event đánh dấu điểm bắt đầu và kết thúc trong chuỗi tác vụ accelerator, đo thời gian accelerator dành cho đoạn công việc đó. Các công cụ profiling (hồ sơ hiệu suất) vẽ các lời gọi host, tính toán của accelerator và việc sao chép trên cùng một trục thời gian, làm lộ rõ từng đoạn việc xếp hàng, chờ phụ thuộc và thực thi sau khi nộp.

Trục thời gian trong Hình 5-3 bắt đầu khi đầu vào đã sẵn sàng. Nếu chương trình chưa được biên dịch, còn một bước chuẩn bị phía trước trục thời gian: lần gọi đầu tiên có thể cần biên dịch mã hoặc cấp phát một workspace. Giả sử lần biên dịch đầu tiên mất 100 ms, và mỗi lần thực thi tiếp theo mất 1 ms; một lần gọi khi đó tổng cộng 101 ms, còn 100 lần gọi tổng cộng 200 ms, trung bình 2 ms mỗi lần gọi. Càng gọi nhiều, phần thời gian biên dịch phân bổ càng nhỏ, trong khi 1 ms cần cho mỗi lần thực thi không đổi. Mục 5.5 sẽ dùng chính phương pháp này để phán đoán cần lặp lại bao nhiêu lần trước khi bucketing (đệm đầu vào về một số ít hình dạng đại diện) và specialization (sinh mã chuyên dụng cho một hình dạng cụ thể) trở nên đáng giá.

5.1.5 Kernel thực thi trên một SM như thế nào

Hình 5-3 vẽ kernel như một khối 20 μs duy nhất; việc nó có rút xuống được 5 μs hay không phụ thuộc vào điều diễn ra trên SM trong khoảng thời gian đó. Mục này cố định một kernel và đi qua năm bước — thread, khả năng cư trú (residency), truy cập bộ nhớ, lệnh ma trận và đồng bộ hóa — để suy ra cách nó thực thi trên một SM. Kernel đó là khối nhân-tích lũy trong Mục 4.3.2: khối đầu ra 128×128, bước 64 dọc theo K; khối A BF16 và khối W mỗi cái 16 KiB, bộ tích lũy FP32 64 KiB, cả ba đều nằm trong shared memory, tạo thành một tập làm việc (working set — dữ liệu phải đồng thời nằm trong bộ nhớ cục bộ suốt quá trình tính toán) tổng cộng 96 KiB. Một thread block xử lý một khối đầu ra và chứa 256 thread. Phần cứng là một SM của H100: nó có thể giữ đồng thời tối đa 2048 thread cư trú (tức 64 warp, sẽ giải thích bên dưới) và 32 thread block, với 64K register (mỗi cái 32 bit) và 228 KB shared memory (theo đúng tài liệu CUDA, tức 233.472 byte). Kernel dùng 128 register mỗi thread, mỗi lệnh nạp lấy 16 byte mỗi thread, mỗi warp có tối đa 4 lệnh nạp đang chờ (outstanding loads), hình dạng lệnh ma trận là 16×8×16, và độ trễ bộ nhớ được lấy bằng 600 ns.

warp: 32 thread được phát lệnh cùng nhau. Một SM không phát lệnh từng thread một; nó gom 32 thread lại và phát chúng như một đơn vị duy nhất, gọi là warp. Cả 32 thread trong một warp thực thi cùng một lệnh, mỗi thread xử lý dữ liệu khác nhau; do đó một lệnh nạp đơn lẻ tạo ra 32 địa chỉ cùng lúc. Một thread block 256 thread chứa 8 warp. 32 “lane” được bàn trong Mục 5.2.3 khi phân tích các bank của shared memory (các phân vùng shared memory có thể phục vụ yêu cầu một cách độc lập) chính là 32 thread của một warp: 32 địa chỉ chúng phát ra đồng thời rơi vào bao nhiêu bank sẽ quyết định lệnh đó chiếm giao diện shared memory trong bao nhiêu lượt.

Khả năng cư trú (residency) và occupancy. Để một thread block cư trú trên một SM — tức luôn ở lại trên SM và sẵn sàng được lập lịch tại mọi thời điểm — nó phải được cấp đồng thời các slot thread (mỗi thread một slot), register và shared memory. Số thread block mỗi loại tài nguyên có thể giữ bằng tổng tài nguyên chia cho nhu cầu mỗi block, làm tròn xuống; số block cư trú là giá trị nhỏ nhất giữa các loại:

Tài nguyênTổng trên SMNhu cầu mỗi blockSố thread block chứa được
Slot thread20482568
Register64K256 × 128 = 32K2
Shared memory228 KB (233.472 byte)Tập làm việc 96 KiB + 1 KiB dành riêng mỗi block = 97 KiB2
Trần thread block32132

Cả register và shared memory đều chỉ cho phép 2 block, nên 2 thread block — tức 16 warp — cư trú trên SM. Occupancy là tỷ lệ giữa số warp cư trú và số warp tối đa cư trú của SM, ở đây là 16/64 = 25%. Hình 5-4 vẽ ba loại tài nguyên thành ba thanh: hai thread block lấp đầy chính xác các register, lấp shared memory chỉ còn dư 34 KiB, trong khi các slot thread chỉ dùng một phần tư.

Hình 5-4 Ba thanh lần lượt biểu thị tổng register, shared memory và slot thread của SM, với màu xanh dương và xanh lá thể hiện mức chiếm dụng của hai thread block cư trú. Register được lấp chính xác; shared memory còn dư 34 KiB, không đủ chỗ cho nhu cầu 97 KiB của block thứ ba — cả hai cùng giới hạn khả năng cư trú; chỉ 512 slot thread được dùng. Mỗi block có 256 thread, 128 register mỗi thread và 96 KiB shared memory, cộng 1 KiB dành riêng mỗi block.

Tài nguyên nào đáng thay đổi phụ thuộc vào tài nguyên nào là ràng buộc chặt nhất. Đưa bộ tích lũy 64 KiB từ shared memory vào register đòi hỏi mỗi thread phải giữ thêm 64 giá trị FP32, tức thêm 64 register. Nhu cầu shared memory mỗi block khi đó giảm xuống 32 KiB, chứa được 6 block; nhưng nếu ban đầu mỗi thread chỉ dùng 64 register, việc thêm bộ tích lũy nâng nó lên 128, và register vẫn chỉ cho phép 2 block. Số block cư trú không đổi; ràng buộc chặn chỉ chuyển từ hai tài nguyên sang chỉ còn register. Để một block thứ ba có thể cư trú, phải đồng thời giảm cả register mỗi thread lẫn shared memory mỗi block, hoặc thu nhỏ tập làm việc mỗi block.

Che độ trễ (latency hiding): các warp cư trú cung cấp các yêu cầu đang bay (in-flight). Occupancy quan trọng vì mọi yêu cầu truy cập bộ nhớ độc lập đều đến từ các warp cư trú: với đủ yêu cầu đang bay, khi một yêu cầu đang chờ kết quả trả về, giao diện vẫn bị các yêu cầu khác chiếm giữ — đó là che độ trễ. Định luật Little ở Mục 4.3.3 cho lượng byte đang bay cần để duy trì băng thông: băng thông × độ trễ. Với băng thông bộ nhớ 3,35 TB/s và độ trễ 600 ns của H100, cả card cần khoảng 2,01 MB đang bay; trải đều ra 132 SM, mỗi SM phải giữ khoảng 15,2 KB đang bay. Sau đó ta tính các warp cư trú cung cấp được bao nhiêu: một lệnh nạp của một warp có 32 thread mỗi thread lấy 16 byte, tổng cộng 512 byte; mỗi warp có tối đa 4 lệnh nạp đang chờ; 16 warp cùng nhau cung cấp 32 KiB.

Bneed=3.35 TB/s×600 ns132≈15.2 KB,Bhave=16×4×512 bytes=32 KiB.B_{\mathrm{need}}=\frac{3.35\ \mathrm{TB/s}\times600\ \mathrm{ns}}{132}\approx15.2\ \mathrm{KB},\qquad B_{\mathrm{have}}=16\times4\times512\ \mathrm{bytes}=32\ \mathrm{KiB}.

Lượng byte đang bay sẵn có gấp 2,15 lần nhu cầu, nên chỉ cần 8 warp cư trú (một thread block) là đủ. Mục 5.2.2 sẽ gặp hiện tượng “khối lớn hơn đọc ít dữ liệu hơn vẫn có thể chậm hơn” — đây chính là dạng định lượng của nó: tập làm việc càng lớn, số block cư trú càng ít, và lượng byte đang bay sẵn có càng giảm; một khi rơi xuống dưới nhu cầu, băng thông hiệu dụng giảm xuống mức byte đang bay sẵn có chia cho độ trễ, và dù giao diện có băng thông bao nhiêu cũng vô ích.

Độ mịn của lệnh ma trận. Mục 4.2.1 đã giải thích rằng đơn vị ma trận thực thi các khối có hình dạng cố định. Ở đây một lệnh nhân-tích lũy ma trận (MMA) có hình dạng 16×8×16: khối A 16×16, khối W 16×8, và một lệnh thực hiện 2×16×8×16=40962\times16\times8\times16=4096 phép toán dấu phẩy động. Một khối 128×128×64 cần (128/16)×(128/8)×(64/16)=8×16×4=512(128/16)\times(128/8)\times(64/16)=8\times16\times4=512 lệnh, khớp với 2.097.152 FLOPs của khối. 512 lệnh này được chia đều cho 8 warp: khối đầu ra chứa 8×16=1288\times16=128 khối đầu ra 16×8, mỗi warp xử lý 16 trong số đó, phát tổng cộng 64 lệnh trong suốt bốn bước dọc theo K. Các lệnh MMA tích lũy trong register; khối trong mục này giữ kết quả tích lũy trong shared memory, nên trước và sau mỗi bước K, 16 khối này phải được di chuyển giữa register và shared memory — phương án đặt bộ tích lũy trong register ở đoạn trên loại bỏ đúng cú di chuyển này.

MMA có thể được phát theo hai cách. Lệnh cấp warp được một warp tự phát lên, với toán hạng được nạp trước vào register của warp đó; lệnh cấp warpgroup của Hopper được 4 warp liền kề (gọi chung là warpgroup) cùng phát, chúng có thể nạp toán hạng trực tiếp từ shared memory và thực thi bất đồng bộ: warp phát lệnh có thể làm việc khác trong lúc chờ kết quả. FlashAttention-3 khai thác chính cách sắp xếp lệnh này để chồng lấp tính toán.

Đồng bộ hóa: sao chép bất đồng bộ và barrier. Mục 4.4.2 dùng hai slot đầu vào để chồng lấp việc nạp với việc tính toán, với tiền đề là biết được hai mốc “slot đã đầy” và “slot đã dùng xong”. Trên một SM, hai mốc này được xác lập bởi barrier. Barrier là một bộ đếm trong shared memory: mỗi người tham gia đến barrier sau khi hoàn thành bước của mình, và khi bộ đếm đạt tới một giá trị định sẵn, các thread đang chờ tại barrier được giải phóng. Các bản sao chép bất đồng bộ (bản sao chép do phần cứng phát hành và hoàn tất trong nền, người phát không cần chờ) cũng đến barrier khi hoàn tất, đếm theo số byte được chuyển vào. Mỗi slot đi kèm hai barrier: barrier “đầy” được kích hoạt khi việc sao chép xong, và phép tính chờ trên nó; barrier “trống” được kích hoạt khi phép tính kết thúc, và bản sao chép tiếp theo chờ trên nó.

Với barrier được cài đặt, các warp có thể phân công lao động: một warp chỉ phát hành việc sao chép (producer — nhà sản xuất), còn các warp còn lại chỉ tính toán (consumer — người tiêu dùng) — cách sắp xếp này gọi là warp specialization (chuyên môn hóa warp). Sau khi warp producer phát hành bản sao chép cho block j+1, nó không cần chờ bản sao chép hoàn tất; thay vào đó nó chờ trên barrier “trống”, chuẩn bị ghi block j+2. Mỗi lần một warp consumer đạt tới barrier “đầy”, nó tiến hành tính block tiếp theo.

Trước tiên hãy tính thời gian mỗi việc sao chép và tính toán. Chia đều băng thông bộ nhớ 3,35 TB/s và đỉnh ma trận BF16 đặc 989,4 TFLOP/s của H100 cho 132 SM cho mỗi SM khoảng 25,4 GB/s và 7,50 TFLOP/s; phép tính dưới đây giả định thread block này dùng riêng phần của SM mình. A và W của một khối K cộng lại 32 KiB, sao chép mất khoảng 1,29 μs; phép nhân khối là 2.097.152 FLOPs, tính mất khoảng 0,28 μs. Việc sao chép chậm hơn tính toán khoảng 4,6 lần, vì khối này chỉ thực hiện 64 phép toán dấu phẩy động mỗi byte được di chuyển, trong khi tỷ lệ đỉnh ma trận so với băng thông bộ nhớ của H100 vào khoảng 295 FLOP/byte. Hình 5-5 vẽ bốn khối K trải trên hai slot dùng hai giá trị thời gian này.

Hình 5-5 Hàng trên là bản sao chép bất đồng bộ do warp producer phát hành, hàng dưới là phép tính do warp consumer thực hiện; màu xanh dương và xanh lá lần lượt là slot A và B. Mũi tên liền là barrier "đầy" — phép tính không thể bắt đầu cho đến khi việc sao chép xong. Mũi tên đứt nét là barrier "trống": block 0 dùng xong slot A ở mốc 1,57 μs, và bản sao chép block 2 không bắt đầu cho đến 2,58 μs, nên producer không bao giờ phải chờ slot. Trên một SM H100, mỗi block sao chép mất 1,29 μs và tính mất 0,28 μs; các bản sao chép chạy liên tục, và consumer chờ barrier "đầy" tiếp theo mỗi lần xong một block.

Sự phân công lao động này chính là hiện thực phần mềm của đường ống trong Mục 4.2.3. Ở đó, ba tài nguyên là đơn vị ma trận, giao diện shared memory và đơn vị lũy thừa; ở đây, warp producer điều khiển việc sao chép, chiếm giữ giao diện shared memory, còn warp consumer điều khiển các phép toán ma trận và lũy thừa, với barrier để mỗi bên chỉ chờ đúng tại điểm bàn giao dữ liệu. FlashAttention-3 được thiết kế chính theo cách này: một warpgroup producer phát hành các bản sao chép khối K và V, trong khi hai warpgroup consumer luân phiên nhân ma trận và softmax, để một bên tính lũy thừa trong khi bên kia nhân ma trận. Mục 5.3.2 sẽ dùng chính các giá trị từ đây — thời gian sao chép 1,29 μs và thời gian tính toán 0,28 μs mỗi block — để suy ra tổng thời gian của đường ống, còn Mục 5.3.3 sẽ trình bày sâu hơn về tiling của FlashAttention.

Sao chép khối lượng lớn bất đồng bộ cũng đặt ra một yêu cầu về bố cục dữ liệu. Mục 5.2.3 sẽ dùng padding (đệm — thêm một word phụ ở cuối mỗi hàng) để trải đều các bank; sao chép khối bất đồng bộ di chuyển cả khối theo một tensor descriptor (cấu trúc ghi lại hình dạng và bước nhảy của tensor) và không thể chèn padding giữa các hàng, nên thay vào đó nó hoán vị các địa chỉ trong một khối theo một quy tắc cố định — kỹ thuật này gọi là swizzle: tổng khối lượng dữ liệu không đổi, nhưng các yêu cầu hướng tới cùng một cột vẫn rơi vào các bank khác nhau.

Ánh xạ sang Ascend. Năm bước được bàn trong mục này — thread, khả năng cư trú, truy cập bộ nhớ, lệnh ma trận và đồng bộ hóa — trên Ascend tương ứng với một bộ thành phần khác từ Mục 4.2.3 và 4.6.2:

Cơ chế trong mục nàyNVIDIA GPUAscend
Phát lệnh và phân công lao độngWarp là đơn vị phát lệnh cho 32 thread; warp producer và consumer phân chia công việc trong cùng một SMAIC và AIV là các core độc lập chứ không phải warp: công việc ma trận thực thi trên AIC, công việc vector thực thi trên AIV, mỗi cái có luồng lệnh riêng
Khả năng cư trú và dung lượngRegister, shared memory và slot thread quyết định số thread block cư trú mỗi SMDung lượng của L0, L1 và unified buffer quyết định mỗi core giữ được bao nhiêu block cùng lúc
Di chuyển và đồng bộ hóaCác bản sao chép bất đồng bộ đến barrier, và các warp chờ trên barrierMTE/NDDMA thực hiện việc di chuyển, và BufferID xác định các buffer và tổ chức bàn giao

Kết luận ở trên — “đủ byte đang bay” — dựa vào độ trễ 600 ns và 4 lệnh nạp đang chờ mỗi warp. Điều kiện lật (flip condition): nếu độ trễ tăng lên 1000 ns và mỗi warp chỉ duy trì được 2 lệnh nạp đang chờ, mỗi SM cần 25,4 KB đang bay, nhưng 16 warp chỉ cung cấp được 16 KiB, kéo tỷ lệ bao phủ (byte đang bay sẵn có chia cho nhu cầu) xuống 0,65; để đóng khoảng trống này cần tối thiểu 25 warp, tức 4 thread block, trong khi register và shared memory chỉ cho phép 2. Lúc này, occupancy 25% không còn đủ: hoặc phải thu nhỏ tập làm việc để có thêm block cư trú, hoặc phải dùng sao chép khối bất đồng bộ để một thread phát yêu cầu cho cả một khối cùng lúc, với lượng byte đang bay được tính theo kích thước khối chứ không phụ thuộc số warp.

Bài tập 5-1 [Mở rộng]: Nếu thu nhỏ khối đầu ra xuống 64×64, tài nguyên nào quyết định occupancy?

Giữ nguyên 256 thread và 128 register mỗi thread, đổi khối đầu ra thành 64×64, giữ bước K là 64, và giữ bộ tích lũy trong shared memory. Tìm nhu cầu shared memory mỗi block, số thread block mỗi loại trong bốn loại tài nguyên cho phép, và occupancy, đồng thời xác định ràng buộc nào chặt nhất. Sau đó tính lại với register mỗi thread hạ xuống 64, và nêu rõ ràng buộc chuyển sang đâu; cuối cùng, với độ trễ 600 ns và 4 lệnh nạp mỗi warp, kiểm tra xem lượng byte đang bay còn đủ để che phủ nhu cầu hay không.

5.2 Tiling, bố cục và truy cập bộ nhớ thực tế của một operator

Một phép tính ma trận lớn có thể được chia khối (tiling) giống như cách cắt đậu phụ: cắt thành các miếng kích thước vừa để xử lý tiện nhất, rồi hoàn tất từng miếng một. Nhưng hình dạng cắt phải khớp với thứ đơn vị tính toán xử lý được và lượng dữ liệu có thể nằm gần đó. Chương 4 đã cho thấy các lệnh ma trận phần cứng chỉ hỗ trợ những hình dạng giới hạn; nhiều đơn vị tính toán trải khắp accelerator có thể cùng làm việc trên nhiều khối cùng lúc, và trong mỗi phần công việc, các bước vẫn phải được tổ chức thành những gia số nhỏ mà phần cứng thực thi được.

Tiling cũng phải khiến dữ liệu được đưa vào được dùng nhiều lần hơn. Theo phép loại suy bàn làm việc của Chương 4, đầu vào hiện tại và các tổng riêng phần chưa xong cần nằm gần đơn vị tính toán, và chỉ khi chúng xong thì chỗ đó mới được giải phóng. Nếu khối đầu ra quá nhỏ, cùng một đầu vào có thể phải di chuyển đi di chuyển lại cho các khối đầu ra khác nhau; nếu khối quá lớn, nó chiếm nhiều bộ nhớ cục bộ và chen lấn các công việc khác có thể chạy đồng thời. Dưới đây, trước tiên ta lần theo một đầu ra đơn cần dữ liệu gì, rồi gom nhiều đầu ra thành các khối, và phân tích kích thước khối, thứ tự truy cập và dung lượng cục bộ phối hợp ảnh hưởng đến việc tái sử dụng ra sao.

5.2.1 Hình dạng ma trận và các lần đọc lặp lại

Gọi A là ma trận đầu vào X ở đầu chương, và viết phép chiếu là C=AWC=AW, với M=1024,K=4096,N=12288M=1024,K=4096,N=12288. A, W và C đều là BF16, lần lượt có kích thước 8, 96 và 24 MiB. Nếu A và W mỗi cái được đọc một lần và C được ghi một lần, chỉ cần 128 MiB truyền tải. Để giải thích vì sao một cách triển khai chia khối tạo ra khoảng 3 GiB truy cập, trước tiên hãy xem một đầu ra đơn được tính thế nào.

Để tính C[i,j], lấy hàng i của A và cột j của W, nhân các phần tử tương ứng, và cộng dọc theo chiều K. Sau đó tính C[i,j+1] ở bên phải, chuyển sang cột liền kề của W trong khi A vẫn giữ nguyên hàng đó. Các phần màu xanh dương của hai phép tính trong hình là giống hệt nhau.

Hình 5-6 Tính một đầu ra, rồi tính đầu ra liền kế. Màu xanh dương đánh dấu hàng của A được dùng, màu cam đánh dấu cột của W được dùng, màu xanh lá đánh dấu phần tử đầu ra thu được ở bước đó. Các ma trận trong hình được vẽ ở kích thước minh họa; trong ví dụ thực tế, K bằng 4096.

Nếu hàng này của A vẫn cư trú trong bộ nhớ cục bộ hoặc cache giữa hai lần dùng, nó có thể được tái sử dụng trực tiếp; nếu nó đã bị đuổi ra, một yêu cầu khác phải đi tới tầng lưu trữ kế tiếp. Khi chuyển sang hàng đầu ra kế tiếp, các phần tử của W cũng được tái sử dụng. Cùng một mẩu dữ liệu cần được dùng lặp đi lặp lại giữa các phép toán. Cách chương trình sắp xếp các phép toán đó quyết định dữ liệu này vượt qua giao diện lưu trữ bao nhiêu lần.

Mục này đếm tại giao diện giữa buffer làm việc và tầng lưu trữ kế tiếp, thống kê từng byte chương trình đọc và ghi mỗi lần. Con số truy cập khoảng 3 GiB dưới đây cũng được tính tại giao diện này.

Viết phép tính elementwise trong hình dưới dạng vòng lặp, ta có:

for i in range(M):
    for j in range(N):
        acc = 0.0
        for k in range(K):
            acc += A[i, k] * W[k, j]
        C[i, j] = acc

Mỗi lượt đi qua vòng lặp k bên trong tạo ra một phần tử đầu ra. Mỗi phép nhân–cộng được tính là hai FLOP, nên tổng khối lượng tính toán cần cho toàn bộ đầu ra là:

Fgemm=2MKN≈103 GFLOPs.F_{\mathrm{gemm}}=2MKN\approx103\ \mathrm{GFLOPs}.

Đường cơ sở để so sánh là khối lượng đọc/ghi của ba ma trận BF16, mỗi ma trận đi qua giao diện đúng một lần:

Vonce=2(MK+KN+MN)=(8+96+24) MiB=128 MiB.V_{\mathrm{once}}=2(MK+KN+MN)=(8+96+24)\ \mathrm{MiB}=128\ \mathrm{MiB}.

Ở đây MiB là 2202^{20} byte. Lấy việc đọc một lần mỗi đầu vào và ghi một lần đầu ra làm đường cơ sở, ta có thể tính ở dưới đây xem tiling làm tăng thêm bao nhiêu truy cập lặp lại.

Thứ tự vòng lặp cũng quyết định kiểu địa chỉ của mỗi lần truy cập. Khi W được lưu theo thứ tự hàng chính (row-major), vòng lặp trong ở trên đọc W[k,j], W[k+1,j] liên tiếp nhau, các địa chỉ cách nhau N phần tử. Chuyển sang thứ tự i, k, j, vòng lặp trong khi đó đọc W[k,j], W[k,j+1] liên tiếp, đồng thời tái sử dụng A[i,k] cho nhiều đầu ra. Như vậy, dữ liệu lấy về trong một lần truyền liên tục có thể được dùng ngay trong tính toán.

Tuy nhiên, thứ tự mới cũng phải giữ một đoạn kết quả đầu ra đang được tích lũy cư trú tại cùng một thời điểm. Nếu bộ nhớ cục bộ không đủ chứa các kết quả này, mỗi vòng cập nhật có thể phát sinh thêm việc đọc và ghi vào tầng lưu trữ kế tiếp. Kế tiếp, ta dùng tiling để giới hạn lượng dữ liệu được giữ đồng thời, cân bằng việc tái sử dụng đầu vào với việc giữ lại đầu ra.

5.2.2 Tái sử dụng theo tile trong giới hạn dung lượng

Trước tiên chọn một khối C nhỏ làm đầu ra hiện tại cần hoàn thành, và cấp một buffer tích lũy cho tile đầu ra m×nm\times n đó. Mỗi lần, đọc vào một cặp tile đầu vào dọc theo chiều K: một tile A kích thước m×km\times k và một tile W kích thước k×nk\times n. Hai tile được nhân với nhau, cập nhật vào cùng một bộ tích lũy đầu ra; khi cặp tile đầu vào kế tiếp được đưa vào, tổng từng phần đã thu được vẫn được giữ lại. Chỉ khi việc tích lũy dọc theo K hoàn tất, kết quả mới được chuyển đổi sang BF16 và ghi ra. Tile ở đây do phần mềm tổ chức phục vụ việc tái sử dụng dữ liệu, và bên trong nó vẫn có thể được phân rã thành nhiều lệnh nhân ma trận phần cứng. Bản thân tiling không đòi hỏi phải sao chép trước ma trận gốc thành nhiều ma trận con tách rời — việc di chuyển thực tế do chương trình đang thực thi sắp xếp.

Hình 5-7 Cố định tile đầu ra m×n hiện tại, lần lượt đưa các tile A và W tương ứng vào dọc theo K, cập nhật lặp lại cùng một tổng từng phần, và chỉ ghi ra sau khi tích lũy hoàn tất. Đầu vào A và W là BF16, bộ tích lũy là FP32. Đây là tiling ở mức phần mềm; một phép nhân tile có thể bao gồm nhiều lệnh nhân ma trận của phần cứng; nhãn kích thước chỉ hình dạng, diện tích ô không theo tỉ lệ số byte.

Mỗi phần tử trong tile A có thể được dùng cho n cột đầu ra, và mỗi phần tử trong tile W có thể được dùng cho m hàng đầu ra. Phóng to tile đầu ra giúp một lần đọc dữ liệu tham gia được vào nhiều phép nhân–cộng hơn, nhưng cũng đòi hỏi phải giữ nhiều đầu vào và nhiều tổng từng phần đầu ra hơn.

Ví dụ 5-3: Phóng to tile đầu ra có thể giảm được bao nhiêu lần đọc đầu vào? Trong hình, tile đầu vào A, tile trọng số W và bộ tích lũy đầu ra lần lượt chiếm 2mk2mk, 2kn2kn và 4mn4mn byte. Chỉ với một bộ buffer đầu vào, yêu cầu bộ nhớ cục bộ là:

S=2mk⏟A tile+2kn⏟W tile+4mn⏟accumulatorbytes.S=\underbrace{2mk}_{A\text{ tile}}+\underbrace{2kn}_{W\text{ tile}}+\underbrace{4mn}_{\text{accumulator}}\quad\mathrm{bytes}.

Lấy k=32k=32, trước tiên dùng tile đầu ra 64×64. Hai tile A và W mỗi tile chiếm 4 KiB, bộ tích lũy chiếm 16 KiB, tổng cộng 24 KiB. Phóng to tile đầu ra lên 128×128, hai đầu vào mỗi cái chiếm 8 KiB, bộ tích lũy chiếm 64 KiB, tổng cộng 80 KiB. Tái sử dụng đầu vào tăng lên, nhưng bộ tích lũy đầu ra tăng nhanh hơn.

Tiếp theo, tính khối lượng đọc/ghi cho toàn bộ ma trận. Giả sử mỗi tile đầu ra tự đọc các đầu vào của riêng nó, không xét việc tái sử dụng cache giữa các tile đầu ra khác nhau. Có N/n tile cột theo hướng N, mỗi tile cột duyệt toàn bộ A; có M/m tile hàng theo hướng M, mỗi tile hàng duyệt toàn bộ W. C chỉ được ghi ra một lần ở cuối, nên:

V=2MKNn⏟repeated reads of A+2KNMm⏟repeated reads of W+2MN⏟C write.V=\underbrace{2MK\frac{N}{n}}_{\text{repeated reads of }A}+\underbrace{2KN\frac{M}{m}}_{\text{repeated reads of }W}+\underbrace{2MN}_{C\text{ write}}.

Tile đầu ra 64×64 chia đầu ra thành 16 tile hàng và 192 tile cột. 192 tile cột mỗi tile đọc 8 MiB của A một lần, tổng cộng 1536 MiB; 16 tile hàng mỗi tile đọc 96 MiB của W một lần, cũng là 1536 MiB. Cộng thêm 24 MiB đầu ra cho tổng 3096 MiB. Đây chính là nguồn gốc của con số khoảng 3 GiB ở đầu mục này.

Chuyển sang 128×128, số tile hàng giảm từ 16 xuống 8, số tile cột giảm từ 192 xuống 96, giảm một nửa số lần đọc cả hai đầu vào. Kích thước đầu ra không đổi, và tổng giảm xuống 768+768+24=1560768+768+24=1560 MiB. Hình 5-8 so sánh lợi ích này với bộ nhớ cục bộ mà nó chiếm dụng.

Hình 5-8 Mỗi điểm đánh dấu một hình dạng tile đầu ra. Phóng to tile đầu ra làm giảm số lần đọc lặp lại qua giao diện nhưng đòi hỏi nhiều bộ nhớ cục bộ hơn. Đường nét đứt đánh dấu mức 128 MiB cho ba ma trận mỗi ma trận đi qua một lần; điều kiện là k=32, một bộ buffer đầu vào, và mỗi tile đầu ra tự đọc các đầu vào của mình.

Các con số đầy đủ được cho dưới đây; cường độ tính toán được tính bằng FLOP của cùng phép nhân ma trận chia cho số byte giao diện ở trên:

Tile đầu raBộ nhớ cục bộ cầnTruy cập giữa buffer làm việc và tầng lưu trữ kế tiếpCường độ tính toán
32×328 KiB6168 MiB15.9 FLOP/byte
64×6424 KiB3096 MiB31.7 FLOP/byte
128×12880 KiB1560 MiB63.0 FLOP/byte

Cường độ tính toán của cả ba tile đầu ra đều thấp xa so với tỉ lệ 295 FLOP/byte giữa đỉnh nhân ma trận của H100 và băng thông bộ nhớ trong Mục 5.1.5. Chỉ nhìn lưu lượng ở mức giao diện này, cả ba cách tiling đều nằm về phía bị băng thông giới hạn, với trần MFU của khối nhân ma trận chỉ lần lượt là 5%, 11% và 21%. Phóng to tile đầu ra đẩy cường độ tính toán về phía điểm giao cắt, nhưng riêng tiling ở mức này không thể vượt qua nó — phần tái sử dụng còn lại phải đến từ các tầng lưu trữ gần khối tính toán hơn.

Khối lượng truy cập gần như giảm một nửa, nhưng thời gian có giảm theo hay không còn tùy thuộc vào lượng công việc accelerator có thể lập lịch chạy đồng thời. Trên GPU, buffer đầu vào nằm trong shared memory, và bộ tích lũy có thể nằm trong shared memory hoặc thanh ghi. Ở đây ta theo cách tiếp cận của Mục 5.1.5: mỗi thread block xử lý một tile đầu ra, cả buffer đầu vào lẫn bộ tích lũy đều trong shared memory. Lấy một SM đơn của RTX PRO 6000 làm phần cứng: card này thuộc compute capability 12.x, giới hạn shared memory mỗi SM là 100 KB (tức 102,400 byte), và mỗi thread block cũng dành riêng 1 KiB.

Hình 5-9 Mọi thanh đều biểu diễn 100 KB shared memory mỗi SM của RTX PRO 6000. Working set 24 KiB, cộng với khoản dành riêng 1 KiB mỗi block, vừa đủ cho bốn bản sao; working set 80 KiB chỉ vừa một bản sao.

Các thread block cư trú đồng thời có thể xen kẽ các yêu cầu bộ nhớ: trong khi một số block chờ dữ liệu, các block khác vẫn tiếp tục làm việc. Khi số bản sao working set cư trú đồng thời giảm từ bốn xuống một, có ít tính toán khác để chạy trong lúc chờ bộ nhớ hơn. Vậy trong khi tile lớn hơn đọc ít dữ liệu hơn, băng thông hiệu dụng thực tế đạt được có thể giảm. Mục 5.1.5 đã tính ra điều này: các warp cư trú cung cấp các yêu cầu bộ nhớ đang trên đường đi, và khi số byte đang bay trên đường không đủ, băng thông hiệu dụng tụt xuống dưới băng thông giao diện.

Giả sử cả hai cách triển khai đều bị băng thông giới hạn ở mức này, với băng thông hiệu dụng lần lượt là B64B_{64} và B128B_{128}:

T128T64=15603096B64B128.\frac{T_{128}}{T_{64}}=\frac{1560}{3096}\frac{B_{64}}{B_{128}}.

Nếu hai băng thông bằng nhau, thời gian giảm khoảng một nửa; nếu tile lớn hơn chỉ đạt 40% băng thông ban đầu, tỉ số thời gian trở thành khoảng 0.504/0.4=1.260.504/0.4=1.26 — thay vào đó thời gian tăng khoảng 26%. Ta có thể trước hết loại bỏ các phương án không vừa trong dung lượng, sau đó thu hẹp phạm vi so sánh bằng khối lượng truy cập, và cuối cùng đo hiệu quả thực thi thực tế. Mục 5.4 sẽ trao quy trình này cho trình biên dịch.

5.2.3 Bố trí dữ liệu: dàn trải các yêu cầu đồng thời qua các bank khác nhau

Trong mục trước, tile lớn hơn đọc ít dữ liệu hơn nhưng có thể chậm hơn do độ song song thấp. Độ song song không chỉ phụ thuộc vào số thread block cư trú, mà còn vào việc các yêu cầu đọc mà các thread này phát ra có thể được phục vụ đồng thời hay không. Hãy xét bố trí bộ nhớ cục bộ trước. Lấy một tile nháp FP32 32×32, bộ nhớ cục bộ của nó được chia thành 32 bank; mỗi bank cung cấp một từ 32-bit mỗi chu kỳ, và số hiệu bank là địa chỉ từ lấy modulo 32. Mục 5.1.5 đã giới thiệu rằng 32 thread của một warp là 32 lane; khi 32 lane mỗi lane đọc một từ trong cùng một hàng, các yêu cầu dàn trải qua 32 bank; khi đọc một cột, bước hàng là 32 từ, nên mọi địa chỉ dồn vào cùng một bank, cần 32 chu kỳ.

Hình 5-10 32 từ khác nhau trong cùng một cột được 32 lane yêu cầu đồng thời. Ở nửa trên, bước hàng là 32 từ, các yêu cầu dồn vào cùng một bank; ở nửa dưới, thêm padding thành 33 từ, các yêu cầu dàn trải qua 32 bank. Hình chỉ vẽ bốn yêu cầu đầu tiên; giả sử mỗi bank cung cấp một từ 32-bit mỗi chu kỳ, bỏ qua broadcast.

Thêm một từ đệm vào cuối mỗi hàng sẽ đổi bước hàng thành 33. Số hiệu bank cho các hàng liền kề trong cùng một cột khi đó tăng thêm một đơn vị mỗi lần, nên 32 yêu cầu dàn trải qua 32 bank và hoàn thành trong một chu kỳ. Không gian chiếm dụng tăng từ 4096 byte lên 4224 byte, chỉ tăng khoảng 3.1%. Cả dữ liệu hiệu dụng lẫn số lần đọc đều không thay đổi, nhưng một phép đọc cột vốn cần 32 chu kỳ nay hoàn thành trong một chu kỳ.

Từ thừa ở cuối mỗi hàng này là padding — nó không thay đổi bản thân phép tính, chỉ dàn trải các yêu cầu đồng thời hiện có qua các bank khác nhau.

Padding dùng thêm byte để đổi lấy sự dàn trải. Ta cũng có thể tránh thêm byte và chỉ thay đổi nơi dữ liệu được lưu: đặt phần tử ở hàng rr, cột cc vào cột c⊕rc\oplus r, trong đó ⊕\oplus là XOR theo bit. Bước hàng vẫn là 32 từ, nhưng khi đọc một cột, các chỉ số cột mà 32 lane lấy ra là 32 giá trị phân biệt của c⊕rc\oplus r khi rr thay đổi, nên các số hiệu bank khác nhau nhau, cũng hoàn thành trong một chu kỳ. Thư viện mẫu nhân ma trận CUTLASS của NVIDIA gọi kiểu thay đổi vị trí lưu trữ này là swizzle. Swizzle tiết kiệm không gian padding, đổi lấy chi phí tính chỉ số cột bằng công thức mỗi lần truy cập.

5.2.4 Tách theo chiều nào: các tác vụ song song và tổng từng phần

Padding giải quyết việc các yêu cầu có thể được phục vụ đồng thời hay không. Nếu chương trình ban đầu chỉ sắp xếp một số ít tác vụ, thay đổi bố trí địa chỉ là chưa đủ — bản thân phép tính phải được tách ra. Lấy RMSNorm làm ví dụ: trước tiên tính một hệ số tỉ lệ dùng chung từ giá trị bình phương trung bình của các phần tử trong một hàng, sau đó chia tỉ lệ hàng đó. Tính hệ số tỉ lệ là một phép rút gọn (reduction): nó đến từ tổng bình phương của cả hàng; sau đó, mỗi phần tử được nhân với hệ số đó và tham số học được gamma tương ứng. Với đầu vào BF16 [1024,4096][1024,4096], nếu một nhóm xử lý một hàng, 1024 nhóm có thể được lập lịch độc lập; nhưng trong giai đoạn decode, một yêu cầu đơn chỉ sinh ra một token tại một thời điểm, nên cách tách như trên chỉ tạo ra một nhóm, khiến hầu hết tài nguyên thực thi của accelerator rảnh rỗi.

Tách mỗi hàng thành tám đoạn 512 phần tử làm tăng số tác vụ song song ở giai đoạn đầu. Phép rút gọn đầy đủ khi đó trở thành ba giai đoạn: mỗi đoạn tạo ra một tổng bình phương cục bộ, tám tổng cục bộ được gộp lại để lấy hệ số tỉ lệ của hàng, và hệ số đó sau đó được dùng để chuẩn hóa đầu vào. Hai giai đoạn đầu chỉ truyền đi một lượng dữ liệu nhỏ, nhưng giai đoạn cuối phải đọc lại dữ liệu gốc.

Hình 5-11 Hàng trên có một nhóm xử lý trọn một hàng, giữ đầu vào cư trú cho đến khi chuẩn hóa xong; hàng dưới tách một hàng thành tám đoạn, tính các tổng cục bộ trước, gộp chúng lại, rồi đọc lại đầu vào để chuẩn hóa. Mỗi hàng có 4096 phần tử BF16; hình chỉ vẽ một hàng, và mở rộng lên 1024 hàng làm tăng số tác vụ giai đoạn đầu từ 1024 lên 8192. Hộp màu cam đánh dấu lần đọc thêm vào đầu vào gốc.

Đếm dọc theo các mũi tên trong Hình 5-11 sẽ tính ra chi phí truy cập bộ nhớ của việc tăng số tác vụ song song. Vẫn lấy 1024 token, với vector đặc trưng của mỗi token chiếm một hàng của ma trận đầu vào. Số tác vụ song song giai đoạn đầu tăng từ 1024 nhóm lên 8192 nhóm; 8192 tổng cục bộ FP32 chiếm 32 KiB, và 1024 hệ số tỉ lệ FP32 chiếm 4 KiB. Chương trình với một nhóm mỗi hàng giữ đầu vào cư trú trong bộ nhớ cục bộ, đọc 8 MiB đầu vào, đọc lại một gamma BF16 cho mỗi hàng tổng cộng 8 MiB, và cuối cùng ghi ra 8 MiB kết quả, tổng khối lượng truy cập là 8+8+8=248+8+8=24 MiB. Sau khi tách, giai đoạn chuẩn hóa cuối đọc lại 8 MiB đầu vào; việc ghi và đọc các tổng cục bộ và hệ số tỉ lệ cộng thêm khoảng 0.1 MiB, tổng cộng 32.1 MiB, và số lần khởi chạy kernel tăng từ một lên ba.

Trong hình, mỗi nhóm chỉ xử lý một phần tám hàng, mang lại một lợi ích khác: cần giữ ít dữ liệu tạm hơn. Lưu cả một hàng ở dạng FP32 cần 16 KiB, trong khi lưu một phần tám hàng chỉ cần 2 KiB. Dữ liệu tạm vốn phải tràn sang bộ nhớ xa hơn do không đủ thanh ghi giờ có thể ở lại trong thanh ghi. Lợi ích của việc tách phép rút gọn đến từ hai nguồn: nhiều tác vụ song song hơn, và ít dữ liệu tạm hơn mỗi nhóm; cái giá phải trả là một lượt đi thêm qua đầu vào, cộng với việc phải truyền dữ liệu giữa các giai đoạn.

Cách tách của RMSNorm là dọc theo chiều rút gọn: tổng bình phương của một hàng được tách thành tám tổng cục bộ, sau đó phải được gộp thêm một lần nữa. Phép nhân ma trận cũng có lựa chọn tương tự. Với CM×N=AM×KWK×NC_{M\times N}=A_{M\times K}W_{K\times N}, công việc có thể được tách theo bất kỳ chiều nào trong ba chiều cho các bộ thực thi khác nhau. Hình 5-12 cho thấy mỗi cách tách cần những đầu vào nào và tạo ra kết quả gì.

Hình 5-12 Tách theo hàng đầu ra M hoặc theo cột đầu ra N cho mỗi bộ thực thi một mảnh kết quả riêng biệt, đầy đủ, được nối ghép lại sau đó khi cần; tách theo chiều rút gọn K cho mỗi bộ thực thi một tổng từng phần của cùng một đầu ra, phải cộng lại với nhau trước khi dùng. Các hộp biểu diễn quyền sở hữu dữ liệu và không vẽ theo tỉ lệ với số phần tử ma trận.

Hai cách tách đầu tiên chỉ cần các đầu vào được sao chép hoặc phân mảnh khi cần, và đầu ra có thể giữ nguyên dạng phân mảnh cho đến khi một operator phía sau cần dữ liệu đầy đủ; cách tách thứ ba tiết kiệm việc sao chép đầu vào, nhưng, giống các tổng cục bộ của RMSNorm, thêm một bước gộp. Tile đầu ra của Mục 5.2.2 tách đồng thời theo cả MM lẫn NN, rồi tích lũy từng khối dọc theo KK: bộ tích lũy ở lại trong một thread block, nên việc gộp tổng từng phần hoàn toàn không vượt qua ranh giới nào.

Thứ tự gộp cũng làm thay đổi kết quả. Phép cộng dấu phẩy động không thỏa mãn tính kết hợp: tính 1+2−24+2−241+2^{-24}+2^{-24} trong FP32, cộng hai số hạng đầu trước thì tổng vẫn là 1, và cộng thêm số hạng thứ ba thì tổng vẫn là 1; cộng hai số hạng cuối trước ta được 2−232^{-23}, rồi cộng với 1 cho 1+2−231+2^{-23}. Hai thứ tự khác nhau một ULP (unit in the last place, khoảng cách giữa hai số dấu phẩy động liền kề ở một số mũ nhất định). Các bit mất đi trong một lần làm tròn không thể phục hồi, và mọi phép cộng tiếp theo đều tiếp tục từ một số đã được làm tròn.

RMSNorm của mục này giúp định lượng được mức chênh lệch này lớn đến đâu. Lấy một hàng gồm 4096 phần tử BF16, trong đó bình phương của mỗi phần tử biểu diễn chính xác trong FP32, nên mọi khác biệt giữa các cách tiếp cận dưới đây chỉ đến từ thứ tự cộng. Nửa trên của Hình 5-13 so sánh các tổng bình phương tính với sáu cách tách khác nhau — 1, 2, 4, 8, 16, 32 đoạn — với trục hoành hiển thị khoảng cách của chúng so với giá trị chính xác. Khi một nhóm xử lý cả hàng, bộ tích lũy tăng dần từng bước tới khoảng 2038, trong khi mỗi số hạng cộng vào chỉ khoảng 0.5, nên các bit thấp liên tục bị làm tròn mất đi, kết cục thấp hơn giá trị chính xác 175.2 ULP. Sau khi tách thành tám đoạn, mỗi bộ tích lũy chỉ tăng tới khoảng 255, và tám tổng cục bộ chỉ được cộng lại một lần ở cuối, cho kết quả cách giá trị chính xác chỉ 9.2 ULP. Hai đường đi khác nhau 166 ULP, chênh lệch tương đối khoảng 9.9×10−69.9\times10^{-6}. Vậy việc tách không chỉ thay đổi số tác vụ song song, mà còn thay đổi cả độ chính xác.

Nửa dưới cố định tám tổng cục bộ này và chỉ thay đổi thứ tự gộp. 40,320 thứ tự khả dĩ chỉ cho ra ba kết quả FP32 phân biệt, mỗi kết quả khác nhau một ULP. Nếu một phép cộng nguyên tử (atomic add) tích lũy các tổng cục bộ vào cùng một địa chỉ, thứ tự đến được quyết định bởi lập lịch; chạy cùng một đầu vào qua cùng một chương trình hai lần khi đó có thể cho các kết quả khác nhau. Thay vào đó, ghi tám tổng cục bộ vào một buffer, chờ cho tất cả được ghi xong, rồi để một bộ thực thi cộng chúng theo một thứ tự cố định sẽ luôn cho cùng một kết quả. Điều này tốn hai thứ: một lượt ghi ra và đọc lại; và việc gộp phải chờ đoạn chậm nhất hoàn thành mới bắt đầu được, trong khi atomic add tích lũy theo thứ tự đến mà không cần chờ tất cả sẵn sàng.

Hình 5-13 Khoảng cách của tổng bình phương của cùng một hàng 4096 phần tử BF16 so với giá trị chính xác, trục hoành theo đơn vị ULP của FP32. Nửa trên so sánh sáu cách tách từ 1 đến 32 đoạn; nửa dưới cố định tám tổng cục bộ và liệt kê kết quả của mọi thứ tự gộp khả dĩ. Giá trị chính xác được tính bằng số hữu tỉ, chỉ dùng làm đường cơ sở so sánh.

Sự khác biệt không dừng ở tổng bình phương. Tiếp tục phép tính với rms_norm_eps của Qwen3-8B, hệ số tỉ lệ của hai đường đi khác nhau 59 ULP, phần tử đầu ra đầu tiên của hàng khác nhau 50 ULP, và sự khác biệt này lan truyền cùng vào lớp kế tiếp.

Vì vậy, lựa chọn cách tách không chỉ là chuyện tốc độ. Như đã nêu ở đầu mục này, tách thành bao nhiêu đoạn phụ thuộc vào số tác vụ song song có thể sắp xếp: với 1024 hàng, một nhóm mỗi hàng là đã đủ; với chỉ một hàng trong decode, bắt buộc phải tách. Cùng một kernel khi đó sẽ thay đổi số đoạn theo kích thước batch, nên cùng một hàng đầu vào sinh ra đầu ra khác nhau dưới các kích thước batch khác nhau. Để đầu ra chỉ phụ thuộc vào đầu vào, cả số đoạn lẫn thứ tự gộp phải được giữ cố định bất kể kích thước batch, đổi lấy việc mất đi độ song song tăng thêm từ việc tách khi batch nhỏ. Mục 10.5.3 dùng đúng điểm này để giải thích vì sao generation và training nạp cùng các trọng số, mà logprob tính từ cùng một tiền tố vẫn có thể khác nhau.

Tiling cũng có thể vượt qua ranh giới của một card đơn. Trong một card, việc tính các tile này được xử lý phối hợp bởi các thread block, shared memory và các phép rút gọn cục bộ; một khi được tách sang nhiều card, các bộ thực thi tách theo một chiều đầu ra mỗi cái phải tự lấy các đầu vào của mình, và các bộ thực thi tách theo chiều rút gọn phải gửi các tổng từng phần đi để cộng lại — toàn bộ sự di chuyển này đi qua kết nối liên card. Sự phụ thuộc toán học không thay đổi; thay đổi là khoảng cách và chi phí di chuyển dữ liệu.

Lập lịch (scheduling), đến lượt nó, quyết định thứ tự thực thi, vị trí đặt, việc đệm và sự chồng lấp sau tiling. Ví dụ, sau khi tách theo KK, việc chờ mọi đầu vào đến rồi mới tính, hay tích lũy dần khi dữ liệu đến, làm thay đổi mức chiếm dụng lưu trữ và thời gian chờ thực tế. Samples, sequences, features, layers và experts trong một model cung cấp thêm các hướng phân chia công việc; trong số đó, layers là các giai đoạn nối tiếp trong đồ thị tính toán, còn experts là các nhánh được chọn theo routing. Mục 6.1.4 sẽ ánh xạ các hướng phân chia công việc này lên ba cách tách ở đây, và Mục 6.2 sẽ lần lượt suy ra lượng truyền thông cho từng cách.

Tiling ma trận, padding bank và tách phép rút gọn đều là các cách sắp xếp lại cùng một khối công việc toán học. Tiling giúp nhiều đầu ra tái sử dụng đầu vào, padding dàn trải các yêu cầu đọc qua các bank khác nhau, và tách phép rút gọn làm tăng số tác vụ chạy song song. Khi phân tích một chương trình, cần tính toán chính xác cả lợi ích lẫn chi phí của từng thay đổi, và cũng kiểm tra xem nó có làm thay đổi kết quả dấu phẩy động hay không. Sự so sánh giữa truy cập liên tục và rút gọn hàng rộng trong các thí nghiệm kèm theo cung cấp các bản ghi thực thi tương ứng.

5.3 Fusion, đệm và thực thi theo pipeline giữa các operator liền kề

Trong một operator đơn, bộ nhớ cục bộ có thể được dùng để tái sử dụng đầu vào. Cơ hội tương tự cũng tồn tại giữa các operator liền kề: đầu ra của một operator thường ngay lập tức trở thành đầu vào của operator kế tiếp. Nếu operator kế tiếp có thể đọc kết quả trực tiếp ra khỏi bộ nhớ cục bộ, nó có thể bỏ qua một lần ghi ra và đọc lại. Để làm được điều này, ta cần xác định operator kế tiếp sẽ dùng dữ liệu nào, và khi nào toàn bộ dữ liệu đó sẵn sàng.

5.3.1 Tensor trung gian và ranh giới fusion

Giờ mở rộng ví dụ đang chạy sang hai phép chiếu của FFN. Đặt G=XWgG=XW_g, U=XWuU=XW_u, với chuỗi kích hoạt tính Z=SiLU⁡(G)⊙UZ=\operatorname{SiLU}(G)\odot U, sau đó chuyển cho phép chiếu down để đưa kết quả trở lại độ rộng ẩn. G, U và Z đều có hình dạng [1024,12288][1024,12288], mỗi cái chiếm 24 MiB khi lưu ở dạng BF16.

Cả SiLU lẫn ⊙\odot đều là các phép toán elementwise. Z[i,j] chỉ phụ thuộc vào G[i,j] và U[i,j], không có phép rút gọn nào qua các vị trí. Vậy một khi hai phần tử ở cùng một vị trí được đọc vào, kết quả cuối tại vị trí đó có thể được tính ngay, không cần tính trước toàn bộ kết quả SiLU.

Trước tiên, so sánh xem kết quả trung gian T=SiLU⁡(G)T=\operatorname{SiLU}(G) đi về đâu. Khi thực thi tách rời, toàn bộ T phải được ghi ra tầng lưu trữ kế tiếp rồi được phép nhân đọc lại; khi fusion, một mẩu nhỏ kết quả SiLU có thể được chuyển thẳng cho phép nhân trong cùng một kernel, qua đó giải phóng không gian cục bộ đã dùng. Điều này giống như chuyển vật liệu vừa gia công trên bàn làm việc thẳng cho bước tiếp theo, bỏ qua vòng khứ hồi gửi đi rồi lấy lại. Tương ứng ở đây, thứ được bỏ qua là việc ghi ra và đọc lại kết quả trung gian; cả hai phép toán vẫn đều phải được hoàn thành.

Hình 5-14 Ở nửa trên, toàn bộ T đi qua một lần ghi ra và một lần đọc lại; ở nửa dưới, mảnh cục bộ t được chuyển thẳng cho phép nhân. Cả hai cách đều vẫn đọc G, U và ghi Z; hình lược bỏ các cạnh vào/ra dùng chung này. Sau fusion, các quy tắc làm tròn gốc cho kết quả trung gian vẫn được giữ nguyên.

Ví dụ 5-4: Fusion các operator kích hoạt giảm được bao nhiêu đọc/ghi tensor trung gian? Khi thực thi tách rời, trước tiên tính T=SiLU⁡(G)T=\operatorname{SiLU}(G), đọc G và ghi T, tổng cộng 48 MiB; sau đó tính Z=T⊙UZ=T\odot U, đọc T và U rồi ghi Z, tổng cộng 72 MiB. Tổng khối lượng đọc/ghi là 120 MiB. Fusion hai bước lại với nhau, mỗi lần đọc một lát G và U, hoàn thành SiLU và phép nhân cục bộ, rồi ghi ra Z, chỉ cần 72 MiB.

48 MiB tiết kiệm được đúng bằng một lần ghi và một lần đọc T. Theo mô hình băng thông của mục trước, giả sử cả hai chương trình có cùng băng thông hiệu dụng, tỉ số thời gian là 72/120=0.672/120=0.6, tăng tốc khoảng 1.7×. Chuỗi kích hoạt đo được giảm từ khoảng 14.2 μs xuống 7.7 μs, tăng tốc đo được khoảng 1.8×. Cả tính toán lý thuyết lẫn kết quả đo được đều chỉ ra cùng một nguồn lợi ích: bỏ qua việc ghi ra và đọc lại toàn bộ T.

Tiếp theo, lượng tử hóa Z với một hệ số tỉ lệ cố định, nên mỗi phần tử chỉ chiếm 1 byte. Chuyển đổi độc lập đòi hỏi đọc 24 MiB Z và ghi ra 12 MiB kết quả lượng tử hóa, thêm 36 MiB. Tổng cho cả ba bước hoàn toàn tách rời là 156 MiB; fusion hai bước đầu cho 108 MiB; fusion cả ba bước cho 24+24+12=6024+24+12=60 MiB. Hình 5-15 vẽ các đọc/ghi đầu vào/đầu ra bắt buộc ở phía trái mỗi thanh, rồi lần lượt cộng thêm đọc/ghi của hai đại lượng trung gian.

Hình 5-15 Mỗi đại lượng trung gian 24 MiB không được giữ lại sẽ tiết kiệm một lần ghi và một lần đọc, tổng cộng 48 MiB. G, U là BF16, đầu ra cuối chiếm 1 byte/phần tử, và hệ số tỉ lệ lượng tử hóa được cho trước. Mọi sơ đồ đều dùng cùng các phép toán và cùng thứ tự làm tròn.

Ở đây, dung lượng và khối lượng truy cập của các tensor trung gian là hai điều khác nhau. Giả sử mỗi bước cấp phát đầu ra trước, rồi giải phóng đầu vào khi nó không còn cần nữa. Chạy SiLU độc lập, G, U và T chiếm bộ nhớ đồng thời, đỉnh 72 MiB; fusion cả ba bước, G, U và đầu ra cuối chiếm bộ nhớ đồng thời, đỉnh 60 MiB. Đọc/ghi giảm 96 MiB, nhưng đỉnh chỉ giảm 12 MiB, vì cùng một khối không gian lưu trữ có thể được đọc và ghi đi ghi lại.

Vì mỗi đầu ra của một phép toán elementwise chỉ phụ thuộc vào đầu vào ở vị trí tương ứng, phép tính có thể xảy ra ngay sau khi đọc, giảm lượng kết quả trung gian cần giữ lại. Để sự tái sử dụng này diễn ra trơn tru, hai bước cũng cần dùng các bố trí dữ liệu tương thích với nhau, và buffer không được ghi đè trước khi lần đọc cuối của nó hoàn thành. Dưới đây, ta theo dõi vòng đời của một khối dữ liệu để phân tích hai điều kiện này.

5.3.2 Chuyển đổi bố trí, tái sử dụng buffer và đệm kép

Các operator liền kề đôi khi dùng các bố trí khác nhau. Operator phía trước có thể tự nhiên ghi ra từng hàng một, trong khi operator phía sau cần đọc vào theo một thứ tự khối khác. Tạo một bản sao sắp xếp lại 24 MiB riêng nghĩa là đọc kết quả gốc và ghi bản mới — tổng cộng 48 MiB. Còn nếu operator phía trước có thể ghi trực tiếp theo bố trí mà phép tính tiếp theo cần, bước sắp xếp lại riêng sẽ được loại bỏ. So sánh tổng thời gian của “tính rồi sắp xếp lại” với “tính đồng thời ghi trực tiếp theo bố trí mới” sẽ cho biết cách nào nhanh hơn.

Buffering cũng quyết định liệu hai giai đoạn có thể chồng lấp hay không. Giả sử một bộ di chuyển (mover) đặt khối 0 vào ngăn A, và khối tính toán sau đó đọc A. Trong lúc khối tính toán đọc A, bộ di chuyển đặt khối 1 vào ngăn B; khi tính toán chuyển sang B, khối 2 có thể được nạp vào A. Với hai ngăn luân phiên như vậy, bộ di chuyển có thể ghi khối kế tiếp trong khi khối tính toán đọc khối hiện tại.

Lưu ý rằng “dữ liệu đã di chuyển xong” và “dữ liệu đã được dùng xong” cho ngăn A là hai thời điểm khác nhau. Tính toán chỉ bắt đầu đọc sau khi việc di chuyển hoàn thành; cũng chỉ sau lần đọc cuối hoàn thành, khối kế tiếp mới có thể được ghi vào ngăn A, ghi đè lên nội dung cũ. Tái dùng các con số SM đơn của H100 trong Mục 5.1.5: mỗi khối mất 1.29 μs để di chuyển và 0.28 μs để tính, với bộ di chuyển và khối tính toán hoạt động độc lập. Hình 5-16 cho thấy trạng thái của mỗi ngăn qua ba khoảng thời gian, và Hình 5-17 vẽ toàn bộ dòng thời gian di chuyển và tính toán dựa trên đó.

Hình 5-16 Khối 0 chiếm ngăn A trong 1.29–1.57 μs, sau đó ngăn A được giải phóng; khối 2 phải chờ cho đến khi bộ di chuyển di chuyển xong khối 1 tại thời điểm 2.58 μs mới bắt đầu ghi vào A. Hình vẽ ba khoảng thời gian; trong 1.57–2.58 μs, không ngăn nào chứa dữ liệu sẵn sàng tính toán, nên khối tính toán ngồi không. Màu sắc cố định danh tính của ngăn A và B; chữ ghi chú khối nào đang được đọc hoặc ghi.

Ví dụ 5-5: Đệm kép chồng lấp di chuyển dữ liệu với tính toán như thế nào? Mỗi khối mất 1.29 μs để di chuyển và 0.28 μs để tính. Bộ di chuyển và khối tính toán hoạt động độc lập, và ta coi chi phí sự kiện là bằng không. Thực thi bốn khối tuần tự mất 4×(1.29+0.28)≈6.284\times(1.29+0.28)\approx6.28 μs.

Với hai ngăn, khối 0 di chuyển vào trong 0–1.29 μs, và được tính trong 1.29–1.57 μs; khối 1 di chuyển vào đồng thời trong 1.29–2.58 μs. Khối 0 tính xong ở 1.57 μs, ngăn A được giải phóng ngay, nhưng bộ di chuyển không thể di chuyển xong khối 1 và bắt đầu ghi khối 2 cho đến 2.58 μs. Sau đó bộ di chuyển làm việc liên tục, chuyển một khối mỗi 1.29 μs, trong khi khối tính toán chỉ bận 0.28 μs mỗi lần. Khối thứ tư hoàn thành ở 5.44 μs — chỉ giảm khoảng 13% thời gian.

Hình 5-17 Hai ngăn đầu vào luân phiên để di chuyển và tính toán chồng lấp lên nhau. Trên một SM của H100, mỗi khối mất 1.29 μs để di chuyển và 0.28 μs để tính, với tài nguyên độc lập và bỏ qua chi phí đồng bộ; cùng một màu biểu thị cùng một ngăn. Sự chồng lấp chỉ che giấu được phép tính — thời điểm hoàn thành bị quyết định bởi chuỗi di chuyển nối tiếp nhau.

Tổng quát hóa từ bốn khối lên n khối: khối đầu tiên phải di chuyển trước, sau đó mỗi khối hoàn thành sau mỗi khoảng thời gian bằng thời gian của giai đoạn chậm hơn, và khối cuối vẫn cần phép tính của nó hoàn thành. Vậy:

Tpipe=tm+(n−1)max⁡(tm,tc)+tc.T_{\mathrm{pipe}}=t_m+(n-1)\max(t_m,t_c)+t_c.

Ở đây tmt_m là thời gian di chuyển mỗi khối và tct_c là thời gian tính mỗi khối. Trong ví dụ vừa đưa, tm>tct_m>t_c, bộ di chuyển bận liên tục từ 0 đến 5.16 μs, và tổng thời gian là 4×1.29+0.28≈5.444\times1.29+0.28\approx5.44 μs — khối tính toán chỉ bận khoảng 20% thời gian. Thêm một ngăn thứ ba không giúp gì: hai ngăn đã giữ cho bộ di chuyển không có thời gian rảnh, nên thêm ngăn không thể tăng tốc được việc di chuyển, và tăng tốc tính toán cũng vô dụng tương tự. Hai ngăn là cấu hình tối thiểu cho phép chồng lấp; mở rộng số ngăn lên n, cho phép di chuyển đi trước n−1 khối, được gọi là n stages trong CUTLASS. Trong ví dụ này, bộ di chuyển đã làm việc liên tục, nên các stage thêm vào chỉ tiêu tốn thêm dung lượng; chỉ khi thời gian di chuyển mỗi khối không ổn định, hoặc khi độ trễ của một lần di chuyển vượt quá thời gian tính một khối, các stage thừa mới lấp được khoảng trống — và dung lượng cần thiết tăng tỉ lệ thuận với số stage. Để rút ngắn thời gian hơn nữa, các lựa chọn duy nhất là giảm số byte di chuyển mỗi khối — ví dụ, bằng cách phóng to khối đầu ra như trong Mục 5.2.2 để mỗi byte đầu vào tham gia vào nhiều phép nhân–cộng hơn, hoặc để các thread block đang thực thi đồng thời dùng chung các khối đầu vào trong L2 — hoặc chuyển sang một card có băng thông bộ nhớ cao hơn.

Quay lại tình huống sắp xếp lại ở đầu mục này: khi hai giai đoạn truy cập dữ liệu theo các thứ tự khác nhau, buffer phải giữ nhiều dữ liệu chưa được dùng hơn. Ba khối 16 KiB xếp hàng chiếm 48 KiB, cộng với 32 KiB không gian sắp xếp lại, tổng cộng 80 KiB. Một thread block trên RTX PRO 6000 có thể dùng tối đa 99 KB shared memory, vừa khít khi cư trú một mình; nhưng nếu một SM phải chứa hai thread block như vậy cùng lúc, mỗi block chỉ nhận tối đa 49 KiB (100 KB trừ đi 1 KiB dành riêng mỗi block, chia đều). Sau khi trừ đi 32 KiB không gian sắp xếp lại, chỉ còn lại một ngăn xếp hàng 16 KiB: trước khi ghi khối thứ hai, bộ di chuyển phải chờ khối tính toán dùng xong và giải phóng ngăn đó. Cơ chế này — khi phía hạ nguồn không theo kịp và buộc phía thượng nguồn phải tạm dừng — được gọi là backpressure. Chọn một bố trí dùng chung làm giảm việc xếp hàng; phóng to buffer cho phép tích lũy nhiều dữ liệu đã được ghi nhưng chưa được dùng hơn.

Đầu vào từ host trong Mục 5.1 có thể được tổ chức theo cùng cách. Một luồng sao chép nạp batch kế tiếp trong khi một luồng tính toán xử lý batch hiện tại; luồng tính toán chờ một sự kiện “sao chép xong” trước khi đọc đầu vào mới, và luồng sao chép chờ một sự kiện “tính xong” trước khi ghi đè buffer cũ. Hai sự kiện này tương ứng đảm bảo rằng đầu vào đã sẵn sàng và dữ liệu cũ đã được tiêu thụ; lợi ích về thời gian đến từ việc hai tài nguyên làm việc đồng thời.

5.3.3 FlashAttention: Tiling và Softmax trực tuyến

Các phép toán elementwise được fusion trong Mục 5.3.1 đều có đầu ra chỉ phụ thuộc vào đầu vào ở cùng vị trí; nhưng Softmax trong attention cần cả một hàng điểm số. Để tính một hàng đầu ra attention, ta có thể trước tiên biến điểm số của mỗi vị trí thành một trọng số chưa chuẩn hóa, dùng các trọng số đó để tính tổng có trọng số của các giá trị, và cuối cùng chia cho tổng các trọng số. Ta có thể xử lý từng khối một, ghi lại đóng góp của nó vào tổng có trọng số và vào tổng trọng số, rồi chuyển sang khối kế tiếp; miễn là các đóng góp gộp lại đúng cách, không cần giữ toàn bộ các điểm số đã xử lý trước đó. Điểm khó là Softmax lấy hàm mũ dựa trên các điểm số, và khi một điểm số lớn hơn xuất hiện ở phía sau, mốc chuẩn số học được dùng cho các kết quả đã tích lũy từ trước phải được điều chỉnh.

FlashAttention sắp xếp việc lưu trữ quanh cách tính từng khối như vậy, giảm các lần đọc/ghi bộ nhớ của các ma trận trung gian attention cho các chuỗi dài. Phương pháp này được đề xuất năm 2022 bởi Tri Dao và các đồng nghiệp tại Đại học Stanford cùng các cộng tác viên tại Đại học Buffalo, SUNY.

Đọc về FlashAttention, tác giả nhớ lại một trải nghiệm với AKG hồi năm 2019. Khi đó, để fusion operator Softmax, tác giả đã tìm kiếm rộng rãi một thuật toán trực tuyến phù hợp, và cuối cùng tìm thấy một công trình nghiên cứu do NVIDIA công bố năm 2018, nhờ đó kết hợp được khả năng tiling và fusion của AKG để xâu chuỗi phép nhân ma trận phía trước với Softmax phía sau. Lúc bấy giờ tác giả thậm chí còn chưa hiểu cơ chế attention, và chỉ sau này mới nhận ra rằng điều đó tình cờ chạm vào một trong những ý tưởng cốt lõi của FlashAttention. Để từ bước đó đi tới FlashAttention đầy đủ vẫn cần gấp cả phép tính tổng có trọng số với V ở sau vào cùng một bản cập nhật trực tuyến, và sắp xếp việc lưu trữ trên chip cùng với việc di chuyển dữ liệu quanh toàn bộ phép tính attention. Suy luận dưới đây cho thấy phép thay đổi thuật toán tưởng như cục bộ này loại bỏ các ma trận trung gian khổng lồ như thế nào.

Quay lại bản thân phép tính attention. Attention chuẩn trước tiên tính S=QKT/dS=QK^\mathsf{T}/\sqrt d, rồi tính P=softmax⁡(S)P=\operatorname{softmax}(S), và cuối cùng tính O=PVO=PV.

Trao ba phép toán này cho các operator đa dụng tách rời cho phép ta dựng chương trình trực tiếp theo công thức, với các operator trao đổi kết quả qua các tensor trung gian. Một khi hiểu được đầy đủ các phụ thuộc dữ liệu của phép tính attention, ta có thể áp dụng một chiến lược biên dịch và thực thi khác: khai thác cấu trúc rút gọn của Softmax để vượt qua ranh giới của các operator độc lập và tính liên tục trên các khối dữ liệu nhỏ. Dưới đây ta trước hết tính chi phí của việc giữ nguyên các tensor trung gian như cách làm gốc, sau đó suy ra trạng thái phải được giữ lại cho việc tính từng khối.

Lấy một head, với độ dài chuỗi L=8192L=8192 và chiều head d=128d=128, tính attention trên toàn bộ token. Q, K, V, O dùng BF16 và mỗi cái chiếm 2 MiB; S và P đầy đủ dùng FP32 và mỗi cái chiếm 256 MiB. Q, K, V mỗi cái được đọc một lần và O được ghi một lần, tổng cộng 8 MiB; S và P mỗi cái được ghi một lần và đọc một lần, nhưng riêng việc này đã chuyển 1 GiB dữ liệu. Các ma trận trung gian tăng theo bình phương độ dài chuỗi, gây ra khối lượng đọc/ghi lớn này.

Hình 5-18 Sơ đồ trên giữ đầy đủ S và P; hai ma trận FP32 mỗi cái chiếm 256 MiB, và tổng viết cộng đọc là 1 GiB. Sơ đồ dưới chỉ truyền tiếp giá trị max m, tổng mũ ℓ và giá trị có trọng số u của phần đã xử lý; sau khi một khối được xử lý, buffer điểm số của nó có thể được tái sử dụng. Các mũi tên tóm tắt thứ tự xử lý; phép tính đầy đủ vẫn cần đọc vào Q, K, V.

Để loại bỏ 1 GiB đọc/ghi trung gian này, ta có thể làm theo cách tiếp cận trong Mục 5.3.1; nhưng áp dụng trực tiếp fusion mức elementwise vấp phải một trở ngại: mẫu số của Softmax là tổng các hàm mũ trên toàn bộ hàng. Mới chỉ đọc khối trước đó, ta chưa biết mẫu số cuối cùng, cũng chưa biết liệu một điểm số lớn hơn có xuất hiện ở phía sau hay không.

Để vứt bỏ các điểm số đã xử lý, ta cần giữ lại ba đại lượng: giá trị max chạy của điểm số m, tổng mũ ℓ\ell tương ứng với m, và tổng có trọng số của các giá trị u dưới cùng một trọng số mũ. Ở đây u thường là một vector; ta sẽ trước tiên làm qua quy trình cập nhật với một ví dụ vô hướng nhỏ.

Ví dụ 5-6: Softmax trực tuyến gộp kết quả các khối trong khi vẫn giữ chuẩn hóa đúng như thế nào? Các điểm số là [0,ln⁡2][0,\ln2], với các giá trị tương ứng là [1,3][1,3], và mỗi khối chứa đúng một phần tử. Khối đầu tiên cho m=0,ℓ=1,u=1m=0,\ell=1,u=1. Sau khi xử lý khối thứ hai, giá trị max tăng lên ln⁡2\ln2, nên ta nhân ℓ\ell và u gốc với r=1/2r=1/2; hàm mũ của phần tử mới là 1, nên ℓ′=1/2+1=1.5\ell'=1/2+1=1.5, u′=1/2+3=3.5u'=1/2+3=3.5, cho kết quả 7/37/3. Xử lý cả hàng trong một lần, hai trọng số tỉ lệ với 1 và 2, cho cùng kết quả (1+2×3)/3=7/3(1+2\times3)/3=7/3.

Hình 5-19 Hai khối có điểm số lần lượt là 0 và ln 2, với giá trị lần lượt là 1 và 3. Sau khi giá trị max tăng lên, tổng mũ cũ và tổng có trọng số cũ đều được nhân với 1/2, rồi đóng góp của khối mới được cộng vào, và chỉ sau đó phép chia mới được thực hiện. Các mũi tên mang các thống kê tiến lên phía trước — các điểm số cũ không cần được giữ lại. m là max của các điểm số đã xử lý, ℓ là tổng mũ tương ứng với m, u là tổng có trọng số tương ứng của các vector giá trị, và r là hệ số co giãn được áp dụng khi mốc max thay đổi.

Trong ví dụ nhỏ, khối thứ hai nâng max từ 0 lên ln 2. Số mũ của số hạng cũ vốn là 1, và chuyển sang mốc mới biến nó thành 1/2; đóng góp cũ ở mẫu số và đóng góp cũ ở tử số phải được co giãn xuống cùng nhau để giữ nguyên trọng số tương đối của chúng.

Tổng quát hóa lên một hàng điểm số đầy đủ, trừ đi max dùng chung sẽ tránh được tràn hàm mũ. Đặt ℓ=∑jexp⁡(sj−m)\ell=\sum_j\exp(s_j-m), u=∑jexp⁡(sj−m)Vju=\sum_j\exp(s_j-m)V_j, với đầu ra cuối là u/ℓu/\ell. Khi một khối mới đến, dịch trạng thái cũ sang mốc max mới, rồi cộng đóng góp của khối mới:

m′=max⁡(m,max⁡(s)),r=exp⁡(m−m′),pj=exp⁡(sj−m′),m'=\max(m,\max(s)),\qquad r=\exp(m-m'),\qquad p_j=\exp(s_j-m'), ℓ′=rℓ+∑jpj,u′=ru+∑jpjVj.\ell'=r\ell+\sum_jp_j,\qquad u'=ru+\sum_jp_jV_j.

Tổng ở vế phải chỉ chạy trên j trong khối mới — mọi đóng góp của các khối cũ đã được chứa trong ℓ\ell và u, không cần đọc lại các điểm số cũ.

FlashAttention khai thác đúng quan hệ truy hồi này để loại bỏ việc lưu trữ toàn bộ S và P. Sau khi mỗi khối được xử lý, m, ℓ\ell và u được cập nhật; điểm số và xác suất của khối đó có thể bị vứt bỏ sau khi dùng, và buffer của chúng có thể được chuyển cho khối kế tiếp. Trong cùng một phạm vi attention, việc đối sánh giữa query và key cũng như việc gán trọng số cho các value vẫn đều phải được tính; thứ thay đổi là thứ tự tính toán và cách giữ các kết quả trung gian, và việc thực thi dấu phẩy động cũng đưa vào các sai khác về làm tròn.

Việc này chuyển bài toán lưu trữ ma trận điểm số đầy đủ thành bài toán sắp xếp khối hiện tại vào trong bộ nhớ nhanh. Tiếp theo ta xác định nên xử lý cùng lúc bao nhiêu hàng của Q và bao nhiêu hàng của K/V, để tận dụng tốt nhất không gian này. Lấy 99 KB shared memory (101,376 byte) có sẵn cho một thread block trên RTX PRO 6000 làm buffer nhanh, giữ a hàng Q và một bộ tích lũy đầu ra FP32 cùng lúc, với một khối điểm số kích thước a×b. K và V tái sử dụng cùng một ngăn đầu vào b×d qua các giai đoạn, điểm số và xác suất tái sử dụng một ngăn a×b, và ba vector FP32 được dành riêng cho các thống kê theo hàng và cho việc cập nhật trung gian. Yêu cầu bộ nhớ cục bộ là:

S=2ad⏟Q+2bd⏟K/V slot+4ab⏟scores/probabilities+4ad⏟output accumulator+12a⏟per-row state≤101376.S=\underbrace{2ad}_{Q}+\underbrace{2bd}_{K/V\text{ slot}}+\underbrace{4ab}_{\text{scores/probabilities}}+\underbrace{4ad}_{\text{output accumulator}}+\underbrace{12a}_{\text{per-row state}}\le101376.

Khối K/V càng nhỏ, không gian mà nó và khối điểm số của nó tiêu thụ càng ít, cho phép chứa được nhiều hàng Q hơn. Nhiều hàng Q dùng chung một lượt đi qua K/V nghĩa là toàn bộ K/V được đọc lại ít lần hơn. Lấy b=64 thì vừa khít a=82 hàng Q, cần ⌈8192/82⌉=100\lceil8192/82\rceil=100 lượt đi đầy đủ. Mỗi lượt đọc 4 MiB K/V, cộng một lần đọc Q và một lần ghi lại O, tổng cộng 100×4+4=404100\times4+4=404 MiB.

Thu nhỏ b xuống 1 thì vừa khít a=128 hàng Q, số lượt giảm xuống 64, với khối lượng truy cập giảm còn 64×4+4=26064\times4+4=260 MiB. Nhưng mỗi lượt giờ phải đi từ 128 khối K/V lên 8192 khối, và trạng thái trực tuyến phải được cập nhật thường xuyên hơn rất nhiều.

Số hàng của khối K/V bSố hàng Q vừa khít aSố lượt đi qua toàn bộ K/VLưu lượng buffer đến tầng kế tiếpSố lần cập nhật trạng thái khối Q–khối K/V
112864260 MiB524288
6482100404 MiB12800
12853155624 MiB9920

Hình 5-20 Buffer nhanh là 99 KB shared memory của một thread block RTX PRO 6000, với độ dài chuỗi 8192, chiều head 128, không có masking; mỗi điểm tương ứng với một kích thước khối K/V trong bảng ở văn bản. Trục hoành là lưu lượng giữa buffer và tầng lưu trữ kế tiếp; trục tung là số lần cập nhật trực tuyến, theo thang log. Di chuyển từ điểm b=64 đến điểm b=1 giảm số lần đọc nhưng tăng số lần cập nhật lên gấp khoảng 41 lần. b kí hiệu số token trong một khối K/V.

Hình 5-20 đặt cả hai chi phí trên cùng một mặt phẳng: càng về trái là càng ít dữ liệu đọc, càng xuống dưới là càng ít cập nhật thực hiện. Thu nhỏ từ khối 64 hàng xuống khối 1 hàng đọc ít hơn 144 MiB, nhưng đẩy số lần cập nhật trạng thái từ 12800 lên 524288 — gấp khoảng 41 lần. Mỗi lần cập nhật phải thực thi điều khiển vòng lặp, một phép so sánh max, và việc co giãn hàm mũ và đầu ra; chỉ xử lý một hàng K/V mỗi lần cũng làm cho một chiều của phép nhân ma trận quá nhỏ để tận dụng tốt khối tính toán ma trận. Khối K/V nhỏ hơn giảm số lần đọc nhưng tăng số lần cập nhật; khối lớn hơn có lợi cho khối tính toán ma trận và giảm chi phí vòng lặp. Vì vậy, chọn kích thước khối đòi hỏi phải cân nhắc cả hai loại chi phí cùng nhau.

Rút gọn trực tuyến loại bỏ các lần đọc/ghi của các ma trận trung gian đầy đủ, nhưng phép nhân ma trận trong khối, phép tính hàm mũ và các lần cập nhật thống kê vẫn phải chạy. Các tối ưu hóa về sau tập trung vào phần công việc còn lại này. FlashAttention-2 cải thiện việc phân chia công việc giữa các thread, giảm việc trao đổi dữ liệu trung gian và các phép toán phi ma trận; FlashAttention-3 dùng thực thi bất đồng bộ để việc di chuyển, nhân ma trận và Softmax chồng lấp lẫn nhau. Sau khi đã loại bỏ các ma trận trung gian khổng lồ, những cải tiến này tiếp tục rút ngắn thời gian tính toán và di chuyển còn lại.

Để xem xét tác động của các cơ chế này trong một chương trình thực tế, một thí nghiệm kèm theo đã so sánh hai cách triển khai đằng sau cùng một giao diện attention trong PyTorch. Ở đây, backend chỉ chương trình tính toán mà framework thực sự lựa chọn nội bộ cho giao diện: với cùng Q, K, V, framework có thể hoàn thành phép tính attention bằng các chương trình khác nhau.

Backend toán học (SDPBackend.MATH) ghép các phép toán đa dụng như nhân ma trận và Softmax theo công thức attention, giữ đầy đủ các ma trận trung gian điểm số và xác suất. Backend FlashAttention (SDPBackend.FLASH_ATTENTION) dùng phương pháp tiling và rút gọn trực tuyến đã mô tả ở trên, tránh lưu trữ các ma trận trung gian đầy đủ. Thí nghiệm chỉ định lần lượt hai backend này và so sánh thời gian cùng bộ nhớ GPU phụ thêm cần thiết để hoàn thành cùng một phép tính attention.

Thí nghiệm dùng một RTX PRO 6000 Blackwell, với đầu vào và đầu ra BF16, xử lý từng chuỗi và từng attention head một lúc, với chiều head 128. Độ dài chuỗi là 8192 token; attention nhân quả (causal) được dùng, nghĩa là mỗi vị trí chỉ chú ý tới chính nó và các vị trí trước nó. Sau khi warmup, một CUDA Graph được dùng để ghi trước và phát lại lặp lại một tập các thao tác GPU, nhằm tránh nhiễu từ việc gửi lệnh từ host theo từng lần gọi. Đo cho mỗi phép tính attention, cách triển khai toán học giữ đầy đủ các ma trận trung gian mất khoảng 2.6 ms, trong khi cách triển khai FlashAttention dùng tiling và rút gọn trực tuyến mất khoảng 80 μs — tăng tốc khoảng 33 lần.

Ta cũng so sánh đỉnh bộ nhớ GPU phụ thêm được dùng bởi một lần gọi thông thường đơn. Với đầu vào Q, K, V đã cư trú trong bộ nhớ GPU, phép đo bao phủ đầu ra và không gian làm việc tạm: khoảng 848 MiB cho cách triển khai toán học, so với khoảng 22 MiB cho cách triển khai FlashAttention.

Sự so sánh song song hai cách triển khai đầy đủ này cho thấy tiling, rút gọn trực tuyến và fusion phối hợp với nhau như thế nào: tiling giúp dữ liệu hiện tại vừa vào bộ nhớ nhanh, rút gọn trực tuyến giúp Softmax được tính từng khối, và fusion loại bỏ việc ghi ra các kết quả trung gian trong khối. Kết hợp lại, cả ba thay đổi cả thời gian thực thi lẫn yêu cầu lưu trữ trung gian.

5.4 Trình biên dịch: Biểu diễn, biến đổi và lựa chọn

Hai mục trước chọn vòng lặp, kích thước khối, buffer và ranh giới operator thông qua phân tích thủ công. Các mô hình thực tế dùng nhiều hình dạng đầu vào và tổ hợp operator, và mỗi accelerator có một phân cấp bộ nhớ cùng tập tài nguyên thực thi khác nhau. Công việc của trình biên dịch là biểu diễn các lựa chọn triển khai này thành chương trình, loại bỏ các phép biến đổi sai dựa trên các phụ thuộc, rồi so sánh chi phí thực thi của các cách triển khai còn lại.

5.4.1 Vì sao phải so sánh nhiều cách triển khai

Xem lại chuỗi elementwise: SiLU áp một phép biến đổi phi tuyến lên mỗi số, số đó rồi được nhân elementwise với một đầu vào khác, và cuối cùng được lượng tử hóa sang một biểu diễn ít bit. Giữa mỗi cặp operator liền kề, ta có thể chọn thực thi tách rời hoặc fusion, cho bốn cách phân chia: cả ba bước tách rời, hai bước đầu fusion, hai bước cuối fusion, hoặc cả ba bước fusion. Nếu một chuỗi chứa mười operator, có 29=5122^9=512 cách phân chia các ranh giới liền kề. Mỗi cách phân chia cũng có thể tự chọn kích thước khối, bố trí cục bộ và độ sâu pipeline.

Các lựa chọn này tương tác với nhau. Gộp hai bước loại bỏ một lần ghi trung gian, nhưng cũng làm tăng lượng dữ liệu phải được giữ sống cùng lúc; nhiều dữ liệu tạm hơn rồi lại thay đổi các kích thước khối khả thi. Vì vậy, trình biên dịch phải biểu diễn đồng thời nội dung tính toán, thứ tự thực thi và vị trí đặt dữ liệu. Chỉ ghi lại một tên operator thì không có cách nào suy ra các chi phí này.

Một trình biên dịch có thể chọn cách triển khai qua ba bước. Thứ nhất, biểu diễn phép toán thành các vòng lặp và truy cập mảng, làm rõ những vòng lặp nào có thể được biến đổi và quan hệ đọc/ghi của chúng là gì. Thứ hai, kiểm tra tính hợp lệ đối chiếu với các phụ thuộc đọc/ghi và các quy tắc số học. Thứ ba, kết hợp các ràng buộc dung lượng và thời gian đo được để so sánh chi phí thực thi. Dưới đây, mã vòng lặp dễ đọc minh họa hai bước đầu, tiếp theo là lời giải thích về cách hệ thống sinh mã operator tự động AKG tổ chức toàn bộ quá trình thông qua biên dịch đa diện (một phương pháp biên dịch mô tả các lần lặp vòng lặp và truy cập mảng bằng các ràng buộc tuyến tính).

5.4.2 Biểu diễn tiling và fusion bằng các phép biến đổi vòng lặp

Mục 5.2 đã xác định cần xử lý bao nhiêu dữ liệu mỗi khối; ta vẫn cần quyết định thứ tự tính các khối này. Dù cùng một dữ liệu được cắt theo cùng cách, nếu một mảnh bị đẩy ra ngay sau khi dùng rồi sau đó phải được chuyển lại vào, vẫn xảy ra các truy cập không cần thiết. Các phép biến đổi vòng lặp cho phép một chương trình biểu diễn những sắp xếp này: split tách một vòng lặp thành “khối nào” và “vị trí nào trong khối”; reorder điều chỉnh thứ tự duyệt sao cho dữ liệu sắp được tái sử dụng có cơ hội ở lại trong bộ nhớ cục bộ. Việc chọn ở mức vòng lặp nào để giữ hoặc dùng một kết quả trung gian khi đó làm rõ nó phải được giữ trong bao lâu.

Lấy Y=SiLU⁡(AW)Y=\operatorname{SiLU}(AW), với M=1024,K=4096,N=12288M=1024,K=4096,N=12288 như trước, phép nhân ma trận tích lũy trong FP32 và kết quả được làm tròn sang BF16 trước phép kích hoạt. Chương trình trực tiếp nhất trước tiên sinh toàn bộ C, rồi duyệt C để sinh Y:

for i in range(M):
    for j in range(N):
        acc = 0.0
        for k in range(K):
            acc += A[i, k] * W[k, j]
        C[i, j] = round_to_bf16(acc)

for i in range(M):
    for j in range(N):
        Y[i, j] = silu(C[i, j])

Phép biến đổi đầu tiên là split. Tách i thành một chỉ số khối ioi_o và một tọa độ trong khối iii_i, với i=64io+iii=64i_o+i_i. 1024 hàng gốc giờ trở thành 16 khối, mỗi khối 64 hàng. Làm tương tự với j cho 192 khối cột. Tách đồng thời hai chiều đầu ra tạo thành các tile 64×64.

Split trước hết thay đổi cách các chỉ số được biểu diễn. Để hoàn thành liên tục một tile, ta còn cần reorder: duyệt các khối đầu ra trong vòng ngoài, và duyệt các phần tử trong một khối ở vòng trong. Rồi tách K thành 128 khối rút gọn, mỗi khối 32. Chương trình chọn một tile đầu ra, nạp dần các phần A, W cần thiết, và tích lũy vào đúng khối đầu ra đó.

Với cấu trúc vòng lặp như vậy, ta giờ có thể xác định các khối đầu vào phải được giữ trong bao lâu. Các khối A và W hiện tại mỗi khối chiếm 4 KiB, và có thể được thay thế ngay sau khi một khối rút gọn tính xong; 4096 bộ tích lũy FP32 chiếm 16 KiB, và phải được giữ cho đến khi cả 128 khối rút gọn hoàn thành. Vì vậy, hai loại buffer này nên được cấp phát và giải phóng ở các mức vòng lặp khác nhau.

Cuối cùng, điều chỉnh vị trí tính toán của phép kích hoạt. Một khi một tile đầu ra đã hoàn thành toàn bộ các phép rút gọn K của nó, cả 4096 giá trị C đều đã được tính, và SiLU có thể được áp dụng rồi Y được ghi ra ngay — điều này không ảnh hưởng đến việc các tile đầu ra khác chưa hoàn thành. Với cách này, chỉ kết quả của khối đầu ra hiện tại cần được giữ cục bộ; toàn bộ 24 MiB của C không còn cần thiết, loại bỏ 48 MiB ghi ra và đọc lại.

for io in range(ceil_div(M, BM)):
    for jo in range(ceil_div(N, BN)):
        acc = zeros_fp32(BM, BN)
        for ko in range(ceil_div(K, BK)):
            a = load_tile(A, io, ko, BM, BK)
            w = load_tile(W, ko, jo, BK, BN)
            matmul_accumulate(acc, a, w)
        c = round_to_bf16(acc)
        store_tile(Y, io, jo, silu(c))

Ở đây load_tile và matmul_accumulate lần lượt biểu diễn việc nạp khối và tích lũy ma trận; tile cuối không đầy đủ được che để đánh dấu các vị trí hợp lệ. Việc thụt lề trong mã giả đánh dấu rõ các quan hệ trước–sau này: các khối đầu vào được cập nhật bên trong vòng ko, bộ tích lũy tồn tại suốt toàn bộ vòng ko, và phép tính kích hoạt xảy ra sau khi vòng ko kết thúc.

Hình 5-21 Các mức vòng lặp quyết định dữ liệu tạm phải được giữ trong bao lâu. Vòng ngoài chọn khối đầu ra và tạo một bộ tích lũy 16 KiB; vòng trong ko lặp đi lặp lại đọc các khối A và W, và chỉ sau khi toàn bộ phép rút gọn hoàn thành, việc làm tròn và hàm kích hoạt mới được áp dụng.

Các hành động lập lịch này có những đối tác trong các hệ lập trình. Halide là một hệ lập trình tách bạch “tính cái gì” khỏi “lập lịch thế nào”, cho phép lập trình viên dùng các hành động như split, tile và reorder để thay đổi cách thực thi tiến hành. Ở đây ta gọi operator sinh ra một kết quả trung gian là producer, và operator tiêu thụ kết quả đó là consumer. compute_at chỉ định ở mức vòng lặp nào của consumer thì producer thực thi. TVM là một hệ biên dịch cho tính toán tensor, cũng dùng schedule để chỉ định vị trí tính toán và việc đọc/ghi cache, cho phép các operator liền kề tính liên tục trên cùng một mẩu dữ liệu.

Vị trí tính toán đồng thời quyết định việc tái sử dụng so với tính lại. Đặt phép tính sinh ra một kết quả trung gian ở mức ngoài cùng sẽ tính trước toàn bộ kết quả, sẵn sàng cho mọi phép tính tiếp theo đọc; đặt nó bên trong vòng lặp khối của một phép tính phía sau chỉ tính phần hiện cần. Nếu nhiều khối đầu ra cần cùng một phần của kết quả trung gian, tính nó trong mỗi khối gây ra sự trùng lặp. Vì vậy, trình biên dịch phải so sánh hai cách tiếp cận: giữ kết quả trung gian cho nhiều lần đọc, hoặc tính lại trước mỗi lần dùng.

fuse trong một vòng lặp cũng có thể gộp nhiều chiều lặp thành một. Ví dụ, tọa độ hai chiều của 16×192 khối đầu ra có thể được chuyển thành một chỉ số tuyến tính duy nhất từ 0 đến 3071, rồi phân phối cho các nhóm làm việc của accelerator. Điều này thay đổi cách các khối đầu ra được duyệt và cách các tác vụ được gán; đặt SiLU ngay sau tile của phép nhân ma trận loại bỏ việc ghi ra và đọc lại ma trận trung gian. Cả hai đều có thể được kết hợp trong cùng một schedule.

5.4.3 Các phụ thuộc và quy tắc làm tròn ràng buộc phép biến đổi như thế nào

Mục trước đã chuyển phép kích hoạt vào bên trong tile đầu ra, nhưng vẫn lập lịch nó sau khi vòng ko kết thúc. Hai con số giải thích vì sao. Giả sử một đầu ra có hai tổng từng phần, 1 và −1; tổng đầy đủ là 0, và SiLU(0) = 0. Nếu SiLU được áp cho từng tổng từng phần trước khi cộng, kết quả là SiLU⁡(1)+SiLU⁡(−1)≈0.4621\operatorname{SiLU}(1)+\operatorname{SiLU}(-1)\approx0.4621. Chuyển phép tính kích hoạt lên trước khi tổng đầy đủ được tính đã làm thay đổi phép tính.

Hình 5-22 Với cùng hai tổng từng phần, cộng chúng trước để được không rồi áp SiLU vẫn cho không; áp SiLU cho từng phần trước rồi mới cộng cho khoảng 0.4621. Hai con số này cho thấy việc chuyển phép kích hoạt lên trước phép cộng làm thay đổi kết quả.

Vì vậy, các phụ thuộc mà trình biên dịch phải tôn trọng có hai mức: các tổng từng phần của cùng một đầu ra phải được kết hợp theo quy tắc rút gọn, và các operator tiếp theo phải chờ các kết quả chúng cần. Các đầu ra khác nhau có thể được tính song song; phép rút gọn cho một đầu ra nhất định và các phép toán tiếp theo của nó phải thực thi theo thứ tự phụ thuộc.

Rút gọn theo khối (chunked reduction) bị ràng buộc bởi chính ràng buộc phụ thuộc này. FlashAttention (Mục 5.3.3) có thể hoàn thành phép tính từng khối vì nó giữ lại m, ℓ\ell và u, và điều chỉnh tổng mũ cùng tổng có trọng số đã tích lũy mỗi khi giá trị max thay đổi. Chỉ giữ đầu ra đã chuẩn hóa của từng khối sẽ mất đi độ lớn tương đối của các mẫu số giữa các khối. Ví dụ, nếu hai khối mỗi khối chứa một giá trị đơn, 1 và 3, thì đầu ra trong khối vẫn là 1 và 3; từ riêng hai đầu ra này, không có cách nào biết softmax của cả hàng nên gán cho chúng tỉ lệ trọng số 1:2 hay 2:1. Vì vậy, rút gọn theo khối phải giữ đủ thống kê để gộp đúng kết quả của từng khối.

Việc làm tròn dấu phẩy động còn ràng buộc phép biến đổi thêm nữa. Chương trình gốc làm tròn một kết quả tích lũy FP32 sang BF16 trước phép kích hoạt; chương trình fusion cũng phải thực thi round_to_bf16 trước phép kích hoạt. Giá trị trung gian có thể đi thẳng từ một thanh ghi đến phép kích hoạt, nhưng bước làm tròn này vẫn xảy ra. Nơi dữ liệu cư trú và định dạng số của nó là hai lựa chọn độc lập.

Ví dụ 5-7: Một hệ số tỉ lệ lượng tử hóa theo cả hàng ràng buộc thứ tự tính theo khối như thế nào? Một hàng chứa 256 phần tử, tách thành hai khối mỗi khối 128 phần tử; phần tử đầu tiên của hai khối lần lượt là 1 và 10, phần còn lại bằng không. Lượng tử hóa dùng định dạng E4M3FN: một biểu diễn dấu phẩy động tám bit với một bit dấu, bốn bit số mũ và ba bit phần định trị, với giá trị hữu hạn tối đa là 448. Hệ số tỉ lệ được xác định bởi giá trị tuyệt đối tối đa trên toàn bộ hàng; tích vô hướng sau đó chỉ dùng trọng số 1 cho phần tử đầu tiên. Kết quả vì vậy phụ thuộc hoàn toàn vào cách phần tử đầu tiên đó được lượng tử hóa.

Hình 5-23 Một hàng được tách thành hai khối; giá trị 10 trong khối thứ hai quyết định hệ số tỉ lệ cho cả hàng. Số hạng đầu tiên trước hết phải được ánh xạ bằng hệ số tỉ lệ này, rồi được làm tròn sang một giá trị mà định dạng cho phép, và cuối cùng được khử lượng tử.

Đọc cả hàng trước sẽ có max là 10. Số hạng đầu tiên được co giãn thành 1×448/10=44.81\times448/10=44.8, làm tròn về giá trị biểu diễn được là 44 trong định dạng này, và khử lượng tử thành 44×10/448=55/5644\times10/448=55/56. Nếu lượng tử hóa xảy ra ngay sau khi đọc khối đầu tiên, max ở thời điểm đó là 1, nên số hạng đầu tiên có thể được biểu diễn chính xác là 448, và kết quả khử lượng tử đúng bằng 1.

Một khi giá trị 10 trong khối thứ hai được đọc, không sự điều chỉnh nào với hệ số tỉ lệ dùng trong các phép tính tiếp theo có thể xóa được việc số hạng đầu tiên đã đi qua phép làm tròn 44.8→44. Hai cách tiếp cận khác nhau 1/561/56, khoảng 1.8%. Vì vậy, khi hệ số tỉ lệ được xác định từ max của cả hàng, toàn bộ hàng phải được đọc trước để tìm max đó, và chỉ sau đó việc lượng tử hóa mới nên tiến hành.

Các ví dụ này cho thấy phép biến đổi nào được phép và thứ tự tính toán nào phải được giữ nguyên. Các đầu ra có thể được thực thi theo bất kỳ thứ tự nào tương đối với nhau, đầu vào có thể được cache theo khối, và các kết quả đã tính có thể được đưa thẳng cho các phép toán tiếp theo; tuy nhiên, cách các phép rút gọn được kết hợp và nơi xảy ra việc làm tròn lượng tử hóa là một phần của định nghĩa phép toán và phải được giữ nguyên qua các phép biến đổi.

5.4.4 AKG: Tổ chức vòng lặp và bộ nhớ bằng biên dịch đa diện

AKG là một hệ sinh mã và tối ưu hóa tự động cho các operator tensor, với kỹ thuật cốt lõi là biên dịch đa diện (polyhedral compilation). Phương pháp này biểu diễn các vòng lặp đều đặn thành ba loại thông tin: những lần lặp nào thực thi, mỗi lần lặp đọc và ghi những phần tử mảng nào, và những lần đọc/ghi nào phải giữ nguyên thứ tự tương đối. Chặn vòng lặp và các chỉ số dưới đều đặn có thể được mô tả bằng các ràng buộc tuyến tính, đây chính là nguồn gốc của cái tên “đa diện”. Với biểu diễn này, trình biên dịch có thể suy ra các phụ thuộc giữa các lần lặp và điều chỉnh thứ tự thực thi tương ứng.

Lấy phép nhân ma trận làm ví dụ: trình biên dịch có thể xác định rằng các cặp i, j khác nhau tương ứng với các đầu ra khác nhau, trong khi các bản cập nhật dọc theo k trỏ về cùng một kết quả tích lũy. Một khi một tile 64×64×32 được chọn, các vùng A và W cần thiết cũng có thể được suy ra từ các truy cập mảng: A là 64×32, W là 32×64. Điều này xác định cả thứ tự thực thi vòng lặp lẫn những dữ liệu nào phải được đọc vào, cũng như dữ liệu đó phải được giữ trong bao lâu.

AKG sinh một biểu diễn đa diện từ các biểu thức tensor, kết hợp tiling với fusion phân cấp: mức ngoài sắp xếp thứ tự thực thi của toàn bộ operator, trong khi mức trong sắp xếp việc tính liên tục của từng khối và việc cache cục bộ của nó. Quản lý lưu trữ cấp phát không gian trên chip và chèn các lệnh di chuyển dữ liệu dựa trên các vùng này; sinh mã sau đó ánh xạ phép tính lên các chế độ thực thi của phần cứng, và giai đoạn tinh chỉnh (tuning) chọn các tham số như hình dạng tile. Quy trình thủ công đã mô tả trước đó — “tách vòng lặp, xác định vùng đọc, giữ kết quả tích lũy, và áp phép kích hoạt sau khi phép rút gọn hoàn thành” — ở đây trở thành một tập các bước biên dịch liên kết với nhau.

Việc phân tích tính toán và lưu trữ cùng nhau cũng giải thích vì sao fusion có thể làm tăng truy cập bộ nhớ. Trong chuỗi elementwise trước đó, fusion giảm số lần đọc các giá trị trung gian bằng cách tránh các vòng khứ hồi; nếu kết quả trung gian thay vào đó dùng một biểu diễn ít bit hơn và phải được các phép tính tiếp theo đọc đi đọc lại, giữ nó lại thực ra có thể làm giảm tổng số lần đọc.

Ví dụ 5-8: Lượng tử hóa thực thi tách rời hay được fusion vào phép nhân ma trận tạo ra ít lưu lượng đọc/ghi hơn? Lấy một đầu vào A FP16 [4096,4096][4096,4096], chiếm 32 MiB, sau khi lượng tử hóa sang FP8 trở thành 16 MiB; các trọng số FP8 là [4096,1536][4096,1536], chiếm 6 MiB, và đầu ra FP16 chiếm 12 MiB. Đầu ra được tiling ở 128×128, cho 32 khối hàng và 12 khối cột. Mỗi khối đầu ra tự đọc dữ liệu nó cần, và hệ số tỉ lệ lượng tử hóa đến từ giá trị max của hàng.

Xét trước cách lượng tử hóa chạy tách rời. Với mỗi hàng được xử lý, 8 KiB không gian cục bộ giữ đầu vào của hàng đó, từ đó hệ số tỉ lệ được tính và hàng được lượng tử hóa. Toàn bộ đầu vào được đọc một lần, tổng cộng 32 MiB, và kết quả lượng tử hóa được ghi ra, tổng cộng 16 MiB. Sau đó, mỗi một trong 12 khối cột đầu ra đọc đầu vào FP8 một lần, tổng cộng 192 MiB. Vậy khối lượng đọc/ghi liên quan đến đầu vào là 32+16+192=24032+16+192=240 MiB.

Một khi lượng tử hóa được fusion vào phép nhân ma trận, đầu vào gốc 32 MiB được đọc trước để tính hệ số tỉ lệ theo hàng; mỗi một trong 12 khối cột sau đó đọc lại đầu vào gốc 32 MiB và lượng tử hóa cục bộ. Truy cập liên quan đến đầu vào tổng cộng 32+12×32=41632+12\times32=416 MiB. Mặc dù cách này tiết kiệm việc ghi ra 16 MiB kết quả lượng tử hóa, lần đọc đầu vào của mỗi khối cột tăng từ 16 MiB lên 32 MiB, và 12 lần đọc cộng thêm 192 MiB.

Cả hai sơ đồ đều có 32 khối hàng, mỗi khối đọc 6 MiB trọng số, rồi ghi 12 MiB đầu ra, tổng cộng 32×6+12=20432\times6+12=204 MiB. Các khối lượng truy cập tổng vì vậy lần lượt là 240+204=444240+204=444 MiB và 416+204=620416+204=620 MiB — fusion làm tăng truy cập khoảng 40%.

Hình 5-24 Cả hai sơ đồ đều trước tiên đọc toàn bộ đầu vào để xác định hệ số tỉ lệ của mỗi hàng. Sau khi lưu kết quả FP8, 12 khối cột cùng nhau đọc lại 192 MiB; sơ đồ fusion đọc lại đầu vào FP16 với 384 MiB. Trọng số và đầu ra cộng thêm cùng 204 MiB trong cả hai trường hợp.

Trong Hình 5-24, thứ quyết định khoảng cách là các lần đọc lặp lại nuôi 12 khối cột: mỗi lần đọc xử lý lượng dữ liệu gấp đôi. Số khối cột đầu ra vì vậy trực tiếp quyết định sự khác biệt về truy cập bộ nhớ giữa hai cách tiếp cận. Mỗi khối cột thêm vào cộng thêm 16 MiB số lần đọc cho sơ đồ giữ lại kết quả lượng tử hóa, và 32 MiB cho sơ đồ fusion. Đưa toàn bộ 1536 cột đầu ra vào cùng một khối loại bỏ sự lặp lại xuyên khối này; bộ tích lũy đầu ra tương ứng khi đó tăng từ 128 cột lên 1536 cột, cần gấp 12 lần không gian. Vậy việc có fusion lượng tử hóa hay không vẫn quy về việc cân bằng giữa diện tích bộ nhớ cục bộ và việc tái sử dụng dữ liệu, như đã thảo luận trong Mục 5.2.

AKG xét fusion và tiling cùng nhau chính là để xử lý loại tương tác này. Một khi vị trí tính kết quả trung gian thay đổi, cả việc mỗi khối đầu ra đọc dữ liệu nào lẫn số lần đọc đều thay đổi; yêu cầu về buffer đổi theo, và điều đó đến lượt lại ảnh hưởng tới các kích thước khối có thể dùng được.

5.4.5 Ước lượng chi phí và lựa chọn thực nghiệm

Ví dụ 5-8 đã so sánh hai cách triển khai theo khối lượng đọc/ghi. Đối mặt với nhiều tổ hợp tiling và fusion hơn, một trình biên dịch thu hẹp không gian tìm kiếm theo cùng cách, từng bước: trước kiểm tra các phụ thuộc và quy tắc số học, rồi loại bỏ các cách triển khai không vừa trong bộ nhớ cục bộ, và cuối cùng ước lượng chi phí thực thi để quyết định cách triển khai nào được thử trước. Mục 5.2 đã đưa ra một phương pháp ước lượng: so sánh hai tile bằng khối lượng truy cập và băng thông hiệu dụng. Trình biên dịch cũng có thể tính thêm tỷ lệ sử dụng lệnh ma trận, yêu cầu về thanh ghi và chi phí khởi chạy để xác định thứ tự thử nghiệm giữa các các cách triển khai.

Tiếp tục với ví dụ ma trận và shared memory của RTX PRO 6000 trong Mục 5.2, quy trình lựa chọn có thể được tổ chức thành bảng dưới đây. Mục đó đặt cả buffer đầu vào lẫn bộ tích lũy trong 100 KB shared memory của mỗi SM; nếu bộ tích lũy thay vào đó được đặt trong thanh ghi, shared memory và thanh ghi sẽ phải được kiểm tra riêng, như trong Mục 5.1.5, vì hai dung lượng không thể vay mượn của nhau. Nếu cả hai ứng viên đều chuyển sang đệm kép, mức dùng đầu vào sẽ tăng đôi, khiến working set lần lượt là 32 KiB và 96 KiB, và số thread block cư trú mỗi SM sẽ đổi từ 4 và 1 thành 3 và 1.

Bước lựa chọnCần điền gì cho ví dụ ma trận của chương nàyCơ sở loại trừ hoặc so sánh
Khắc định mục tiêuHình dạng đầu vào, độ chính xác, dung sai lỗi, tần suất gọiSo sánh cùng một công việc toán học và cùng một khối lượng công việc
Liệt kê ứng viênTile, thứ tự vòng lặp, số lượng buffer, phạm vi fusionTrình biên dịch và phần cứng có hỗ trợ chúng hay không
Kiểm tra khả thiShared memory, thanh ghi, thread, canh lề, phụ thuộcBáo cáo tài nguyên của trình biên dịch và các kiểm tra đúng-sai
Ước lượng thời gianKhối lượng tính toán, khối lượng đọc/ghi thực tế, chi phí khởi chạy và biến đổi bố tríThời gian phục vụ của từng tài nguyên và đường tới hạn
Đo và lựa chọnThời gian cho các hình dạng đại diện, độ biến thiên, và toàn bộ chuỗi operatorTối thiểu hóa chi phí khối lượng công việc trong khi giữ lại các ứng viên sát sao

Các báo cáo tài nguyên được dùng trước hết để loại các ứng viên không vừa trong dung lượng có sẵn; các ứng viên còn lại sau đó được đo về băng thông hiệu dụng, thời gian tính toán và hành vi trên toàn bộ chuỗi operator. Nếu việc tràn thanh ghi, padding cạnh ma trận hoặc độ đồng thời giảm triệt tiêu các lợi ích của tái sử dụng, quy trình quay lại bảng ứng viên để thay đổi hình dạng khối hay số lượng buffer. Chương 6 đi theo cùng quy trình, nhưng mở rộng vị trí đặt tile sang nhiều card và dịch các phụ thuộc liên card thành phép truyền thông tập hợp (các phép toán truyền thông mà nhiều card trao đổi hoặc gộp dữ liệu theo một mẫu cố định).

Bước đo có thể được minh họa bằng chương trình chuyển vị ma trận trong Thí nghiệm 5-5 kèm theo: cùng một chương trình trên RTX PRO 6000 Blackwell mất khoảng 7.9 μs với các khối 16×16 và khoảng 2.9 μs với các khối 32×32, tăng tốc khoảng 2.7×. Việc nhân đôi mỗi chiều của khối làm tăng gấp bốn số phần tử bên trong, và bố trí lưu trữ cùng sự hợp tác giữa các thread đều đổi theo.

Tuning tự động của TVM tìm kiếm trên tiling, thứ tự vòng lặp và các scheme lập lịch khác dựa trên kết quả đo được của phần cứng; việc để một agent có khả năng gọi các công cụ biên dịch và kiểm thử sửa đổi schedule hoặc mã kernel cũng có thể khám phá các cách triển khai ngoài các template lập lịch ban đầu. Cả hai phương pháp tìm kiếm đều đánh giá các chương trình đã biên dịch, lọc ra các cách triển khai đúng-sai qua các kiểm tra số học rồi đánh giá schedule theo thời gian thực thi. Việc đánh giá nhắm vào chương trình đã biên dịch chính xác vì trình biên dịch có thể fusion các operator được viết tách biệt trong mã nguồn thành một kernel duy nhất. Phương pháp tìm kiếm cũng phải giải quyết một câu hỏi khác: dùng khối lượng công việc nào để đánh giá các cách triển khai này.

Ví dụ 5-9: Tỉ lệ gọi giữa các hình dạng đầu vào thay đổi lựa chọn cách triển khai operator như thế nào? Giả sử một operator có hai hình dạng đầu vào A và B, và cách triển khai gốc mất 10 μs cho cả hai. Một cách triển khai mới mất 5 μs cho A và 20 μs cho B: thời gian của A giảm một nửa, thời gian của B tăng đôi. Nếu mỗi hình dạng xuất hiện một lần, tổng thời gian tăng từ 20 μs lên 25 μs, tức tăng 25% cho cách triển khai mới; nếu A được gọi 100 lần và B một lần, tổng thời gian giảm từ 1010 μs xuống 520 μs, tăng tốc khoảng 1.9× cho cách triển khai mới.

Gọi p là tỉ lệ các lần gọi có hình dạng đầu vào A. Thời gian thực thi trung bình của cách triển khai gốc là 10 μs, và thời gian thực thi trung bình của cách triển khai mới là:

T‾new=5p+20(1−p)=20−15pμs.\overline T_{\mathrm{new}}=5p+20(1-p)=20-15p\quad\mathrm{\mu s}.

Để cách triển khai mới có thời gian trung bình ngắn hơn, ta cần 20−15p<1020-15p<10, tức p>2/3p>2/3. Điểm giao cắt trong Hình 5-25 cho thấy cách triển khai mới nhanh hơn chỉ khi A chiếm quá hai phần ba tổng số lần gọi.

Hình 5-25 Khi tỉ lệ của hình dạng A vượt quá 2/3, thời gian thực thi tổng của cách triển khai mới ngắn hơn. Cả A lẫn B ban đầu đều mất 10 μs; cách triển khai mới mất lần lượt 5 và 20 μs. Các lần gọi là tuần tự; hình so sánh phép thực thi ở chế độ ổn định. Điểm giao cắt là nơi thời gian trung bình của hai cách triển khai bằng nhau.

Vì vậy, tuning có thể chọn một cách triển khai riêng cho từng hình dạng, hoặc chọn một scheme duy nhất có tổng chi phí thấp hơn dựa trên tần suất gọi thực tế. Thời gian biên dịch và đo đạc mà bản thân quá trình tìm kiếm tiêu tốn được tính là chi phí chuẩn bị. Thí nghiệm về agent kèm theo lưu lại các bản ghi sinh mã, gửi lặp lại và đo đạc; mục sau kết hợp chúng với số lần gọi để tính xem phải mất bao lâu thì thời gian thực thi tiết kiệm được mới bù đắp chi phí tìm kiếm.

5.5 Runtime: gửi lệnh, phát lại và hình dạng động

Một khi trình biên dịch đã chọn một chương trình cho accelerator, runtime vẫn phải lặp đi lặp lại chuẩn bị đầu vào, chọn chương trình, gửi tác vụ và quản lý buffer. Dù dùng cùng mã kernel, việc thay đổi cách gửi nó, hoặc tái sử dụng các kết quả chuẩn bị đã có, đều làm thay đổi tổng thời gian. Mục này tiếp tục dòng thời gian từ Mục 5.1, phân tích phần công việc phía host này.

5.5.1 Chồng lấp việc chuẩn bị phía host với việc tính toán trên accelerator

Mục 5.1 đã phân biệt thời gian gửi lệnh của CPU với thời gian thực thi của accelerator, và Mục 5.3 đã cho thấy hai loại tài nguyên có thể chồng lấp. Áp dụng cả hai điểm này cho một lượt forward của một model: CPU kiểm tra hình dạng và chuẩn bị tham số hoặc metadata cho giai đoạn kế tiếp trong khi GPU thực thi giai đoạn hiện tại. Nếu việc chuẩn bị cho giai đoạn kế tiếp có thể xảy ra sớm, hai phía tạo thành một pipeline; nếu CPU chờ đến khi GPU xong mới bắt đầu chuẩn bị, accelerator bị đình trệ giữa các giai đoạn.

Ví dụ 5-10: Vì sao host trở thành điểm nghẽp sau khi accelerator tăng tốc? Giả sử có 100 đoạn công việc, mỗi đoạn cần 20 μs chuẩn bị phía host và 20 μs tính toán phía accelerator. Dữ liệu cần để chuẩn bị đoạn kế tiếp đã được biết trước, buffer đủ dùng, và hai loại tài nguyên độc lập với nhau. Thực thi tuần tự chặt chẽ mất 100×(20+20)=4000100\times(20+20)=4000 μs; với pipeline, việc chuẩn bị đoạn đầu tiên hoàn thành trước, và sau đó mỗi đoạn hoàn thành sau mỗi 20 μs, cho tổng thời gian 20+100×20=202020+100\times20=2020 μs.

Rút ngắn thời gian tính toán trên accelerator xuống 5 μs, host vẫn mất 20 μs để chuẩn bị mỗi đoạn. Pipeline trở thành 100×20+5=2005100\times20+5=2005 μs, chỉ ít hơn 15 μs so với trước. Tốc độ gửi tác vụ của CPU không được cải thiện, nên GPU tính càng nhanh thì nó chờ tác vụ kế tiếp càng lâu. Rút ngắn cả thời gian chuẩn bị phía host xuống 5 μs cho phép hai phía mỗi bên hoàn thành một đoạn sau mỗi 5 μs, đưa thời gian hoàn thành xuống 5+100×5=5055+100\times5=505 μs.

Hình 5-26 vẽ dòng thời gian cho bốn đoạn đầu tiên trong cả ba tình huống.

Hình 5-26 Việc chuẩn bị phía host cho mỗi đoạn mất 20 μs. Trong thực thi tuần tự, hai loại tài nguyên luân phiên; sau khi có pipeline, host chuẩn bị đoạn kế tiếp trong khi accelerator tính đoạn hiện tại, hoàn thành một đoạn sau mỗi 20 μs; một khi accelerator tăng tốc lên 5 μs, mỗi đoạn vẫn phải chờ host chuẩn bị, và accelerator rảnh rỗi phần lớn thời gian. Hình chỉ vẽ bốn đoạn đầu tiên; ví dụ trong văn bản bao gồm 100 đoạn.

Ví dụ này chỉ ra hai hướng tối ưu phía host. Một là giảm công việc chuẩn bị — ví dụ, bằng cách tái sử dụng một hình dạng và một kế hoạch thực thi đã được xác định (việc chọn kernel, phân chia tác vụ và sắp xếp workspace được xác định trước dựa trên các điều kiện như độ dài yêu cầu và bố trí cache). Hướng còn lại là đưa việc chuẩn bị lên sớm hơn, vào trong cửa sổ thực thi của accelerator cho đoạn trước — ví dụ, xử lý bất đồng bộ kết quả của vòng trước. Hai hướng tương ứng giảm lượng công việc trong pipeline và giảm các khoảng hở giữa các giai đoạn.

Trong việc sinh tự động (autoregressive generation), việc chuẩn bị có thể xảy ra sớm hay không phụ thuộc vào việc dữ liệu cần đã được tính chưa. Phép tính bước tiếp theo cần token đã chọn ở vòng trước; việc ghép batch, cập nhật trạng thái ngữ pháp (khi định dạng đầu ra bị ràng buộc, ghi lại mức nào mà nội dung sinh ra đã khớp với ngữ pháp), và việc xác định có dừng hay không mỗi cái đều cần kết quả tương ứng của riêng nó. Runtime có thể thực hiện sớm các công việc không phụ thuộc vào các kết quả này, giảm thời gian chờ giữa các vòng. Việc thực thi và lập lịch bất đồng bộ trong các engine như vLLM xoay quay đúng vấn đề này.

Một khi một accelerator chuyên dụng rút ngắn thời gian thực thi model xuống dưới hẳn thời gian chuẩn bị mỗi vòng của host, cận dưới của khoảng cách bước trong pipeline ổn định bị xác định bởi việc chuẩn bị phía host. Việc rút ngắn khoảng cách bước hơn nữa khi đó đòi hỏi tối ưu đồng thời việc chuẩn bị batch động (một batch có thể được thêm hoặc bớt yêu cầu ở mỗi bước), cập nhật địa chỉ, và việc gửi lệnh.

Ra mắt vLLM v0.6.0 vào tháng 9/2024 là một trường hợp mà chi phí phía host chi phối khoảng cách bước. Profiling trước khi phát hiện cho thấy, khi chạy Llama 3 8B trên một H100 đơn, trong một bước, API server chiếm 33%, lập lịch chiếm 29%, và thực thi trên accelerator chỉ chiếm 38%: API server và inference engine chạy trong cùng một tiến trình Python, và các coroutine của chúng tranh chấp GIL (global interpreter lock), khiến accelerator phải chờ host. v0.6.0 chuyển API server sang một tiến trình riêng, giao tiếp với engine qua một message queue và tránh tranh chấp GIL; nó cũng giới thiệu lập lịch đa bước (multi-step scheduling), thực thi nhiều bước liên tục sau một lượt lập lịch, tránh việc chuẩn bị đầu vào lặp lại; và nó làm cho việc xử lý đầu ra trở nên bất đồng bộ, chồng lấp nó với phép tính của bước kế tiếp. Kết hợp lại, ba thay đổi này nâng thông lượng Llama 3 8B lên 2.7 lần so với v0.5.3, và thông lượng 70B lên 1.8 lần. Không thay đổi nào trong số đó chạm vào bất kỳ kernel nào, cũng không liên quan tới việc thay phần cứng. Theo tiêu chí trong Mục 1.3.4, một khi chi phí trừu tượng hóa phần mềm phía host đạt 60% khoảng cách bước, việc gỡ bỏ chi phí đó hiệu quả hơn nhiều so với việc chuyển sang một accelerator nhanh hơn.

Ngoài hai vòng liền kề của cùng một yêu cầu, còn có công việc chuẩn bị có thể tái sử dụng giữa các lần gọi khác nhau. Hai lần gọi có thể lần lượt yêu cầu model giải thích mã và kiểm tra test — nội dung đầu vào khác nhau, nhưng đi qua cùng một tập lớp model. Khi độ dài và kích thước batch nằm trong một khoảng đã được một kế hoạch thực thi hiện có hỗ trợ, runtime có thể tái sử dụng chương trình, chỉ cập nhật token, chỉ số vị trí và địa chỉ trạng thái; khi gặp một hình dạng mới hay một đường điều khiển khác, nó chọn một kế hoạch thực thi khác. Các ứng dụng biểu diễn sự thay đổi qua ngữ cảnh, và runtime khai thác cấu trúc thực thi ổn định để giảm việc chuẩn bị lặp lại. Phần sau tính lần lượt lợi ích của hai hình thức tái sử dụng này — phát lại đồ thị và chuyên biệt hóa theo hình dạng.

5.5.2 CUDA Graph: lợi ích gửi lệnh so với chi phí sao chép đầu vào

Trong mục trước, GPU đã tăng tốc nhưng vẫn bị bỏ chờ host. Khi một model thực thi lặp lại các chuỗi tác vụ tương tự, phần lớn công việc gửi lệnh này có thể được tái sử dụng: ghi một lần các tác vụ accelerator cùng phụ thuộc của chúng thành một đồ thị, rồi gửi lại nó lặp đi lặp lại. Với CUDA Graph, mỗi lần host chỉ cần phát một lệnh phát lại đồ thị duy nhất, thay vì chuẩn bị và gửi từng kernel một cách riêng lẻ. Accelerator vẫn thực thi mọi nút trong đồ thị — thứ được giảm là công việc lặp lại của host. Hình 5-27 minh họa sự khác biệt giữa hai chế độ gửi lệnh bằng sáu kernel của một FFN đơn.

Hình 5-27 Với gửi lệnh thông thường, host phát một lệnh khởi chạy cho mỗi kernel; với phát lại đồ thị, host phát một lệnh khởi chạy đồ thị duy nhất, và accelerator vẫn thực thi sáu kernel được ghi trong đồ thị. Việc giảm số kernel đòi hỏi fusion, không phải phát lại đồ thị.

Hình 5-28 lấy ba lần thực thi FFN làm phép so sánh. Gửi lệnh thông thường tạo ra 18 lần host khởi chạy kernel và 18 kernel trên accelerator; phát lại đồ thị giảm con số này xuống 3 lần khởi chạy đồ thị, trong khi accelerator vẫn thực thi 18 kernel. Một khi chuỗi kích hoạt được fusion, số kernel trên accelerator giảm xuống 15; cộng thêm phát lại đồ thị, host vẫn chỉ cần 3 lần khởi chạy đồ thị. Fusion gộp phép tính của nhiều operator thành một kernel duy nhất, trong khi phát lại đồ thị giảm công việc chuẩn bị và gửi tác vụ từng cái một của host.

Hình 5-28 Dấu vết đo được của 3 lần thực thi FFN với gửi lệnh thông thường. Hàng trên là các lệnh gọi API khởi chạy kernel phía host, hàng dưới là các kernel trên accelerator; trục hoánh được tính thời gian từ mốc bắt đầu của đoạn này. 18 lần khởi chạy kernel tương ứng 18 kernel trên accelerator.

Hình 5-29 Sau khi fusion chuỗi kích hoạt trên 3 lần thực thi FFN, host phát 15 lần khởi chạy kernel, và accelerator tương ứng thực thi 15 kernel. Dữ liệu đến từ một cửa sổ capture được đánh dấu riêng trong cùng thí nghiệm với hình trước.

Hình 5-30 3 lần khởi chạy đồ thị tương ứng 18 kernel trên accelerator. Phát lại đồ thị giảm số lần gửi lệnh phía host, trong khi accelerator vẫn thực thi mọi nút trong đồ thị gốc; hình này đánh dấu thời gian thực tế trong cửa sổ capture của chính nó.

Hình 5-31 Fusion trước rồi phát lại, host phát 3 lần thực thi đồ thị, và accelerator thực thi 15 kernel. Fusion giảm số kernel trên accelerator; phát lại đồ thị giảm số lần gửi lệnh phía host.

Trước khi phát lại đồ thị, đầu vào mới vẫn cần được đặt vào các địa chỉ mà đồ thị đã ghi. Một cách triển khai điển hình dành sẵn các buffer cố định cho đồ thị; mỗi lần, phía thượng nguồn ghi dữ liệu mới vào các buffer này, và đồ thị đọc từ các địa chỉ đã ghi. Nếu kết quả đầu ra của phía thượng nguồn nằm ở một tensor khác, cần thêm một phép sao chép trước khi thực thi đồ thị; nếu phía thượng nguồn ghi thẳng vào buffer cố định, phép sao chép này có thể được bỏ qua.

Hình 5-32 Thực thi đồ thị đọc từ địa chỉ đã ghi G. Khi đầu vào mới nằm ở X, nó được sao chép sang G trước; khi phía thượng nguồn ghi thẳng vào G, cùng một buffer được tái sử dụng, loại bỏ phép sao chép trung gian.

Ví dụ 5-11: Khi nào việc sao chép đầu vào triệt tiêu lợi ích của phát lại đồ thị? Thực thi thông thường mất 20 μs tính toán trên accelerator cộng 20 μs accelerator chờ host gửi tác vụ, tổng cộng 40 μs. Phát lại đồ thị cộng chuẩn bị metadata mất 5 μs, với thời gian tính toán trên accelerator vẫn là 20 μs. Đồ thị vốn sẽ tiết kiệm 15 μs.

Lấy một đầu vào BF16 [256,4096][256,4096] cần cho thực thi đồ thị, với lượng dữ liệu là 2 MiB. Phép sao chép vừa đọc nguồn vừa ghi đích, cho khối lượng đọc/ghi cộng dồn là 4 MiB; với băng thông bộ nhớ 1792 GB/s của RTX PRO 6000 (dùng chung cho đọc và ghi), phép sao chép mất khoảng 2.3 μs, khiến tổng thời gian của đồ thị khoảng 27.3 μs — ít hơn khoảng 12.7 μs so với thực thi thông thường.

Giữ thời gian tính toán và chuẩn bị cố định, mở rộng đầu ra lên 2048 hàng, nâng lượng dữ liệu lên 16 MiB. Khối lượng đọc/ghi của phép sao chép tăng lên 32 MiB, mất khoảng 18.7 μs, và tổng thời gian của đồ thị trở thành khoảng 43.7 μs. Phép sao chép này giờ vượt quá 15 μs tiết kiệm được ban đầu.

Hình 5-33 Màu cam là chuẩn bị, màu xanh là lần sao chép đầu vào thêm, màu xanh lá là tính toán trên accelerator. Thực thi thông thường mất 40 μs; đồ thị với đầu vào 2 MiB mất khoảng 27 μs, và với đầu vào 16 MiB mất khoảng 44 μs. Phép sao chép được tính theo băng thông bộ nhớ 1792 GB/s của RTX PRO 6000.

Trong Hình 5-33, thời gian mà phát lại đồ thị tiết kiệm ở bước chuẩn bị là một độ dài cố định, trong khi thanh biểu diễn thời gian sao chép dài ra khi đầu vào lớn lên. Khi hai cái bằng nhau, phát lại đồ thị không còn tiết kiệm được thời gian nào nữa. Gọi lượng dữ liệu đầu vào là X và băng thông sao chép là B; điều kiện để đồ thị nhanh hơn là 2X/B<15 μs2X/B<15\ \mathrm{\mu s}. Thay B=1792 GB/s vào, ta có:

X<15×10−6×1.792×10122 bytes≈12.8 MiB.X<\frac{15\times10^{-6}\times1.792\times10^{12}}{2}\ \mathrm{bytes}\approx12.8\ \mathrm{MiB}.

2 MiB nằm dưới ngưỡng này, trong khi 16 MiB vượt qua nó. Việc cho operator phía trước ghi kết quả của nó thẳng vào buffer của đồ thị, hoặc bắt đầu capture từ một điểm mà tensor đầu vào nhỏ hơn, đều có thể giảm chi phí sao chép. Vậy việc chọn các phép toán nào để gộp vào một đồ thị đòi hỏi cân nhắc giữa thời gian gửi lệnh phía host được tiết kiệm với thời gian sao chép đầu vào bị cộng thêm.

Phát lại đồ thị tái sử dụng chuỗi gửi lệnh; cùng ý tưởng này cũng có thể áp cho công việc chuẩn bị đứng trước các tác vụ này. Ví dụ, một kế hoạch thực thi có thể được dùng chung cho nhiều lớp model. Thư viện toán học cho inference FlashInfer tách việc sinh kế hoạch (plan) khỏi việc thực thi (run): plan chọn các kernel, phân chia tác vụ và sắp xếp workspace dựa trên độ dài theo từng yêu cầu của bước này cùng nơi KV nằm trong bộ nhớ GPU, rồi sao chép metadata này sang GPU; run sau đó thực thi attention theo kế hoạch. Nếu, trong cùng một bước, cả 36 lớp đều có cùng độ dài yêu cầu và cùng bố trí KV, chúng có thể dùng chung một kế hoạch duy nhất; việc chạy lại plan cho mỗi lớp nghĩa là làm cùng một công việc chuẩn bị 36 lần. Trong hai phép so sánh trên RTX PRO 6000, việc chuyển sang một lần gọi plan duy nhất giảm thời gian cho các lần gọi attention của 36 lớp từ khoảng 2.0/2.4 ms xuống 0.72/0.75 ms, và giảm số lần sao chép từ host sang GPU từ 144 xuống 4, trong khi số kernel trên accelerator giữ nguyên. Ngược lại, một khi độ dài yêu cầu thay đổi, plan phải được sinh lại — tái sử dụng kế hoạch cũ sẽ tính theo các độ dài cũ và cho ra kết quả sai.

5.5.3 Hình dạng động và chi phí biên dịch

Việc nhiều lớp model dùng chung một kế hoạch dựa trên cấu trúc tính toán giống hệt nhau; việc các lần gọi khác nhau dùng chung một chương trình đã biên dịch cũng phải xử lý các thay đổi về hình dạng đầu vào. Khi số hàng đầu vào thay đổi, runtime có thể chọn giữa ba cách tiếp cận. Chương trình đa dụng tính các chỉ số và chặn từ hình dạng hiện tại; bucketing đệm đầu vào về một số ít hình dạng đại diện và tái sử dụng chương trình cho các hình dạng đó; specialization khai thác các điều kiện đầu vào đã biết để sinh một cách triển khai riêng — ở đây, biên dịch một chương trình riêng cho mỗi hình dạng thường gặp. Hình dạng được khai thác càng đầy đủ, càng có nhiều cơ hội để đơn giản hóa công việc trên accelerator, nhưng điều này cũng đòi hỏi biên dịch và lưu trữ nhiều chương trình hơn.

Ví dụ 5-12: Bao nhiêu lần thực thi lặp lại mới khiến specialization có lợi? Lấy một nhóm 10 lần gọi FFN: 8 lần với 256 hàng, 1 lần với 1536 hàng, và 1 lần với 2048 hàng, tổng cộng 8×256+1536+2048=56328\times256+1536+2048=5632 hàng. Dùng hai bucket, 512 và 2048, tổng thực tế được xử lý là 8×512+2048+2048=81928\times512+2048+2048=8192 hàng, nhiều hơn khoảng 45% tính toán. Hình 5-34 cho thấy số hàng thực tế nào được đệm về bucket nào.

Hình 5-34 Các lần gọi với 256 hàng được đệm về bucket 512 hàng, các lần gọi với 1536 hàng được đệm về bucket 2048 hàng, và 2048 hàng rơi đúng vào ranh giới bucket. Phần màu xám là lượng tính toán thừa do padding; qua 10 lần gọi, tổng thực tế là 5632 hàng, nhưng 8192 hàng được thực thi.

Ba phép chiếu của một FFN đầy đủ tổng cộng cần 6HF6HF FLOP mỗi hàng. Trên RTX PRO 6000, lấy các cách triển khai đa dụng, theo bucket và chuyên biệt lần lượt đạt khoảng 20%, 40% và 50% đỉnh BF16 dense 503.8 TFLOP/s — tức MFU 20%, 40% và 50% — ta được các tốc độ xử lý ma trận hiệu dụng lần lượt là 100, 200 và 250 TFLOP/s, nên mỗi nhóm mất tương ứng khoảng 17.0, 12.4 và 6.8 ms. Mặc dù bucketing tính nhiều hơn 45% công việc, tốc độ xử lý của nó tăng gấp đôi, nên nó vẫn nhanh hơn cách triển khai đa dụng. Giả sử thời gian chuẩn bị của chúng lần lượt là 100, 400 và 900 ms; ta có bảng dưới đây.

Chiến lượcChuẩn bị một lần PThực thi mỗi nhóm TXử lý hình dạng
Đa dụng100 ms17.0 msThực thi đúng 5632 hàng
Theo bucket400 ms12.4 msĐệm về hai bucket, thực thi 8192 hàng
Chuyên biệt theo hình dạng900 ms6.8 msSinh chương trình riêng cho mỗi ba hình dạng

Tổng thời gian chạy r nhóm là Ttotal=P+rTT_{\mathrm{total}}=P+rT. Trong Hình 5-35, độ cắt của mỗi đường biểu diễn thời gian chuẩn bị, và độ dốc biểu diễn thời gian thực thi mỗi nhóm. Trong ví dụ này, scheme thực thi mỗi nhóm nhanh hơn có đường phẳng hơn nhưng điểm xuất phát cao hơn, vì nó đòi hỏi thời gian chuẩn bị lâu hơn.

Hình 5-35 Khi cùng một nhóm gọi được thực thi lặp lại, chiến lược tiết kiệm thời gian nhất thay đổi theo số lần tái sử dụng. Các đường cong dùng hình dạng, tốc độ xử lý và thời gian chuẩn bị từ Ví dụ 5-12; các khoảng số lần gọi nguyên mà mỗi chiến lược áp dụng được tính bằng giá trị chưa làm tròn. "Đa dụng" dùng cùng scheme thực thi cho mọi hình dạng; "theo bucket" gộp các hình dạng tương tự lại; "chuyên biệt" chọn một scheme riêng cho mỗi hình dạng cụ thể.

Bucketing tốn thêm 300 ms chuẩn bị so với scheme đa dụng, nhưng tiết kiệm khoảng 4.6 ms mỗi nhóm, nên sau khoảng 65 nhóm thì khoản chuẩn bị thừa được bù đắp. Specialization tốn thêm 500 ms chuẩn bị so với bucketing, tiết kiệm thêm khoảng 5.6 ms mỗi nhóm, cần khoảng 90 nhóm để bù đắp. Tính từ các giá trị chính xác, đa dụng rẻ nhất cho 1–64 nhóm, theo bucket rẻ nhất cho 65–89 nhóm, và chuyên biệt rẻ nhất từ 90 nhóm trở đi.

Nếu một instance chỉ chạy 70 nhóm, tổng thời gian lần lượt xấp xỉ 1.29, 1.27 và 1.38 s. Bucketing mất ít thời gian nhất vì phần tiết kiệm thời gian thực thi của nó đã vượt quá chi phí chuẩn bị thừa; mặc dù specialization thực thi mỗi nhóm nhanh hơn, phần tiết kiệm thời gian tích lũy của nó chưa đủ để bù chi phí chuẩn bị thừa. Nếu cả ba chương trình đã được cache, chi phí chuẩn bị cho lần chạy này là bằng không, và specialization mất ít thời gian nhất ngay từ nhóm đầu tiên. Vì vậy, lựa chọn của runtime phụ thuộc vào cả tần suất của từng hình dạng lẫn việc chương trình đã được biên dịch và cache chưa.

Ta cũng có thể tính cần bao nhiêu lần gọi để tuning tự động bù đắp chi phí chuẩn bị của nó. Nếu một cuộc tìm kiếm mất 10 phút và mỗi lần gọi dùng cách triển khai mới tiết kiệm 5 μs, thì để bù thời gian tìm kiếm bằng phần tiết kiệm tích lũy cần 600/(5×10−6)=1.2×108600/(5\times10^{-6})=1.2\times10^8 lần gọi. Với 10,000 lần gọi mỗi giây, việc này mất khoảng 3.3 giờ; với 100 lần mỗi giây, khoảng 14 ngày. Quá trình tuning càng mất nhiều thời gian, càng nên ưu tiên tối ưu các hình dạng được gọi thường xuyên và dùng trong thời gian dài.

Chương 8 sẽ kết hợp kiểu mẫu yêu cầu đến với batching động để thảo luận thêm về việc chọn và cache chương trình như thế nào. Nếu một hình dạng nhất định lặp lại trong thời gian ngắn, giữ lại kết quả biên dịch tương ứng giảm việc chuẩn bị lặp lại; nếu nó không còn được dùng về lâu dài, mục cache đó có thể bị xóa để giải phóng không gian cho các hình dạng khác.

5.5.4 Persistent kernel: lập lịch tác vụ theo khối dữ liệu

Phát lại đồ thị và cache chương trình giảm việc chuẩn bị phía host lặp lại, nhưng còn một loại chờ khác trên accelerator: với hai kernel có quan hệ phụ thuộc, kernel sau thường vẫn phải chờ kernel trước hoàn thành hoàn toàn mới bắt đầu. Nếu operator sau chỉ cần một khối kết quả của operator trước, điều kiện chờ có thể được thu hẹp về khối dữ liệu: một khi khối đầu tiên sẵn sàng, việc tính toán phía sau có thể bắt đầu ngay, trong khi các khối còn lại tiếp tục được tính.

Một persistent kernel cư trú trên accelerator, rút các tác vụ từ một hàng đợi và dùng các sự kiện (event) để xác nhận liệu dữ liệu cần đã sẵn sàng chưa. Cách này vừa giảm công việc khởi chạy kernel lặp lại của host, vừa cho phép các operator liền kề chồng lấp thực thi theo khối. Mirage Persistent Kernel/MPK là một hệ thống nghiên cứu áp dụng cách lập lịch tác vụ này.

Lại lấy phép chiếu up rồi kích hoạt SiLU trong Qwen3-8B, lấy một vector đặc trưng gồm 512 token, gom thành tám khối mỗi khối 64 token. Phép chiếu của mỗi khối là 2×64×4096×12288≈6.442\times64\times4096\times12288\approx6.44 GFLOP, mất khoảng 12.8 μs ở đỉnh BF16 dense 503.8 TFLOP/s của RTX PRO 6000; phép kích hoạt của mỗi khối xử lý 64×1228864\times12288 phần tử BF16, mỗi phần tử đọc 2 byte và ghi 2 byte, mất khoảng 1.76 μs ở băng thông bộ nhớ 1792 GB/s. Giả sử tài nguyên ma trận và vector độc lập với đệm dồi dào, và mỗi lần khởi chạy kernel mất 5 μs, việc chờ mọi phép chiếu hoàn thành rồi mới bắt đầu phép kích hoạt mất tổng cộng 2×5+8×(12.8+1.76)≈1262\times5+8\times(12.8+1.76)\approx126 μs.

Giờ tách phép tính thành tám tác vụ chiếu và tám tác vụ kích hoạt, mỗi tác vụ thêm 0.5 μs thời gian lập lịch và 0.2 μs thời gian thông báo hoàn thành. Phép chiếu mỗi khối trở thành 13.5 μs, và phép kích hoạt trở thành 2.46 μs. Với một lần khởi chạy kernel, tám khối chiếu chạy liên tiếp nhau, phép kích hoạt của mỗi khối bắt đầu ngay sau khi phép chiếu của nó xong, và phép kích hoạt cuối hoàn thành sau khi mọi phép chiếu xong, cho thời gian hoàn thành 5+8×13.5+2.46≈1155+8\times13.5+2.46\approx115 μs, tăng tốc khoảng 1.1.

Ở đây, giai đoạn chậm hơn là phép chiếu. Khoảng 11.0 μs tiết kiệm được đến từ hai nguồn: một lần khởi chạy kernel 5 μs ít đi, và khoảng 12.3 μs của bảy khối kích hoạt bị che khuất trong thời gian phép chiếu; đồng thời, chi phí tác vụ cộng lại 6.3 μs trên đường tới hạn. Vì mỗi khối kích hoạt chỉ mất 1.76 μs, việc bắt đầu nó sớm theo từng khối chỉ có thể che được tối đa ngần ấy thời gian. Hai giai đoạn có thời lượng càng gần nhau, việc chồng lấp theo khối càng tiết kiệm nhiều thời gian; khi một giai đoạn dài hơn hẳn giai đoạn còn lại, lợi ích về cơ bản chỉ giới hạn ở một lần khởi chạy kernel ít đi.

Hình 5-36 Thực thi thô (coarse-grained) trước hết hoàn thành cả tám khối chiếu, rồi bắt đầu cả tám khối kích hoạt. Ở đỉnh ma trận và băng thông bộ nhớ của RTX PRO 6000, mỗi khối chiếu mất khoảng 12.8 μs, mỗi khối kích hoạt khoảng 1.76 μs, và hai lần khởi chạy kernel mỗi lần 5 μs; đường thẳng đứng đánh dấu điểm bắt đầu kích hoạt.

Hình 5-37 Tài nguyên ma trận và vector độc lập với đệm dồi dào. Mỗi tác vụ còn phải chịu thêm 0.7 μs lập lịch và thông báo; phép kích hoạt có thể bắt đầu ngay khi khối chiếu đầu tiên xong. Các phép chiếu chạy liên tiếp nhau, và phép kích hoạt chỉ làm việc thoáng qua sau mỗi khối chiếu; pipeline tám khối hoàn thành trong khoảng 115 μs.

Theo dòng thời gian phép chiếu trong Hình 5-37, tám khối chiếu chạy liên tiếp nhau không có khoảng hở. Điều này theo cùng mẫu với đệm kép: khối đầu tiên hoàn thành giai đoạn trước, và sau đó mỗi khối hoàn thành sau mỗi khoảng bằng thời gian mà giai đoạn chậm hơn cần — ở đây, phép chiếu là giai đoạn chậm hơn, đúng như việc di chuyển là giai đoạn chậm hơn trong Mục 5.3.2, và nhịp của pipeline bị nó xác định. Mỗi ngăn buffer bị chiếm từ lúc bắt đầu phép chiếu cho đến lúc kết thúc phép kích hoạt; khi không còn ngăn trống, tác vụ phép chiếu phải chờ. Hàng đợi tác vụ, các sự kiện hoàn thành và các buffer trống vì vậy cùng nhau xác định khi nào khối kế tiếp có thể bắt đầu.

5.6 Từ tối ưu cục bộ đến các yêu cầu đầy đủ

Mỗi tối ưu ở trên đều có mục tiêu được xác định rõ ràng: tiling giảm số lần đọc đầu vào lặp lại, fusion giảm đọc/ghi trung gian, và các tối ưu runtime giảm việc chuẩn bị và chờ đợi. Nhưng một yêu cầu đầy đủ cũng bao gồm các công việc khác, và thời gian tiết kiệm được ở cục bộ chỉ là một phần của tổng thể. Mục này trước giải thích cách quy đổi thời gian yêu cầu từ các tăng tốc cục bộ, rồi phân tích một yêu cầu thực tế cho Qwen3-8B.

5.6.1 Tỉ lệ hotspot quyết định dư địa tăng tốc

Giả sử một yêu cầu gồm các giai đoạn tuần tự, thời gian của các giai đoạn khác không đổi, hotspot chiếm tỉ lệ f của thời gian gốc, và mức tăng tốc của hotspot là s. Chuẩn hóa thời gian yêu cầu gốc về 1, thời gian mới là (1−f)+f/s(1-f)+f/s, và mức tăng tốc tổng thể là:

Srequest=1(1−f)+f/s.S_{\mathrm{request}}=\frac{1}{(1-f)+f/s}.

Nếu hotspot chiếm 20% và thời gian của nó giảm một nửa, tổng thời gian trở thành 0.8+0.2/2=0.90.8+0.2/2=0.9, tăng tốc khoảng 1.11. Tăng tốc hotspot thêm lên mười lần cho tổng thời gian 0.82; ngay cả khi thời gian của hotspot tụt về bằng không, các giai đoạn khác vẫn chiếm 0.8 tổng thời gian gốc, nên mức tăng tốc tổng thể bị chặn ở 1.25×.

Đây là định luật Amdahl, đã được giới thiệu trong Chương 1. Định luật giải thích vì sao lợi ích từ việc tối ưu dần thu hẹp: hotspot mất càng ít thời gian, các giai đoạn khác chiếm càng lớn tỉ trọng trong tổng thời gian.

5.6.2 Các nhánh song song và việc chuyển đường tới hạn

Trong quan hệ Amdahl, các thời gian cộng dồn theo các giai đoạn tuần tự. Nếu một yêu cầu có hai nhánh thực thi đồng thời, rút ngắn một nhánh vẫn có thể khiến nó phải chờ nhánh còn lại; khi đó, nên truy vết đồ thị phụ thuộc để tìm đường dài nhất quyết định thời gian hoàn thành. Giả sử sau 10 μs chuẩn bị, hotspot A và nhánh B bắt đầu đồng thời, lấy lần lượt 60 và 40 μs; bước kết thúc mất 10 μs và chờ cả hai nhánh hoàn thành. Thời gian yêu cầu gốc là:

T=10+max⁡(60,40)+10=80 μs.T=10+\max(60,40)+10=80\ \mathrm{\mu s}.

Tăng tốc A bốn lần xuống 15 μs, với B vẫn mất 40 μs, cho thời gian mới 10+max⁡(15,40)+10=6010+\max(15,40)+10=60 μs. A tiết kiệm 45 μs, nhưng yêu cầu chỉ tiết kiệm 20 μs, vì B thế chỗ A thành đường tới hạn.

Hình 5-38 Bước kết thúc chờ cả nhánh A lẫn nhánh B. Sau khi A rút từ 60 μs xuống 15 μs, nhánh chậm hơn chuyển từ A sang B, và yêu cầu giảm từ 80 μs xuống 60 μs. Bước chuẩn bị và bước kết thúc mỗi cái 10 μs, và hai nhánh dùng các tài nguyên độc lập.

Khi A co xuống 40 μs, nó hoàn thành đúng vào lúc B hoàn thành; việc tiếp tục chỉ tối ưu A khiến thời gian hoàn thành mắc kẹt ở 60 μs. Điểm giao cắt này đánh dấu lợi ích tối đa có được từ việc chỉ tối ưu A. Để rút ngắn yêu cầu hơn nữa, nên tối ưu B, hoặc giảm bước chuẩn bị và bước kết thúc.

Các tài nguyên dùng chung cũng có thể thay đổi bản thân thời gian của các nhánh. Giả sử cách triển khai mới của A tiêu tốn nhiều băng thông, khiến B đang chạy đồng thời mất tới 90 μs; yêu cầu khi đó trở thành 10+max⁡(15,90)+10=11010+\max(15,90)+10=110 μs. Thời gian thừa mà B mất vượt quá thời gian tính toán mà A tiết kiệm, và cả yêu cầu rốt cuộc lại chậm hơn.

Đường tới hạn cho thấy thời gian tiết kiệm của một nhánh phải được đưa trở lại toàn bộ đồ thị phụ thuộc để biết mức thực sự nhanh hơn của yêu cầu. Việc tiết kiệm lưu trữ cũng phải được đưa trở lại toàn bộ quá trình thực thi theo cùng cách: mức giảm mà một mẩu trạng thái chiếm dung lượng không nhất thiết bằng mức giảm số byte đọc mỗi bước, vì lượng đọc mỗi bước phụ thuộc vào việc bao nhiêu lớp đọc nó và mỗi lớp đọc những mục nào. Phiên bản V4.1 trong Mục 2.3.6 đã đưa ra một tập số liệu dưới điều kiện 8K: V4.1 có 38 lớp global attention dùng chung 4 bản sao global KV, cắt mức cư trú xuống khoảng một phần tư của V4-Flash, nhưng mức đọc logic mỗi lớp chỉ giảm khoảng 9%, và mức giảm đó nhỏ hơn mức giảm cư trú. Việc dùng chung trạng thái không đòi hỏi mỗi lớp phải lưu một bản sao riêng, nhưng mọi lớp dùng nó vẫn phải đọc nó, nên mức giảm dung lượng không chuyển hóa trực tiếp thành mức giảm đọc.

Khi ngữ cảnh dài thêm, phép quét chỉ số vẫn phải chạy cũng sinh thêm chi phí. Ở 128K, khối lượng trạng thái lịch sử toàn cục và khối lượng đọc chỉ số logic của hai thế hệ lần lượt xấp xỉ 65.085 và 33.422 MiB. Trong V4.1, việc truy cập chỉ số của các lớp phía sau chỉ chạm tới nhóm ứng viên (candidate pool), trong khi chỉ số ban đầu vẫn phải duyệt toàn bộ trạng thái toàn cục. Vậy dù phạm vi truy cập chỉ số của các lớp sau bị cố định, chi phí của toàn bộ bước decode vẫn tăng theo độ dài ngữ cảnh.

Điều làm thay đổi chi phí không chỉ là dữ liệu đến từ đâu, mà còn là chương trình có thực sự tránh được phép tính hay không. Nhóm ứng viên thu hẹp phạm vi tìm kiếm cho các chỉ số phía sau, giảm khối lượng công việc tương ứng. Bản tham chiếu công khai trước tiên tính tích vô hướng trên toàn bộ chỉ số đã cache, rồi che các vị trí ngoài nhóm ứng viên; bản triển khai production trong bài báo trực tiếp thu hẹp số mục mà các chỉ số phía sau quét. Bản trước đã thực hiện tích vô hướng ngoài nhóm, trong khi bản sau bỏ qua công việc đó tại thời điểm phép tính được phát. Cùng một quy tắc lựa chọn vì vậy tương ứng với hai chi phí thực thi khác nhau.

Hình 5-39 Mức cư trú toàn cục và mức đọc logic theo từng lớp được so sánh riêng. Phần trên chỉ đếm dung lượng KV toàn cục; phần dưới bao gồm các mục toàn cục, cửa sổ cục bộ và chỉ số; cả hai biểu đồ dùng cùng một thang trục hoành. Các giá trị được tính dưới bố trí production.

5.6.3 Ca tổng hợp: các lựa chọn tối ưu cho một yêu cầu Qwen3

Áp dụng phương pháp này cho Qwen3-8B trên RTX PRO 6000. Đầu vào của yêu cầu là 7239 token, sinh ra 32 token đầu ra, độ đồng thời 1, dùng BF16 và thực thi eager; ở đây, eager có nghĩa là host gửi từng mục công việc cho accelerator theo thứ tự chương trình. Model có 36 lớp, trước tiên thực thi một lượt prefill để sinh đầu ra đầu tiên, rồi 31 lượt decode. Mỗi lớp thực thi một phép kích hoạt SwiGLU mỗi bước, tức Z=SiLU⁡(G)⊙UZ=\operatorname{SiLU}(G)\odot U của Mục 5.3.1, cho tổng cộng 36×32=115236\times32=1152 lần gọi, trong đó 36 lần là prefill và 1116 lần là decode.

Hai loại gọi mang khối lượng công việc khác nhau. Trong cùng một yêu cầu, prefill thực thi operator này trên 7239 token đầu vào qua từng một trong 36 lớp model, xử lý tổng cộng 36×7239=26060436\times7239=260604 biểu diễn token; 31 bước decode mỗi bước xử lý một token mới cho mỗi lớp, thực thi tổng cộng 36×31=111636\times31=1116 lần gọi operator đơn token. Ở đây ta thống kê công việc operator theo “số lớp × số token”, nên cùng một token tham gia nhiều lần bị đếm nhiều lần. Decode có số lần gọi gấp 31 lần prefill, nhưng xử lý ít hàng hơn rất nhiều. Đây chính xác là khác biệt hình dạng mà chương này đã gặp đi gặp lại: ma trận lớn mang lại nhiều công việc song song hơn, ma trận nhỏ mất ít thời gian tính toán hơn, và chi phí khởi chạy kernel cùng phân phối tác vụ trở nên nổi bật hơn và khó che giấu hơn, khiến việc tận dụng hết accelerator trở nên khó khăn hơn.

Trong cùng một engine, việc thay thế phép SwiGLU trên cả 36 lớp bằng một kernel dùng scheme lập lịch khác, và xác nhận từ các dấu vết thực thi rằng cả 1152 lần gọi đều dùng kernel thay thế, cho các thời gian theo giai đoạn sau.

Giai đoạn SwiGLUHình dạng và số lần gọiKernel gốcKernel thay thế
Prefill7239 hàng, 36 lần gọi~11.2 ms~11.2 ms
Decode1 hàng, 1116 lần gọi~2.4 ms~0.92 ms

Kernel thay thế chủ yếu rút ngắn thời gian kích hoạt decode với đầu vào một hàng, tăng tốc khoảng 2.6×, tiết kiệm tổng cộng khoảng 1.45 ms cho phần decode. Thời gian kích hoạt prefill về cơ bản không đổi, thực ra lâu hơn khoảng 0.03 ms sau khi thay thế. Các phép kích hoạt này thực thi tuần tự với các công việc accelerator khác trong dấu vết. Tổng thời gian tính toán kích hoạt của kernel gốc khoảng 13.6 ms, khoảng 1.6% của cửa sổ capture khoảng 851 ms; giữ mọi công việc khác cố định, ngay cả khi loại bỏ toàn bộ thời gian tính toán kích hoạt thì cũng chỉ tiết kiệm tối đa 13.6 ms. Mới chỉ phần decode được rút ngắn cho đến nay, đưa tổng thời gian tính toán kích hoạt từ 13.6 ms xuống 12.2 ms, nên cả yêu cầu dự kiến giảm khoảng 1.4 ms.

Đo các yêu cầu đan xen trước và sau khi thay thế phía client cho 11 cặp kết quả. Thời gian token đầu tiên trước và sau khi thay thế đều khoảng 436 ms, và thời gian yêu cầu đầy đủ trung vị lần lượt khoảng 817 ms và 815 ms. Tính “thời gian trước khi thay thế trừ thời gian sau khi thay thế” cho mỗi vòng, trung vị của các cặp chênh lệch khoảng 1.3 ms, khoảng 0.16% thời gian yêu cầu gốc, khớp với dự kiến ~1.4 ms. Hình 5-40 liệt kê chênh lệch từng vòng, với 9 cặp cho thời gian ngắn hơn sau khi thay thế và 2 cặp cho thời gian dài hơn.

Hình 5-40 Thời gian yêu cầu trước khi thay thế trừ thời gian yêu cầu sau, trong cùng một vòng; giá trị dương nghĩa là cách thay thế nhanh hơn. Qwen3-8B trên RTX PRO 6000, đầu vào 7239 token, đầu ra cưỡng bức 32 token, độ đồng thời 1, BF16, eager, tắt prefix cache; 11 cặp đo đan xen, thứ tự ngẫu nhiên trong mỗi vòng, với các dịch vụ cư trú khác chạy đồng thời. Đường nét đứt đánh dấu trung vị cặp chênh lệch khoảng 1.3 ms. Bảng theo giai đoạn đến từ một bản ghi profiling được capture riêng.

Cải thiện decode xảy ra sau đầu ra đầu tiên, nên với yêu cầu đầu vào dài hiện tại, thời gian token đầu tiên vẫn bị quyết định bởi phép tính và thời gian chờ bên trong prefill. Đầu ra càng dài, phép tính kích hoạt một hàng càng lặp lại nhiều lần, và phần tiết kiệm theo từng bước càng tích lũy nhiều hơn.

Khi phân tích một yêu cầu đầy đủ, trước hết thống kê số lần gọi cho từng hình dạng, rồi tính thời gian tiết kiệm mỗi lần gọi, và cuối cùng dùng cấu trúc phụ thuộc để phán đoán liệu các phần tiết kiệm đó có rút ngắn cả yêu cầu hay không. Điều này nối mức tăng tốc của một kernel, thời gian mà mỗi giai đoạn mất, và thời gian mà người dùng thực sự phải chờ.

Các cạm bẫy thường gặp

Cạm bẫy: Phán đoán chi phí thực thi chỉ bằng số kernel. Phát lại đồ thị đã giảm số lần gửi lệnh phía host cho 3 lần gọi FFN từ 18 xuống 3, trong khi accelerator vẫn thực thi 18 kernel; chỉ fusion mới giảm số kernel trên accelerator xuống 15. Việc gửi lệnh của host và việc tính toán trên accelerator nên được đọc trên các dòng thời gian riêng của chúng.

Cạm bẫy: Cách triển khai có khối lượng truy cập nhỏ nhất luôn nhanh nhất. Thu nhỏ khối K/V của attention từ 64 hàng xuống 1 hàng làm giảm khối lượng truy cập 144 MiB, nhưng số lần cập nhật trạng thái tăng khoảng 41×. Các lần đọc, các lần cập nhật trạng thái và tốc độ xử lý ma trận cùng nhau quyết định chi phí thực thi.

Cạm bẫy: Thêm buffer luôn cải thiện thông lượng. Trên một SM của H100, mỗi khối mất 1.29 μs để di chuyển và 0.28 μs để tính; hai ngăn đã giữ cho bộ di chuyển làm việc liên tục, hoàn thành bốn khối trong 5.44 μs. Một ngăn thứ ba không thể làm cho việc di chuyển nhanh hơn, cũng như việc tăng tốc tính toán không thể; để rút ngắn thời gian, phải giảm số byte di chuyển mỗi khối hoặc tăng băng thông.

Cạm bẫy: Thực thi chế độ ổn định nhanh nhất nghĩa là tổng thời gian ngắn nhất. Chạy 70 nhóm gọi, bucketing mất khoảng 1.27 s trong khi chuyên biệt theo hình dạng mất khoảng 1.38 s. Chuyên biệt thực thi mỗi nhóm nhanh hơn, nhưng phần tiết kiệm được chưa đủ để bù chi phí chuẩn bị thừa.

Bài tập và thí nghiệm kèm theo

Các câu hỏi sau dần dần thay đổi các điều kiện của các ví dụ trong chương này. Trước hết hoàn thành suy luận trên giấy, rồi dùng mã kèm theo hoặc các dấu vết thực thi để giải thích sự khác biệt; các mục thí nghiệm giữ nguyên đánh số gốc, với cấu hình, quy trình và hồ sơ đầy đủ nằm trong tài liệu tương ứng.

Thí nghiệm 5-1 [Mở rộng]: Kích thước khối ma trận thay đổi bộ nhớ cục bộ và khối lượng đọc như thế nào

Đổi khối đầu ra của Ví dụ 5-3 thành 64×128, giữ khối rút gọn ở 32, và tìm yêu cầu bộ nhớ cục bộ, số lần đọc riêng của A và W, và tổng khối lượng truy cập. So với scheme 64×64, băng thông hiệu dụng của scheme 64×128 phải cao hơn bao nhiêu thì thời gian truy cập bộ nhớ của nó mới ngắn hơn? Rồi so sánh các dấu vết thực thi cho hai thứ tự vòng lặp, và giải thích việc truy cập liên tục ảnh hưởng tới tốc độ xử lý ra sao.

Thí nghiệm 5-2 [Cốt lõi]: Fusion các operator kích hoạt giảm được bao nhiêu truy cập bộ nhớ và mức chiếm dụng GPU memory

Với một chuỗi kích hoạt mà G và U mỗi cái 24 MiB, vẽ thứ tự cấp phát, sử dụng và giải phóng của từng tensor khi ba operator thực thi độc lập so với khi được fusion hoàn toàn, và suy ra khối lượng truy cập 156 MiB và 60 MiB cùng mức đỉnh chiếm dụng GPU memory của hai scheme. Nếu cách triển khai fusion thêm chi phí cố định 5 μs mỗi lần gọi, và băng thông giao diện lấy là băng thông bộ nhớ của RTX PRO 6000 là 1792 GB/s, cần tối thiểu bao nhiêu hàng đầu vào thì thời gian truy cập bộ nhớ tiết kiệm được mới vượt quá chi phí này? Rồi, kết hợp với kết quả thí nghiệm kèm theo, giải thích mối quan hệ giữa tỉ lệ giảm khối lượng truy cập và tỉ lệ giảm thời gian.

Thí nghiệm 5-3 [Mở rộng]: Rút gọn theo khối và prefetching thay đổi phép tính attention và truy cập bộ nhớ như thế nào

Thêm một phần tử nữa vào đầu vào của Ví dụ 5-6, với điểm số ln⁡4\ln4 và giá trị tương ứng 5, cập nhật m, ℓ\ell và u, và so sánh với việc xử lý ba phần tử cùng lúc. Giữ tổng dung lượng buffer attention ở 99 KB (101,376 byte), khoét thêm một ngăn prefetch BF16 b×d, lấy b=64, và tìm số hàng Q tối đa mới, số lượt quét, và khối lượng truy cập. Giải thích vì sao buffer prefetch làm thay đổi số lần dữ liệu được đọc lặp lại.

Thí nghiệm 5-4 [Mở rộng]: Fusion operator làm thay đổi kết quả tính toán hay chi phí thực thi không

Trong mã giả của Mục 5.4.2, lần lượt chuyển SiLU vào trong vòng ko, loại bỏ phép làm tròn BF16 trung gian, và chuyển phép tính lượng tử hóa vào trong vòng lặp theo từng khối cột đầu ra. Nêu ra mỗi thay đổi làm thay đổi phép toán toán học, phép làm tròn số học, hay chi phí thực thi, và đưa ra một đầu vào có thể phân biệt các kết quả hoặc khối lượng truy cập trước và sau mỗi thay đổi.

Thí nghiệm 5-5 [Mở rộng]: Các vòng lặp theo khối thay đổi thứ tự truy cập và bộ nhớ tạm như thế nào

Viết vòng lặp i-k-j theo khối dùng chỉ số khối và tọa độ trong khối, đánh dấu nơi bộ cache đầu vào và bộ tích lũy được tạo, dùng lần cuối, và được giải phóng trong chương trình. Rồi, đối chiếu với phép thực thi của chương trình chuyển vị kèm theo dưới hai kích thước khối, so sánh tính liên tục của truy cập, các thao tác đồng bộ, và lượng dữ liệu tạm, và dùng điều này để giải thích khác biệt về thời gian thực thi.

Thí nghiệm 5-6 [Mở rộng]: Tỉ lệ gọi theo hình dạng đầu vào ảnh hưởng lợi ích tuning ra sao

Dùng hai hình dạng đầu vào của Ví dụ 5-9, cho tỉ lệ gọi của hình dạng A lần lượt là 1/2, 2/3 và 9/10, và tính thời gian trung bình của cách triển khai gốc và mới. Nếu việc chọn cách triển khai theo hình dạng tốn thêm 1 μs mỗi lần, dùng cách triển khai mới cho A và cách cũ cho B, khi nào cách này thắng cách luôn dùng cách triển khai gốc? Rồi, dựa trên các hồ sơ tìm kiếm kèm theo, thống kê tổng thời gian tìm kiếm, và tính cách triển khai được tìm ra tiết kiệm bao nhiêu thời gian thực thi qua các lần gọi lặp lại.

Thí nghiệm 5-7 [Mở rộng]: Specialization cần bao nhiêu lần thực thi lặp lại mới có lời

Vẽ toàn bộ dòng thời gian chuẩn bị và thực thi trong Ví dụ 5-12. Đổi số lần gọi của hình dạng nhỏ mỗi nhóm từ 8 xuống 4, giữ hai hình dạng còn lại ở mỗi cái một lần gọi. So sánh từng cặp scheme và tìm số lần lặp lại mà tổng thời gian của chúng bằng nhau. Rồi, giả sử hai chương trình theo bucket đã được biên dịch và cache, tìm khoảng số lần lặp lại mà mỗi scheme có tổng thời gian ngắn nhất sau khi các thời gian chuẩn bị thay đổi.

Thí nghiệm 5-8 [Cốt lõi]: Phát lại đồ thị và fusion operator lần lượt giảm những chi phí nào

Đọc Hình 5-28 và thống kê số lần gửi lệnh của host và số kernel của accelerator cho bốn cách tiếp cận. Chuyển Ví dụ 5-11 sang RTX 4090, lấy băng thông sao chép là băng thông bộ nhớ 1008 GB/s của nó, và tìm kích thước đầu vào mà chi phí sao chép triệt tiêu đúng thời gian tiết kiệm được nhờ phát lại đồ thị; rồi cho operator phía trước ghi thẳng vào buffer của đồ thị, và so sánh thời gian hoàn thành của gửi lệnh thông thường so với phát lại đồ thị ở các kích thước đầu vào 2 MiB và 16 MiB.

Thí nghiệm 5-9 [Cốt lõi]: Tăng tốc kernel làm thay đổi thời gian yêu cầu và đường tới hạn như thế nào

Dùng 11 cặp hồ sơ yêu cầu kèm theo, trước hết tính chênh lệch thời gian cho từng cặp, rồi lấy trung vị của các chênh lệch này. Giải thích điều này khác với việc lấy trung vị mỗi nhóm thời gian yêu cầu rồi mới trừ cho nhau ra sao. Dùng các điều kiện của ca tổng hợp trong chương này, giữ nguyên prefill và tăng số bước decode từ 31 lên 127; giả sử phần tiết kiệm mỗi bước không đổi, ước tính tổng thời gian tiết kiệm trong việc tính toán kích hoạt. Rồi rút dần thời gian thực thi của A trong Hình 5-38, và tìm thời gian thực thi của A mà tại đó thời gian yêu cầu đầy đủ không rút ngắn thêm nữa.

Lịch sử và tài liệu đọc thêm

Halide tách bạch thuật toán khỏi schedule, biến việc tiling, sắp xếp lại và đặt vị trí tính toán thành các hành động có thể tổ hợp cho lập trình viên; TVM kết nối tính toán tensor, sinh mã backend và tuning tự động; AKG kết hợp lập lịch đa diện, fusion phân cấp và sắp xếp bộ nhớ trên chip để sinh mã cho các bộ xử lý thần kinh. Việc đọc tư liệu gốc đằng sau cả ba thứ cho phép truy vết sâu hơn cách một biểu diễn nhất định định hình những phép biến đổi nào mà trình biên dịch có thể tự động thực hiện.

Một FPGA là một thiết bị logic cấu hình được; tổng hợp cấp cao (HLS) chuyển các chương trình cấp cao thành các đường dữ liệu phần cứng. Công việc FPGA/HLS sớm của tác giả cũng trải qua một bước chuyển từ việc bọc giao diện sang biến đổi mã: việc hình thành pipeline yêu cầu nhận diện các primitive và các phụ thuộc đọc/ghi, và thay đổi đường dữ liệu được sinh ra. Trải nghiệm này minh họa cách trọng tâm công việc đã dịch chuyển từ việc đơn giản hóa việc lập trình sang cải thiện việc thực thi.

Hệ thống tối ưu điều phối operator Korch cung cấp một ví dụ đáng học hỏi khác: trong một trường hợp lịch sử, một đồ thị con vốn thực thi dưới dạng ba kernel mất tổng khoảng 91 μs; sau khi được tái cấu trúc thành bốn kernel, tổng giảm xuống khoảng 69 μs. Việc chia lại và tổ hợp lại các operator rút ngắn thời gian thực thi trên accelerator đủ để bù chi phí khởi chạy kernel thừa và chi phí trao tay dữ liệu. Cùng với phản ví dụ về fusion trong chương này, hướng này cho thấy việc tối ưu nên so sánh chi phí của toàn bộ quá trình tính toán.

Chú thích nguồn

Tóm tắt chương

Chương này bắt đầu từ một phép chiếu ma trận đơn và đi qua truy cập bộ nhớ, fusion operator, biên dịch và runtime, kết thúc bằng việc phân tích một yêu cầu đầy đủ. Các mục đã bám vào ba phán đoán cơ bản lặp đi lặp lại.

Thứ nhất, tái sử dụng cần không gian. Một tile đầu ra lớn hơn cho phép cùng một đầu vào tham gia vào nhiều phép nhân–cộng hơn, nhưng nó cũng đòi hỏi nhiều không gian tích lũy hơn; giữ lại một kết quả lượng tử hóa cho phép nhiều khối cột đọc dữ liệu ở độ rộng bit thấp hơn, nhưng nó cũng đòi hỏi lưu trữ trung gian. Việc phán đoán khoản lưu trữ này có đáng không nghĩa là nhìn vào mức đọc lại hay tính lại mà nó tránh được. Trên một SM, không gian này cũng quyết định bao nhiêu thread block có thể cư trú, và điều đó đến lượt quyết định bao nhiêu yêu cầu bộ nhớ đang bay trên đường có thể che giấu độ trễ.

Thứ hai, phụ thuộc dữ liệu quyết định thứ tự tính toán. Một phép toán elementwise có thể hoàn thành ngay khi đầu vào ở một vị trí sẵn sàng, trong khi một phép rút gọn phải giữ các tổng từng phần hoặc các thống kê trực tuyến. Trình biên dịch di chuyển phép tính dựa trên các phụ thuộc này, và runtime dùng các event để xác nhận liệu dữ liệu đã sẵn sàng chưa và liệu một buffer có thể được giải phóng chưa. Fusion và pipeline đều được xây dựng trên các quan hệ phụ thuộc này.

Thứ ba, tổng thời gian phụ thuộc vào thứ tự thực thi của từng mẩu công việc. Băng thông quyết định việc chuyển một loạt dữ liệu mất bao lâu, giai đoạn chậm hơn quyết định khoảng cách hoàn thành giữa các khối liền kề trong một pipeline, số lần tái sử dụng quyết định liệu phần tiết kiệm thời gian thực thi có đủ để bù thời gian chuẩn bị ban đầu, và đường tới hạn quyết định liệu một cải thiện cục bộ có rút ngắn được yêu cầu hay không. Vậy cần vừa tính xem bao nhiêu công việc đã được giảm, vừa xác định xem chính công việc đó có ảnh hưởng tới thời gian hoàn thành cuối cùng ngay từ đầu hay không.

Chương tiếp theo mở rộng các quan hệ này sang nhiều accelerator. Khi một mẩu kết quả sẵn sàng, ai cần đọc nó, và nó phải được giữ trong bao lâu sẽ quyết định khi nào việc truyền thông có thể bắt đầu và mỗi đầu cần bao nhiêu bộ đệm.