Phân tích hiệu năng trong PyTorch (Phần 3): Attention là tất cả những gì bạn cần phân tích

Đây là bài viết thứ ba trong series Phân tích hiệu năng trong PyTorch (Profiling in PyTorch), một chuỗi bài viết giúp bạn xây dựng kỹ năng đọc các bản vết (trace) của trình phân tích hiệu năng (profiler) để tối ưu hóa mô hình:

1. Phân tích hiệu năng trong PyTorch (Phần 1): Hướng dẫn cho người mới bắt đầu về torch.profiler
2. Phân tích hiệu năng trong PyTorch (Phần 2): Từ nn.Linear đến Fused MLP
3. Phân tích hiệu năng trong PyTorch (Phần 3): Attention là tất cả những gì bạn cần phân tích (Bài hiện tại)

Series này nhằm giúp bạn làm quen với việc đọc các bảng và vết thực thi từ profiler. Trong Phần 1, chúng ta đã phân tích các phép toán cơ bản như cộng và nhân. Chúng ta đã thấy cách bảng profiler chỉ ra các “điểm nóng” (hotspots) và cách bản vết cho thấy thứ tự chạy của một thuật toán theo thời gian.

Trong Phần 2, chúng ta đã gói các phép cộng và nhân đó vào một lớp linear của torch, sau đó xếp chồng nhiều lớp linear lên nhau (một mạng perceptron đa lớp - MLP) và phân tích nó. Trong quá trình đó, chúng ta cũng đã phân tích các kernel hợp nhất (fused kernels) và các kernel được tinh chỉnh thủ công.

Xét từ góc độ kiến trúc Transformer, bước logic tiếp theo là phân tích một thuật toán nền tảng khác: Attention. Mặc dù nổi tiếng với độ phức tạp thời gian bậc hai, nhưng có nhiều thủ thuật thông minh để giảm thiểu vấn đề này và tăng tốc độ. Mục tiêu của chúng ta ở đây không phải là đi sâu vào chi tiết mọi thủ thuật, mà là xem mỗi thủ thuật trông như thế nào dưới góc nhìn của profiler.

Các mã nguồn cho bài viết này nằm tại đây: 04_a_naive_attention.py, 04_b_inplace_ops_attention.py, 04_c_sdpa_attention.py, và 04_d_kernels_attention.py. Bạn nên mở chúng trong một tab riêng để vừa đọc vừa theo dõi mã nguồn. Chúng tôi sử dụng GPU NVIDIA A100-SXM4-80GB để chạy các script này.

Attention cơ bản (Naive attention)

Attention hoạt động với Queries (q), Keys (k), và Values (v). Sự tương tác giữa chúng có thể được viết thành một chuỗi các bước ngắn:

1. Xây dựng điểm số attention scores: matmul(q, k.T)
2. Thu nhỏ (scale) điểm số: scores * scale
3. Áp dụng mặt nạ nhân quả (causal mask) cho điểm số: scores.masked_fill(mask, "-inf")
4. Chuẩn hóa điểm số bằng softmax để lấy trọng số attention attn: softmax(scores)
5. Tính toán lại giá trị với các trọng số đó: matmul(attn, v)

Vì vậy, attention thực chất là một tập hợp các phép toán nguyên thủy. Một số phép toán chúng ta đã biết (matmul), và những phép còn lại cũng dễ dàng nhận diện. Hãy cùng viết một module attention cơ bản trong PyTorch và phân tích nó.

class NaiveCausalAttention(nn.Module):
    def __init__(self, head_dim):
        super().__init__()
        self.scale = 1.0 / math.sqrt(head_dim)

    def forward(self, q, k, v, mask):
        scores = torch.matmul(q, k.transpose(-2, -1))
        scores = scores * self.scale
        scores = scores.masked_fill(mask, float("-inf"))
        attn = torch.softmax(scores, dim=-1)
        out = torch.matmul(attn, v)
        return out

Trước khi mở bản vết, hãy cùng dự đoán xem chúng ta sẽ thấy gì. Khi theo dõi hàm forward của module này, chúng ta kỳ vọng thấy:

• Một kernel matmul (q . k.T)
• Một kernel mul (phép scale)
• Một thao tác cho masking
• Một kernel softmax
• Một kernel matmul (attn . v)

uv run 04_a_naive_attention.py
uvx trace-util -f traces/ -b <hf_uname>/traces

Hình 1 cho thấy luồng CPU (luồng GPU được thu gọn để tránh gây nhiễu). Bên trong attn_fwd (lời gọi forward được đánh dấu), chúng ta có thể thấy chính xác các thao tác đã dự đoán:

• mul: phép scale
• masked_fill: causal masking
• softmax: kernel softmax

Bây giờ, hãy mở luồng GPU để xem các kernel nào thực sự được khởi chạy.

Hình 2 cho thấy luồng GPU bên cạnh luồng CPU. Hãy phóng to một khối attn_fwd trên luồng GPU để xem chi tiết từng kernel.

Hình 3 cho phép chúng ta đọc các kernel riêng lẻ cho một bước:

1. matmul (query và key)
2. mul (scaling)
3. copy bộ nhớ (memory copy) 🤔
4. causal masking
5. softmax (tạo trọng số attention)
6. matmul (trọng số attention và values)

Năm trong số này là đúng như dự kiến. Phép copy bộ nhớ là điểm bất thường, vậy nó từ đâu ra? Manh mối là PyTorch có các thao tác “tại chỗ” (in-place). Khi bạn thao tác trên một tensor theo cách thông thường (out-of-place), PyTorch thường tạo một bản sao, áp dụng thao tác lên đó và trả về bản sao này. Theo trình tự các thao tác, “thủ phạm” ở đây chính là masked_fill.

Điều gì sẽ xảy ra nếu chúng ta thay thế nó bằng một thao tác tại chỗ?

Attention cơ bản với Causal Masking tại chỗ (In-place)

Tất cả những gì chúng ta thay đổi là chuyển masked_fill thành masked_fill_ (lưu ý dấu gạch dưới ở cuối, quy ước của PyTorch cho các thao tác in-place).

def forward(self, q, k, v, mask):
    # q, k, v: [batch, heads, seq, head_dim]
    scores = torch.matmul(q, k.transpose(-2, -1))  # [batch, heads, seq, seq]
    scores = torch.mul(scores, self.scale)
-    scores = scores.masked_fill(mask, float("-inf"))
+    scores.masked_fill_(mask, float("-inf"))
    attn = torch.softmax(scores, dim=-1)
    out = torch.matmul(attn, v)  # [batch, heads, seq, head_dim]
    return out

Hãy xem bản vết để thấy sự thay đổi.

uv run 04_b_inplace_ops_attention.py
uvx trace-util -f traces/ -b <hf_uname>/traces

Phiên bản tại chỗ (Hình 5) bao gồm ít thao tác CPU hơn nhiều trong bước masking so với phiên bản thông thường (Hình 4). Đây là một tín hiệu đáng khích lệ. Hãy mở luồng GPU để xác nhận.

Trên luồng GPU, kernel Memcpy đã biến mất hoàn toàn (Hình 6 và 7). Chỉ với một thay đổi nhỏ, chúng ta đã loại bỏ được một kernel trong mỗi lần forward. Điều này có vẻ không nhiều, nhưng hãy nhớ đây chỉ là một thao tác attention đơn lẻ. Trong một mô hình lớn dựa trên transformer (LLM, mô hình Diffusion…), thao tác này lặp lại ở mỗi lớp, và có rất nhiều lớp, vì vậy mức tiết kiệm sẽ tích lũy nhanh chóng.

Việc mặc định dùng out-of-place trong PyTorch là có lý do. Để tính gradient, autograd phải ghi nhớ các giá trị tensor trong lượt forward vì nhiều công thức backward cần dùng lại chúng. Thao tác in-place sẽ ghi đè các giá trị này trong bộ nhớ, khiến lượt backward đọc sai số. Vì chúng ta chạy forward trong torch.no_grad, việc dùng in-place là an toàn. Thao tác in-place không chỉ tiết kiệm thời gian mà còn tiết kiệm bộ nhớ vì không cần bản sao thêm.

Scaled Dot Product Attention (SDPA)

Chúng ta vừa xây dựng attention từ các phép toán cơ bản và thậm chí loại bỏ được Memcpy. Tin tốt là đội ngũ PyTorch đã làm điều này cho chúng ta và đóng gói toàn bộ quy trình vào một hàm duy nhất:

from torch.nn import functional as F

F.scaled_dot_product_attention(q, k, v, is_causal=True)

Một dòng này thay thế toàn bộ module tự viết, và is_causal=True giúp chúng ta không cần tự tạo mặt nạ (mask). Scaled Dot Product Attention (SDPA) không chỉ có một triển khai duy nhất. Bên dưới, nó sẽ điều phối (dispatch) đến một trong số nhiều backend và chọn cái nhanh nhất hỗ trợ đầu vào của chúng ta (kiểu dữ liệu, chiều head, mask, phần cứng, v.v.).

Các backend được liệt kê trong enum torch.nn.attention.SDPBackend:

• MATH
• FLASH_ATTENTION
• EFFICIENT_ATTENTION
• CUDNN_ATTENTION

Chúng ta có thể cố định một backend cụ thể bằng context manager torch.nn.attention.sdpa_kernel để phân tích từng cái một.

Math backend

uv run 04_c_sdpa_attention.py --backend math
uvx trace-util -f traces/ -b <hf_uname>/traces

Chúng ta kỳ vọng bản vết sẽ đơn giản hơn và nhanh hơn vì dùng một dòng lệnh thay vì tự viết. Hãy kiểm tra bảng profiler trước:

Bất ngờ đầu tiên: dòng lệnh ngắn gọn này lại chậm hơn 3.7 lần.

Bản vết (Hình 9) cho thấy lý do: math backend khởi chạy 20 kernel GPU cho mỗi lượt forward, trong khi bản naive chỉ chạy 5 (Hình 8). Tại sao lại như vậy?

Tensor cores bị bỏ trống

Hãy nhìn vào tên kernel:

Lượt chạy	Kernel matmul
Hình 10: Attention Naive
Hình 11: SDPA math backend

GPU A100 có Tensor Cores, phần cứng chuyên dụng để tăng tốc matmul, nhanh hơn nhiều so với CUDA cores thông thường.

• Ký hiệu s16816 trong kernel naive (Hình 10) là chữ ký của matmul Tensor Core bfloat16, nghĩa là bản naive đang đi trên “đường cao tốc”.
• Ký hiệu sgemm (Hình 11) là matmul độ chính xác đơn (FP32) chạy trên CUDA cores thông thường.

Nói cách khác, math backend không hề chạm vào Tensor Cores: để đổi lấy độ chính xác số học, nó nâng kiểu dữ liệu lên FP32 (gấp đôi lượng dữ liệu cần di chuyển) và quay về dùng CUDA cores chậm hơn.

Mặt nạ nhân quả được xây dựng lại

Trong bản naive, chúng ta xây dựng mặt nạ một lần và tái sử dụng. Ở đây, với is_causal=True, math backend tạo mới mặt nạ trong mọi lần gọi. Bạn có thể thấy điều này trên luồng CPU (Hình 12):

Thứ tự thao tác: aten::ones $\rightarrow$ aten::tril (tạo ma trận tam giác dưới) $\rightarrow$ aten::scalar_tensor $\rightarrow$ aten::fill_ (tạo giá trị -inf) $\rightarrow$ aten::where (chuyển thành bias cộng).

Safe softmax

Bản tự viết dùng aten::softmax đơn giản. Math backend dùng aten::_safe_softmax, dẫn đến nhiều kernel hơn (Hình 13):

Nếu một hàng bị mask hoàn toàn (tất cả là -inf), softmax thông thường sẽ tính exp(-inf)/sum(exp(-inf)) = 0/0 = NaN. _safe_softmax ngăn chặn điều này.

Vậy math backend dùng để làm gì?

Math backend là bản triển khai tham chiếu (reference implementation). Nó là một sự phân rã attention thành các op ATen cơ bản, an toàn về kiểu dữ liệu và an toàn với NaN. Nó cơ bản là bản attention naive chúng ta viết, nhưng cẩn thận hơn. Chính sự cẩn thận này khiến nó cực kỳ chậm.

Nhiệm vụ của nó không phải là nhanh, mà là luôn luôn hoạt động đúng. Điều này khiến nó trở thành mốc so sánh (baseline) hoàn hảo.

Efficient backend

uv run 04_c_sdpa_attention.py --backend efficient
uvx trace-util -f traces -b <hf_uname>/traces

Thay vì 20 kernel, efficient backend chỉ khởi chạy một kernel duy nhất: fmha_cutlassF_bf16_aligned_64x64_rf_sm80 (Hình 14).

Giải mã tên kernel:

• fmha (fused multi-head attention): Tất cả các op nguyên thủy trong attention giờ được “hợp nhất” (fused) vào một op.
• cutlassF: Xây dựng trên CUTLASS (mẫu open-source của NVIDIA cho Tensor-core GEMMs), F là forward.
• bf16_aligned: Chạy ở bfloat16 (không nâng lên FP32 như bản math).
• 64x64: Kích thước tile.
• rf (register file): Tập làm việc được giữ trong thanh ghi, bộ nhớ nhanh nhất trên chip.
• sm80: Biên dịch cho Ampere (A100).

Đây là kernel attention tiết kiệm bộ nhớ phát triển từ thư viện xformers của Meta.

Flash backend

uv run 04_c_sdpa_attention.py --backend flash
uvx trace-util -f traces -b <hf_uname>/traces

Kernel void pytorch_flash (Hình 15) chính là FlashAttention-2, được tích hợp vào PyTorch.

Tại sao Flash Attention lại tồn tại?

Vấn đề thực sự của math backend không phải là số lượng 20 kernel, mà là những gì các kernel đó chuyển cho nhau.

Bước 1 tạo ra ma trận điểm số đầy đủ attn = q . k.T với kích thước [seq, seq] cho mỗi head. Với chiều dài chuỗi 4096, đó là $4096 \times 4096 \approx 16$ triệu số. Ma trận này được ghi ra HBM (bộ nhớ chính của GPU). Sau đó, nó được đọc lại để scale, ghi lại cho mask, đọc lại cho softmax, v.v. Chi phí của Attention bị chi phối bởi lưu thông qua lại với HBM, không phải bởi chính các phép matmul.

FlashAttention giải quyết chính xác điều này. Thay vì tính toàn bộ ma trận s rồi mới thu gọn, nó duyệt qua k và v theo từng tile, giữ một softmax chạy (online softmax) và tích lũy kết quả theo từng tile. Ma trận điểm số đầy đủ [seq, seq] không bao giờ được ghi ra HBM, nó chỉ tồn tại trên chip.

Tại sao Flash trông có vẻ “sai” dưới profiler?

Flash là backend nhanh nhất, nhưng profiler báo cáo tỷ lệ lấp đầy (occupancy) rất thấp (~13% ở Hình 16). Tại sao?

Một kernel GPU chạy nhiều luồng (threads) được nhóm thành các khối (blocks). Các block được lập lịch lên Streaming Multiprocessors (SMs). Một SM có thể chứa nhiều block nếu có đủ tài nguyên (thanh ghi, bộ nhớ chia sẻ). Khi nói occupancy thấp, nghĩa là mỗi SM có ít warp (nhóm 32 luồng) hơn mức lý thuyết có thể hỗ trợ.

Flash sử dụng rất nhiều thanh ghi trên mỗi luồng và lượng lớn bộ nhớ chia sẻ trên mỗi block (Hình 17). Việc sử dụng tài nguyên “nặng” này là có chủ đích để giữ các tile attention trên chip, tái sử dụng dữ liệu quyết liệt và tránh việc ghi ma trận attention đầy đủ ra bộ nhớ toàn cục.

cuDNN backend

uv run 04_c_sdpa_attention.py --backend cudnn
uvx trace-util -f traces -b <hf_uname>/traces

cuDNN cũng cung cấp một kernel hợp nhất kiểu flash (Hình 18). Vậy tại sao lại cần thêm cuDNN nếu đã có Flash? Sự khác biệt nằm ở cách kernel được viết và xây dựng.

Flash và Efficient là các kernel cố định, được biên dịch sẵn. cuDNN là thư viện của NVIDIA, kernel attention của nó được tạo ra và tinh chỉnh cho bài toán cụ thể hiện tại.

Tên kernel cuDNN rất dài: cudnn_generated_fort_native_sdpa_sm80_flash_fprop_wmma_f16_knob_6_128x64x64_4x1x1_cga1x1x1_kernel0_0. Nó cho thấy:

• cudnn_generated: Được tạo ra bởi cuDNN, không phải binary sẵn có.
• flash_fprop: Kiểu Flash Attention forward pass.
• wmma_f16: Dùng API WMMA cho Tensor-core pipeline 16-bit.
• knob_6: Chọn từ một tập hợp các cấu hình đã tinh chỉnh (“knobs”).

Điều này dẫn đến những điểm bất thường trong bản vết:

1. Không có phép transpose: Luồng CPU đi thẳng từ _cudnn_attention_forward đến cấp phát bộ nhớ rồi chạy kernel, không có aten::transpose (Hình 19, 20, 21). cuDNN tiêu thụ trực tiếp layout [B, H, S, D].
2. Khởi chạy qua cuLaunchKernelEx: Thay vì cudaLaunchKernel thông thường, cuDNN dùng driver-level extended launch (Hình 22).
3. Báo cáo 0% occupancy: Đây là lỗ hổng đo lường. CUPTI không thể gán occupancy cho cuLaunchKernelEx như cách nó làm với cudaLaunchKernel. Thực tế, nó vẫn dùng rất nhiều thanh ghi (Hình 23), tương đương với Flash.

Chi phí chuyển sang CPU

Dù không có phép transpose, cuDNN không phải là backend “nhẹ” nhất trên CPU.

cuDNN tốn khoảng 214 µs mỗi lượt forward trên CPU, nhiều hơn Flash và Efficient. Phần lớn thời gian nằm ở việc engine runtime của cuDNN chọn và chuẩn bị kế hoạch chạy (tìm kiếm “knob”) trong mỗi lần gọi. Khi một bản vết đột nhiên trông “sạch” hơn (ít op hơn), không có nghĩa là công việc biến mất, đôi khi nó chỉ chuyển vào sâu trong thư viện nơi profiler không thể phân tách.

Tổng kết mọi biến thể

Lời kết cho series

Nếu bạn chỉ rút ra một điều từ toàn bộ series này, hãy chọn thói quen: Dự đoán trước, rồi mới xem.

Hãy nói thành lời những gì bạn kỳ vọng bản vết sẽ chứa, mở nó ra, và coi mọi điểm không khớp là điều thú vị nhất trên màn hình. Mọi hiểu biết sâu sắc trong ba bài viết này — từ Memcpy ẩn, đến 20 kernel của math backend, occupancy “sai” của Flash hay thanh CPU “béo” của cuDNN — đều đến từ một dự đoán không khớp với thực tế.

Phân tích hiệu năng không phải là một kỹ năng đáng sợ chỉ dành cho chuyên gia GPU. Đó chỉ là kỷ luật của việc quan sát kỹ và hỏi “đợi đã, tại sao điều này lại xảy ra?” cho đến khi tìm ra câu trả lời. Giờ đây bạn đã có từ vựng và phản xạ để làm điều đó với mô hình của chính mình. Hãy mở một bản vết, đưa ra dự đoán, và đi tìm điểm không khớp.

Cảm ơn bạn đã đọc series Profiling in PyTorch. Bây giờ, hãy đi phân tích hiệu năng một thứ gì đó đi! 🤗

Link bài viết gốc

• Tags:
• Ai
• 10 July 2026
• Huggingface.co