Viết kernel GPU bằng tay đòi hỏi sự kiên nhẫn, nỗ lực và cả sự tỉnh táo của bạn. Tôi đã thử nghiệm với các LLM hiện đại (Claude Opus 5, GPT 5.6 Sol) để có thể "in" kernel ra cho tôi theo đúng nghĩa đen. Tôi sẽ nói về những gì tôi học được từ việc đó.
Thành thật mà nói, viết kernel GPU là một tác vụ có thể kiểm chứng một cách dễ dàng đến bất ngờ.
Đầu tiên, bạn quyết định (các) phép toán mà bạn muốn viết kernel. Thứ hai, bạn viết phiên bản đầu tiên và đảm bảo nó biên dịch mà không có lỗi rõ ràng nào. Thứ ba, bạn kiểm tra tính đúng đắn bằng một triển khai tham chiếu chậm. Nếu nó không khớp, bạn cố gắng sửa lỗi sai. Khi đã xong, bạn benchmark thời gian thực thi của kernel. Các phiên bản tiếp theo sẽ xây dựng dựa trên kết quả này, và bạn tiếp tục tối ưu cho đến khi đạt được các chỉ số roofline hoặc đến khi bạn hài lòng.
Vòng lặp có thể kiểm chứng: Phát triển kernel GPU
Hình ảnh trên thể hiện nó như một vòng lặp có thể kiểm chứng:
- Kiểm tra biên dịch là một vòng lặp cục bộ khép kín (B <-> C) trước khi bạn nghĩ đến tính đúng đắn.
- Kiểm tra tính đúng đắn (D <-> E <-> F) là vòng lặp cốt lõi mang lại phần thưởng có thể kiểm chứng. Đây là phần khiến việc viết kernel trở thành một bài toán tốt cho việc xác minh tự động, vì bạn có một tham chiếu chuẩn (ground-truth) để đối chiếu.
- Tối ưu hóa (G -> H -> quay lại D) tái sử dụng cùng một vòng lặp kiểm tra tính đúng đắn cho mỗi phiên bản mới, vì một kernel nhanh nhưng sai là vô giá trị, nên tính đúng đắn phải được xác minh cho từng phiên bản.
- Kiểm tra roofline/sự hài lòng (I) là vòng lặp bên ngoài quyết định nên tiếp tục tối ưu hay dừng lại.
Phiên bản đầu tiên
Hình ảnh trên vẫn chỉ là góc nhìn cấp cao và sự phức tạp nằm trong chi tiết. Chúng ta cần đảm bảo rằng agent LLM của chúng ta có đầy đủ ngữ cảnh cần thiết để bắt đầu viết một phiên bản đầu tiên tốt.
Đây là lúc các DSL CUDA phát huy tác dụng. Triton, CuTeDSL và Tilelang là những công cụ rất dễ bắt đầu sử dụng trong Python. Đường cong học tập ít dốc hơn so với CUDA C++, tuy nhiên, các lớp trừu tượng (abstraction) trong các DSL này có thể khiến agent của chúng ta bối rối hơn. Chúng ta cần một cách để truyền ngữ cảnh của các lớp trừu tượng đó cho agent.
Các LLM hiện đại đã biết cách viết Triton "tốt". Chúng có thể làm việc tốt với các lớp trừu tượng của Triton ngay cả khi không có ngữ cảnh. Tuy nhiên, đối với các DSL khác như CuTeDSL (cung cấp nhiều quyền kiểm soát hơn nhiều so với Triton), tôi nhận thấy rằng việc có một thư mục ngữ cảnh nơi agent có thể tìm hiểu các lớp trừu tượng của DSL giúp ích rất nhiều.
Ví dụ, việc clone
NVIDIA cutlass
repository vào thư mục ngữ cảnh là một cách tốt để cho phép agent tìm kiếm các lớp trừu tượng liên quan đến
Layout Algebra, Copy/GEMM atoms, phân cấp bộ nhớ, kernel mẫu
, v.v. khi viết kernel bằng CuTeDSL.
Theo kinh nghiệm của tôi, một phiên bản đầu tiên tốt của kernel sẽ biên dịch mà không có lỗi rõ ràng và vượt qua bài kiểm tra tính đúng đắn mà tôi sẽ đề cập bên dưới.
Kiểm thử, Benchmark và Profile
Khi đã cung cấp đủ ngữ cảnh cho agent, điểm nghẽn thực sự lúc này chuyển sang việc xác minh. Triển khai tham chiếu và việc xác minh dựa trên nó ngày càng trở nên quan trọng hơn. Tôi gọi giai đoạn này là kiểm thử tính đúng đắn hay đơn giản là kiểm thử. Tốc độ của triển khai tham chiếu không quan trọng bằng mục đích của nó. Điều bạn định đo lường và xác minh chính là thứ mà agent của bạn sẽ tối ưu hóa.
Thông thường, khi phép tính không được thực hiện ở độ chính xác thấp hơn (lower precision), tôi đo Sai số tuyệt đối/tương đối tối đa (MAE), Sai số bình phương trung bình (MSE/RMSE) và PSNR (Tỷ lệ tín hiệu trên nhiễu đỉnh). Khi có sự tham gia của độ chính xác thấp, tôi có xu hướng đo PSNR và Độ tương tự cosine (cossim).
Cách bạn làm cho các phiên bản kernel thực sự chạy trên GPU phụ thuộc vào việc GPU có sẵn cục bộ hay qua cloud. Dù vậy, agent của chúng ta nên có khả năng truy cập đầu ra của nó bằng cách này hay cách khác.
Tôi nhận thấy phương pháp rung dưới đây là một cách tốt để có N hàm kiểm thử:
1def rung(name):2 def deco(fn):3 try:4 out = fn()5 results[name] = {"ok": True, **(out or {})}6 print(f"[{name}] ok " + " ".join(f"{k}={v}" for k, v in (out or {}).items()))7 except Exception as e:8 results[name] = {"ok": False, "err": f"{type(e).__name__}: {e}"}9 print(f"[{name}] FAILED {type(e).__name__}: {e}")10 traceback.print_exc()11 return fn12 return deco
mà bạn có thể gọi như sau:
1out = {}23@rung("pre-checks")4def _():5 run_pure_checks()6 run_dsl_checks()78@rung("run")9def _():10 out["o"] = custom_kernel(*inputs)11 torch.cuda.synchronize()12 return {"shape": tuple(out["o"].shape),13 "finite": bool(torch.isfinite(out["o"]).all())}
Đối với bậc benchmark, bạn có thể làm nhiều việc:
- Benchmark thời gian thực thi kernel từ đầu đến cuối cho tổng thời gian chạy
- Sử dụng tính năng theo dõi nội bộ kernel (intra-kernel tracing) để benchmark các phần bên trong kernel và xuất chúng ra đầu ra (bằng một trình theo dõi tùy chỉnh hoặc CUPTI)
- Kết xuất IR, PTX, SASS và CUBIN được tạo ra vào một thư mục dumps và để agent đọc qua
Điểm cuối cùng có thể được mở rộng thêm một chút. Đôi khi, DSL có thể hạ cấp (lower) để tạo ra PTX không tối ưu (cuối cùng là SASS) và bạn tìm thấy (các) lệnh hoặc hình dạng (shape) tốt hơn để sử dụng thay thế. Agent của chúng ta có thể đọc qua các tệp văn bản PTX/SASS và chèn trực tiếp mã cấp thấp hơn thay vì để DSL xử lý phần không tối ưu đó. Một lần nữa, việc truyền tài liệu PTX "có thể tìm kiếm được" làm ngữ cảnh sẽ rất hữu ích ở đây.
Điều cuối cùng để kết nối tất cả lại với nhau là Profiling. Nếu agent của bạn có thể truy cập CLI của NCU (Nsight Compute Systems), bạn có thể yêu cầu nó profile và tạo báo cáo cho kernel của bạn như một phần của vòng lặp phản hồi có thể kiểm chứng ở trên.
Lời kết
"Vậy việc phát triển kernel GPU đã chết rồi sao?"
"Ừ thì đúng, nhưng thực ra là không"
Đúng, vì phần khó về layout, lập chỉ mục (indexing), các lớp trừu tượng và cấu trúc tổng thể phần lớn có thể được giải quyết bởi các agent có đủ ngữ cảnh. Bạn có thể dễ dàng giảm khối lượng công việc 2-3 tuần xuống còn 1-2 ngày. Không, vì điểm nghẽn thực sự đã chuyển từ kernel sang việc xác minh. Giờ đây không có một cách làm duy nhất nào: ngữ cảnh và bộ khung kiểm thử (harness) của bạn càng tốt, quy trình càng nhanh. Các trường hợp đặc thù sẽ được hưởng lợi nhiều hơn nữa, và tất cả những gì bạn cần chỉ là một bộ khung kiểm thử (harness) tốt.
Cuối cùng, thay vì coi các agent là các hệ thống tự động, hãy coi chúng như những trợ lý cực kỳ thông minh mà bạn có thể dẫn dắt. Đây là lúc kiến thức nền tảng của bạn về GPU và kernel trở nên hữu ích. Phần con người (bạn) vẫn cần thiết ở đây.
Đó là một cảm giác vừa ngọt ngào vừa cay đắng, tôi biết mà :)
