Đừng vội viết kernel Python mới

cuTile Python rất đáng thử, nhưng quyết định đúng không phải là “có dùng không”, mà là đặt nó vào đúng lớp thử nghiệm, fallback và đo đạc.

Bụi Wire

2026-06-12

Có một niềm tin đang lan khá nhanh trong mấy team build AI: hễ NVIDIA đưa thêm cách viết GPU kernel bằng Python thì chắc nên nhảy vào ngay, vì Python dễ hơn CUDA C++ và GPU đang là chỗ nghẽn.

Nghe hợp lý. Ai từng debug một đoạn CUDA kernel lúc nửa đêm đều hiểu cảm giác muốn có thứ gì đó gần Python hơn, ít nghi thức hơn, chạy được trong notebook hơn. Nhưng mình muốn kéo phanh nhẹ: cuTile Python không nên được xem là đường tắt để mọi developer viết kernel custom. Nó nên được xem là một lớp thử nghiệm có kiểm soát để quyết định xem team bạn có thật sự cần custom kernel hay chỉ đang vá víu performance bằng mũi chỉ sai chỗ.

Sau bài này, điều mình muốn bạn nghĩ khác là: đừng hỏi “tool mới này có nhanh không?”, hãy hỏi “nó giúp mình ra quyết định triển khai GPU kernel an toàn hơn không?”

Sơ đồ minh họa cho bài Đừng vội viết kernel Python mới

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

Mục tiêu: may một đường thử, không thay cả áo

cuTile Python là một giao diện lập trình GPU theo kiểu tile-based — chia dữ liệu thành các ô nhỏ để load, tính toán và store hiệu quả trên GPU. Trong tutorial chính, workflow đi từ kiểm tra Colab runtime, GPU, driver, CUDA, cuTile, rồi mới chạy các kernel cho vector addition, matrix addition và matrix multiplication. Quan trọng hơn: notebook vẫn giữ PyTorch fallback, tức nếu runtime không đủ điều kiện thì quay về PyTorch thay vì chết giữa đường.

Đây mới là phần đáng học.

Không phải vì vector addition hay matrix addition là bài toán quá mới. Những ví dụ đó giống mẫu vải thử: đủ đơn giản để bạn nhìn ra sợi ngang sợi dọc của hệ thống — môi trường, correctness, benchmark, fallback — trước khi cắt vào phần vải thật là inference pipeline hoặc training workload.

Với developer hoặc tech lead, mục tiêu trong một buổi không phải là “viết kernel xịn”. Mục tiêu là trả lời 3 câu:

Runtime của mình có đủ ổn để chạy cuTile không?
Custom kernel có cho kết quả đúng so với PyTorch không?
Nếu cuTile không chạy được, production path có còn sống không?

Nếu chưa trả lời được 3 câu này, bàn chuyện tối ưu sâu là hơi sớm.

Checklist trước khi đụng vào kernel

Trước khi viết một dòng kernel, mình sẽ yêu cầu team có checklist này:

| Câu hỏi | Vì sao cần hỏi | Nếu câu trả lời là “không” |
|---|---|---|
| GPU, driver, CUDA có đúng yêu cầu không? | cuTile phụ thuộc runtime khá sát | Chạy PyTorch fallback, không ép demo |
| Có baseline PyTorch chưa? | Không có chuẩn thì không biết custom kernel đúng hay sai | Viết baseline trước |
| Có correctness check chưa? | GPU bug có thể sai âm thầm | So sánh output bằng tolerance rõ ràng |
| Có benchmark có synchronize chưa? | GPU chạy async, đo sai rất dễ | Thêm synchronize trước/sau đo |
| Có đường rollback không? | Notebook chạy được không có nghĩa prod chịu được | Bọc API fallback ngay từ đầu |

Nói thẳng ra thì: cái mới nhất không tự động là cái nên đưa vào production. Với GPU programming, thứ nguy hiểm nhất không phải kernel chậm, mà là kernel sai nhưng nhìn có vẻ chạy ngon.

Một buổi chiều đủ để ra quyết định sơ bộ

Bạn có thể dựng một spike nhỏ trong khoảng một buổi. Không cần biến nó thành dự án nghiên cứu.

Bước 1: kiểm tra runtime trước

Trong Colab hoặc notebook nội bộ, chạy nhóm lệnh kiểm tra môi trường:

python --version
nvidia-smi
nvcc --version || true

Trong Python, kiểm tra PyTorch thấy GPU chưa:

import torch

print(torch.cuda.is_available())
if torch.cuda.is_available():
    print(torch.cuda.get_device_name(0))

Nếu chỗ này đã không ổn, đừng cố viết kernel. Ghi lại trạng thái runtime và chuyển sang fallback. Đây là khác biệt giữa demo có trách nhiệm và demo kiểu “máy mình chạy được”.

Bước 2: viết wrapper trước, kernel sau

Thay vì gọi cuTile trực tiếp khắp notebook, hãy bọc thành một hàm có fallback:

def vector_add(a, b, use_cutile=False):
    if use_cutile and cutile_available():
        return vector_add_cutile(a, b)
    return a + b

Tên hàm cutile_available() có thể kiểm tra import, CUDA, device capability, hoặc flag cấu hình. Điểm chính là call site không cần biết backend nào đang chạy.

Đây là thói quen rất đáng giữ nếu sau này bạn chuyển từ notebook sang service. Backend GPU giống lớp lót bên trong áo: người dùng không cần thấy, nhưng nếu rách thì cả trải nghiệm bị cộm.

Bước 3: correctness check không được để cuối

Ví dụ cụ thể: với matrix multiplication, hãy so output custom kernel với PyTorch:

def assert_close(name, got, expected, atol=1e-3, rtol=1e-3):
    ok = torch.allclose(got, expected, atol=atol, rtol=rtol)
    max_diff = (got - expected).abs().max().item()
    print(name, ok, max_diff)
    assert ok

Với float32, tolerance có thể chặt hơn. Với float16, bạn cần nới hơn vì độ chính xác thấp hơn. Đừng dùng cùng một ngưỡng cho mọi dtype rồi tự kết luận kernel sai hoặc đúng.

Bước 4: benchmark phải có synchronize

GPU execution thường là async — lệnh được xếp hàng rồi trả về trước khi GPU làm xong. Nếu bạn đo bằng time.time() mà không synchronize, kết quả có thể đẹp một cách đáng nghi.

import time

def bench(fn, repeats=20):
    torch.cuda.synchronize()
    start = time.perf_counter()
    for _ in range(repeats):
        out = fn()
    torch.cuda.synchronize()
    end = time.perf_counter()
    return (end - start) / repeats, out

Ở bước này, bạn chưa cần tuyên bố thắng thua. Bạn chỉ cần biết custom path có chạy ổn, đúng, và đáng điều tra tiếp không.

Khung quyết định: khi nào đáng đi tiếp?

Mình thích dùng ma trận 2 trục:

Độ nóng của bottleneck: phần này có thật sự chiếm nhiều thời gian trong workload không?
Độ ổn định của shape: kích thước tensor có lặp lại đủ nhiều để tối ưu theo tile không?

Từ đó có 4 trường hợp:

| Bottleneck | Shape | Quyết định |
|---|---|---|
| Thấp | Biến động | Ở lại PyTorch |
| Thấp | Ổn định | Chưa cần custom, chỉ ghi nhận |
| Cao | Biến động | Tối ưu pipeline trước, cân nhắc batching |
| Cao | Ổn định | cuTile spike đáng làm sâu hơn |

Đây là chỗ nhiều team hiểu sai. Họ thấy custom kernel là “cấp cao hơn” nên muốn dùng cho phần lõi. Nhưng nếu workload của bạn thay shape liên tục, hoặc bottleneck nằm ở data loading, network, orchestration, thì kernel mới chỉ là một miếng vá đẹp đặt nhầm chỗ.

Nguồn liên quan về Isaac Lab trên SageMaker cũng gợi ý một bài học tương tự ở cấp hạ tầng: workload GPU nặng cần tách pha thử nghiệm ngắn và pha chạy dài. Với cuTile, mình cũng sẽ tách như vậy: notebook spike để học và đo, rồi mới tính đến packaging, CI, monitoring và rollout.

Bẫy thường gặp khi đưa vào hệ thống thật

Bẫy 1: Colab chạy được, production chưa chắc chạy

Colab là môi trường tiện để học, không phải hợp đồng vận hành. Driver, CUDA, GPU type, package version có thể khác production. Nếu team bạn dùng container, hãy chốt image sớm. Nếu chạy trên managed training hoặc inference, hãy kiểm tra quyền cài package và version CUDA.

Bẫy 2: fallback có nhưng không được test

Fallback không phải để trang trí. Hãy ép CI chạy cả hai path:

def test_vector_add_fallback():
    y = vector_add(a, b, use_cutile=False)
    assert_close('fallback', y, a + b)

def test_vector_add_cutile_or_skip():
    if not cutile_available():
        return
    y = vector_add(a, b, use_cutile=True)
    assert_close('cutile', y, a + b)

Nếu không có GPU trong CI, ít nhất phải có test cho wrapper và fallback. Đừng để production lần đầu phát hiện đường lui bị đứt chỉ.

Bẫy 3: benchmark thiếu workload thật

Vector addition và matrix addition là bài tập tốt để hiểu API. Nhưng quyết định triển khai phải dựa trên workload gần thật: kích thước tensor, dtype, batch size, memory layout, tần suất gọi. Nếu bạn build model serving, hãy đo trong request path. Nếu bạn build training loop, hãy đo trong step thật, không đo một kernel cô lập rồi suy rộng quá tay.

Bẫy 4: tối ưu kernel nhưng bỏ qua chi phí vận hành

Custom kernel kéo theo chi phí bảo trì: version CUDA, GPU compatibility, test correctness, debug numerical issue, fallback logic. Với team nhỏ, chi phí này có thể lớn hơn lợi ích. Đây là lý do mình xem cuTile như công cụ ra quyết định, không phải mặc định là dependency mới.

Nếu là mình, mình sẽ triển khai thế này

Với một team AI ở Việt Nam, giả sử 4-6 engineer, đang có inference hoặc training workload bị nghi là nghẽn GPU, mình sẽ làm theo thứ tự:

Chọn một operation nghi ngờ thật sự nóng, không chọn vì nó dễ demo.
Viết PyTorch baseline rõ ràng, có correctness check.
Dựng cuTile notebook có fallback, kiểm tra runtime trước khi chạy.
Benchmark với synchronize và nhiều shape đại diện, không chỉ một case đẹp.
Ra quyết định bằng ma trận bottleneck/shape, không bằng cảm giác tool mới.
Nếu đi tiếp, đóng gói wrapper, test fallback, pin environment, rồi mới nghĩ tới production rollout.

SkillOpt trong nguồn liên quan có một điểm rất đáng mượn: mọi vòng tối ưu đều được instrument — có baseline, lịch sử thay đổi, token usage, validation gating. Với GPU kernel cũng vậy. Tối ưu mà không đo, không gate, không rollback thì chỉ là hy vọng được bọc bằng code.

Takeaway của mình: cuTile Python đáng thử nhất khi bạn dùng nó để may một đường đo đạc chắc tay, không phải để thay toàn bộ tủ đồ GPU chỉ sau một notebook chạy xanh. Kernel nhanh là tốt; kernel đúng, đo được, và rút lui được mới là thứ sống lâu trong production.

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