- AI द्वारा जनरेट किए गए CUDA-C kernels ने PyTorch के विशेषज्ञ-ऑप्टिमाइज़्ड kernels के बराबर या उससे बेहतर प्रदर्शन दिखाया
- एक अकेले LLM (large language model) ने प्राकृतिक भाषा में optimization ideas बनाना और अलग-अलग code branching को दोहराते हुए, मौजूदा तरीकों की तुलना में optimization diversity और parallel search में बेहतर प्रदर्शन हासिल किया
- प्रमुख benchmark नतीजों में Matmul(101%), Conv2D(179.9%), Softmax(111.8%), LayerNorm(484.4%), Conv2D+ReLU+MaxPool(290.1%) आदि में PyTorch की तुलना में भारी बढ़त दिखी
- मौजूदा “sequential kernel improvement” की सीमाओं को पार करने के लिए natural language reasoning + branching structure लागू किया गया, जिससे तेज़ी से high-performance kernels बनाए गए
- FP16, Flash Attention जैसे आधुनिक ML workloads में भी प्रगति जारी है, और यह संकेत मिलता है कि भविष्य में AI खुद ही तेज़ kernels खोजने और सुधारने वाला paradigm मुख्यधारा बन सकता है
TL;DR मुख्य उपलब्धियां
- Stanford CRFM शोध टीम ने पाया कि AI द्वारा जनरेट किए गए high-performance CUDA-C kernels कुछ मामलों में PyTorch के विशेषज्ञ-डिज़ाइन किए गए kernels जैसी या उससे बेहतर गति दे सकते हैं
- मूल लक्ष्य synthetic data के ज़रिए बेहतर kernel auto-generation model को train करना था, लेकिन केवल synthetic data generation के दौरान ही आश्चर्यजनक रूप से तेज़ kernels अपने-आप बनते दिखे
- इसी वजह से टीम ने method, performance benchmarks, optimization strategy और आगे की दिशा को जल्दी साझा किया
- Benchmark: Nvidia L40S GPU के आधार पर. प्रदर्शन(%): PyTorch execution time ÷ generated kernel execution time
- Matmul (FP32): PyTorch की तुलना में 101.3% (4096x4096 matrix)
- Conv2D: 179.9% (input: 100, 3, 224, 224; AlexNet Conv1 spec)
- Softmax: 111.8% (4096x65536 tensor)
- LayerNorm: 484.4% (16, 64, 256, 256 tensor)
- Conv2D + ReLU + MaxPool: PyTorch reference की तुलना में 290.1%, torch.compile() की तुलना में 189.0%
शोध की प्रेरणा और तरीका
- मूल उद्देश्य kernel generation model training के लिए synthetic data बनाना था, लेकिन प्रयोग के दौरान खुद जनरेट किए गए kernels ने अपेक्षा से कहीं बेहतर high performance हासिल कर ली
- KernelBench (Stanford का सार्वजनिक benchmark, दिसंबर 2024 में जारी) का उपयोग किया गया
- दिए गए torch code के लिए LLM सबसे तेज़ kernel को नए सिरे से लिखता है
- input/output परिणामों की numerical equivalence से correctness verify की जाती है
- पुराना तरीका: kernel को चरण-दर-चरण थोड़ा-थोड़ा बदलते हुए सुधारने वाला sequential modification loop
- कमी: ideas की विविधता कम, एक जैसी optimization बार-बार, local optimum पर फंसने की संभावना
- सुधार का विचार
- optimization ideas को प्राकृतिक भाषा में सोचना और दर्ज करना, फिर उन ideas की कई code implementation branches एक साथ जनरेट करना
- हर round में कई optimization प्रयास parallel चलाना → सबसे तेज़ kernel को अगले round का seed बनाना
- इससे सीमित search iterations के भीतर विभिन्न optimization strategies की खोज और parallel search संभव हुई
optimization ideas के उदाहरण
- memory access optimization: global/shared/register memory hierarchy की efficiency बेहतर करना
- asynchronous processing और latency hiding: computation और data movement को overlap करना
- data type / precision optimization: FP16/BF16 का उपयोग और hardware-specific operations
- computation और instruction optimization: compute amount, instruction count, और hardware special instructions का बेहतर उपयोग
- parallelism और occupancy: SM (Streaming Multiprocessors) का अधिकतम उपयोग
- loop / branching / indexing optimization: loops को कम करना, अनावश्यक index operations हटाना
kernel optimization की वास्तविक प्रक्रिया (Conv2D उदाहरण)
- हर round के optimization ideas और performance improvement का flow
- शुरुआती चरण (round 0): साधारण CUDA porting (PyTorch की तुलना में 20% गति)
- अगले rounds: → read cache उपयोग → FP16 Tensor Core operations (GEMM conversion) → double buffering / pipeline → precomputed indices / shared memory → vectorization → K-loop concurrent buffering जैसे उन्नत GPU architecture उपयोग
- अंतिम चरण (round 13): PyTorch की तुलना में 179.9% प्रदर्शन हासिल
- वास्तविक code में advanced CUDA programming techniques का व्यापक उपयोग हुआ
- हर round में नए ideas प्राकृतिक भाषा में बनाए गए, कई implementations को parallel आज़माया गया, और सबसे अच्छे code को चुना गया
Takeaways और संकेत
- AI-आधारित kernel generation, मज़बूत search loop और प्राकृतिक भाषा पर आधारित विविध optimization ideas के संयोजन के साथ, मानव विशेषज्ञ स्तर को पार कर सकती है
- कुछ आधुनिक operators (FP16 matmul, Flash Attention) में अभी इस समय PyTorch की तुलना में प्रदर्शन कम है
- हाल के hardware में FP32 computation, FP16/BF16 की तुलना में अपेक्षाकृत कम optimized है → इसलिए इस क्षेत्र में प्रदर्शन बढ़त संभव है
- सीमित search token budget (input/output मिलाकर 7 million tokens) में भी लगातार performance improvement देखा गया
- AlphaEvolve, Gemini 2.5 Pro जैसी हालिया research भी यह संकेत देती है कि massive branching + verification-based search बिना retraining के भी बड़ा performance improvement दे सकती है
- आगे चलकर इसी तरीके से synthetic data और high-performance kernels का बड़े पैमाने पर generation होगा, और AI self-improving loop की दिशा में विकास संभव है
निष्कर्ष
- AI-आधारित kernel auto-generation और optimization अब विशेषज्ञ hand-coded स्तर तक पहुँच चुकी है, और निकट भविष्य में model + branching search + verification के संयोजन से self-improving AI systems संभव हो सकते हैं
- यह संभावना उभर रही है कि AI, PyTorch और TensorFlow जैसे frameworks की performance limits को भी अपने-आप पार कर सकेगी
परिशिष्ट: Conv2D kernel का पूरा code
- इसमें वास्तविक function और kernel का पूरा source code शामिल है (विस्तृत code यहाँ छोड़ा गया है)
- code की मुख्य विशेषताएं:
- shared memory vectorization, K-dim hierarchical double-buffering, CUDA WMMA, warp-level output buffer आदि
- real-time dynamic index calculation, bias handling, K loop के भीतर delayed data की concurrent loading आदि
- पूरा sample code और example kernels GitHub repository में देखा जा सकता है
2 टिप्पणियां
इसे एक तरह की ensemble तकनीक कहना चाहिए शायद। काफ़ी शानदार है।
Hacker News राय
मुझे लगता है कि इस लेख के लेखक "AI एजेंट" को जिस नज़रिए से देखते हैं, वह काफ़ी दिलचस्प है
ज़्यादातर लोग एजेंट को इंसानी कर्मचारी की तरह सोचते हैं
वे कुछ ही एजेंट्स को parallel में सेट करते हैं, अक्सर सिर्फ़ एक को, और हर एक को loop में चलाते हुए एक बार में केवल एक काम करने देते हैं
यानी वे अब भी ऐसी दुनिया में अटके हुए हैं जहाँ कर्मचारियों की संख्या तय है, वे एक समय में एक ही काम कर सकते हैं, और context switching धीमा है
लेकिन LLM ऐसा नहीं है
आप व्यावहारिक रूप से अनंत एजेंट जब चाहें तब बना सकते हैं
LLM requests को serial में प्रोसेस करें या parallel में, लागत में कोई फ़र्क नहीं होता
जब यह बात समझ में आती है, तो यह पैटर्न स्वाभाविक लगता है कि हर एजेंट ज़रूरत पड़ने पर खुद को कई sub-agents में बाँट ले
लेखकों ने ठीक यही तरीका आज़माया
मुझे लगता है कि एजेंट को "task" या "job" मानकर, और Celery या Sidekiq से सीखी बातों को लागू करने वाला नज़रिया ज़्यादा उपयुक्त है
"FP32 का इस्तेमाल आधुनिक ML workloads में काफ़ी कम हो गया है, और आधुनिक hardware पर यह FP16 या BF16 की तुलना में कम optimized भी है। इसलिए FP32 kernels में PyTorch से बेहतर प्रदर्शन हासिल करना अपेक्षाकृत आसान रहा होगा"
डेवलपर्स ने वर्षों तक fp32 version kernels को optimize करने में लगभग कोई समय नहीं लगाया
असल में सचमुच दिलचस्प बात तब होगी, जब वे उन kernels की performance बढ़ा सकें जिन पर सक्रिय रूप से काम हुआ है और जिन्हें लोग वास्तव में इस्तेमाल करते हैं
मुझे लगता है कि NVIDIA GPU के लिए पर्याप्त विस्तार से documentation नहीं देता, इसलिए ये अच्छे नतीजे आंशिक रूप से उसी से समझे जा सकते हैं
अगर processor की microarchitecture अच्छी तरह documented हो, तो programmer या compiler निर्णायक रूप से optimal program लिख सकते हैं, इसलिए ML/AI लगाकर पहले से ज्ञात समाधान ढूँढने से आगे बढ़कर बड़ा सुधार करना आसान नहीं होता
लेकिन NVIDIA GPU जैसी कम documented चीज़ों में optimal program ढूँढने के लिए अक्सर पहले से optimized programs के उदाहरणों के आधार पर random search करनी पड़ती है, या फिर अलग-अलग स्थितियों में असली hardware कैसे व्यवहार करता है, इसका reverse engineering करना पड़ता है
ऐसी स्थिति में ML/AI संभावना दिखा सकता है, और अच्छे programs को data की तरह सीखकर वह undocumented behavior भी पकड़ सकता है जिसे इंसानी programmer के लिए ढूँढना मुश्किल हो
मुझे जिज्ञासा है कि कहीं fp16/bf16 kernels में पहले से ज्ञात सुधारों को बस fp32 पर port तो नहीं कर दिया गया
"क्या आप कह रहे हैं कि लोगों ने वर्षों तक fp32 kernels को छुआ ही नहीं?"
वाह, अगर ऐसा है, तो यह सचमुच कमाल की बात है कि AI ने ऐसे क्षेत्र में नया algorithm बना दिया जहाँ पहले से कोई तैयार समाधान नहीं था
मेरा निष्कर्ष यह है कि इस आर्टिकल, Google के AlphaEvolve(यहाँ), और हाल की उस ख़बर से कि o3 ने Linux kernel में zero day ढूँढ निकाला(यहाँ)
मुझे ऐसा लगा कि खासकर Gemini Pro 2.5 और o3 अचानक उस नए capability level पर पहुँच गए हैं जहाँ वे ऐसे ideas भी निकाल रहे हैं जो पुराने models से नहीं हो पाते थे
मुझे नहीं लगता कि अचानक कुछ होने लगा है जो पहले नहीं हो रहा था
असल में बात यह है कि अब iteration और testing इंसानों की तुलना में बहुत तेज़ी से हो सकती है
और बड़ी मात्रा में जानकारी का तुरंत उपयोग संभव हो गया है
यानी हम उस बिंदु पर पहुँच गए हैं जहाँ विशाल जानकारी, प्रगति, और बुद्धिमानी से लागू किया गया brute force का संयोजन कुछ applications में सफल हो रहा है
मुझे लगता है कि जिन उदाहरणों का तुमने ज़िक्र किया और यह नतीजा, उनमें वास्तव में बहुत समानताएँ हैं
मुख्य लेख में भी कहा गया है कि "test-time iterative loop एक interactive chatbot जैसा नहीं है जो बारी-बारी से code modifications के नतीजे देखता है, बल्कि यह साफ़ optimization hypotheses के आधार पर आक्रामक parallel evaluation करने वाली structured search के ज़्यादा करीब है"
मेरा निष्कर्ष यह है कि अब LLM की क्षमताओं के आधार पर
हमने ऐसे solution spaces को काफ़ी घटाने का तरीका सीख लिया है जहाँ evaluation function स्पष्ट है या patterns मिलते-जुलते हैं
कौन-सा model किसे पीछे छोड़ता है, या कौन-सा model ही किसी समस्या को हल कर सकता है... यह models के बीच तुलना का मुद्दा नहीं है
मुझे ज़्यादा महत्वपूर्ण यह लगता है कि इस तरह के applications अब साफ़ तौर पर सामने आ रहे हैं
Gemini Pro 2.5 मुझे पहला ऐसा AI लगा जिसे इंसान वास्तव में इस्तेमाल कर सकता है, लेकिन सख़्ती से कहें तो उसने बस मुश्किल से वह threshold पार किया है
कई बार success rate 20% से नीचे भी गिर जाती है
अगर 3.0 आ गया... तो शायद वह थोड़ा डरावना भी लग सकता है
एक मिनट, मतलब क्या? इसका Linux kernel से कोई लेना-देना नहीं है, यहाँ GPU programming के "kernel" की बात हो रही है
क्या तुमने शायद पूरा लेख ग़लत समझ लिया?
दिलचस्प तो है, लेकिन क्या इससे मज़बूत सबूत भी है? सिर्फ़ एक बार का नतीजा काफ़ी convincing नहीं है
"baseline code default FP32 है, और tolerance 1e-02 है"
इतनी tolerance हो तो fp16 operations से किसी "fp32" kernel को आसानी से replace किया जा सकता है
एक कदम और आगे सोचें तो लगता है कि शायद इस काम का असली सार यही था कि उस kernel में जितने संभव हों उतने fp32 operations को fp16 में बदल दिया जाए
ऐसा porting काम व्यवहार में कितना कठिन है, यह जाँचने की ज़रूरत है
लेकिन सहज रूप से देखें तो यह इतना प्रभावशाली नहीं लगता
अगर मेरी सोच ग़लत है, तो अच्छा होगा कि लेखक इस बिंदु को साफ़ तौर पर संबोधित करें
अगर ऐसा है तो नतीजा अर्थहीन हो जाता है
पता नहीं उन्होंने relative error चेक किया भी या नहीं
float32 को float16 से बदलना कोई मायने नहीं रखता
precision ही शायद float32 चुनने की एकमात्र वजह होती है
अगर वही गुण खो गया, तो versions के बीच फ़र्क भी नहीं बचता
क्या सिर्फ़ मैं ही ऐसा था जिसने यह शीर्षक देखकर समझ लिया कि AI ने OS kernel बना दिया है?
400% speedup अपने-आप में प्रभावशाली है
लेकिन उससे भी ज़्यादा दिलचस्प बात यह है: हर iteration में सिर्फ़ arithmetic optimization नहीं, बल्कि search diversity लाने के लिए language reasoning step को ज़बरदस्ती शामिल करने वाली methodology
यह तरीका वास्तव में अच्छी तरह काम करता दिखा, इसलिए यह बहुत रोचक है
मुझे तो यह बस साधारण memetic evolution लगता है
यह सचमुच बहुत दिलचस्प नतीजा है
लगता है ये लोग इतने उत्साहित हो गए कि इसे blog पर ही डाल दिया
शायद उनका इरादा publication से पहले feedback लेने का भी रहा हो
पता नहीं यह "self improvement" की वह दंतकथाओं वाली राह है या नहीं
लेकिन ऐसा लगता है कि इस तरह के नतीजे उसी राह में दिखने वाली चीज़ों जैसे हैं
ये तरीके तभी काम करते हैं जब evaluation function बेहद स्पष्ट हो
अपना अनुभव साझा करूँ: replication की कोशिश में LayerNorm kernel numerical रूप से stable नहीं था, इसलिए मुझे वह verifiable नहीं लगा
उसे सिर्फ़ mean 0, standard deviation 1 पर टेस्ट किया गया था, इसलिए numerical cancellation तुरंत सामने नहीं आया
साथ ही, बाद में यह भी पुष्टि हुई कि उन्होंने एक नया numerically stable version बनाया
मुझे यह अच्छी बात लगी
यह सचमुच शानदार नतीजा है
o3 और Gemini 2.5 Pro का इस्तेमाल किया गया
लेकिन अफ़सोस, यह नहीं बताया गया कि किसने बेहतर kernel बनाया
यह देखना एक दिलचस्प बिंदु होगा कि AI-generated code fused kernel जैसे बड़े दायरे को कैसे संभालता है
जैसे gemm + relu + gemm + normalization वगैरह
इसे tuner से पूरी तरह explore करना या इंसानों से हाथ से लिखवाना काफ़ी मेहनत का काम है
बस इतना तय है कि यह OS kernel नहीं है