Làm chủ CUDA và Điện toán hiệu năng cao, Phần VIII

Nơi Phần VII đã để lại chúng ta

Phần VII đã kết thúc với một lời hứa và một nút thắt kiến trúc đầy tò mò.

Lời hứa: trên kiến trúc Hopper, tỷ lệ giữa lệnh tính toán và lệnh tải trong một vòng lặp nội bộ GEMM tiến dần đến vô cực từ góc nhìn của SMSP.

Nút thắt: một lệnh duy nhất di chuyển một khối dữ liệu (tile) 128×128 BF16, đơn vị TMA tự tạo tất cả các địa chỉ và một thứ gọi là mbarrier thay thế cho __syncthreads() mà bạn đã viết từ khi bắt đầu học chương trình "hello world" trong CUDA.

Hãy cùng phân tích chính xác điều đó có nghĩa là gì, tại sao NVIDIA lại đưa ra những lựa chọn đó, và bạn cần phải hiểu những gì để viết, đọc hoặc debug các nhân (kernel) CUTLASS 3.x mà không cảm thấy như mình đang đọc một ngôn ngữ xa lạ.

Chúng ta sẽ đi sâu vào chi tiết. Không còn cách nào khác.

Vấn đề mà cp.async chưa hoàn toàn giải quyết được

Phần VII đã xác nhận rằng cp.async vượt trội hơn con đường truyền thống LDG → STS vì nó loại bỏ các thanh ghi đích khỏi bảng điểm (scoreboard). SMSP phát lệnh sao chép, bàn giao nó cho Công cụ Sao chép Bất đồng bộ (Async Copy Engine) và ngay lập tức rảnh tay để thực hiện lệnh tiếp theo.

Điều này thực sự rất tuyệt vời. Nhưng nó có một chi phí ẩn chỉ trở nên rõ ràng khi bạn nhìn vào luồng lệnh SMSP của một nhân GEMM thực tế.

Hãy xem xét một khối BF16 128×128×32. Việc tải khối đó đòi hỏi 128 × 32 phần tử BF16 = 4096 BF16 = 8 KB. Với 16 byte trên mỗi lệnh cp.async, đó là 512 lệnh CP.ASYNC.CA.SHARED.GLOBAL riêng biệt.

512 lệnh đó phải được lấy từ bộ nhớ đệm lệnh (instruction cache), giải mã, gửi qua đơn vị MIO và được theo dõi bởi phần cứng. Chúng tiêu tốn băng thông lệnh SMSP mặc dù không tạo ra kết quả nào trong thanh ghi.

Trên kiến trúc Ampere, SMSP có thể phát khoảng một lệnh cp.async 128-bit mỗi 4 chu kỳ trên mỗi SMSP. Đối với 512 lệnh, đó là khoảng 2048 chu kỳ SMSP cho mỗi lần tải khối, chỉ tính riêng chi phí quản lý lệnh. Việc di chuyển dữ liệu thực tế diễn ra bất đồng bộ, nhưng luồng lệnh không miễn phí.

Đối với các khối lớn, điều này vẫn trong tầm kiểm soát. Nhưng với các khối nhỏ hơn, hoặc với các kiến trúc mà bạn muốn SMSP dành mọi chu kỳ cho các lệnh nhân tensor (tensor core), thì đây là một giới hạn.

Hopper (SM90, H100) được thiết kế để loại bỏ hoàn toàn giới hạn đó. Câu trả lời chính là Bộ tăng tốc bộ nhớ Tensor (Tensor Memory Accelerator).

Bộ tăng tốc bộ nhớ Tensor

TMA là một đơn vị phần cứng được giới thiệu trong Hopper, thực hiện việc sao chép tensor đa chiều giữa bộ nhớ toàn cục (global memory) và bộ nhớ chia sẻ (shared memory) (hoặc bộ nhớ chia sẻ phân tán trong một cụm, nhưng chúng ta sẽ nói về các cụm sau).

Nó chấp nhận một tensor descriptor được tính toán trên host và một tập hợp các tọa độ được tính toán trên device, và nó xử lý mọi thứ khác: tính toán địa chỉ, striding, chuyển đổi kiểu dữ liệu, clamping vượt giới hạn, chính sách cache và báo hiệu hoàn tất giao dịch.

Hãy cụ thể hơn về ý nghĩa của cụm từ “mọi thứ khác”.

Trong một tiled GEMM thông thường, với mỗi tile bạn tải, mọi thread trong warp phải tính toán phần địa chỉ bộ nhớ toàn cục (global memory) của riêng mình.

Việc tính toán địa chỉ đó bao gồm chỉ số block, chỉ số thread, kích thước tile, matrix stride và kích thước phần tử. Đó hoàn toàn là các phép toán xác định tạo ra kết quả giống hệt nhau mỗi khi bạn thực hiện cùng một vòng lặp tile.

Đó cũng là các phép toán mà SMSP phải thực thi. Trên kiến trúc Ampere với cp.async, các phép toán đó vẫn diễn ra trong SMSP mặc dù giao dịch bộ nhớ sau đó là không đồng bộ.

TMA loại bỏ các phép toán đó khỏi SMSP. Một thread phát ra một lệnh với handle của tensor descriptor và một cặp tọa độ (y, x).

TMA unit sử dụng các tọa độ đó và metadata của descriptor để tính toán mọi địa chỉ cần thiết cho toàn bộ quá trình truyền tile, scatter hoặc gather dữ liệu, và ghi nó vào shared memory. SMSP chỉ cần phát ra một lệnh. Một lệnh duy nhất.

Đây không phải là một tối ưu hóa nhỏ. Đó là một sự thay đổi về chất trong những gì SMSP thực hiện trong một GEMM kernel. Trên Hopper, công việc của SMSP là chạy các lệnh WGMMA.MMA_ASYNC.

Công việc của TMA là di chuyển dữ liệu. Hai công việc này diễn ra đồng thời, trên các đơn vị phần cứng riêng biệt, và sự giao tiếp duy nhất giữa chúng là một đối tượng đồng bộ hóa mbarrier.

Subscribe now

Tensor descriptor

Trước khi một Hopper kernel chạy, host phải tạo một tensor descriptor bằng cách sử dụng cuTensorMapEncodeIm2col hoặc, phổ biến hơn cho GEMM, cuTensorMapEncodeTiled. Đây là một cấu trúc mờ 128-byte được lưu trữ trong constant memory (hoặc truyền qua thanh ghi và tải vào L1).

Descriptor mã hóa:

Base pointer: địa chỉ bộ nhớ toàn cục của phần tử tensor [0, 0, 0, ...].

Global dimensions: kích thước thực của mỗi chiều trong tensor đầy đủ, tính theo số phần tử. Với ma trận M×K A, đây là {M, K} (hoặc {K, M} nếu là column-major).

Global strides: byte stride giữa các phần tử liên tiếp trong mỗi chiều. Với ma trận row-major có K cột và các phần tử BF16, stride giữa hàng i và hàng i+1 là K × 2 byte. Các stride này cho phép các tensor không liên tục tùy ý.

Box dimensions: kích thước của tile cần truyền trong mỗi chiều. Với tile 128×32 BF16, đây là {128, 32}.

Interleave and swizzle mode: cách dữ liệu được sắp xếp lại trong quá trình truyền để tạo ra bố cục shared memory tránh được xung đột bank (bank conflicts). Đây là phần thay thế cho tất cả các phép toán đệm (padding arithmetic) từ Phần VII.

Element stride and data type: cách thông dịch các byte thô.

Descriptor được tạo một lần trên CPU và truyền đến kernel. Trên device, một warp đơn lẻ hoặc thậm chí một thread đơn lẻ sau đó có thể sử dụng descriptor này để khởi tạo một quá trình truyền tile hoàn chỉnh chỉ với một lệnh, vì tất cả thông tin bất biến theo tile đều đã được mã hóa.

Đây là một lựa chọn thiết kế có chủ đích: chuyển các phép tính tốn kém (tạo descriptor) sang host, nơi mà độ trễ không quan trọng so với chi phí khởi chạy kernel, để lệnh ở phía device có thể rẻ nhất có thể.

Lệnh TMA

PTX cho một lệnh load TMA 2D trông như sau:

cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes
    [smem_dst], [gmem_desc, {coord_y, coord_x}], [mbar];

Hãy phân tích từng token.

cp.async.bulk có nghĩa đây là một quá trình sao chép khối không đồng bộ; “bulk” phân biệt nó với cp.async vô hướng. Kích thước truyền được xác định bởi descriptor, không phải mã hóa trong lệnh.

tensor.2d nghĩa là TMA sẽ thông dịch các tọa độ như một truy cập tensor 2D. Có các biến thể cho tensor từ 1D đến 5D.

shared::cluster là phạm vi đích: shared memory có thể truy cập được bởi toàn bộ cluster của thread block (sẽ nói thêm về cluster sau). Đối với các kernel đơn CTA, đây chỉ đơn giản là shared memory.

global là nguồn: global memory, được lập chỉ mục thông qua descriptor.

mbarrier::complete_tx::bytes là cơ chế báo hiệu hoàn tất. Khi quá trình truyền hoàn tất, TMA sẽ báo hiệu cho một đối tượng mbarrier, giảm số lượng giao dịch của nó.

Khi số lượng về 0, các thread đang chờ tại barrier sẽ được giải phóng. Điều này thay thế cho consumer_wait() và __syncthreads() theo nghĩa là bản thân barrier theo dõi cả việc dữ liệu đến và đồng bộ hóa thread trong một primitive duy nhất.

[smem_dst] là địa chỉ đích trong shared memory.

[gmem_desc, {coord_y, coord_x}] là descriptor cộng với tọa độ. TMA trích xuất base pointer, stride và box dimensions từ descriptor, áp dụng các tọa độ và tạo ra toàn bộ dải địa chỉ.

[mbar] là con trỏ tới đối tượng mbarrier trong shared memory.

Trong CUDA C++, API cuda::experimental::tma:: (hoặc __pipeline_memcpy_async cho các trường hợp đơn giản hơn) sẽ tạo ra lệnh này. Đường dẫn sản xuất chính tắc là thông qua cute::copy của CUTLASS 3.x với một TMA copy atom, mà chúng ta sẽ xem xét trong phần CUTLASS.

Subscribe now

Một primitive đồng bộ hóa bạn chưa từng thấy trước đây

__syncthreads() là một barrier toàn bộ thread block. Mọi thread trong block phải đến nơi trước khi bất kỳ thread nào được tiếp tục.

Nó được thực hiện thông qua một bộ đếm chia sẻ, bộ đếm này bị giảm bởi mỗi thread khi đến nơi và được kiểm tra bởi một cơ chế barrier phần cứng. Chi phí của nó tỉ lệ thuận với kích thước thread block, và nó không thể phân biệt giữa “Tôi đã tính toán xong” và “dữ liệu của tôi đã đến từ DMA engine”.

mbarrier (memory barrier, hay chính xác hơn là đối tượng barrier của Hopper) giải quyết cả hai vấn đề đó.

Một đối tượng mbarrier là một giá trị 64-bit được lưu trữ trong shared memory. Nó có hai pha, expect và arrive, và nó theo dõi hai số lượng riêng biệt:

Arrival count được giảm bởi các thread gọi mbarrier.arrive hoặc mbarrier.arrive_drop. Khi số lượng này về 0, pha của barrier sẽ lật ngược.

Transaction count (số lượng giao dịch) được giảm bởi chính engine TMA khi quá trình sao chép khối hoàn tất. Đây chính là complete_tx::bytes trong lệnh PTX ở trên. Lập trình viên khởi tạo số lượng này bằng số byte dự kiến mà TMA sẽ truyền tải.

Rào cản (barrier) được coi là “hoàn tất” khi cả hai bộ đếm đều về không: tất cả các luồng (thread) tham gia đã đến nơi và tất cả các giao dịch TMA dự kiến đã hoàn thành.

Điều này có nghĩa là bạn có thể để một consumer (bên tiêu thụ) đợi trên một rào cản được báo hiệu một phần bởi các luồng và một phần bởi các engine DMA phần cứng, mà không cần vòng lặp polling, không cần các thao tác atomic trong đường dẫn quan trọng (critical path) và không cần __syncthreads() để tuần tự hóa toàn bộ 128 luồng trong block.

Thiết lập này trông như sau trong CUDA C++:

__shared__ cuda::barrier<cuda::thread_scope_block> mbar;

// One thread initializes the barrier for N_THREADS participants
if (thread_rank == 0) {
    init(&mbar, N_THREADS);
    // Tell the barrier to also expect TMA_BYTES bytes of async data
    cuda::device::barrier_native_handle(mbar).arrive_tx(TMA_BYTES);
}
__syncthreads();  // This syncthreads is to publish the initialized mbar

// Producer thread issues TMA
if (thread_rank == 0) {
    tma_load(&mbar, smem_A, gmem_desc_A, tile_coord_m, tile_coord_k);
}

// All threads arrive at the barrier (decrement arrival count)
auto token = cuda::device::barrier_native_handle(mbar).arrive();

// Wait for both arrival count and transaction count to reach zero
cuda::device::barrier_native_handle(mbar).wait(std::move(token));

Lưu ý về sự bất đối xứng: một luồng phát lệnh TMA, tất cả các luồng tham gia vào việc đồng bộ hóa rào cản. Đây không phải là lỗi; đó là thiết kế.

TMA là một thao tác singleton mà một luồng khởi tạo, nhưng dữ liệu mà nó chuyển đến được tiêu thụ bởi tất cả các luồng, vì vậy tất cả các luồng phải đồng bộ hóa dựa trên việc hoàn thành của nó.

Lệnh gọi arrive_tx thông báo cho rào cản rằng các byte TMA đang được chờ đợi. Nếu không có nó, rào cản sẽ hoàn thành ngay khi tất cả các luồng đến nơi, bất kể dữ liệu DMA đã nằm trong bộ nhớ chia sẻ hay chưa. Đó sẽ là một race condition (tình trạng tranh chấp).

token được trả về bởi arrive là một phase token (mã thông báo giai đoạn). mbarrier hoạt động theo các giai đoạn luân phiên (giống như double buffer ở cấp độ đồng bộ hóa), và token đảm bảo rằng wait đợi đúng giai đoạn.

Đây là cách Hopper tránh vấn đề ABA khi tái sử dụng rào cản: bạn không thể vô tình đợi trên một giai đoạn rào cản đã hoàn thành trong lần lặp trước đó.

Đăng ký ngay

Warpgroup MMA

Phần VII đã không đề cập sâu về phía tính toán của Hopper vì phía bộ nhớ đã đủ để tiêu hóa rồi. Bây giờ chúng ta cần nói về WGMMA, và nó cũng cấp tiến không kém.

Trên kiến trúc Ampere, các lệnh tensor core được phát ra theo từng warp: HMMA.1688 hoặc mma.sync.aligned trong PTX hoạt động trên các tile 16×8×16 với 32 luồng tham gia. Mỗi warp thực thi độc lập tile nhân ma trận của riêng nó.

Các lệnh tensor core cấp độ warp vốn đã là một sự khác biệt đáng kể so với SIMT, vì tất cả 32 luồng trong một warp cùng hợp tác để tạo ra một tile đầu ra 16×8 duy nhất. Nhưng warp vẫn là đơn vị lập lịch và là đơn vị thực thi tensor core.

Trên kiến trúc Hopper, lệnh tensor core là cấp độ warpgroup: WGMMA.MMA_ASYNC hoạt động trên một nhóm gồm 4 warp (128 luồng) cùng lúc. Các kích thước tile đầu vào cho BF16 là:

• A: 64×16 trên mỗi warpgroup (đóng góp từ thanh ghi hoặc bộ nhớ chia sẻ)

• B: 16×256 trên mỗi warpgroup (luôn lấy từ bộ nhớ chia sẻ)

• C/D: 64×256 bộ tích lũy (trong thanh ghi, chia nhỏ trên 128 luồng)

Một lệnh WGMMA.MMA_ASYNC duy nhất tính toán một BFGEMM 64×256×16, tạo ra 64×256 = 16.384 phần tử đầu ra trong một lệnh.

Để so sánh, một lệnh mma.sync.aligned trên Ampere với hình dạng BF16 lớn nhất tạo ra BFGEMM 16×8×16, với 128 phần tử đầu ra.

Tỷ lệ khối lượng đầu ra là 128:1. Đây chính là ý nghĩa thực tế của việc “tiệm cận tỷ lệ tính toán-trên-tải vô hạn”.

Hậu tố _ASYNC là rất quan trọng: WGMMA.MMA_ASYNC không hoàn thành một cách đồng bộ. 4 warp phát lệnh và kết quả không được đảm bảo nằm trong các thanh ghi tích lũy cho đến khi một lệnh WGMMA.WAIT_GROUP được thực thi.

Phần cứng có thể chồng lấp nhiều thao tác WGMMA đang chạy cùng lúc, và lập trình viên phải chèn các lệnh đợi rõ ràng trước khi đọc các bộ tích lũy.

Do đó, mô hình lập trình ở cấp độ lệnh trông như sau:

WGMMA.MMA_ASYNC D, A, B   ; issue tile multiply k=0
WGMMA.MMA_ASYNC D, A, B   ; issue tile multiply k=1
WGMMA.MMA_ASYNC D, A, B   ; issue tile multiply k=2
...
WGMMA.WAIT_GROUP 0         ; wait for all outstanding WGMMAs
; D accumulator registers now hold valid results

Trong CUDA C++, điều này được thể hiện thông qua các trừu tượng cute::wgmma trong CUTLASS 3.x, hoặc thông qua namespace cuda::wgmma:: cấp thấp hơn. PTX trực tiếp cũng khả thi nhưng hoàn toàn không được khuyến khích ngoài các bối cảnh nghiên cứu.

Lý do B luôn phải đến từ bộ nhớ chia sẻ (không phải thanh ghi) là do ràng buộc phần cứng. Các đơn vị tensor core trên Hopper được kết nối trực tiếp với các mảng bộ nhớ chia sẻ.

Toán hạng B được nạp trực tiếp từ các bank bộ nhớ chia sẻ bởi đường dẫn dữ liệu tensor core mà không cần đi qua file thanh ghi.

Đây là lý do tại sao TMA truyền B vào bộ nhớ chia sẻ nằm trên đường dẫn quan trọng, nhưng không có bước “nạp B từ bộ nhớ chia sẻ vào thanh ghi”. Tensor core đọc trực tiếp bộ nhớ chia sẻ.

A có thể đến từ thanh ghi hoặc bộ nhớ chia sẻ. Đối với các kernel hiệu năng cao nhất, A cũng đến từ bộ nhớ chia sẻ, nghĩa là cả hai toán hạng đều bỏ qua file thanh ghi hoàn toàn ở phía tính toán. File thanh ghi chỉ giữ bộ tích lũy C/D.

Đăng ký ngay

Thread Block Clusters

Hopper giới thiệu một cấp độ mới trong phân cấp GPU nằm giữa thread block và grid: thread block cluster.

Một cluster là một nhóm gồm tối đa 8 thread block được đảm bảo lên lịch cùng nhau trên cùng một GPC (Graphics Processing Context, một nhóm các SM chia sẻ chung một lát L2).

Các thread block bên trong một cluster có thể truy cập bộ nhớ chia sẻ của nhau thông qua cơ chế Distributed Shared Memory (DSMEM), sử dụng TMA để di chuyển dữ liệu giữa các SM mà không cần đi qua L2.

Lệnh PTX cho việc truyền TMA liên SM là:

cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes
    [smem_dst], [gmem_desc, {coord_y, coord_x}], [mbar];

Đây là cùng một hướng dẫn như lệnh nạp TMA thông thường, với phạm vi shared::cluster cho biết rằng đích đến hiển thị trên toàn bộ cluster. Đơn vị TMA quản lý việc di chuyển dữ liệu giữa các SM một cách trong suốt.

Tại sao điều này quan trọng đối với GEMM? Hãy xem xét một cluster gồm 2 CTA, mỗi CTA chịu trách nhiệm cho một khối hàng khác nhau của C. Cả hai đều cần truy cập vào cùng các cột của B.

Với các cluster, CTA 0 nạp B vào shared memory của nó thông qua TMA, và CTA 1 có thể đọc trực tiếp shared memory của CTA 0 thông qua DSMEM. B được nạp một lần và được tiêu thụ bởi hai CTA. Điều này thực sự nhân đôi khả năng tái sử dụng B mà không cần nhân đôi shared memory trên mỗi CTA.

Đối với một cluster N=8, 8 CTA chia sẻ việc nạp tile B, phân bổ băng thông HBM cho B trên lượng tính toán lớn gấp 8 lần.

Đây là cơ chế mà qua đó các nhân Hopper GEMM đạt tới hiệu năng đỉnh của phần cứng trên các kích thước bài toán lớn: kiến trúc cluster cho phép tập hợp làm việc (working set) của toàn bộ phép tính được lưu giữ trong shared memory phân tán, với HBM chỉ được truy cập một lần cho mỗi phần tử.

Kích thước cluster được chỉ định tại thời điểm khởi chạy nhân (kernel launch):

cudaLaunchConfig_t config = {};
config.gridDim = grid;
config.blockDim = block;
cudaLaunchAttribute attr;
attr.id = cudaLaunchAttributeClusterDimension;
attr.val.clusterDim.x = 2; // 2 CTAs per cluster
attr.val.clusterDim.y = 1;
attr.val.clusterDim.z = 1;
config.attrs = &attr;
config.numAttrs = 1;
cudaLaunchKernelEx(&config, my_kernel, args...);

Lập lịch cluster mang tính hợp tác: phần cứng sẽ cố gắng đặt các CTA của một cluster vào cùng một GPC, nhưng đây chỉ là một gợi ý, không phải là sự đảm bảo đối với các cluster lớn hơn khả năng chứa của một GPC.

Trên H100 SXM5 với 132 SM được tổ chức thành 7 GPC, các cluster lên đến 8 luôn được thỏa mãn trong một GPC duy nhất.

Mô hình Nhân bền bỉ (Persistent Kernel Model)

Trên Ampere, một nhân GEMM điển hình là một “grid kernel”: mỗi thread block xử lý một tile đầu ra (M_TILE, N_TILE) và kết thúc. CUDA runtime lập lịch các thread block mới ngay khi có dung lượng SM trống.

Đối với các ma trận lớn thì điều này hoàn toàn ổn: có đủ số lượng tile để bộ lập lịch SM luôn bận rộn.

Đối với các ma trận nhỏ hơn, chi phí khởi chạy và kết thúc các thread block chiếm ưu thế. Mỗi thread block phải nạp các tile A và B của nó từ đầu, ghi tile C của nó vào global memory và kết thúc. Trạng thái shared memory không được tái sử dụng giữa các thread block.

Hệ thống phân cấp bộ nhớ và mô hình cluster của Hopper làm cho một phương pháp khác trở nên hấp dẫn: persistent kernels (các nhân bền bỉ).

Trong một persistent kernel, một thread block (hoặc warpgroup) không kết thúc sau khi xử lý xong một tile.

Thay vào đó, nó lặp qua nhiều tile đầu ra, duy trì các tile A và B trong shared memory giữa các lần lặp nơi tile được tái sử dụng, và chỉ nạp các tile mới qua TMA khi cần thiết. Nhân chỉ kết thúc sau khi tất cả các tile đầu ra trong phân vùng được giao của nó đã hoàn thành.

CUTLASS 3.x triển khai điều này thông qua Tile Scheduler, một thành phần phía device giúp quản lý việc phân công các tile đầu ra cho các CTA bền bỉ.

Bộ lập lịch tăng nguyên tử một biến đếm công việc (work counter) được lưu trong global memory, gán cặp (m_tile, n_tile) khả dụng tiếp theo cho CTA đang yêu cầu. Khi tất cả các tile đã được gán, bộ lập lịch sẽ báo hiệu hoàn thành và CTA thoát khỏi vòng lặp công việc.

Các lợi ích là rất cụ thể:

Khả năng tái sử dụng L2 được cải thiện vì cùng một CTA xử lý nhiều tile liền kề, và các tile A hoặc B mà chúng chia sẻ vẫn nằm trong L2 (hoặc thậm chí trong shared memory) giữa các lần lặp.

Chi phí khởi chạy thread block được phân bổ: GPU khởi chạy một làn sóng các CTA bền bỉ và chúng chạy cho đến khi hoàn thành, thay vì khởi chạy hàng ngàn block tạm thời.

Các kích thước bài toán bất thường được xử lý linh hoạt hơn: tile dư cuối cùng được xử lý bởi bất kỳ CTA nào nhận lấy nó, mà không cần khởi chạy nhân epilogue riêng biệt.

Nhược điểm là độ phức tạp lập trình: bạn đang viết một bộ lập lịch phần mềm bên trong một CUDA kernel, với tất cả các mối lo ngại kèm theo về tính đúng đắn khi truy cập đồng thời và cân bằng tải giữa các công việc tile không đồng nhất.

CUTLASS xử lý điều này giúp bạn, đó là một lý do tại sao thư viện này tồn tại.

Đăng ký ngay

Kiến trúc CUTLASS 3.x

CUTLASS 3.x là phiên bản viết lại hoàn toàn của CUTLASS 2.x, được xây dựng trên một lớp trừu tượng mới gọi là CuTe (CUDA Template library).

Việc hiểu CUTLASS 3.x đòi hỏi phải hiểu CuTe, bởi vì CUTLASS 3.x về cơ bản là CuTe cộng với một tập hợp các template nhân sử dụng nó.

CuTe: Layout là các Đối tượng Hạng nhất

Ý tưởng trung tâm của CuTe là một layout là một hàm từ không gian tọa độ logic sang độ lệch vật lý trong bộ nhớ. Một layout mã hóa cả hình dạng (phạm vi của mỗi chiều) và bước nhảy (khoảng cách giữa các phần tử liên tiếp dọc theo mỗi chiều).

Trong CuTe, một layout được viết dưới dạng Shape:Stride. Ví dụ, một ma trận 4×8 row-major với các phần tử có kích thước 2 byte có layout (4,8):(8,1), nghĩa là: chiều ngoài cùng (hàng) có bước nhảy 8 (mỗi hàng cách nhau 8 phần tử), và chiều trong cùng (cột) có bước nhảy 1. Một phiên bản column-major của cùng ma trận đó sẽ là (4,8):(1,4).

Sức mạnh của biểu diễn này là nó có thể kết hợp. Một thao tác tiling chỉ đơn giản là sự kết hợp layout. Một swizzle (hoán vị bit của các địa chỉ để tránh xung đột bank) là một phép biến đổi layout hoán vị các bit địa chỉ theo một mẫu cụ thể.

Toàn bộ phép tính địa chỉ cho một tensor đã được tiling, swizzled và chuyển vị được biểu thị dưới dạng một chuỗi các phép kết hợp layout mà trình biên dịch đánh giá tại thời điểm biên dịch, tạo ra một công thức địa chỉ duy nhất.

Đây là lý do tại sao CUTLASS 3.x có thể biểu diễn các mô hình truy cập phức tạp mà không cần bất kỳ phân nhánh thời gian chạy (runtime branching) nào trong tính toán địa chỉ.

using LayoutA = Layout<Shape<_128, _32>, Stride<_32, _1>>;  // 128x32 row-major
using LayoutA_Swizzled = ComposedLayout<Swizzle<3,3,3>, LayoutA>;

Template Swizzle<B,M,S> mã hóa một phép hoán vị địa chỉ dựa trên XOR cụ thể. Các bit B được hoán vị với các bit S, lệch bởi M bit.

Đối với BF16 với 32 bank, mỗi bank 4 byte, swizzle chính xác sẽ loại bỏ tất cả các xung đột bank mà không cần padding. CUTLASS đi kèm với các tham số swizzle chính xác cho mọi loại phần tử và kích thước tile mà nó hỗ trợ.

MMA Atom và Copy Atom

Trong CUTLASS 3.x, một lệnh tensor core là một MMA atom: một đối tượng có kiểu mô tả hình dạng đầu vào/đầu ra, ánh xạ thread-đến-dữ liệu và lệnh cần phát ra. MMA atom chuẩn của Hopper cho BF16 là:

using MMA_Atom = MMA_Atom<SM90_64x256x16_F32BF16BF16F32_SS>;

Tên gọi mã hóa: SM90 (Hopper), kích thước tile 64×256×16, bộ tích lũy F32, đầu vào A và B là BF16, đầu ra F32, SS nghĩa là cả A và B đều đến từ shared memory.

Một bản sao TMA là một copy atom:

using Copy_Atom_A = Copy_Atom<SM90_TMA_LOAD, bfloat16_t>;

Mẫu kernel CUTLASS kết hợp các atom này với kích thước tile, hình dạng cluster và các giai đoạn pipeline thành một kernel hoàn chỉnh:

using CollectiveMainloop = cutlass::gemm::collective::CollectiveMma<
    cutlass::gemm::MainloopSm90TmaGmmaRmemAAccumulator<3>,  // 3-stage pipeline
    Shape<_128, _256, _64>,                                   // tile MxNxK
    bfloat16_t, LayoutA,
    bfloat16_t, LayoutB,
    TiledMma,
    GmemTiledCopyA,
    SmemLayoutA,
    SmemCopyAtomA,
    cute::identity,
    GmemTiledCopyB,
    SmemLayoutB,
    SmemCopyAtomB,
    cute::identity
>;

Cách này khá dài dòng, nhưng mỗi tham số template đều ánh xạ tới một cơ chế phần cứng cụ thể: MainloopSm90TmaGmmaRmemAAccumulator<3> có nghĩa là “sử dụng TMA để load, sử dụng WGMMA để tính toán, giữ bộ tích lũy (accumulator) trong thanh ghi, với 3 giai đoạn pipeline”.

Trình biên dịch giải quyết tất cả những điều này thành một kernel mà trong đó phần thân vòng lặp chính là một chuỗi các lệnh WGMMA.MMA_ASYNC liên tục, chỉ bị ngắt quãng bởi các lệnh chờ mbarrier do TMA khởi tạo tại các ranh giới giai đoạn.

Việc tính toán địa chỉ cho các lệnh load TMA gần như không xuất hiện trong mã nguồn thiết bị, vì đã được chuyển sang phần xây dựng descriptor trên host.

Mô hình Warpgroup Producer-Consumer

CUTLASS 3.x trên kiến trúc Hopper áp dụng mô hình chuyên biệt hóa warpgroup trong mỗi CTA. Một block thread gồm 128 thread (một warpgroup) được chia tại thời điểm biên dịch thành một producer warpgroup và một hoặc nhiều consumer warpgroup.

Producer warpgroup chịu trách nhiệm thực hiện các lệnh TMA load (mỗi thread phụ trách một load, các thread khác sẽ chờ tại barrier). Các consumer warpgroup chịu trách nhiệm thực hiện các lệnh WGMMA.MMA_ASYNC và chạy phần epilogue (ghi dữ liệu C vào bộ nhớ global thông qua TMA store đầu ra).

Sự chuyên biệt hóa này được thể hiện rõ ràng:

if (warpgroup_id == 0) {
    // Producer: thực hiện TMA load vào các giai đoạn shared memory
    collective_mainloop.load(params, smem_tensors, pipeline, pipeline_state, k_tile_count);
} else {
    // Consumer: thực hiện các lệnh WGMMA, chạy epilogue
    collective_mainloop.mma(params, smem_tensors, accumulators, pipeline, pipeline_state, k_tile_count);
    collective_epilogue.store(params, accumulators, ...);
}

Các producer và consumer warpgroup giao tiếp độc quyền thông qua pipeline bộ nhớ chia sẻ được bảo vệ bằng mbarrier. Không có __syncthreads() nào giữa chúng trong trạng thái ổn định. Các barrier là đủ dùng.

Điều này rất quan trọng về mặt kiến trúc: __syncthreads() là một barrier toàn cục cho CTA. Trong mô hình producer-consumer, nơi mà producer và consumer có khối lượng công việc khác nhau trong mỗi lần lặp, một barrier CTA toàn cục sẽ buộc nhóm nhanh hơn phải chờ nhóm chậm hơn ở mỗi lần lặp.

Nguyên thủy mbarrier cho phép đồng bộ hóa bất đối xứng: consumer chỉ chờ dữ liệu mà nó cần, thay vì chờ producer đạt đến một điểm cụ thể trong luồng điều khiển của nó.

Đăng ký ngay

Pipeline N-Giai đoạn trên Hopper

Phần VII đã mô tả cơ chế double buffering (2 giai đoạn) trên Ampere. Trên Hopper, CUTLASS sử dụng mặc định từ 3 đến 8 giai đoạn, với số lượng giai đoạn tối ưu phụ thuộc vào kích thước tile, kích thước bài toán và mục tiêu chiếm dụng (occupancy).

Máy trạng thái pipeline trên Hopper quản lý N giai đoạn bộ nhớ chia sẻ, N mbarrier cho producer (mỗi giai đoạn một cái, báo hiệu dữ liệu đã đến) và N mbarrier cho consumer (mỗi giai đoạn một cái, báo hiệu rằng consumer đã đọc xong và giai đoạn đó có thể được tái sử dụng).

Vòng lặp ở trạng thái ổn định trông như sau về mặt khái niệm:

Giai đoạn 0: [TMA load A0, B0] → [mbar_full[0] được báo hiệu] → [WGMMA trên A0,B0] → [mbar_empty[0] được báo hiệu]
Giai đoạn 1: [TMA load A1, B1] → [mbar_full[1] được báo hiệu] → [WGMMA trên A1,B1] → [mbar_empty[1] được báo hiệu]
Giai đoạn 2: [TMA load A2, B2] → [mbar_full[2] được báo hiệu] → [WGMMA trên A2,B2] → [mbar_empty[2] được báo hiệu]
Giai đoạn 0: [TMA load A3, B3] → ...

Producer thực hiện TMA load vào giai đoạn i và báo hiệu mbar_full[i]. Consumer đợi mbar_full[i], chạy WGMMA, báo hiệu mbar_empty[i] và chuyển sang giai đoạn (i+1) % N.

Producer đợi mbar_empty[i] trước khi tái sử dụng giai đoạn đó cho lần load tiếp theo. Bộ đệm vòng trong bộ nhớ chia sẻ này, được quản lý bởi các cặp mbarrier, là cấu trúc dữ liệu cơ bản của một Hopper GEMM kernel.

Phần prologue thực hiện load N-1 tile trước khi vòng lặp chính bắt đầu (tương tự như invariant của prologue double buffer trong Phần VII, chỉ là với nhiều giai đoạn hơn). Phần epilogue sẽ xử lý các tile còn đang thực thi sau khi vòng lặp k kết thúc.

Với 3 giai đoạn trên một card H100 có 228 KB bộ nhớ chia sẻ mỗi SM (tăng từ 192 KB trên Ampere), một cặp tile 128×256 BF16 tiêu tốn khoảng:

• Tile A: 128 × 64 × 2 bytes = 16 KB

• Tile B: 64 × 256 × 2 bytes = 32 KB

• Mỗi giai đoạn: 48 KB

• 3 giai đoạn: 144 KB

• Dành cho mbarrier và các accumulator spill: 84 KB

Ở mức 3 giai đoạn và tile 128×256, mỗi SM chỉ chạy được 1 CTA là khả thi. Hai CTA sẽ yêu cầu 288 KB, vượt quá giới hạn 228 KB bộ nhớ chia sẻ.

Do đó, độ chiếm dụng (occupancy) là 1 CTA mỗi SM, điều này hoàn toàn ổn trên Hopper vì một CTA đơn lẻ đó đã lấp đầy SM với các lệnh WGMMA và đơn vị TMA cũng đã được tận dụng tối đa.

Đây là một triết lý về độ chiếm dụng hoàn toàn khác so với Ampere. Trên Ampere, bạn thường cần 2-4 CTA mỗi SM để ẩn độ trễ bộ nhớ thông qua việc chuyển đổi warp.

Trên Hopper, một CTA với TMA và WGMMA đã đạt được thông lượng gần mức đỉnh trên các tile lớn, bởi vì các đơn vị phần cứng quan trọng (TMA, tensor core) đều đã hoạt động hết công suất.

Đăng ký ngay

Trình phân tích (Profiler) cho thấy gì trên Hopper

Các chỉ số trong Nsight Compute thay đổi đáng kể so với Ampere.

smsp__warp_issue_stalled_long_scoreboard tiến gần về 0. Không phải vì bộ nhớ nhanh, mà vì các lệnh TMA load hoàn toàn không liên quan đến scoreboard. SMSP không chờ đợi bộ nhớ; nó không phải là đơn vị đưa ra yêu cầu bộ nhớ đó.

smsp__warp_issue_stalled_mio_throttle cũng thấp. Một lệnh TMA duy nhất cho mỗi tile hầu như không gây áp lực lên đơn vị MIO.

smsp__warp_issue_stalled_wgmma_global_wait trở thành điểm nghẽn (stall) chính mới: đây là lúc SMSP chờ lệnh WGMMA.WAIT_GROUP hoàn thành để có thể đọc các thanh ghi accumulator.

Điểm nghẽn này là không thể tránh khỏi đối với các kernel đọc accumulator giữa các nhóm WGMMA (ví dụ: cho các phép giảm thiểu một phần split-K). Đối với các kernel có chiều K lớn, pipeline WGMMA sẽ đầy và điểm nghẽn này sẽ biến mất.

sm__pipe_tensor_op_hmma_cycles_active nên nằm trong khoảng 80-95% đối với một nhân GEMM được tối ưu hóa tốt trên kiến trúc Hopper. Bất kỳ giá trị nào dưới 70% đều cho thấy vấn đề về độ sâu pipeline (quá ít giai đoạn) hoặc vấn đề lập lịch cluster (GPC không lập lịch các CTA của cluster cùng nhau).

l1tex__data_pipe_lsu_wavefronts_mem_shared_op_ld đếm các thao tác đọc bộ nhớ chia sẻ (shared memory). Đối với một kernel mà cả A và B đều được đọc từ shared memory bởi WGMMA (biến thể SS), số liệu này phản ánh thông lượng của nhân tensor (tensor core), chứ không phải các lệnh tải mà lập trình viên có thể thấy. Các tensor core đang đọc trực tiếp từ shared memory, và điều này hiển thị trong các số liệu LSU.

Các số liệu về thông lượng TMA nằm trong namespace tma: tma__read_bytes và tma__read_transactions. Một kernel đạt được thông lượng TMA đỉnh điểm sẽ cho thấy băng thông TMA gần với băng thông HBM lý thuyết, bởi vì TMA là thành phần duy nhất truy cập vào HBM.

Thông tin chẩn đoán then chốt trên Hopper: nếu mức sử dụng WGMMA của bạn cao và băng thông TMA cao, thì kernel đó đang hoạt động tốt. Hai đơn vị phần cứng này là các điểm nghẽn theo thiết kế. Mọi thứ khác nên ở trạng thái nhàn rỗi hoặc gần như nhàn rỗi.

Đăng ký ngay

Xem xét lại biểu đồ Roofline trên Hopper

Phần VII đã giới thiệu mô hình roofline và lưu ý rằng việc chẩn đoán hữu ích phải mang tính phân cấp: không chỉ là “giới hạn bộ nhớ” (memory-bound) mà là “giới hạn bộ nhớ ở mức L2, đạt 60% mức đỉnh của L2”. Trên Hopper, hệ thống phân cấp có các cấp độ tương tự (L1, L2, HBM) nhưng với các độ dốc mới.

Các thông số roofline của H100 SXM5:

• Băng thông đỉnh HBM3: 3.35 TB/s

• Băng thông đỉnh L2: khoảng 12 TB/s (trên 50 MB L2, hai slice)

• Băng thông đỉnh Shared memory: khoảng 33 TB/s tổng cộng (SM-local)

• Đỉnh Tensor core (dense BF16): 494 TFLOP/s

Các điểm Ridge:

• HBM ridge: 494 / 3.35 ≈ 147 FLOP/byte

• L2 ridge: 494 / 12 ≈ 41 FLOP/byte

• Shared memory ridge: 494 / 33 ≈ 15 FLOP/byte

Đối với một GEMM có cường độ tính toán từ 147 FLOP/byte trở lên, kernel sẽ bị giới hạn bởi khả năng tính toán (compute-bound) với giả định rằng hệ thống phân cấp bộ nhớ được sử dụng đúng cách. Dưới 147 FLOP/byte, nó bị giới hạn bởi băng thông HBM.

Dưới 41, ngay cả một tỷ lệ trúng L2 hoàn hảo cũng không thể cứu vãn hiệu năng. Dưới 15, thông lượng của tensor core bị giới hạn bởi băng thông shared memory, điều này có nghĩa là xảy ra xung đột bank hoặc kích thước tile không làm bão hòa được đường dẫn dữ liệu WGMMA.

Thông tin mới quan trọng trên Hopper: TMA thay đổi hình thái đóng góp của hệ thống phân cấp bộ nhớ. Băng thông lệnh SMSP, vốn là một điểm nghẽn thứ cấp trên Ampere (và là điểm nghẽn chính đối với các tile nhỏ), đã bị loại bỏ khỏi tính toán băng thông HBM một cách hiệu quả.

Băng thông thô tới shared memory hiện chỉ bị giới hạn bởi thông lượng của đơn vị TMA, mà tài liệu H100 liệt kê là khoảng 900 GB/s tổng cộng (trên tất cả các đơn vị TMA của tất cả SM).

Con số này thấp hơn băng thông HBM là 3.35 TB/s, vì vậy đối với các kernel hoàn toàn bị giới hạn bởi băng thông (không phải compute-bound), TMA không phải là điểm hạn chế; HBM mới là điểm đó.

Đối với các kernel compute-bound với tile lớn, việc giảm tải lệnh của TMA chính là yếu tố cho phép SMSP chạy WGMMA ở thông lượng tối đa.

Một cái nhìn ngắn gọn về Blackwell

Blackwell (SM100, B100/B200) đã được công bố vào tháng 3 năm 2024 và bắt đầu giao hàng cho các nhà cung cấp hyperscale vào cuối năm 2024. Quỹ đạo kiến trúc được thiết lập bởi Hopper vẫn tiếp tục và tăng tốc.

Tensor core trên Blackwell giới thiệu thế hệ MMA thứ 5 với hỗ trợ FP4 (các định dạng MXFP4 và NF4), cho phép đạt mức đỉnh 20 PFLOP/s ở cấp độ hệ thống B200 đầy đủ (dual-die). Thông lượng dense FP8 là khoảng 9 PFLOP/s mỗi chip.

TMA trên Blackwell đạt được hỗ trợ gốc cho các phép biến đổi mô hình im2col (liên quan đến các phép tích chập) và lệnh lưu chuyển vị (transposed stores), giảm nhu cầu cho các kernel chuyển vị riêng biệt.

Giới hạn kích thước cluster tăng lên 16 CTA (từ 8 trên Hopper), giúp giảm bớt tải cho các B-tile trên nhiều tài nguyên tính toán hơn.

Một NVLink thế hệ thứ năm mới cung cấp băng thông hai chiều 1.8 TB/s cho mỗi GPU trong các hệ thống kết nối NVLink (rack NVL72), cho phép các kernel đa GPU mà tại đó “bộ nhớ toàn cục” được nhìn thấy bởi một thao tác TMA sẽ phân bổ trên 72 GPU. Đây là cấp độ mà sự khác biệt giữa một kernel đơn GPU và một biểu đồ tính toán phân tán bắt đầu mờ nhạt.

CUTLASS 3.x hỗ trợ Blackwell thông qua các collective template SM100 mới. Mô hình lập trình vẫn giữ nguyên; chỉ là các con số lớn hơn.

Đăng ký ngay

Kết luận

Quỹ đạo từ Volta qua Ampere đến Hopper là một câu chuyện mạch lạc: mỗi thế hệ đều đẩy nhiều công cụ di chuyển dữ liệu ra khỏi SMSP và chuyển sang phần cứng chuyên dụng.

Volta cung cấp cho bạn tensor core, vì vậy SMSP ngừng thực hiện các phép tính số học. Ampere cung cấp cho bạn cp.async, vì vậy SMSP ngừng chờ đợi các lệnh tải. Hopper cung cấp cho bạn TMA, vì vậy SMSP đã ngừng hoàn toàn việc phát hành các lệnh tải.

SMSP trên một nhân GEMM Hopper được tối ưu hóa tốt là một cỗ máy chỉ làm một việc: phát hành lệnh WGMMA.MMA_ASYNC. Mọi thứ khác đã được ủy quyền.

Đây không phải là sự ngẫu nhiên. Đó là điểm kết thúc logic của quan sát rằng nhân matrix multiply là nhân quan trọng nhất đối với các khối lượng công việc ML hiện đại, và phần cứng hiệu quả nhất cho matrix multiply là phần cứng nơi các đơn vị tính toán không bao giờ nhàn rỗi.

Mọi đổi mới kiến trúc từ năm 2017 trở đi đều là cuộc tấn công vào các lý do khác nhau khiến các đơn vị tính toán bị nhàn rỗi: độ trễ số học (tensor core), độ trễ bộ nhớ (cp.async), băng thông lệnh (TMA), băng thông liên SM (clusters, NVLink).

Mbarrier, tensor descriptor, chuyên môn hóa warpgroup, pipeline producer-consumer, trình lập lịch tile: đây không phải là sự phức tạp mang tính trang trí.

Chúng là các cơ chế giúp một GPU 2024 chạy kernel 2024 đạt được 80-90% đỉnh lý thuyết trên phép nhân ma trận, một con số mà những người viết các routine BLAS thủ công cách đây một thập kỷ có vẻ là điều không tưởng.

Phần IX sẽ tạm lùi lại từ bức tranh đơn GPU để nhìn vào sự song song hóa đa GPU: song song hóa tensor, pipeline parallelism, NCCL, và câu hỏi về cách băng thông NVLink tương tác với hiệu suất tính toán trên mỗi GPU mà chúng ta đã dành tám phần để xây dựng.

Công cụ thay đổi; nguyên tắc thì không: tìm điểm nghẽn, xử lý xung quanh nó, và đo lường lại.