Stanford CRFM: Kernel CUDA do AI tạo ra vượt hiệu năng mã tối ưu hóa PyTorch

 (crfm.stanford.edu)
7 điểm bởi GN⁺ 2025-05-31 | 2 bình luận | Chia sẻ qua WhatsApp
  • Các kernel CUDA-C do AI tạo ra cho thấy hiệu năng tương đương hoặc tốt hơn các kernel tối ưu hóa do chuyên gia của PyTorch viết
  • Một LLM (mô hình ngôn ngữ lớn) duy nhất lặp lại việc tạo ý tưởng tối ưu hóa bằng ngôn ngữ tự nhiên và nhiều nhánh mã nguồn khác nhau, đạt hiệu quả vượt trội hơn các phương pháp trước đây về độ đa dạng tối ưu hóa và khả năng tìm kiếm song song
  • Kết quả benchmark tiêu biểu cho thấy vượt trội rõ rệt so với PyTorch ở các tác vụ như Matmul (101%), Conv2D (179.9%), Softmax (111.8%), LayerNorm (484.4%), Conv2D+ReLU+MaxPool (290.1%)
  • Để vượt qua giới hạn của cách tiếp cận “cải tiến kernel tuần tự” truyền thống, nhóm áp dụng suy luận ngôn ngữ tự nhiên + cấu trúc phân nhánh, từ đó tạo ra các kernel hiệu năng cao với tốc độ nhanh
  • Phương pháp này cũng đang tiến bộ trên các workload ML hiện đại như FP16, Flash Attention, và cho thấy trong tương lai AI có thể tự tìm kiếm và cải tiến các kernel nhanh hơn sẽ trở thành một mô hình chủ đạo

TL;DR Thành tựu chính

  • Nhóm nghiên cứu Stanford CRFM phát hiện rằng các kernel CUDA-C hiệu năng cao do AI tạo ra có thể đạt tốc độ tương đương hoặc tốt hơn các kernel do chuyên gia thiết kế trong PyTorch
  • Mục tiêu ban đầu là huấn luyện một mô hình tự động sinh kernel tốt hơn bằng dữ liệu tổng hợp, nhưng họ quan sát thấy chỉ riêng quá trình tạo dữ liệu tổng hợp đã tự động sinh ra các kernel nhanh ở mức đáng kinh ngạc
  • Vì vậy, nhóm đã công bố sớm phương pháp, benchmark hiệu năng, chiến lược tối ưu hóa và định hướng tiếp theo
  • Benchmark: dựa trên GPU Nvidia L40S. Hiệu năng (%): thời gian chạy của PyTorch ÷ thời gian chạy của kernel được tạo ra
    • Matmul (FP32): 101.3% so với PyTorch (ma trận 4096x4096)
    • Conv2D: 179.9% (đầu vào: 100, 3, 224, 224; cấu hình AlexNet Conv1)
    • Softmax: 111.8% (tensor 4096x65536)
    • LayerNorm: 484.4% (tensor 16, 64, 256, 256)
    • Conv2D + ReLU + MaxPool: 290.1% so với PyTorch reference, 189.0% so với torch.compile()

Động lực nghiên cứu và phương pháp

  • Mục tiêu ban đầu là tạo dữ liệu tổng hợp để huấn luyện mô hình sinh kernel, nhưng trong quá trình thử nghiệm, các kernel do hệ thống tự tạo đã đạt hiệu năng cao vượt xa kỳ vọng
  • Sử dụng KernelBench (benchmark công khai của Stanford, phát hành tháng 12/2024)
    • Với đoạn mã torch cho trước, LLM sẽ viết lại một kernel tối ưu tốc độ mới
    • Tính đúng đắn được kiểm chứng bằng việc đối chiếu tương đương số học giữa đầu vào/đầu ra
  • Cách làm trước đây: một vòng lặp chỉnh sửa tuần tự, trong đó kernel được sửa từng bước nhỏ để cải tiến dần
    • Nhược điểm: thiếu đa dạng ý tưởng, lặp lại cùng một kiểu tối ưu hóa, dễ hội tụ vào cực trị cục bộ
  • Ý tưởng cải tiến
    1. Trước tiên phát sinh và ghi lại ý tưởng tối ưu hóa bằng ngôn ngữ tự nhiên, sau đó tạo đồng thời nhiều nhánh hiện thực mã nguồn của các ý tưởng đó
    2. Mỗi vòng thực hiện song song nhiều thử nghiệm tối ưu hóa → chọn kernel có hiệu năng tốt nhất làm hạt giống (seed) cho vòng tiếp theo
    • Cách này cho phép khám phá nhiều chiến lược tối ưu hóa và tìm kiếm song song trong số vòng lặp tìm kiếm hữu hạn

Ví dụ về các ý tưởng tối ưu hóa

  • Tối ưu truy cập bộ nhớ: cải thiện hiệu quả các tầng bộ nhớ global/shared/register
  • Xử lý bất đồng bộ và che giấu độ trễ: chồng lấp tính toán với di chuyển dữ liệu
  • Tối ưu kiểu dữ liệu/độ chính xác: tận dụng FP16/BF16 và các phép toán chuyên biệt theo phần cứng
  • Tối ưu tính toán và lệnh: giảm lượng phép tính, số lệnh, tận dụng tối đa các lệnh phần cứng đặc thù
  • Song song hóa và occupancy: tối đa hóa việc sử dụng SM (Streaming Multiprocessors)
  • Tối ưu vòng lặp/rẽ nhánh/chỉ số hóa: giảm vòng lặp, loại bỏ phép tính chỉ số không cần thiết

Quy trình tối ưu hóa kernel thực tế (ví dụ Conv2D)

  • Luồng ý tưởng tối ưu hóa và cải thiện hiệu năng theo từng vòng
    • Ban đầu (vòng 0): port CUDA đơn giản (đạt 20% tốc độ so với PyTorch)
    • Các vòng tiếp theo: → tận dụng read cache → phép toán FP16 Tensor Core (chuyển sang GEMM) → double buffering/pipeline → tiền tính chỉ số/shared memory → vectorization → đồng đệm vòng lặp K, v.v. để khai thác sâu kiến trúc GPU
    • Cuối cùng (vòng 13): đạt 179.9% hiệu năng so với PyTorch
  • Mã thực tế tận dụng rất nhiều kỹ thuật lập trình CUDA nâng cao
    • Ở mỗi vòng, hệ thống tạo ý tưởng mới bằng ngôn ngữ tự nhiên, thử song song nhiều cách hiện thực và chọn mã tối ưu nhất

Takeaways và hàm ý

  • Sinh kernel bằng AI có thể vượt trình độ chuyên gia con người nhờ vòng lặp tìm kiếm mạnh cùng tổ hợp đa dạng các ý tưởng tối ưu hóa dựa trên ngôn ngữ tự nhiên
  • Một số toán tử hiện đại (FP16 matmul, Flash Attention) ở thời điểm hiện tại vẫn có hiệu năng thấp hơn PyTorch
  • Tính toán FP32 trên phần cứng gần đây được tối ưu kém hơn tương đối so với FP16/BF16 → có thể tạo lợi thế hiệu năng trong khu vực này
  • Ngay cả khi bị giới hạn ngân sách tìm kiếm token (tổng vào/ra 7 triệu token), hiệu năng vẫn tiếp tục được cải thiện
  • Các nghiên cứu gần đây như AlphaEvolve, Gemini 2.5 Pro cũng cho thấy tìm kiếm dựa trên phân nhánh quy mô lớn + kiểm chứng có thể đem lại cải thiện đột phá mà không cần tái huấn luyện
  • Trong tương lai, cách tiếp cận này có thể sản xuất hàng loạt dữ liệu tổng hợp và kernel hiệu năng cao, đồng thời phát triển thành vòng lặp AI tự cải tiến

Kết luận

  • Tự động sinh và tối ưu hóa kernel bằng AI đã đạt tới trình độ tương đương kỹ thuật hand-coding của chuyên gia, và trong tương lai gần có thể mở ra các hệ thống AI tự cải tiến dựa trên tổ hợp mô hình + tìm kiếm phân nhánh + kiểm chứng
  • Khả năng AI tự vượt qua giới hạn hiệu năng của các framework như PyTorch và TensorFlow đang dần trở nên hiện thực

Phụ lục: Toàn bộ mã kernel Conv2D

  • Bao gồm toàn bộ mã nguồn hàm và kernel thực tế (lược bỏ mã chi tiết)
  • Các đặc điểm chính trong mã:
    • vector hóa shared memory, double-buffering phân cấp theo K-dim, CUDA WMMA, bộ đệm đầu ra cấp warp, v.v.
    • tính chỉ số động theo thời gian thực, xử lý bias, đồng thời nạp dữ liệu trễ trong vòng lặp K, v.v.
  • Có thể xem mã mẫu đầy đủ và các kernel ví dụ trong kho GitHub

Bài viết liên quan

2 bình luận

 
aer0700 2025-06-02

Có thể xem đây như một dạng kỹ thuật ensemble chăng. Thật đáng kinh ngạc.
 
GN⁺ 2025-05-31
Ý kiến Hacker News

  • Tôi thấy cách các tác giả của bài này nhìn nhận về "AI agent" khá thú vị
    Đa số mọi người có xu hướng nghĩ về agent như nhân viên con người
    Họ thiết lập một vài agent chạy song song với nhau (thường chỉ một), rồi để mỗi agent lặp vòng và mỗi lần chỉ xử lý một việc
    Vẫn đang mắc kẹt trong thế giới có số lượng nhân sự cố định, mỗi người chỉ làm được một việc tại một thời điểm, và chuyển đổi tác vụ thì chậm
    Nhưng LLM không như vậy
    Bạn có thể gần như tạo ra vô hạn agent bất cứ lúc nào tùy ý
    Chi phí cho một yêu cầu LLM là như nhau dù xử lý tuần tự hay song song
    Khi nhận ra điều này, tự nhiên sẽ nảy ra mô hình mỗi agent tự phân nhánh thành nhiều sub-agent khi cần
    Đó chính là cách mà các tác giả đã thử
    Tôi cảm thấy sẽ phù hợp hơn nếu xem agent như một "task" hoặc "job", rồi áp dụng những điều đã học từ Celery hay Sidekiq

  • "FP32 hiện được dùng ít hơn hẳn trong các workload ML hiện đại, và trên phần cứng mới cũng được tối ưu kém hơn FP16 hay BF16. Vì vậy đây có thể là một trong những lý do khiến việc cải thiện hiệu năng của kernel FP32 so với PyTorch trở nên dễ dàng hơn"
    Các lập trình viên hầu như đã không đầu tư thời gian tối ưu các kernel phiên bản fp32 trong nhiều năm
    Điều thực sự thú vị, theo tôi, là khi có thể nâng hiệu năng ở phía các kernel đang được phát triển tích cực và thực sự được mọi người sử dụng

    • Tôi nghĩ kết quả tốt như vậy một phần được giải thích bởi việc NVIDIA không cung cấp tài liệu đủ chi tiết về GPU của họ
      Với những bộ xử lý có vi kiến trúc được tài liệu hóa tốt, lập trình viên hay compiler có thể viết ra chương trình tối ưu một cách tất định, nên rất khó để áp dụng ML/AI và tạo ra cải tiến lớn vượt ra ngoài mức tìm lại lời giải đã biết
      Ngược lại, với trường hợp tài liệu ít đầy đủ hơn như GPU NVIDIA, để tìm ra chương trình tối ưu thường phải tìm kiếm ngẫu nhiên dựa trên các ví dụ chương trình đã được tối ưu trước đó, hoặc cần phân tích kiểu reverse engineering để hiểu phần cứng thực sự hoạt động ra sao trong từng tình huống
      Trong bối cảnh đó, ML/AI có thể cho thấy tiềm năng, và bằng cách học từ dữ liệu là các chương trình được viết tốt, nó có thể nắm bắt được những undocumented behavior mà ngay cả lập trình viên con người cũng khó tìm ra

    • Tôi tự hỏi liệu những cải tiến đã được biết đến ở kernel fp16/bf16 có đơn giản chỉ được chuyển sang fp32 hay không

    • "Ý là trong nhiều năm người ta hoàn toàn không đụng vào các kernel fp32 sao?"
      Wow, nếu đúng vậy thì tôi thấy thật tuyệt khi AI đã tạo ra một thuật toán mới ở lĩnh vực chưa có sẵn lời giải trước đó

  • Kết luận của tôi là, như có thể thấy từ bài này, AlphaEvolve của Google(ở đây), và tin gần đây về việc o3 tìm ra zero day trong Linux kernel(ở đây)
    Tôi cảm thấy đặc biệt là Gemini Pro 2.5 và o3 đã đạt tới một mức năng lực mới, nơi chúng đột nhiên làm được cả những ý tưởng mà các model trước đây không làm nổi

    • Tôi không nghĩ là đột nhiên cái trước đây không làm được bỗng nhiên làm được
      Mà là việc lặp lại và thử nghiệm với tốc độ vượt xa con người về thực chất đã trở nên khả thi
      Cùng với khả năng tận dụng ngay lập tức một lượng lớn thông tin
      Nói cách khác, chúng ta đã đi tới thời điểm mà sự kết hợp giữa lượng thông tin khổng lồ, tiến bộ, và brute force được áp dụng một cách thông minh đang thành công trong một số lĩnh vực ứng dụng

    • Tôi nghĩ các ví dụ bạn nhắc tới và kết quả lần này thực ra có khá nhiều điểm tương đồng
      Trong bài cũng có đoạn nói rằng "vòng lặp lặp lại tại thời điểm kiểm thử không giống một chatbot tương tác lần lượt xác nhận kết quả chỉnh sửa code, mà gần hơn với một kiểu tìm kiếm có cấu trúc, chủ động đánh giá song song dựa trên các giả thuyết tối ưu hóa rõ ràng"
      Kết luận của tôi là giờ đây dựa trên năng lực của LLM
      Chúng ta đã học được cách làm rõ hàm đánh giá, hoặc thu hẹp đáng kể không gian lời giải có dạng mẫu tương tự
      Không phải chuyện model này vượt model kia, hay chỉ model nào đó mới giải được... mới là trọng tâm
      Điều có ý nghĩa hơn là thực tế những ứng dụng kiểu này giờ đã bộc lộ đủ rõ

    • Tôi cảm thấy Gemini Pro 2.5 là AI đầu tiên mà con người thực sự dùng được, nhưng nói chính xác thì nó mới chỉ vừa vượt qua ngưỡng đó
      Vẫn có khá nhiều trường hợp tỷ lệ thành công rơi xuống dưới 20%
      Khi 3.0 ra mắt... tôi nghĩ có lẽ lúc đó sẽ hơi đáng sợ

    • Khoan, ý là sao? Cái này không liên quan gì đến Linux kernel, mà là "kernel" theo nghĩa trong lập trình GPU
      Có phải bạn đã hiểu nhầm cả bài không?

    • Khá thú vị, nhưng có bằng chứng mạnh hơn không? Chỉ một kết quả đơn lẻ thì chưa đủ sức thuyết phục

  • "Code chuẩn cơ sở là FP32 mặc định, và sai số cho phép là 1e-02"
    Với mức sai số này thì có thể dễ dàng thay một kernel "fp32" bằng phép tính fp16

    • Nghĩ thêm một bước thì tôi tự hỏi liệu cốt lõi của công việc lần này có phải đơn giản là thay càng nhiều phép toán fp32 trong kernel đó sang fp16 càng tốt hay không
      Việc chuyển đổi kiểu này thực sự khó đến mức nào thì cần kiểm tra thêm
      Nhưng về trực giác tôi không thấy quá ấn tượng
      Nếu tôi hiểu sai, tôi mong các tác giả làm rõ điểm này

    • Nếu vậy thì kết quả sẽ trở nên vô nghĩa
      Tôi không rõ họ có kiểm tra sai số tương đối hay không
      Thay float32 bằng float16 thì không có ý nghĩa gì
      Độ chính xác gần như là lý do duy nhất để chọn float32
      Nếu đánh mất chính đặc tính này thì cũng không còn khác biệt giữa các phiên bản nữa

  • Có phải chỉ mình tôi đọc tiêu đề này rồi tưởng AI đã tạo ra OS kernel không?

    • Tôi cũng vậy

  • Cải thiện tốc độ 400% đã rất ghê rồi
    Nhưng điều thú vị hơn cả là: phương pháp luận đã cưỡng bức đưa bước suy luận ngôn ngữ vào mỗi vòng lặp, thay vì chỉ tối ưu các phép toán đơn thuần, để tạo ra sự đa dạng trong tìm kiếm
    Việc cách làm này thực sự hoạt động tốt khiến nó đặc biệt thú vị

    • Tôi đã nghĩ có thể họ dùng thứ gì đó như map-elites hay kỹ thuật island, nhưng có vẻ tôi đã bỏ lỡ phần đó
      Tôi nghĩ đây chỉ là một dạng tiến hóa meme thông thường

  • Kết quả thực sự rất thú vị
    Có vẻ những người này quá phấn khích nên đã công bố lên blog
    Cũng có thể họ muốn nhận phản hồi trước khi xuất bản chính thức
    Tôi không biết đây có phải con đường huyền thoại dẫn tới "tự cải thiện" hay không
    Nhưng tôi cảm thấy những kết quả như thế này đúng là kiểu ví dụ mà ta kỳ vọng sẽ thấy trên con đường đó

    • "Tôi không biết đây có phải con đường dẫn tới tự cải thiện thật sự hay không"
      Các phương pháp này chỉ hiệu quả khi có một hàm đánh giá cực kỳ rõ ràng

  • Chia sẻ trải nghiệm của tôi: trong lần thử replication, kernel LayerNorm không ổn định về mặt số học nên tôi xem như không thể xác thực
    Họ chỉ kiểm thử với trung bình 0, độ lệch chuẩn 1 nên hiện tượng numerical cancellation không lộ ra ngay
    Ngoài ra, sau đó tôi có xác nhận rằng họ đã làm lại một phiên bản ổn định về mặt số học
    Tôi nghĩ đó là một điểm đáng khen

  • Kết quả thực sự rất ngầu
    Họ có dùng o3 và Gemini 2.5 Pro
    Nhưng tiếc là không nói rõ bên nào tạo ra kernel tốt hơn

  • Điểm đáng quan sát là code do AI tạo ra xử lý các vùng rộng hơn như fused kernel như thế nào
    Ví dụ có thể là gemm + relu + gemm + normalization, v.v.
    Nếu duyệt bằng tuner hoặc để con người tự viết thủ công từng cái thì sẽ là công việc khá vất vả

    • Nhưng tôi vẫn chưa rõ chính xác "kernel" ở đây nghĩa là gì trong bối cảnh AI
      Chỉ chắc chắn một điều là nó không phải OS kernel