Đừng viết GPU kernel vì thấy ngầu

Một playbook cho builder: khi nào nên đụng TileLang, khi nào cứ để PyTorch/cuBLAS làm việc, và cách kiểm chứng trước khi đốt GPU.

“Hay mình tự viết kernel cho nhanh hơn?” — câu này thường xuất hiện lúc 5 giờ chiều, ngay sau khi dashboard latency đỏ như đèn xe máy hết bình.

Ở một team giả định tên là Xưởng Số, lead ML tên Minh đang có một pipeline inference: RAG lấy context, model sinh câu trả lời, thêm một lớp rerank, rồi trả về UI. Latency cao. Một bạn trong team đọc thấy TileLang có thể viết GPU kernel bằng Python DSL, compile qua TVM, đụng tới tensor core GEMM, fused softmax, FlashAttention, autotuning… thế là mắt sáng lên.

Mình hiểu cảm giác đó. Có những ngày đọc release notes xong thấy tay ngứa như thợ mộc nhìn miếng gỗ đẹp: muốn đục ngay một đường cho đã. Nhưng production không thưởng cho đường đục đẹp. Production thưởng cho đúng chỗ cần đục.

Luận điểm của bài này: custom GPU kernel không phải bước nâng cấp mặc định; nó là quyết định kiến trúc khi bạn đã chứng minh bottleneck nằm ở operator cụ thể, có baseline tử tế, và có đường rollback.

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

Mục tiêu thật: chọn nơi tối ưu, không chọn món đồ chơi

TileLang đáng chú ý vì nó cho phép viết GPU kernel ở mức cao hơn CUDA thuần, nhưng vẫn đủ gần phần cứng để điều khiển các thứ như:

• shared memory — vùng nhớ nhanh trong một block GPU, dùng để tái sử dụng dữ liệu gần compute.
• register fragment — mảnh dữ liệu nằm trong register, rất nhanh nhưng dung lượng ít.
• tensor core — phần cứng chuyên nhân ma trận nhanh trên GPU NVIDIA.
• fused epilogue — gộp phần xử lý sau GEMM, ví dụ bias hoặc activation, vào cùng kernel để giảm đọc/ghi bộ nhớ.
• autotuning — thử nhiều cấu hình kernel để tìm cấu hình hợp với GPU đang chạy.

Điểm hay: bạn có thể mô tả tiled GEMM, row-wise softmax, FlashAttention, rồi để compiler lo mapping thread, layout bộ nhớ, synchronization, vectorization và sinh CUDA code.

Điểm nguy hiểm: vì nó dễ tiếp cận hơn CUDA thuần, bạn dễ tưởng mình nên dùng nó sớm hơn.

Ở Xưởng Số, Minh không cho team nhảy thẳng vào TileLang. Anh bắt team trả lời ba câu trước:

1. Operator nào đang chiếm latency nhiều nhất?
2. Baseline hiện tại là gì: PyTorch, cuBLAS, Triton có sẵn, hay kernel từ framework serving?
3. Nếu kernel tự viết sai số hoặc chậm hơn trên GPU khác, rollback bằng gì?

Nói thẳng ra thì: đừng bào phẳng cả cái bàn khi chỉ có một góc bị vênh.

Checklist trước khi chạm TileLang

Nếu bạn là developer/tech lead, đây là checklist mình sẽ dùng trước khi cho team tự viết kernel.

1. Có profile theo operator, không chỉ đo end-to-end

End-to-end latency nói hệ thống chậm. Nó không nói chậm ở đâu.

Bạn cần tách ít nhất:

• retrieval latency nếu là RAG;
• prefill/decode latency nếu là LLM serving;
• thời gian GEMM, attention, softmax, normalization;
• CPU overhead: serialize, network, queue, post-processing.

Ví dụ cụ thể: team bạn thấy request mất lâu, nhưng profile cho thấy phần tốn nhất là vector search do query và document lệch ngôn ngữ. Lúc đó tối ưu GPU kernel không giúp nhiều. Một hướng như HyDE — tạo “tài liệu giả định” từ câu hỏi để retrieval dễ khớp hơn — có thể đáng thử trước, vì bottleneck nằm ở truy xuất, không phải compute.

2. Có baseline đáng tin

Với kernel compute-heavy, baseline tối thiểu nên gồm:

• PyTorch implementation;
• cuBLAS hoặc operator framework đang dùng;
• kiểm tra numerical correctness bằng sai số tương đối/chấp nhận được;
• đo latency bằng warmup và nhiều lần chạy, không lấy một lần đẹp nhất.

Nếu bạn chưa có baseline, mọi “nhanh hơn” đều giống đo ván gỗ bằng mắt: nhìn có vẻ thẳng, ráp vào khung mới biết lệch.

3. Có workload đại diện

Autotuning chỉ có ý nghĩa khi search space phản ánh tải thật.

Đừng tune GEMM trên shape đẹp rồi đem chạy workload nhiều batch size lẻ, sequence length dao động, hoặc attention mask phức tạp. Với LLM serving, workload còn thay đổi theo tỉ lệ prefill/decode. Đây là lúc tư duy từ benchmark phân tán như srt-slurm hữu ích: khai báo recipe, sweep tham số, nhìn throughput-versus-latency theo Pareto frontier — tức tập lựa chọn không bị lựa chọn khác vừa nhanh hơn vừa rẻ hơn đánh bại.

Playbook một buổi: thử mà không tự lừa mình

Hình dung thế này: bạn có một operator nghi ngờ đang chậm, ví dụ row-wise softmax trong pipeline attention riêng, hoặc GEMM kèm bias + activation bị gọi quá nhiều lần. Bạn muốn biết có đáng đầu tư custom kernel không.

Một buổi chiều đủ để làm bản nháp kiểm chứng, không đủ để tuyên bố chiến thắng.

Bước 1: đóng khung ca thử nghiệm

Chọn một operator, một shape chính, một GPU target.

Ví dụ minh họa:

operator: fused GEMM + bias + activation
shape: M x K nhân K x N, dùng shape phổ biến nhất trong production log
dtype: fp16 hoặc bf16
baseline: torch.matmul + bias + activation
metric: latency, throughput, relative error
rollback: dùng lại PyTorch path

Đừng bắt bản thử nghiệm giải quyết mọi shape. Bản thử tốt là bản trả lời được: “có tín hiệu đáng theo tiếp không?”

Bước 2: viết kernel nhỏ nhất có kiểm chứng

Với TileLang, hãy bắt đầu từ kernel đơn giản: vector add hoặc GEMM tiled. Mục tiêu không phải khoe code, mà là dựng được vòng lặp:

# pseudo-code minh họa, không phải API cố định
# 1. tạo input bằng torch
# 2. chạy baseline PyTorch/cuBLAS
# 3. chạy TileLang kernel
# 4. so relative error
# 5. benchmark nhiều lần sau warmup

Điều quan trọng là utility đo và validate phải dùng lại được. Nếu mỗi kernel lại có một cách đo khác nhau, bạn đang so những miếng gỗ chưa cùng mặt phẳng.

Bước 3: chỉ fuse khi đã thấy memory traffic là vấn đề

Fusing hấp dẫn vì giảm round-trip ra global memory. Ví dụ GEMM xong phải cộng bias, chạy activation, rồi ghi lại tensor. Nếu tách thành nhiều kernel, mỗi bước có thể phải đọc/ghi thêm.

Nhưng fuse cũng làm kernel khó maintain hơn, tăng pressure lên register, và có thể giảm occupancy — mức độ tận dụng tài nguyên GPU. Vì vậy, hãy đo cả hai:

• GEMM riêng + epilogue riêng;
• GEMM với fused epilogue.

Nếu chênh lệch không rõ, ưu tiên đường dễ bảo trì.

Bước 4: autotune có giới hạn

Autotuning không phải thả lưới vô tận. Bạn nên giới hạn theo vài biến có ý nghĩa:

• tile size;
• số pipeline stages;
• số thread;
• swizzling để cải thiện locality ở L2 cache;
• dtype và layout.

Ghi lại cấu hình thắng theo GPU. Một config tốt trên A100 không mặc định tốt trên L4, H100 hay GPU cloud khác. Nếu production của bạn chạy nhiều loại GPU, artifact benchmark phải đi kèm kernel.

Bẫy hay gặp: benchmark đẹp nhưng production vẫn cau có

Có bốn bẫy mình thấy team builder dễ dính.

Bẫy 1: tối ưu sai tầng.
Nếu request chậm vì queue, network, hoặc retrieval, custom kernel chỉ làm đẹp một dòng trong profile. Với hệ RAG, đôi khi cải thiện embedding query bằng HyDE hoặc đổi chiến lược retrieval cho tác động rõ hơn việc viết kernel.

Bẫy 2: so với baseline yếu.
Thắng một implementation PyTorch ngây thơ không có nghĩa thắng operator đã tối ưu trong stack production. Hãy so với đường thật bạn đang deploy.

Bẫy 3: quên numerical behavior.
Softmax và attention nhạy với stability. FlashAttention không chỉ là “attention nhanh”; nó là cách tính attention tối ưu memory bằng tiling, tránh vật liệu trung gian quá lớn. Sai một chút ở accumulation hoặc mask có thể thành lỗi khó debug.

Bẫy 4: không có owner sau khi merge.
Kernel custom là code hạ tầng. Nó cần người theo dõi khi đổi driver, đổi CUDA, đổi GPU, đổi shape workload. Nếu team chưa có owner, đừng biến pull request thành món đồ thủ công không ai biết sửa.

Quyết định sau cùng: khi nào bật đèn xanh?

Mình sẽ cho team Xưởng Số đi tiếp với TileLang nếu đạt đủ bốn điều kiện:

| Câu hỏi | Bật xanh khi nào? |
|---|---|
| Bottleneck có rõ không? | Profile chỉ ra operator cụ thể chiếm phần đáng kể latency hoặc memory traffic |
| Baseline có công bằng không? | So với PyTorch/cuBLAS/operator production, có warmup và validate sai số |
| Workload có thật không? | Shape, dtype, batch, sequence length lấy từ log hoặc kịch bản gần production |
| Vận hành có đường lui không? | Có feature flag, fallback path, artifact benchmark, owner bảo trì |

Nếu thiếu một trong bốn, mình sẽ dừng ở prototype. Không phải vì TileLang không đáng dùng, mà vì bạn chưa chứng minh đúng bài toán.

Độc giả nên nghĩ khác điều gì sau bài này? Tối ưu GPU không bắt đầu bằng công cụ mới; nó bắt đầu bằng quyền được nói “chưa cần” sau khi profile đủ kỹ.

Nếu là mình, mình sẽ để TileLang nằm trong hộp dụng cụ của team performance, không đặt nó lên bàn thờ kiến trúc. Dao đục sắc đến đâu cũng phải gặp đúng mộng gỗ — đục nhầm thì chỉ có thêm mùn cưa để dọn.

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