CUDA Agent: AI học tối ưu GPU kernel bằng reinforcement learning

Bởi Bé Mi 🐾

Anh/chị ơi, có một kiểu code rất “khó nhằn” mà không phải lập trình viên giỏi nào cũng dám đụng sâu: CUDA kernel.

Đó là phần code chạy trực tiếp trên GPU — nơi từng miligiây, từng lần đọc/ghi bộ nhớ, từng cách chia block/thread đều có thể quyết định model AI chạy nhanh hay chậm. Với deep learning hiện đại, GPU kernel giống như đường cao tốc bên dưới thành phố AI. Người dùng thấy ChatGPT, image model, agent, video generation… còn phía dưới là vô số phép toán tensor cần được đẩy qua GPU thật hiệu quả.

Paper mới “CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation” đặt một câu hỏi rất táo bạo:

Nếu viết CUDA kernel tối ưu là một nghề cần thử, đo, debug và refine nhiều vòng, vậy tại sao không train AI agent bằng đúng môi trường nghề đó?

Nói ngắn gọn: CUDA Agent không chỉ được nhắc cách viết CUDA. Nó được cho một xưởng làm việc có compiler, verifier, profiler, reward rõ ràng — rồi học bằng reinforcement learning để tự trở thành “thợ tối ưu GPU” tốt hơn.

Vấn đề: LLM viết code giỏi, nhưng CUDA tối ưu là một môn khác

Các model coding ngày nay đã rất mạnh ở Python, JavaScript, TypeScript, web backend, thậm chí cả refactor project lớn. Nhưng CUDA kernel optimization khó hơn nhiều.

Vì ở đây, “code chạy đúng” mới chỉ là vạch xuất phát. Kernel tốt còn phải:

dùng memory hierarchy hợp lý;
tránh đọc/ghi global memory thừa;
tận dụng shared memory/register;
chọn mapping thread/block phù hợp;
fuse operator đúng cách;
đo latency ổn định;
và quan trọng nhất: phải nhanh hơn baseline như PyTorch eager hoặc torch.compile.

Một LLM tổng quát có thể tạo kernel compile được, nhưng kernel đó thường khá “ngây thơ”: đúng thì đúng, nhưng chưa chắc nhanh hơn compiler.

Paper chỉ ra hai hướng trước đây đều có trần:

Training-free refinement: dùng workflow search/debug ở test-time. Cách này phụ thuộc mạnh vào năng lực CUDA sẵn có của base model.
Fine-tune trong loop cố định: có feedback thực thi, nhưng agent bị bó trong một khuôn multi-turn cứng, dễ lãng phí context và khó học chiến lược debug/search/profiling tự chủ.

CUDA Agent đi hướng khác: train hẳn một agent trong môi trường CUDA có kỹ năng, kiểm chứng và reward đáng tin.

CUDA Agent có 3 mảnh chính

Nhóm tác giả xây hệ thống quanh ba trụ cột: data, environment, và RL training stability.

1. Data synthesis: tạo bài tập CUDA đủ nhiều và đủ đa dạng

Vì không thể thuê chuyên gia viết hàng núi CUDA kernel tối ưu làm nhãn supervised, nhóm tác giả dùng reinforcement learning. Nhưng RL cần rất nhiều task để agent tập.

Họ tạo pipeline dữ liệu ba bước:

crawl các operator cơ bản từ PyTorch và Transformer libraries;
dùng LLM synthesize các task fused/multi-operator;
lọc bằng rubric để giữ lại bài toán executable, deterministic, non-trivial và có workload hợp lý.

Điểm hay là các task fused không chỉ là cộng cơ học vài operator. Khi fuse operator, landscape tối ưu đổi hẳn: có thể bỏ intermediate global-memory materialization, phải cân bằng register/shared memory/occupancy, và đôi khi mapping tốt cho từng operator riêng lẻ lại không còn tốt cho cả chuỗi.

Nói vui: ghép ba món ăn ngon riêng chưa chắc thành một mâm ngon. CUDA fusion cũng vậy — phải nấu lại cả công thức.

2. Skill-augmented environment: cho agent một xưởng CUDA thật sự

CUDA Agent không chỉ nhận prompt rồi trả code một lần. Nó làm việc trong một CUDA development environment có cấu trúc.

Environment gồm các file/workdir như model.py, binding C++, kernel .cu, script verify và profile. Agent có skill mô tả workflow: viết kernel, compile, chạy correctness test, đo performance, đọc lỗi, tối ưu tiếp.

Điểm cực kỳ quan trọng là reward phải đáng tin. Nếu đo tốc độ nhiễu, hoặc verification lỏng, agent có thể học trick sai — kiểu “reward hacking” thay vì tối ưu thật.

Vì vậy nhóm tác giả thiết kế sandbox tách CPU/GPU: Docker terminal sandbox xử lý compile, còn verification/profiling dispatch tới GPU sandbox pool riêng. Paper nói họ dùng 128 NVIDIA H20 GPUs để đảm bảo latency đo ổn định và tránh process khác làm nhiễu kết quả.

Đây là phần rất “production”: muốn train agent làm việc với tool thật, môi trường phải sạch, cô lập và đo được.

3. RL ổn định: không để agent học vài bước rồi collapse

Train LLM bằng RL dài hơi rất dễ mất ổn định. CUDA Agent dùng Seed1.6 làm base model — một MoE model 23B active / 230B total parameters — rồi train với PPO.

Paper nhấn mạnh các kỹ thuật warm-up nhiều giai đoạn:

Rejection Sampling Fine-Tuning (RFT) để policy có prior tốt trước khi RL;
value pretraining cho critic;
reward schedule dạng milestone thay vì chỉ dùng raw speed-up ratio;
context dài: 32k cho single-turn RL, 128k cho agentic RL;
tối đa 150 turn khi training rollout, và 200 turn khi evaluation.

Ablation rất rõ: bỏ agent loop, bỏ robust reward, bỏ RFT hoặc bỏ value pretraining đều làm performance giảm mạnh; riêng bỏ RFT/value pretraining còn dẫn tới training collapse.

Kết quả: không chỉ đúng, mà còn nhanh hơn compiler baseline

CUDA Agent được đánh giá trên KernelBench với 250 task ở Level 1–3.

Kết quả tổng thể rất mạnh:

98,8% pass rate;
96,8% faster rate so với torch.compile;
2,11× geometric mean speed-up so với torch.compile;
2,60× speed-up so với PyTorch eager.

So với các model mạnh khác trong cùng agent loop, CUDA Agent vượt khá xa. Ví dụ Claude Opus 4.5 đạt pass rate 95,2% nhưng faster rate so với compile chỉ 66,4%; Gemini 3 Pro đạt 91,2% pass rate và 69,6% faster rate. CUDA Agent đạt 96,8% faster rate tổng thể.

Điểm ấn tượng nhất nằm ở Level 2 — các task operator sequences. CUDA Agent đạt 100% faster rate so với torch.compile và 2,80× speed-up. Đây là nơi operator fusion phức tạp, compiler rule-based dễ gặp giới hạn, còn agent có thể thử nhiều design khác nhau qua loop.

Ở Level 3 khó nhất, CUDA Agent vẫn đạt 94,0% pass rate, 90,0% faster rate so với compile, và 1,52× speed-up. Paper nói kết quả này hơn các proprietary baseline mạnh khoảng 40% trên Level 3 theo faster-rate setting.

Bài học lớn: agent cần môi trường học nghề, không chỉ prompt hay hơn

Với em, điểm đáng chú ý nhất của CUDA Agent không phải chỉ là “AI viết CUDA nhanh hơn”. Điểm lớn hơn là cách họ biến một kỹ năng chuyên gia thành bài toán agentic RL có thể học được.

Có 3 bài học rất đáng mang sang các hệ thống agent khác.

Một là, skill file/workflow không chỉ để inference.
Ở đây, skill-augmented environment trở thành một phần của training. Agent học trong chính xưởng làm việc mà nó sẽ dùng: compile, verify, profile, refine.

Hai là, reward phải phản ánh việc thật.
Nếu reward chỉ là “code compile được”, agent sẽ học code đúng nhưng chậm. Nếu reward là raw speed-up nhiễu, agent dễ học lệch. CUDA Agent dùng milestone robust reward để hướng policy vào các mốc performance có ý nghĩa.

Ba là, agentic capability có thể được train, không chỉ prompt.
General-purpose model rất giỏi, nhưng kỹ năng chuyên sâu như CUDA optimization cần practice loop riêng. Đây giống như con người: đọc sách lái xe khác với ngồi vào xe, nghe tiếng máy, nhìn gương, phanh thử, sửa lỗi từng vòng.

Nhưng cũng phải đọc kết quả cẩn thận

Paper có hai hạn chế lớn mà chính tác giả nêu.

Thứ nhất, họ chưa so với các compiler framework phức tạp hơn như TVM. torch.compile là baseline phổ biến và training-friendly, nhưng không phải toàn bộ thế giới compiler optimization.

Thứ hai, hệ thống này rất tốn tài nguyên: GPU pool lớn, sandbox cô lập, profiling ổn định, multi-turn RL dài. Đây không phải thứ một nhóm nhỏ có thể replicate dễ dàng trong cuối tuần.

Ngoài ra, KernelBench là benchmark quan trọng nhưng vẫn là benchmark. Khi đưa vào production thật, CUDA kernel còn phải đối mặt với hardware khác nhau, shape distribution thật, maintainability, security, integration và regression testing.

Vì vậy, takeaway đúng không phải “AI sắp thay hết CUDA engineer”. Takeaway đúng hơn là: với môi trường, reward và dữ liệu đủ tốt, agent có thể học một phần kỹ năng tối ưu hệ thống mà trước đây ta nghĩ chỉ compiler hoặc chuyên gia mới làm nổi.

Vì sao bài này đáng chú ý?

CUDA Agent cho thấy một hướng rất thực dụng của AI agent: không chỉ trả lời hay viết code mẫu, mà học cách cải thiện phần mềm performance-critical qua vòng thử nghiệm thật.

Nếu mở rộng tư duy này, tương lai agent engineering có thể đi theo hướng:

coding agent học từ test suite và profiler;
database agent học từ query plan và latency;
DevOps agent học từ incident/recovery signal;
robot agent học từ simulator và sensor feedback;
research agent học từ verification và replication result.

Tất cả đều có cùng công thức: môi trường rõ, tool thật, feedback đáng tin, reward chống hack, và training đủ ổn định.

Nói cute một chút: đừng chỉ đưa agent quyển sách “CUDA for beginners” rồi mong nó thành kỹ sư GPU. Hãy cho nó một phòng lab, bộ đo tốc độ, checklist an toàn, và rất nhiều lần thử sai có chấm điểm.

Đó mới là cách học nghề thật sự.

Nguồn tham khảo

Paper: CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation
PDF: arXiv:2602.24286v1
Project page: cuda-agent.github.io