Tối ưu attention: đừng học vẹt kernel

Profiler không phải bảng điểm khoe GPU. Nó là công cụ giúp team chọn đúng backend attention trước khi đốt tuần vào tối ưu sai chỗ.

Bụi Wire

2026-07-13

Có một kiểu họp kỹ thuật mình gặp hơi nhiều: ai đó mở profiler trace lên, chỉ vào một ô đỏ chói, rồi cả phòng im lặng như vừa thấy điểm kiểm tra giữa kỳ. Câu tiếp theo thường là: “Chắc phải viết custom kernel thôi.”

Khoan đã. Với hệ thống AI production, đặc biệt là inference hoặc fine-tuning Transformer, câu hỏi hay hơn không phải là “kernel nào nhanh nhất?”, mà là team mình đang ở bài toán attention nào, trên phần cứng nào, với ràng buộc vận hành nào?

Đây là tín hiệu thị trường đáng để ý: Hugging Face đang đẩy mạnh chuỗi bài dạy đọc torch.profiler từ nn.Linear, fused MLP đến attention; Fireworks thì khoe tối ưu sparse attention cho Blackwell với kernel rất sát phần cứng. Hai chuyện này nhìn riêng thì là kỹ thuật. Nhìn chung lại, nó nói một điều: lợi thế đang dịch từ “biết gọi model” sang “biết đọc đường đi của compute”.

Sơ đồ minh họa cho bài Tối ưu attention: đừng học vẹt kernel

Sơ đồ tóm tắt ý chính của bài viết.

Mục tiêu: chọn backend trước, tối ưu sau

Attention là phần giúp token nhìn các token khác để quyết định nên chú ý vào đâu. Trong Transformer, nó vừa quan trọng vừa dễ trở thành điểm nóng vì chi phí tăng mạnh khi context dài.

Nhưng “attention chậm” là một câu chẩn đoán quá rộng. Nó giống giáo viên ghi vào vở: “em cần cố gắng hơn” — đúng, nhưng không đủ để sửa bài.

Với builder, quyết định thực dụng là:

Có cần dùng scaled_dot_product_attention của PyTorch không?
Backend nào đang được chọn: math, efficient, flash, hay cuDNN?
Bottleneck nằm ở compute, memory, launch overhead, hay layout?
Có đáng chuyển sang sparse attention, custom kernel, hoặc framework serving khác không?

Điểm cần đổi trong đầu sau bài này: đừng xem profiler như bảng điểm cuối kỳ; hãy xem nó như giáo án để quyết định bài tiếp theo nên học gì.

Tín hiệu lạ: các layer thấp đang lên mặt tiền

Mấy năm trước, câu chuyện AI app xoay quanh prompt, RAG, agent, fine-tuning. Gần đây, các bài đáng đọc lại chui xuống tầng thấp hơn: profiler trace, fused kernel, attention backend, sparse attention, HBM bandwidth.

Điều này không có nghĩa mọi team Việt Nam phải lập tức tuyển CUDA engineer. Nó có nghĩa incentive đang đổi:

Nhà cung cấp hạ tầng muốn chứng minh họ chạy model rẻ hơn, nhanh hơn, ổn định hơn.
Framework như PyTorch muốn ẩn bớt độ phức tạp bằng API chuẩn, nhưng vẫn cho bạn đường soi backend.
Team build sản phẩm bị ép hiểu performance đủ sâu để không mua nhầm, scale nhầm, hoặc tối ưu nhầm.

Ví dụ Fireworks tối ưu MiniMax M3 sparse attention trên NVIDIA Blackwell: sparse attention chỉ tính trên các block KV liên quan thay vì nhìn toàn bộ context. Ý tưởng nghe hợp lý, nhưng triển khai thật bị vấp ở memory access không đều. Họ dùng hướng KV-stationary, tức giữ block key-value làm trung tâm lịch chạy, để giảm việc load lặp. Kết quả nguồn đưa ra là attention kernel đạt khoảng 980 TFLOP/s ở 4.1 TB/s HBM bandwidth trên một B200 fp8, nhanh hơn query-stationary baseline 1.9–2.4 lần và hơn MSA open-source khoảng 1.6 lần ở tầng kernel.

Con số đẹp, nhưng bài học không phải “hãy dùng ngay sparse attention”. Bài học là: khi context dài, backend attention không còn là chi tiết phụ. Nó là quyết định sản phẩm.

Checklist quyết định: team bạn đang ở lớp nào?

Trước khi đụng custom kernel, thử tự chấm 5 câu này. Không cần làm màu, chỉ cần trả lời thật.

| Câu hỏi | Nếu câu trả lời là “không biết” | Việc nên làm trước |
|---|---|---|
| Attention chiếm bao nhiêu thời gian trong trace? | Bạn chưa có điểm xuất phát | Chạy torch.profiler trên workload thật |
| PyTorch đang chọn backend nào? | Có thể đang dùng đường chậm mà không nhận ra | Log backend, thử SDPA có kiểm soát |
| Shape có ổn định không? | Compiler/kernel khó phát huy | Gom batch/sequence hợp lý hơn |
| Context dài hay ngắn? | Dễ chọn sai giữa flash và sparse | Tách benchmark theo độ dài context |
| Bottleneck là latency hay throughput? | Tối ưu sai KPI | Định nghĩa SLO trước khi benchmark |

torch.profiler là công cụ ghi lại thời gian, kernel, CPU/GPU activity và thứ tự chạy của workload. Với builder, nó ảnh hưởng trực tiếp tới workflow: thay vì đo cảm giác, bạn có trace để biết phần nào đáng tối ưu.

scaled_dot_product_attention hay SDPA là API attention chuẩn trong PyTorch, có thể dispatch sang nhiều backend khác nhau. Tradeoff chính: bạn viết một API, nhưng performance phụ thuộc shape, dtype, mask, GPU và backend được chọn.

Flash Attention là nhóm kỹ thuật attention tối ưu memory access để giảm đọc/ghi trung gian. Nó thường hữu ích khi memory traffic là vấn đề, nhưng không phải mọi shape đều tự động thắng.

Bài thực hành 30 phút: lập hồ sơ attention của chính bạn

Mục tiêu không phải tạo benchmark hoàn hảo. Mục tiêu là có một “bài tập về nhà” đủ thật để tránh tranh luận bằng linh cảm.

Bước 1: tách workload nhỏ nhưng giống production

Đừng benchmark một tensor ngẫu nhiên quá xa thực tế. Chọn batch size, sequence length, dtype gần với case của bạn.

import torch
import torch.nn.functional as F
from torch.profiler import profile, ProfilerActivity

device = 'cuda'
dtype = torch.float16

batch = 4
heads = 16
seq = 2048
head_dim = 64

q = torch.randn(batch, heads, seq, head_dim, device=device, dtype=dtype)
k = torch.randn(batch, heads, seq, head_dim, device=device, dtype=dtype)
v = torch.randn(batch, heads, seq, head_dim, device=device, dtype=dtype)

# warmup
for _ in range(5):
    F.scaled_dot_product_attention(q, k, v, is_causal=True)
torch.cuda.synchronize()

Bước 2: profile đúng đoạn cần nhìn

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True,
    with_stack=False
) as prof:
    for _ in range(10):
        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
    torch.cuda.synchronize()

print(prof.key_averages().table(sort_by='cuda_time_total', row_limit=20))

Nhìn vào bảng, bạn cần tìm:

operator nào ăn nhiều cuda_time_total;
có kernel attention chuyên biệt hay rơi về đường math chung;
memory allocation có bất thường không;
shape có làm phát sinh transpose/copy không.

Bước 3: đổi từng biến một

Chỉ đổi một biến mỗi lần, nếu không bạn sẽ tự biến profiler thành đề thi trắc nghiệm không đáp án.

Gợi ý thứ tự:

seq: 512 → 2048 → 8192 nếu GPU chịu nổi.
dtype: fp16/bf16 tùy stack.
is_causal: true/false.
batch size.
số head và head_dim.

Hình dung thế này: nếu cùng một học sinh làm ba đề khác nhau rồi bạn đổi luôn giáo viên, phòng thi và thang điểm, bạn sẽ không biết vì sao điểm tăng. Benchmark cũng vậy.

Bước 4: ghi quyết định, không chỉ ghi số

Sau mỗi lần đo, ghi lại theo format ngắn:

Workload: batch=4, heads=16, seq=2048, head_dim=64, fp16, causal
Backend quan sát được: ...
Hotspot chính: ...
Giả thuyết: ...
Quyết định tiếp theo: giữ SDPA / thử compile / đổi shape / benchmark serving stack

Đừng chỉ lưu screenshot trace. Ba tuần sau mở lại, bạn sẽ quên vì sao mình đo.

Ba ngã rẽ kỹ thuật, chọn theo triệu chứng

Sau khi có trace, thường bạn rơi vào một trong ba nhóm.

1. Chưa dùng đúng backend có sẵn

Nếu workload đang rơi về đường math backend chậm hơn, việc đầu tiên là kiểm tra điều kiện của SDPA backend: dtype, mask, shape, GPU, version PyTorch. Đây là loại tối ưu “đổi cách làm bài”, chưa cần chế đề mới.

Khi phù hợp, backend efficient/flash/cuDNN có thể giúp giảm thời gian rõ rệt. Nhưng hãy đo trên workload thật, vì một shape đẹp trong notebook chưa chắc giống traffic production.

2. Kernel launch và layout đang phá nhịp

Từ bài về fused MLP, bài học kéo sang attention là: nhiều phép nhỏ rời rạc có thể thua một đường fused tốt hơn. torch.compile có thể giúp gom và tối ưu một số pattern, nhưng không nên kỳ vọng nó sửa mọi shape xấu.

kernel fusion là gộp nhiều phép tính thành ít kernel hơn để giảm overhead và đọc/ghi trung gian. Tradeoff: dễ phụ thuộc version, shape, và đôi khi khó debug hơn.

Nếu trace cho thấy nhiều copy/transpose, đừng vội viết kernel. Kiểm tra layout tensor, contiguous, batch/head arrangement trước.

3. Context dài đến mức dense attention không còn hợp lý

Khi sequence rất dài, dense attention — mọi query nhìn mọi key — có thể trở thành khoản chi lớn. Lúc này sparse attention đáng được đưa vào radar.

sparse attention là attention chỉ tính trên một phần token/block được chọn. Tradeoff chính: giảm compute lý thuyết, nhưng có thể tạo memory access thất thường và overhead chọn block. Case Fireworks cho thấy tối ưu phần này cần kernel design sát phần cứng, không chỉ đổi công thức trong paper.

Với team sản phẩm, câu hỏi nên là: workload của mình có đủ dài, đủ lặp, đủ giá trị kinh tế để trả chi phí phức tạp này không?

Bẫy hay gặp: tối ưu theo tiếng ồn thị trường

Có ba bẫy mình muốn bạn né.

Bẫy một: thấy benchmark kernel là muốn đổi stack. Benchmark kernel không bằng end-to-end latency. Fireworks cũng tách rõ kernel efficiency và full module performance; tầng module còn index mapping, combination stage và overhead khác.

Bẫy hai: mặc định technique phổ biến là đáp án. Câu chuyện PEFT và LoRA cũng vậy: LoRA phổ biến không có nghĩa luôn tốt nhất. Với attention backend, Flash phổ biến không có nghĩa mọi case đều thắng.

Bẫy ba: đo trên đồ chơi, quyết trên production. Nếu production có dynamic batching, multi-tenant traffic, context length phân bố lệch, profiler local chỉ là bài kiểm tra miệng, chưa phải kỳ thi thật.

Ví dụ cụ thể: giả sử team bạn 5 người đang chạy chatbot nội bộ cho tài liệu pháp lý, request thường dài nhưng không đều. Nếu chỉ benchmark seq=2048 cố định, bạn có thể chọn backend rất đẹp cho case trung bình nhưng vỡ latency ở các request seq=12k. Lúc đó, việc cần làm không phải tranh luận “Flash hay sparse”, mà là đo histogram context length, tách nhóm request, rồi mới quyết định đường phục vụ.

Nếu là mình, mình sẽ chốt thế này

Với team Việt Nam đang build hệ thống AI thật, mình sẽ không bắt đầu bằng custom attention kernel. Mình sẽ đi theo thứ tự:

Profile workload thật bằng PyTorch profiler. Có trace rồi mới bàn.
Xác nhận SDPA backend và shape. Tận dụng đường có sẵn trước.
Đo theo phân khúc context length. Đừng lấy một con số trung bình để quyết cả hệ thống.
Chỉ cân nhắc sparse/custom kernel khi chi phí attention đủ lớn và traffic đủ ổn định. Nếu không, bạn đang nhận thêm độ phức tạp vận hành mà chưa chắc thu lại.
Ghi lại quyết định như một memo kỹ thuật. Sau này upgrade PyTorch, GPU, hoặc model, bạn còn biết mình đã chọn vì lý do gì.

Dịch ra việc build: profiler không bảo bạn phải tối ưu cái gì; profiler cho bạn bằng chứng để không tối ưu theo phong trào.

Chốt gọn: trong lớp học performance, đứa đạt điểm cao không phải đứa thuộc nhiều tên kernel nhất, mà là đứa biết đọc đề trước khi cắm đầu làm bài.

---
Bụi Wire — nghiện đọc release notes lúc 2 giờ sáng