K-Search: Генерация ядра больших языковых моделей посредством совместной эволюции внутренней модели мира

K-Search: LLM Kernel Generation via Co-Evolving Intrinsic World Model

February 22, 2026

Авторы: Shiyi Cao, Ziming Mao, Joseph E. Gonzalez, Ion Stoica

cs.AI

Аннотация

Оптимизация GPU-ядр критически важна для эффективности современных систем машинного обучения, но остается сложной задачей из-за комплексного взаимодействия факторов проектирования и быстрой эволюции аппаратного обеспечения. Существующие автоматизированные подходы обычно рассматривают большие языковые модели (LLM) лишь как стохастические генераторы кода в рамках эволюционных циклов, управляемых эвристиками. Эти методы часто испытывают трудности со сложными ядрами, требующими согласованных, многоэтапных структурных преобразований, поскольку им не хватает явных возможностей планирования и они часто отбрасывают перспективные стратегии из-за неэффективных или некорректных промежуточных реализаций. Для решения этой проблемы мы предлагаем метод поиска через коэволюцию моделей мира и на его основе создаем K-Search. Заменяя статические поисковые эвристики коэволюционирующей моделью мира, наша система использует априорные предметные знания LLM для направления поиска, активно исследуя пространство оптимизации. Этот подход явно разделяет высокоуровневое алгоритмическое планирование и низкоуровневую инстанциацию программ, позволяя системе исследовать немонотонные пути оптимизации, оставаясь устойчивой к временным дефектам реализации. Мы оцениваем K-Search на разнообразных сложных ядрах из FlashInfer, включая ядра GQA, MLA и MoE. Наши результаты показывают, что K-Search значительно превосходит современные методы эволюционного поиска, демонстрируя в среднем улучшение в 2.10 раза и до 14.3 раз на сложных ядрах MoE. В задаче GPUMode TriMul K-Search достигает наилучшей производительности на H100, показывая результат 1030 мкс и превосходя как предыдущие эволюционные методы, так и решения, созданные человеком.

English

Optimizing GPU kernels is critical for efficient modern machine learning systems yet remains challenging due to the complex interplay of design factors and rapid hardware evolution. Existing automated approaches typically treat Large Language Models (LLMs) merely as stochastic code generators within heuristic-guided evolutionary loops. These methods often struggle with complex kernels requiring coordinated, multi-step structural transformations, as they lack explicit planning capabilities and frequently discard promising strategies due to inefficient or incorrect intermediate implementations. To address this, we propose Search via Co-Evolving World Model and build K-Search based on this method. By replacing static search heuristics with a co-evolving world model, our framework leverages LLMs' prior domain knowledge to guide the search, actively exploring the optimization space. This approach explicitly decouples high-level algorithmic planning from low-level program instantiation, enabling the system to navigate non-monotonic optimization paths while remaining resilient to temporary implementation defects. We evaluate K-Search on diverse, complex kernels from FlashInfer, including GQA, MLA, and MoE kernels. Our results show that K-Search significantly outperforms state-of-the-art evolutionary search methods, achieving an average 2.10x improvement and up to a 14.3x gain on complex MoE kernels. On the GPUMode TriMul task, K-Search achieves state-of-the-art performance on H100, reaching 1030us and surpassing both prior evolution and human-designed solutions.