Kỹ thuật GPU thực hành: Metal shader, phân tích hiệu năng và triển khai edge

Kỹ thuật GPU thực hành trải dài từ lập trình shader đa nền tảng, chẩn đoán bằng bộ đếm phần cứng, đến triển khai edge trên thiết bị. Dưới đây là phần phân tích sâu về cả ba quy trình.

1. Shader nhân ma trận theo tile tùy chỉnh bằng Apple Metal (MSL)

Trên Apple Silicon (M1–M4 / dòng A), các compute kernel được viết bằng Metal Shading Language (MSL) và tận dụng Threadgroup Memory (tương đương CUDA Shared Memory của Apple) để tránh các lượt đọc global memory chậm.

Dưới đây là một shader nhân ma trận theo tile (C = A × B) bằng MSL sử dụng các tile threadgroup 16 × 16.

#include <metal_stdlib>
using namespace metal;

constant uint TILE_SIZE = 16;

struct MatrixDims {
    uint M; // Rows of A and C
    uint N; // Cols of B and C
    uint K; // Cols of A and Rows of B
};

kernel void matmul_tiled(
    device const float* A        [[buffer(0)]],
    device const float* B        [[buffer(1)]],
    device float* C              [[buffer(2)]],
    constant MatrixDims& dims    [[buffer(3)]],
    uint2 gid                    [[thread_position_in_grid]],
    uint2 tid                    [[thread_position_in_threadgroup]]
) {
    // Fast on-chip Threadgroup (SRAM) tiles
    threadgroup float tileA[TILE_SIZE][TILE_SIZE];
    threadgroup float tileB[TILE_SIZE][TILE_SIZE];

    float acc = 0.0f;
    uint numTiles = (dims.K + TILE_SIZE - 1) / TILE_SIZE;

    for (uint t = 0; t < numTiles; ++t) {
        uint tiledCol = t * TILE_SIZE + tid.x;
        uint tiledRow = t * TILE_SIZE + tid.y;

        // Collaboratively load 16x16 tiles into Threadgroup memory
        tileA[tid.y][tid.x] = (gid.y < dims.M && tiledCol < dims.K)
                              ? A[gid.y * dims.K + tiledCol] : 0.0f;
        tileB[tid.y][tid.x] = (tiledRow < dims.K && gid.x < dims.N)
                              ? B[tiledRow * dims.N + gid.x] : 0.0f;

        // Synchronize all threads in the threadgroup before reading tiles
        threadgroup_barrier(mem_flags::mem_threadgroup);

        // Compute partial dot product entirely inside registers/SRAM
        for (uint k = 0; k < TILE_SIZE; ++k) {
            acc += tileA[tid.y][k] * tileB[k][tid.x];
        }

        // Wait before overwriting tiles in the next loop iteration
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    if (gid.y < dims.M && gid.x < dims.N) {
        C[gid.y * dims.N + gid.x] = acc;
    }
}

• Ánh xạ phần cứng then chốt: [[thread_position_in_grid]] ánh xạ trực tiếp tới blockIdx * blockDim + threadIdx của CUDA, còn threadgroup_barrier(mem_flags::mem_threadgroup) ánh xạ tới __syncthreads() của CUDA. Với các workload FP16 trong sản xuất, Apple Silicon còn cung cấp các intrinsic phần cứng simdgroup_float8x8 để truy cập trực tiếp vào các đơn vị nhân ma trận của Apple.

2. Hướng dẫn từng bước tối ưu hóa Roofline với ncu

Giả sử bạn viết một kernel CUDA hoặc Triton ngây thơ cho Row-Wise LayerNorm trên tensor FP16 (4096, 4096), và nó chạy chậm hơn đáng kể so với triển khai gốc của PyTorch.

Bước 1: Thu thập các bộ đếm phần cứng

Chạy NVIDIA Nsight Compute (ncu) chỉ nhắm vào tên kernel tùy chỉnh của bạn để tránh overhead phân tích trên code khởi tạo:

ncu --kernel-name regex:layernorm_naive --metrics \
sm__throughput.avg.pct_of_peak_sustained_elapsed,\
dram__throughput.avg.pct_of_peak_sustained_elapsed,\
l1tex__t_sectors_pipe_lsu_mem_global_op_ld.sum,\
smsp__warp_issue_stalled_long_scoreboard_per_warp_active.pct \
python run_layernorm.py

Bước 2: Chẩn đoán nút thắt trên Roofline

Kiểm tra các chỉ số đầu ra:

• Thông lượng tính toán (SM): 11.4% mức đỉnh.

• Thông lượng bộ nhớ (DRAM/HBM): 84.2% mức đỉnh.

• Lý do Warp Stall (long_scoreboard): 68.5% (các warp bị đình trệ chờ kết quả đọc bộ nhớ HBM toàn cục trả về).

• Chẩn đoán: Kernel nằm ngay dưới đường Roof băng thông bộ nhớ nghiêng trên biểu đồ Roofline. Hơn nữa, việc kiểm tra l1tex__t_sectors cho thấy các lượt đọc bộ nhớ không gộp (uncoalesced) - các thread trong cùng một warp nhảy bước theo cột thay vì đọc các địa chỉ hàng liên tục, lãng phí 75% của mỗi cache line HBM 128 byte được nạp.

Bước 3: Áp dụng các bản sửa phần cứng có mục tiêu

1. Sửa việc gộp bộ nhớ (Memory Coalescing): Đảm bảo các thread liền kề (threadIdx.x) đọc các phần tử liền kề dọc theo chiều trong liên tục (n_cols) để 32 lượt đọc FP16 hợp thành một giao dịch gộp 64 byte duy nhất.

2. Tải dạng vector (float4 / block tl.load): Tải 128 bit (8 giá trị FP16) mỗi lệnh thay vì 16 bit, giảm tổng số lệnh bộ nhớ đi 8×.

3. Phép rút gọn Warp-Shuffle: Thay các phép rút gọn Shared Memory chậm bằng warp shuffle thanh ghi-sang-thanh ghi (__shfl_xor_sync trong CUDA hoặc tl.sum trong Triton) để tính trung bình và phương sai theo hàng mà không cần barrier đồng bộ.

Bước 4: Kiểm chứng trên Roofline

Chạy lại ncu xác nhận:

• Global Load Sectors: Giảm 4× (mức sử dụng cache line đạt 100%).

• Long Scoreboard Stalls: Giảm từ 68.5% xuống < 15%.

• Độ trễ Wall-Clock: Cải thiện 3.8×, đẩy băng thông hiệu dụng vượt 85% giới hạn HBM vật lý của H100.

3. Quy trình triển khai edge: Lượng tử hóa & chạy LLM trên di động/laptop

Triển khai một LLM 8B tham số lên laptop hoặc điện thoại thông minh (8 GB–16 GB RAM dùng chung) đòi hỏi nén trọng số xuống số nguyên 4-bit và thực thi qua một runtime C++ không overhead như llama.cpp hoặc Apple MLX.

[HF Safetensors (FP16, 16 GB)] 
       │  1. Convert & Apply Calibration Data (Imatrix)
       ▼
[GGUF Quantized Binary (Q4_K_M, 4.9 GB)]
       │  2. Memory-Map (mmap) into Unified RAM
       ▼
[Metal / Vulkan Compute Shaders (On-the-Fly Dequant + Matrix Mul)]

Bước 1: Chuyển đổi & Hiệu chỉnh (Importance Matrix)

Việc làm tròn 4-bit thô làm giảm độ chính xác suy luận và lập trình. Với llama.cpp, trước tiên hãy chuyển trọng số Hugging Face sang file GGUF không nén, sau đó tính Importance Matrix (imatrix) trên một mẫu văn bản hiệu chỉnh nhỏ để đo xem các kênh trọng số nào có tác động kích hoạt lớn nhất:

python convert_hf_to_gguf.py ./Llama-3.1-8B --outfile llama-8b-f16.gguf
./llama-imatrix -m llama-8b-f16.gguf -f calibration.txt -o imatrix.dat

Bước 2: K-Quantization hỗn hợp bit (Q4_K_M)

Lượng tử hóa mô hình bằng Q4_K_M (K-quant 4-bit mức Medium), giữ các layer value/output attention nhạy cảm ở 6-bit (Q6_K) trong khi nén các layer feed-forward dày đặc xuống 4-bit (Q4_K) theo các super-block gồm 256 trọng số:

./llama-quantize --imatrix imatrix.dat llama-8b-f16.gguf llama-8b-q4_k_m.gguf Q4_K_M

Điều này thu nhỏ mô hình từ 16.0 GB xuống còn ~4.9 GB với mức suy giảm perplexity dưới 1%.

Bước 3: Thực thi Zero-Copy trên bộ nhớ thống nhất của thiết bị

Khi chạy trên iPhone, thiết bị Android hoặc MacBook:

• Tải tức thì qua mmap: Runtime ánh xạ bộ nhớ (mmap) file .gguf 4.9 GB trực tiếp từ bộ nhớ flash vào Unified Memory - loại bỏ việc nhân bản RAM và tải mô hình trong vòng chưa đầy 1 giây.

• Shader khử lượng tử hóa tích hợp: Trong quá trình sinh token, các compute shader Metal tùy chỉnh (ggml-metal.metal) hoặc Vulkan (mul_mat_q4_K.comp) nạp các block trọng số 4-bit đóng gói trực tiếp vào thanh ghi GPU, giải nén chúng bằng phép dịch bit, và nhân chúng với các kích hoạt FP16 ngay lập tức - đạt 40–70+ token/giây trên GPU laptop hiện đại và 15–25 token/giây trên điện thoại di động flagship.