Tự động nghiên cứu với Codex: Hành trình đạt tốc độ nhanh hơn 232 lần cho Kernel QR

Trong cuộc thi auto-research do GPU Mode phối hợp với Core Automation tổ chức, tác giả Sankalp đã sử dụng Codex (mô hình GPT-5.5) để tối ưu kernel QR factorization trên GPU B200, đạt tốc độ nhanh hơn 232 lần so với baseline và xếp hạng 12/183. Đây là một ví dụ điển hình về sức mạnh của agent AI trong việc tối ưu hiệu năng kernel GPU, đồng thời hé lộ những chiến lược quan trọng để điều hướng các agent nghiên cứu tự động.

Cuộc thi trong ngắn gọn

GPU Mode, hợp tác với Core Automation, đã tổ chức một cuộc thi mang chủ đề auto-research. Bài toán đặt ra là triển khai QR factorization dạng batched square compact-Householder, tức là phân rã QR. Tác giả đã xếp hạng 12/183 với tốc độ nhanh hơn baseline 232 lần.

Lưu ý: Bạn không cần hiểu chi tiết toán học hay bài toán để theo dõi bài viết này. Trọng tâm là cách tiếp cận và chiến lược điều hướng agent.

Cuộc thi này thuộc chuỗi Linear Algebra Kernels in the Age of Research của GPU Mode.

Giới thiệu bài toán

Chúng ta được cung cấp một batch ma trận vuông FP32 CUDA với kích thước batch x n x n, và phải trả về biểu diễn compact Householder QR giống như torch.geqrf(A). Kết quả gồm:

• Ma trận H: tam giác trên là R, tam giác dưới lưu trữ Householder vectors
• Vector tau: hệ số phản xạ cho mỗi cột

Trình chấm điểm sẽ tái tạo Q bằng torch.linalg.householder_product(H, tau), lấy R = triu(H), và kiểm tra:

• A ≈ QR
• QᵀQ ≈ I
• QᵀA ≈ R

Các kích thước quan trọng gồm ma trận 512×512 (được đánh giá cao nhất), cùng các case lớn hơn 1024, 2048, và 4096. Được phép dùng FP16, FP8 hoặc NVFP4 nội bộ, nhưng kết quả trả về phải đạt chuẩn FP32.

Vì sao bài toán này phù hợp với auto-research?

GPU Mode cung cấp cho người tham gia popcorn CLI — công cụ giúp agent dễ dàng test, benchmark và submit lên leaderboard. Trình chấm cung cấp phản hồi chi tiết theo từng kích thước cùng thời gian trung bình hình học.

Điều này tạo ra một vòng lặp phản hồi khép kín — điều mà các agent AI cực kỳ ưa thích. Trong 14 ngày, tác giả đã thực hiện hơn 1500 lần submit.

Học đủ để đặt câu hỏi tốt hơn

Điểm mấu chốt: bạn càng hiểu rõ vấn đề, bạn càng có thể prompt LLM tốt hơn, vì bạn chuyển hóa "unknown unknowns" thành "known unknowns".

Tác giả đã dành thời gian đầu để học về QR decomposition là gì và có những cách nào để thực hiện. Có nhiều phương pháp như Gram-Schmidt và Householder reflections. Cuộc thi bắt buộc dùng Householder reflections.

Sau khi thảo luận với Claude và xem video YouTube, rõ ràng cần dùng blocked Householder algorithm với trailing WY-update làm kiến trúc chính.

Blocked Householder: Biến công việc tuần tự thành GEMM

Householder QR chuẩn zeroes phần dưới đường chéo từng cột một. Vấn đề là reflector j+1 được xây dựng từ ma trận sau khi reflector j đã tác động — tạo ra chuỗi phụ thuộc tuần tự không thể tái sắp xếp hay hợp nhất.

Thuật toán blocked là giải pháp cổ điển: chọn một panel hẹp b cột (32 hoặc 64), làm toàn bộ công việc tuần tự bên trong panel. Sau đó, nén b reflectors thành một rank-b update duy nhất (WY representation) và áp dụng lên toàn bộ khối trailing trong một lần với ba phép nhân ma trận liên tiếp.

Cụ thể, WY representation nén b reflectors:

H1·H2·...·Hb = I − VTVᵀ

Và cập nhật trailing-block thành ba bước dạng GEMM:

• W = VᵀA_trail
• Z = TᵀW
• A_trail ← A_trail − VZ

Công việc tuần tự chỉ giới hạn trong panel (chỉ b cột), còn mọi thứ khác trở thành GEMM — đúng hình dạng mà tensor cores mong muốn.

Chiến lược với Codex

Vì sao chọn Codex?

• Có gói ChatGPT Pro ($200/tháng)
• Kinh nghiệm cho thấy OpenAI models tốt hơn với Triton
• Lệnh /compaction hoạt động rất tốt

Thiết lập harness

Tác giả yêu cầu Codex:

• Thêm problem_statement.md
• Ghi chi tiết cách submit và dùng popcorn CLI vào AGENTS.md
• Duy trì log.md để ghi nhận các lần submit (accept/reject kèm timing theo shape)

Logs là bằng chứng về những ý tưởng đã hiệu quả và không hiệu quả. Các phiên agent sau có thể đọc logs và kiểm tra nhanh xem ý tưởng nào đã được thử.

Điều hướng với /goal

Để model lặp cho đến khi đạt mục tiêu, dùng lệnh /goal. Cần đưa ra mục tiêu định lượng cụ thể:

"Chỉ dùng Triton hoặc CUDA và đánh bại thời gian n=512 của best hiện tại. Thử nhiều ý tưởng bằng cách submit trực tiếp hoặc dùng Modal profiling. Loại bỏ cuSolver hoàn toàn trong các thử nghiệm mới."

Một lần, mục tiêu này chạy liên tục hơn một ngày. Tác giả chỉ kiểm tra mỗi 2-3 giờ để điều hướng.

Kiểm tra mà không dừng vòng lặp

Dùng lệnh /btw hoặc /side để hỏi model mà không làm gián đoạn vòng lặp chính. Các câu hỏi điển hình:

• Thuật toán/thay đổi nào trong submission hiện tại?
• Những ý tưởng đang thử là gì?
• Có thể chuyển ý tưởng này sang shape khác không?

Các bước đột phá về hiệu năng

Hành trình từ baseline 108,803 µs → 1,805 µs (nhanh hơn 98.34%) trải qua 10 thay đổi cấu trúc chính:

#	Thay đổi cấu trúc	Loại	Geomean
1	torch.geqrf cho mọi shape	Điểm khởi đầu	>108.8k µs
2	Blocked WY QR trên n512	Algorithm	108.8k µs
3	Blocked route trên mọi shape	Algorithm	10.2k µs
4	Triton panels + grouped WY	Kernel	4.3k µs
5	Cholesky-ORHR cho n4096	Algorithm	4.0k µs
6	CUDA graph replay	Runtime	3.4k µs
7	Fused V/T layout assembly	Kernel	2.75k µs
8	split16 panels + tail-Gram	Algorithm	2.5k µs
9	Fixed-shape kernel specialization	Kernel	2.0k µs
10	Composed superpanels, custom Cholesky	Kernel	1.80k µs

Vấn đề lớn nhất là launch overhead và panel overhead — không phải memory hay compute. Phần lớn nỗ lực tập trung giảm panel overhead và làm WY-update GEMM-able hơn.

Đưa sự đa dạng ý tưởng để thoát local maxima

Thách thức lớn nhất trong giai đoạn 3000 → 1800 µs là model bị kẹt trong local maxima — chỉ tinh chỉnh tham số và các biến thể nhỏ của cùng một ý tưởng.

Chiến lược "beam of candidates"

Thay vì chỉ giữ một candidate tốt nhất, duy trì 3-5 beam ý tưởng song song:

• Một beam exploit gần best hiện tại
• Một beam near-miss
• Một beam structural/risk cao từ profiling hoặc ý tưởng bên ngoài
• Một beam cleanup/compile-time nếu chưa chiếm hết tài nguyên

Nếu hai ý tưởng riêng lẻ trung tính hoặc hơi chậm nhưng tác động lên các chi phí độc lập, hãy thử kết hợp chúng trước khi loại bỏ.

Human-in-the-loop

Tác giả đóng vai trò "secret sauce" để điều hướng khi model kẹt lâu. Khuyến khích model:

• Chấp nhận rủi ro với ý tưởng táo bạo — điều này thực sự hiệu quả!
• Dùng advisor model mạnh hơn (như GPT-5.6 Sol hoặc Fable) để tạo ý tưởng đa dạng
• Dùng sub-agents thử ý tưởng tham vọng, tìm kiếm blog và papers trên web

Lời khuyên về AGENTS.md

Tác giả đã viết AGENTS.md kỹ lưỡng với các nguyên tắc:

• Submission discipline: profile sau thay đổi lớn, dùng advisor model khi kẹt, space các lần submit
• Beam search discipline: không hill-climb với single incumbent, duy trì beam đa dạng
• Promotion hygiene: commit sau khi promote candidate mới
• Profiling habits: ưu tiên bằng chứng hiện tại, lưu logs/profile dưới submit_logs/

Những gì có thể làm tốt hơn

Tác giả rút ra bốn bài học từ việc xem xét top 10 submissions:

1. Data detectors: Các kernel nhanh hơn viết bộ phát hiện phân bố dữ liệu (dense, clustered, rank-deficient, mixed) và khai thác chúng. Ví dụ: case low-rank có nhiều số 0.

2. Loại bỏ thư viện mạnh mẽ hơn: Top 2 và 5 dùng custom triangular inverse thay vì PyTorch triangular solve. Giải pháp của tác giả vẫn qua lại nhiều giữa PyTorch và Triton.

3. Giữ trailing matrix trong fp16: Không nên chuyển đổi qua lại giữa các biểu diễn. Đây là một "unknown unknown" thuần túy về domain expertise.

4. Beam of candidates từ đầu: Nên có beam đa dạng ngay từ khởi đầu.

5. Không dùng được tcgen05: Chưa khai thác được lệnh tensor-core thế hệ thứ năm trên Blackwell (B200).

Kết luận

Tác giả xếp hạng 12/183 với tốc độ nhanh hơn baseline 232 lần. Quan trọng hơn, anh học được nhiều về tối ưu GPU kernel, cơ bản về auto-research (hay "loop engineering"?) và chi tiết riêng của B200.

Có một phổ về cách các kỹ sư muốn xây dựng: người thì muốn harness tổng quát tự cải thiện, người thì muốn harness chuyên biệt cho từng vấn đề. Tôi nghiêng về harness chuyên biệt cho vấn đề.

Nếu bạn quan tâm, cuộc thi thứ hai trong chuỗi — eigen decomposition — đang diễn ra. Hẹn gặp trên leaderboard!

---

Bài viết dựa trên blog của Sankalp tại sankalp.bearblog.dev. Lưu ý: một số chi tiết về mô hình AI (GPT-5.5, Claude Opus 4.8, GPT-5.6) và mốc thời gian trong tương lai là theo bài gốc.