Phân tích hiệu năng trong PyTorch (Phần 2): Từ nn.Linear đến một Fused MLP

Ngày xuất bản: 11 tháng 6, 2026

Trong phần đầu tiên của loạt bài “Phân tích hiệu năng trong PyTorch”, chúng ta đã sử dụng torch.add(torch.matmul(x, w), b) để học cách đọc các dấu vết (traces) từ PyTorch profiler. Chúng ta cũng đã thảo luận về một số chủ đề khác như chuỗi điều phối (dispatch chain) của CPU, chi phí khởi chạy (launch overhead), sự khác biệt giữa chế độ bị giới hạn bởi chi phí quản lý (overhead-bound) và chế độ bị giới hạn bởi tính toán (compute-bound), cũng như một số cơ chế nội bộ của torch.compile.

Trong phần thứ hai này, chúng ta sẽ tiến thêm một bước nữa. Chúng ta thay thế cặp matmul-add viết tay bằng một lớp nn.Linear (với bias=True). Đây là khối xây dựng cơ bản mà mọi mô hình deep learning đều sử dụng. Sau đó, chúng ta sẽ xếp chồng ba lớp này (trong ví dụ của chúng ta), với một hàm kích hoạt ở giữa, để tạo thành một khối Perceptron đa lớp (MLP).

Các mã nguồn cho bài viết này nằm tại đây: 02_linear.py, 03_simple_mlp.py, và 03_kernels_mlp.py. Tương tự như trước, bạn nên mở chúng trong một tab riêng và theo dõi mã nguồn trong khi đọc. Chúng tôi sử dụng GPU NVIDIA A100-SXM4-80GB để chạy các kịch bản này.

Trước khi bắt đầu, hãy nhắc lại nhanh hai khái niệm mà chúng ta sẽ xuyên suốt sử dụng:

1. Một kernel GPU là một chương trình chạy song song trên nhiều luồng (threads) của GPU.
2. CPU thực hiện lập lịch và khởi chạy (schedule and launch) các kernel này. Hầu hết các chi phí quản lý (overhead) của PyTorch mà bạn thấy trong trace profiler chính là công việc lập lịch này.

Từ matmul-add đến Linear

nn.Linear thực chất là một lớp bao bọc (wrapper) quanh phép nhân ma trận và phép cộng mà chúng ta đã phân tích ở Phần 1. Điểm khác biệt duy nhất là nó sở hữu trọng số (weight) và độ chệch (bias) dưới dạng các tham số và cung cấp phương thức forward quen thuộc.

# bias=True sẽ mô phỏng chính xác các thao tác nhân và cộng
# mà chúng ta đã thấy trong phần 1 của loạt bài
linear_layer = nn.Linear(in_dim, out_dim, bias=True)
y = linear_layer(x)

Thao tác này có thể được viết là: y = x @ w.T + b

Trong đó x là đầu vào, w là trọng số và b là độ chệch. Hãy chạy 02_linear.py và kiểm tra profile.

uv run 02_linear.py --batch 1024 --in_dim 32 --out_dim 64
uvx trace-util traces -b traces

Hình 1 cho thấy trace profiler của một lần gọi forward của lớp linear. Chúng ta thấy ba bước Profile (Profile Steps) ở cả luồng CPU và GPU do thiết lập wait=1, warmup=1 và active=3.

Phép chuyển vị (transpose) đang làm gì?

Nếu phóng to trace profiler như trong Hình 2, chúng ta sẽ thấy một toán tử aten::t (chuyển vị) xuất hiện trước toán tử aten::addmm (nhân và cộng). Điều này cho thấy nn.Linear thực hiện chuyển vị tham số trọng số trước khi nhân nó với đầu vào.

Một điểm quan trọng cần lưu ý là aten::t không thực sự sao chép hay tổ chức lại dữ liệu: nó chỉ viết lại siêu dữ liệu (metadata) của tensor (hình dạng - shape và bước nhảy - stride) trên CPU để đại diện cho ma trận đã chuyển vị. Nó không khởi chạy bất kỳ kernel nào trên GPU. Bạn có thể kiểm tra điều này bằng cách nhìn vào luồng GPU hoặc xem bảng profiler, thời gian thực thi trên CUDA của aten::t sẽ bằng 0.

Tại sao không có các kernel mul và add riêng biệt?

Không có aten::add (phép cộng bias) trong chuỗi điều phối của lớp linear (Hình 3). Điều này là do phép cộng bias đã được gộp (folded) vào kernel nhân ma trận, sử dụng cái gọi là epilogue.

Epilogue là một tính toán nhỏ mà kernel GEMM (Nhân ma trận tổng quát) thực hiện ở bước cuối cùng, ngay trước khi ghi kết quả trở lại HBM (Bộ nhớ băng thông cao - bộ nhớ chính của GPU). Việc cộng bias, áp dụng hàm kích hoạt hoặc nhân với một hằng số đều là các epilogue điển hình. Mục đích của epilogue là tránh việc phải đọc/ghi vào HBM lần thứ hai, vì lưu lượng bộ nhớ khiến thao tác trở nên tốn kém.

nn.Linear gọi torch.nn.functional.linear, sau đó gọi aten::linear. aten::linear nhận thấy có bias được truyền vào, nên nó điều phối aten::addmm(bias, x, weight) thay vì thực hiện matmul và add riêng biệt.

Kernel GEMM của cuBLAS chạy trên GPU có sẵn một biến thể cộng bias, và đó là kernel mà aten::addmm chọn. Phép cộng không bao giờ xuất hiện dưới dạng một kernel riêng biệt vì nó là một phần của quá trình ghi kết quả của kernel matmul.

Liệu –compile có giúp ích cho một lớp Linear duy nhất?

Hãy thử biên dịch lời gọi forward và xem trace profiler.

uv run 02_linear.py --batch 1024 --in_dim 32 --out_dim 64 --compile
uvx trace-util traces -b traces

Nếu so sánh trace giữa chế độ Eager và Compiled cho một lớp nn.Linear, bạn sẽ thấy:

• Cùng một kernel cuBLAS GEMM trên GPU.
• Cùng một toán tử aten::addmm trên CPU.
• Một vài hàng bổ sung trên luồng CPU đặc thù cho quá trình biên dịch.

Điều này rất đáng lưu ý. Một phản xạ phổ biến là sử dụng torch.compile bất cứ khi nào mô hình cảm thấy chậm. Nhưng đối với một phép GEMM-với-bias đơn lẻ, trình biên dịch có rất ít việc để làm. Điều này không phải là lỗi, mà vì trình biên dịch cần nhiều hơn một thao tác để có thể thực hiện việc hợp nhất (fusing).

Phép chuyển vị biến đi đâu? Bố cục Kernel và Pre-ops

Một người đọc kỹ hai trace (Eager vs Compile) sẽ nhận thấy chuỗi điều phối CPU ở chế độ Eager có nhiều thành phần hơn.

Để hiểu aten::t thực sự làm gì, chúng ta cần nói về strides (bước nhảy) và views (góc nhìn).

Một tensor lưu trữ dữ liệu dưới dạng một dải số phẳng, liên tục trong bộ nhớ. shape (hình dạng) và stride là siêu dữ liệu nằm trên dải đó để chỉ cho PyTorch cách đọc: stride (s0, s1) có nghĩa là “nhảy s0 phần tử để sang hàng mới, nhảy s1 để sang cột mới”. Khi thay đổi siêu dữ liệu này, bạn có một view khác của cùng một dữ liệu thô mà không cần sao chép:

>>> M = torch.tensor([[0, 1],
...                   [2, 3],
...                   [4, 5]])
>>> M.shape, M.stride()
(torch.Size([3, 2]), (2, 1))   # hai bước mỗi hàng, một bước mỗi cột

>>> T = M.t()                  # chuyển vị
>>> T.shape, T.stride()
(torch.Size([2, 3]), (1, 2))   # shape và stride bị hoán đổi, dữ liệu giữ nguyên
>>> T
tensor([[0, 2, 4],
        [1, 3, 5]])
>>> T.flatten()                # buộc phải hiện thực hóa, dữ liệu được sắp xếp lại
tensor([0, 2, 4, 1, 3, 5])

M.t() không di chuyển một con số nào. Nó trả về một view mới với các stride bị hoán đổi. Đây chính xác là những gì aten::t làm: nó tạo ra một view của trọng số với stride được viết lại.

Như thấy ở Hình 5, torch.compile không loại bỏ kernel GPU; nó loại bỏ chi phí điều phối trên CPU để tạo view đó. Inductor đã theo dõi chuỗi view tại thời điểm biên dịch, tính toán stride kết quả một lần và phát ra một lời gọi aten::addmm trực tiếp với các stride đã được mã hóa cứng.

Nếu nhìn vào luồng GPU, cả hai lần chạy đều sử dụng cùng một kernel: cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8

Hậu tố tn trong tên kernel chính là mô tả bố cục (layout). cuBLAS và CUTLASS biên dịch sẵn các kernel riêng biệt cho mỗi tổ hợp bố cục đầu vào (n là không chuyển vị, t là có chuyển vị).

Xếp chồng ba lớp Linear: Tạo thành MLP

Chúng ta sẽ phân tích hiệu năng của một Perceptron đa lớp (MLP) sử dụng biến thể kích hoạt GeGLU.

class SimpleGeGLUMLP(nn.Module):
    def __init__(self, dim, hidden):
        super().__init__()
        self.gate_proj = nn.Linear(dim, hidden, bias=False)
        self.up_proj = nn.Linear(dim, hidden, bias=False)
        self.down_proj = nn.Linear(hidden, dim, bias=False)

    def forward(self, x):
        g = self.gate_proj(x)
        u = self.up_proj(x)
        h = F.gelu(g, approximate="tanh")
        m = h * u
        y = self.down_proj(m)
        return y

Hãy chạy 03_simple_mlp.py và xem trace. Chúng ta kỳ vọng thấy ba lần điều phối aten::linear và hai lần khởi chạy kernel pointwise (một cho GeLU và một cho phép nhân).

Trong mỗi lượt forward, GPU chạy chính xác 5 kernel. Ba phép GEMM thực hiện một lời gọi cudaOccupancyMaxActiveBlocksPerMultiprocessor trước khi khởi chạy (để cuBLAS tính toán kích thước grid), trong khi các thao tác pointwise (GeLU và mul) khởi chạy trực tiếp.

Các thao tác như aten::t, aten::transpose, aten::reshape, aten::view hiển thị 0.000us thời gian CUDA vì chúng chỉ thay đổi siêu dữ liệu trên CPU.

Tại sao có hai loại kernel GEMM?

Trong trace, chúng ta thấy:

• gate_proj và up_proj sử dụng kernel ...128x128...stages_32x5_tn (0.19ms).
• down_proj sử dụng kernel ...128x256...stages_64x3_tn (0.17ms).

Mặc dù cả ba đều có cùng số lượng phép tính (FLOPs), nhưng down_proj nhanh hơn khoảng 10%. Lý do là vì hình dạng ma trận khác nhau ($N=768$ thay vì $3072$), cuBLAS chọn một “tile” (ô gạch) khác giúp tái sử dụng dữ liệu tốt hơn.

torch.compile làm gì trong trường hợp này?

Khi biên dịch phương thức forward:

Trong chế độ Eager, mỗi nn.Linear bị mở rộng thành một chuỗi các toán tử điều phối. torch.compile loại bỏ chuỗi này. Chúng ta chỉ còn thấy ba lời gọi aten::mm thuần túy.

Kernel Triton được hợp nhất (Fused)

Đây là điểm mấu chốt của bài học biên dịch. Hai kernel pointwise (GeLU và mul) cùng với một phép reshape đã bị gộp thành một kernel duy nhất: triton_poi_fused__unsafe_view_gelu_mul_0.

Tại sao điều này lại hiệu quả? Trong chế độ Eager, kết quả trung gian h = gelu(g) là một tensor lớn (khoảng 50 MB) được ghi vào HBM rồi ngay lập tức được đọc lại bởi kernel nhân. Việc hợp nhất giữ dữ liệu này trong các thanh ghi (registers) (bộ nhớ nằm ngay trong chip, nhanh hơn nhiều so với HBM). Một vòng lặp đọc/ghi dữ liệu trung gian qua bộ nhớ toàn cục đã bị loại bỏ.

Sử dụng các kernel được tinh chỉnh thủ công (Hand-tuned)

Thay vì để PyTorch hoặc trình biên dịch chọn kernel, chúng ta sử dụng một kernel được chuyên gia viết và tinh chỉnh thủ công thông qua thư viện kernels từ Hugging Face Hub (ví dụ: LigerGEGLUMLP).

from kernels import get_kernel

kernels_layers = get_kernel("kernels-community/liger-kernels", version=1).layers
kernels_geglu_mlp = kernels_layers.LigerGEGLUMLP(Config()).to(device, dtype=torch.bfloat16).eval()

Tại sao nên dùng thư viện kernels?

Viết kernel trong Triton hoặc CUDA là một chuyện, nhưng triển khai chúng lại là chuyện khác. Kernel phải được biên dịch chính xác cho kiến trúc GPU, phiên bản CUDA và phiên bản PyTorch của bạn. Thư viện kernels giải quyết vấn đề này bằng cách cung cấp các gói kernel đã được biên dịch sẵn và cố định phiên bản.

Tại sao kernel tinh chỉnh lại tốt hơn?

1. Hợp nhất được tích hợp sẵn: LigerGEGLUMLP chạy một kernel Triton duy nhất để tính gelu(gate) * up trong một lần quét. Chúng ta đạt được điều này mà không cần trình biên dịch, do đó không có chi phí từ Dynamo guards hay rủi ro biên dịch lại (Hình 13 và 14).
2. Tham số khởi chạy được tối ưu cho phần cứng: Thay vì dự đoán ngẫu nhiên, Liger chọn kích thước block dựa trên số lượng cột.

Sự đánh đổi: Kernel của Liger chạy trong 92.8 µs, trong khi kernel hợp nhất của Inductor (biên dịch) chạy trong 89.4 µs. Kernel biên dịch nhanh hơn vì nó được tối ưu riêng cho hình dạng chính xác [8192, 3072]. Nếu bạn thay đổi kích thước batch hoặc sequence, bạn sẽ phải trả chi phí biên dịch lại từ đầu.

Vì vậy, lựa chọn thực sự là giữa một kernel tổng quát nhanh (Liger) và một kernel chuyên biệt cho một hình dạng đầu vào cụ thể (torch.compile).

Kết luận

Bảng dưới đây tổng hợp những gì mỗi bước thay đổi trên GPU và những gì giữ nguyên.

Nếu có một thói quen cần ghi nhớ, đó là: đoán trước, rồi mới xem. Hãy xác định những gì bạn kỳ vọng trace sẽ chứa, mở nó ra, và coi bất kỳ sự sai lệch nào là điều thú vị nhất trên màn hình.

Link bài viết gốc

• Tags:
• Ai
• 11 June 2026
• Huggingface.co