K-Search là gì? AI tự tối ưu GPU kernel

Trả lời nhanh: K-Search là hệ thống dùng LLM để lập kế hoạch, tạo, benchmark và lặp lại quá trình tối ưu GPU kernel. Khác với vòng tiến hóa chỉ sửa code ngẫu nhiên, K-Search duy trì một world model dạng cây giả thuyết, tách chiến lược tối ưu khỏi implementation và cập nhật niềm tin bằng kết quả chạy trên phần cứng thật.

Cập nhật lần cuối: 21/08/2026

Viết code chạy đúng và viết GPU kernel chạy nhanh là hai bài toán khác nhau. Một thay đổi tưởng hợp lý có thể làm tăng register pressure, giảm occupancy hoặc phá pattern truy cập bộ nhớ. Một candidate trung gian có thể chậm hơn nhưng lại là bước bắt buộc để mở đường cho tối ưu lớn ở vòng sau.

Các hệ thống code generation một lượt thường thất bại ở đây. Chúng có thể tạo code hợp lệ nhưng không quản lý được chuỗi giả thuyết, không nhớ tại sao một hướng thất bại và dễ chọn những thay đổi nhỏ đem lại lợi ích tức thời.

K-Search biến công việc này thành một vòng tìm kiếm có trạng thái. LLM không chỉ viết code; nó duy trì một mô hình về bottleneck, chiến lược, rủi ro và bằng chứng đã quan sát được.

K-Search là gì?

K-Search là framework tự động tối ưu GPU kernel bằng LLM và co-evolving intrinsic world model. Hệ thống duy trì cây tìm kiếm gồm các giả thuyết tối ưu, tạo implementation tương ứng, benchmark candidate trên phần cứng thật rồi dùng kết quả để thêm, sửa hoặc loại nhánh.

K-Search được giới thiệu bởi Shiyi Cao, Ziming Mao, Joseph E. Gonzalez và Ion Stoica trong nghiên cứu của UC Berkeley Sky Computing Lab.

Thành phần	Chức năng
Spec	Mô tả bài toán, phần cứng, constraint và pattern tối ưu
World model	Cây giả thuyết về bottleneck và hướng cải tiến
Action selection	Chọn nhánh có tiềm năng để thử tiếp
Code policy	Sinh implementation cụ thể cho một action
Compiler và correctness test	Loại code không build hoặc sai kết quả
Hardware benchmark	Đo latency hoặc throughput trên thiết bị thật
World model update	Insert, Update hoặc Prune dựa trên evidence
Persistence	Lưu solution, report và trạng thái để tiếp tục search

Tên “world model” trong K-Search không mang nghĩa mô phỏng toàn bộ thế giới vật lý. Nó là mô hình nội tại có cấu trúc về không gian tối ưu kernel: điều gì đã thử, bottleneck nào được nghi ngờ, nhánh nào có tiềm năng và độ tin cậy hiện tại là bao nhiêu.

GPU kernel là gì và vì sao khó tối ưu?

GPU kernel là hàm cấp thấp chạy song song trên GPU để thực hiện một phép tính cụ thể. Attention, matrix multiplication, normalization, routing trong Mixture of Experts và state-space scan đều phụ thuộc vào kernel hiệu quả.

Hiệu năng kernel chịu tác động đồng thời bởi:

cách chia tile;
memory layout;
coalesced access;
shared memory;
register pressure;
occupancy;
synchronization;
vectorization;
warp hoặc SIMD primitive;
kích thước input;
kiến trúc GPU và compiler.

Một tối ưu tốt cho workload này có thể tệ cho workload khác. Tăng tile size có thể giảm số lượt đọc memory nhưng lại làm tràn register. Fusing operation có thể giảm memory traffic nhưng tăng độ phức tạp và làm compiler khó tối ưu.

Vì vậy, kernel optimization là bài toán search dưới nhiều constraint, không phải bài toán dịch code cú pháp.

K-Search khác vòng lặp sinh code thông thường thế nào?

Một vòng lặp đơn giản thường làm như sau:

LLM tạo candidate.
Hệ thống compile và benchmark.
Giữ candidate nhanh nhất.
Yêu cầu LLM sửa tiếp.

Quy trình này dễ mắc kẹt ở local optimum. Nếu một thay đổi cấu trúc cần hai bước, candidate sau bước đầu có thể chậm hoặc lỗi. Greedy search sẽ loại nó trước khi bước thứ hai xuất hiện.

K-Search tách hai tầng:

Tầng kế hoạch: mô tả hành động tối ưu, tác động dự kiến và quan hệ giữa các bước.
Tầng hiện thực: thử viết code cụ thể để kiểm chứng một hành động.

Một implementation lỗi không đồng nghĩa giả thuyết sai. Hệ thống có thể thử lại implementation khác hoặc giữ nhánh để quay lại sau.

Tiêu chí	Vòng lặp code đơn giản	K-Search
Bộ nhớ	Candidate tốt nhất gần nhất	Cây gồm giả thuyết, code và evidence
Planning	Thường ẩn trong prompt	Tách thành action node có cấu trúc
Thất bại trung gian	Dễ bị loại cùng ý tưởng	Có thể phân biệt lỗi ý tưởng và lỗi implementation
Khám phá	Mutation hoặc prompt ngẫu nhiên	Chọn nhánh theo rating và confidence
Cập nhật	Giữ hoặc bỏ candidate	Insert, Update, Prune world model
Evidence	Score cuối	Benchmark, bottleneck và trajectory

World model của K-Search hoạt động ra sao?

World model là một prefix tree. Mỗi đường từ root tới leaf thể hiện một kế hoạch tối ưu gồm nhiều action liên tiếp. Các sibling là những hướng cạnh tranh.

Một node có thể chứa:

hành động cụ thể;
độ khó triển khai từ 1 đến 5;
tác động dự kiến lên memory bandwidth;
rủi ro register pressure;
mức phù hợp với compute hardware;
rating tổng từ 0 đến 10;
confidence từ 0 đến 1;
liên kết tới implementation và benchmark đã thử.

Cấu trúc này buộc LLM nói rõ “tại sao thử” trước khi “viết gì”. Sau khi có kết quả, LLM cập nhật cây:

Insert: thêm hướng mới được suy ra từ evidence.
Update: sửa rating hoặc confidence của node cũ.
Prune: loại nhánh kém tiềm năng hoặc vi phạm constraint.

Khi search stagnate, hệ thống có thể lùi lại và mở một nhánh khác thay vì tiếp tục vắt kiệt một candidate.

Vòng lặp K-Search gồm những bước nào?

Đọc Spec để hiểu phép tính, hardware, correctness và constraint.
Phân loại kernel như reduction, scan, attention, softmax hoặc matrix operation.
Chuẩn hóa phép tính tham chiếu để tách semantics khỏi implementation hiện tại.
Phân tích data layout và bottleneck theo từng runtime regime.
Chọn action node tiềm năng từ frontier của world model.
Sinh candidate code thể hiện một thay đổi có thể kiểm chứng.
Compile, test correctness và benchmark trên workload cùng phần cứng mục tiêu.
Gắn evidence vào trajectory gồm latency, lỗi và profiler signal.
Cập nhật world model bằng Insert, Update hoặc Prune.
Lặp cho đến budget hoặc điều kiện dừng rồi xuất kernel tốt nhất cùng lịch sử tìm kiếm.

Điểm quan trọng là hardware benchmark nằm trong loop. LLM chỉ tạo giả thuyết; phần cứng quyết định giả thuyết có tạo hiệu năng thật hay không.

Kết quả nghiên cứu K-Search nói lên điều gì?

Nhóm tác giả đánh giá K-Search trên các kernel phức tạp của FlashInfer gồm GQA, MLA và MoE.

Theo bài báo gốc:

K-Search đạt mức cải thiện trung bình 2,10 lần so với OpenEvolve;
mức tăng cao nhất đạt 14,3 lần trên kernel MoE phức tạp;
trên GPUMode TriMul với H100, nghiên cứu báo cáo latency khoảng 1.030 microsecond và vượt các lời giải tự động lẫn thủ công được so sánh.

Repository chính thức cập nhật benchmark TriMul với geometric mean 1,028 ms trên 7 workload, 3 lần chạy, H100, PyTorch 2.8.0 và Triton 3.4.0.

Những con số này cho thấy planning có trạng thái giúp search tốt hơn trong benchmark được chọn. Chúng không có nghĩa K-Search tạo kernel tối ưu cho mọi model, hardware hoặc production workload.

Vì sao kết quả trên MoE đáng chú ý?

MoE kernel chứa nhiều yếu tố tương tác như routing, top-k, quantization và memory access. Tối ưu thường cần thay đổi cấu trúc phối hợp, không chỉ chỉnh một tham số. Đây chính là trường hợp mà việc tách high-level plan khỏi code instantiation có thể tạo lợi thế.

K-Search có phải AI tự viết code hoàn toàn không?

Không theo nghĩa không cần con người hoặc hạ tầng chuyên môn. Hệ thống cần:

Spec đúng;
reference implementation;
test correctness;
workload đại diện;
compiler và runtime ổn định;
phần cứng benchmark;
budget inference;
rule bảo vệ tính an toàn;
review trước khi tích hợp production.

LLM giúp sinh và quản lý hypothesis. Nó không thay thế benchmark harness và domain constraint. Nếu Spec sai, hệ thống có thể tối ưu rất giỏi cho một mục tiêu không còn ý nghĩa.

K-Search mang lại bài học gì cho AI Agent?

Giá trị của K-Search vượt khỏi GPU kernel ở cấp độ pattern thiết kế.

1. Tách kế hoạch khỏi thực thi

Một agent nên lưu hypothesis hoặc plan như object có version. Tool execution là phép thử, không phải toàn bộ reasoning.

2. Lưu cả thất bại có cấu trúc

Chỉ lưu “best result” khiến agent lặp lại ngõ cụt. Cần lưu action, assumption, environment, error và evidence.

3. Dùng feedback từ môi trường thật

Model không nên tự đánh giá code mình viết. Unit test, benchmark, linter và policy engine phải cung cấp tín hiệu độc lập.

4. Cho phép đường đi không đơn điệu

Một bước trung gian có thể chưa tốt nhưng cần thiết cho thay đổi cấu trúc lớn. Workflow phải phân biệt regression tạm thời trong sandbox với regression được đưa production.

5. Duy trì search budget

Không để agent thử vô hạn. Giới hạn số vòng, chi phí, thời gian, candidate và hardware minute.

Pattern này gần với cách tối ưu agent loop được phân tích trong bài tối ưu Agent Loop và các nguyên tắc vận hành tại xây dựng AI Agent production.

Khi nào K-Search đáng để thử?

K-Search phù hợp hơn khi:

kernel chiếm tỷ trọng lớn trong end-to-end latency;
đã có correctness harness đáng tin;
workload và hardware mục tiêu ổn định;
mỗi phần trăm hiệu năng có giá trị vận hành lớn;
không gian tối ưu chứa nhiều thay đổi cấu trúc;
có ngân sách benchmark và model call.

Không nên bắt đầu từ K-Search khi hệ thống còn bottleneck ở I/O, network, batching hoặc database. Tối ưu kernel không giúp nhiều nếu kernel chỉ chiếm phần nhỏ tổng latency.

Cách đánh giá K-Search trong thực tế

Bước 1: Chọn kernel bằng profiling

Đo end-to-end trước. Ưu tiên kernel có đóng góp lớn theo Amdahl’s law, không chọn vì nó thú vị về kỹ thuật.

Bước 2: Khóa correctness contract

Xác định shape, dtype, tolerance, determinism, edge case và numerical stability. Candidate nhanh nhưng sai không được tính.

Bước 3: Xây workload suite

Không benchmark một shape duy nhất. Dùng phân phối workload production, gồm p50, p95 và case hiếm nhưng quan trọng.

Bước 4: Chọn baseline công bằng

So với native library, compiler output, human kernel và phương pháp search khác trên cùng phần cứng.

Bước 5: Đặt search budget giống nhau

So sánh số iteration, model call, token, hardware minute và số candidate hợp lệ. Nếu budget khác, kết luận về thuật toán search có thể sai.

Bước 6: Tái lập nhiều lần

LLM generation có tính ngẫu nhiên. Báo best, median, variance và failure rate qua nhiều seed.

Bước 7: Kiểm tra portability

Chạy trên driver, compiler, shape và GPU khác. Kernel quá chuyên biệt có thể không giữ lợi ích khi môi trường đổi.

Bước 8: Canary trước production

Đóng gói version, giữ fallback, theo dõi numerical error và latency distribution. Không thay kernel trực tiếp chỉ vì một benchmark tốt.

Những giới hạn và rủi ro nào cần lưu ý?

Benchmark overfitting

Search có thể tối ưu đúng suite nhưng không tốt cho traffic thật. Workload phải đại diện và test set cần tách khỏi vòng search.

Chi phí tìm kiếm

Nhiều model call, compile và hardware benchmark tạo chi phí đáng kể. Cần tính thời gian hoàn vốn từ inference saving.

Correctness tinh vi

Sai số số học có thể chỉ xuất hiện ở shape hoặc giá trị hiếm. Test cần kiểm tra precision, overflow, race condition và non-determinism.

World model cũng có thể sai

Rating và confidence do LLM tạo không phải xác suất đã calibration. Chúng chỉ là heuristic để ưu tiên search và phải bị evidence thực tế kiểm soát.

Phụ thuộc Spec

Constraint thiếu có thể dẫn tới code compile được nhưng dùng pattern không phù hợp phần cứng. Context engineering là một phần của hệ thống, không phải phụ kiện.

Rủi ro supply chain

Generated code, dependency và script benchmark cần được review. Không chạy code không tin cậy trên hạ tầng có credential hoặc dữ liệu nhạy cảm.

FAQ - Câu hỏi thường gặp về K-Search

K-Search có thay thế kỹ sư GPU không?

Chưa. Nó tự động hóa một phần hypothesis search và code generation, nhưng vẫn cần chuyên gia xây Spec, test, benchmark, phân tích correctness và quyết định triển khai.

K-Search khác evolutionary search ở điểm nào?

Nó dùng một world model dạng cây để lập kế hoạch và cập nhật niềm tin, thay vì chỉ mutation code theo heuristic tĩnh. High-level action được tách khỏi low-level implementation.

K-Search có chỉ dùng cho CUDA không?

Nghiên cứu gốc tập trung vào CUDA và FlashInfer. Repository hiện có nhiều backend, gồm MLX cho Apple Silicon, nhưng mỗi backend cần Spec, compiler, benchmark và constraint riêng.

World model trong K-Search có giống world model của robot không?

Không hoàn toàn. Ở đây nó là biểu diễn nội tại về không gian tìm kiếm kernel, hypothesis và evidence, không phải mô hình dự đoán thế giới vật lý đầy đủ.

Kết quả nhanh hơn 14,3 lần có áp dụng cho mọi kernel không?

Không. Đây là mức cao nhất được báo cáo trên kernel MoE cụ thể trong benchmark nghiên cứu. Kết quả thực tế phụ thuộc kernel, baseline, workload, hardware và search budget.

Có thể áp dụng pattern K-Search cho automation marketing không?

Có thể học pattern hypothesis, experiment, evidence và memory. Tuy nhiên, metric marketing nhiễu và có độ trễ cao hơn benchmark kernel; cần thiết kế causal test, guardrail và approval riêng.

Tóm lại

K-Search biến tối ưu GPU kernel thành vòng lặp search có world model và evidence.
Hệ thống tách high-level plan khỏi code implementation để xử lý thay đổi nhiều bước.
Bài báo báo cáo cải thiện trung bình 2,10 lần và tối đa 14,3 lần trong benchmark cụ thể.
Hardware benchmark và correctness harness mới là nguồn sự thật, không phải đánh giá của LLM.
Bài học có thể chuyển sang AI Agent là lưu hypothesis, thất bại, evidence và giới hạn search budget.

Nguồn tham khảo

K-Search: LLM Kernel Generation via Co-Evolving Intrinsic World Model, Cao và cộng sự, 2026.
Trang dự án K-Search tại UC Berkeley Sky Computing Lab.
Repository mã nguồn K-Search.

Lưu ý: benchmark kernel nhạy với hardware, driver, compiler và workload. Chỉ so sánh kết quả khi các điều kiện đo được kiểm soát tương đương.