Vibecoding GPU Kernels

GPU कर्नेल (kernels) हाथ से लिखने में धैर्य, मेहनत और साथ ही आपकी सुध-बुध भी लगती है। मैं आधुनिक LLM (Claude Opus 5, GPT 5.6 Sol) के साथ प्रयोग कर रहा हूँ ताकि वे सचमुच मेरे लिए कर्नेल लिखकर दे सकें। मैं इस दौरान जो कुछ सीखा, उसके बारे में बात करूँगा।

सच कहूँ तो, GPU कर्नेल लिखना एक ऐसा काम है जिसे सत्यापित करना बेहद आसान (embarrassingly verifiable) है।

पहले, आप तय करते हैं कि आप किस ऑपरेशन (operations) के लिए कर्नेल लिखना चाहते हैं। दूसरा, आप पहला वर्ज़न लिखते हैं और सुनिश्चित करते हैं कि यह बिना किसी स्पष्ट त्रुटि के कंपाइल हो जाए। तीसरा, आप एक धीमी रेफ़रेंस इम्प्लीमेंटेशन (reference implementation) से सटीकता (correctness) की जाँच करते हैं। अगर यह मेल नहीं खाता, तो आप सटीकता की समस्या को ठीक करने की कोशिश करते हैं। ऐसा हो जाने के बाद, आप कर्नेल के निष्पादन समय (execution time) का बेंचमार्क (benchmark) लेते हैं। अगले वर्ज़न इसी के ऊपर आधारित होते हैं, और आप तब तक ऑप्टिमाइज़ करते रहते हैं जब तक कि आप रूफलाइन मेट्रिक्स (roofline metrics) तक नहीं पहुँच जाते या संतुष्ट नहीं हो जाते।

सत्यापन योग्य लूप: GPU कर्नेल डेवलपमेंट

ऊपर दी गई छवि इसे एक सत्यापन योग्य लूप के रूप में दिखाती है:

• कंपाइल जाँच सटीकता के बारे में सोचने से पहले का एक तंग स्थानीय लूप (B <-> C) है।
• सटीकता जाँच (D <-> E <-> F) मुख्य सत्यापन-पुरस्कार लूप है। यही वह हिस्सा है जो कर्नेल लेखन को स्वचालित सत्यापन के लिए एक अच्छी समस्या बनाता है, क्योंकि आपके पास जाँच के लिए एक ग्राउंड-ट्रुथ रेफ़रेंस (ground-truth reference) होता है।
• ऑप्टिमाइज़ेशन (G -> H -> वापस D) हर नए वर्ज़न के लिए उसी सटीकता लूप का पुन: उपयोग करता है, क्योंकि तेज़ लेकिन गलत कर्नेल बेकार है; हर वर्ज़न के लिए सटीकता सत्यापित करना ज़रूरी है।
• रूफलाइन/संतुष्टि जाँच (I) बाहरी लूप है जो तय करता है कि ऑप्टिमाइज़ करते रहना है या रुक जाना है।

पहला वर्ज़न

ऊपर की छवि अभी भी एक उच्च-स्तरीय दृश्य है और असली मुश्किल विवरणों में छिपी होती है। हमें यह सुनिश्चित करना होगा कि हमारे LLM एजेंट के पास एक अच्छा पहला वर्ज़न लिखना शुरू करने के लिए सभी आवश्यक संदर्भ (context) हों।

यहाँ CUDA DSL काम में आते हैं। Triton, CuTeDSL और Tilelang वे DSL हैं जिनके साथ Python में शुरुआत करना बहुत आसान है। CUDA C++ की तुलना में इनका सीखने का क्रम (learning curve) कम कठिन है, हालाँकि, उन DSL के एब्स्ट्रैक्शन (abstractions) हमारे एजेंट को और अधिक भ्रमित कर सकते हैं। हमें उन एब्स्ट्रैक्शन का संदर्भ एजेंट तक पहुँचाने का एक तरीका चाहिए।

आधुनिक LLM पहले से ही "अच्छा" Triton लिखना जानते हैं। वे बिना किसी संदर्भ के भी Triton एब्स्ट्रैक्शन के साथ अच्छी तरह काम कर सकते हैं। हालाँकि, CuTeDSL जैसे अन्य DSL के लिए (जो Triton की तुलना में कहीं अधिक नियंत्रण देता है), मैंने पाया है कि एक संदर्भ निर्देशिका (context directory) रखना, जहाँ एजेंट DSL एब्स्ट्रैक्शन को समझने के लिए देख सके, बहुत मददगार होता है।

उदाहरण के लिए, संदर्भ निर्देशिका में

NVIDIA cutlass

रिपॉज़िटरी को क्लोन करना एजेंट को CuTeDSL में कर्नेल लिखते समय

Layout Algebra, Copy/GEMM atoms, memory hierarchy, example kernels

जैसे एब्स्ट्रैक्शन और इसी तरह की अन्य चीज़ें खोजने का एक अच्छा तरीका है।

मेरे अनुभव में, कर्नेल का एक अच्छा पहला वर्ज़न बिना किसी स्पष्ट त्रुटि के कंपाइल हो जाता है और सटीकता परीक्षण (correctness test) पास कर लेता है, जिसके बारे में मैं नीचे बात करूँगा।

परीक्षण, बेंचमार्क और प्रोफाइलिंग

एजेंट को पर्याप्त संदर्भ देने के बाद, असली बाधा (bottleneck) अब सत्यापन (validation) पर आ जाती है। रेफ़रेंस इम्प्लीमेंटेशन स्वयं, और उसके विरुद्ध सत्यापन करना, और अधिक महत्वपूर्ण हो जाता है। मैं इस चरण को सटीकता परीक्षण या केवल परीक्षण कहता हूँ। रेफ़रेंस इम्प्लीमेंटेशन की गति उतनी मायने नहीं रखती जितना उसका उद्देश्य (intention)। आप जो मापना और सत्यापित करना चाहते हैं, आपका एजेंट उसी के लिए ऑप्टिमाइज़ करेगा।

आमतौर पर, जब गणना कम परिशुद्धता (lower precision) में नहीं होनी होती, तो मैं Max Absolute/Relative Error (MAE), Mean Squared Error (MSE/RMSE), और PSNR (Peak Signal to Noise Ratio) मापता हूँ। जब कम परिशुद्धता शामिल होती है, तो मैं PSNR और Cosine similarity (cossim) मापता हूँ।

आप कर्नेल वर्ज़न को वास्तव में GPU पर चलाने का तरीका इस बात पर निर्भर करता है कि GPU स्थानीय रूप से उपलब्ध है या क्लाउड के माध्यम से। फिर भी, हमारे एजेंट के पास किसी न किसी तरह अपने आउटपुट तक पहुँचने की क्षमता होनी चाहिए।

मुझे नीचे दी गई rung पद्धति N संख्या में परीक्षण फ़ंक्शन रखने का एक अच्छा तरीका लगती है:

1def rung(name):
2    def deco(fn):
3        try:
4            out = fn()
5            results[name] = {"ok": True, **(out or {})}
6            print(f"[{name}] ok  " + "  ".join(f"{k}={v}" for k, v in (out or {}).items()))
7        except Exception as e:
8            results[name] = {"ok": False, "err": f"{type(e).__name__}: {e}"}
9            print(f"[{name}] FAILED  {type(e).__name__}: {e}")
10            traceback.print_exc()
11        return fn
12    return deco

जिसे आप इस तरह कॉल कर सकते हैं:

1out = {}
2
3@rung("pre-checks")
4def _():
5    run_pure_checks()
6    run_dsl_checks()
7
8@rung("run")
9def _():
10    out["o"] = custom_kernel(*inputs)
11    torch.cuda.synchronize()
12    return {"shape": tuple(out["o"].shape),
13            "finite": bool(torch.isfinite(out["o"]).all())}

बेंचमार्क rung के लिए, आप कई चीज़ें कर सकते हैं:

• कुल खर्च हुए समय के लिए एंड-टू-एंड कर्नेल निष्पादन समय का बेंचमार्क लें
• इंट्रा-कर्नेल ट्रेसिंग (intra-kernel tracing) का उपयोग करके कर्नेल के भीतर के अनुभागों का बेंचमार्क लें और उन्हें आउटपुट में डंप करें (कस्टम ट्रेसर या CUPTI का उपयोग करके)
• जनरेट किए गए IR, PTX, SASS और CUBIN को dumps निर्देशिका में डंप करें और एजेंट को उसकी समीक्षा करने दें

आखिरी बिंदु को थोड़ा और विस्तार से समझा जा सकता है। कभी-कभी ऐसा हो सकता है कि DSL कम-इष्टतम (suboptimal) PTX (और अंततः SASS) जनरेट करे और आपको इसके बजाय उपयोग किए जा सकने वाले बेहतर इंस्ट्रक्शन या शेप (shapes) मिल जाएँ। हमारा एजेंट PTX/SASS टेक्स्ट फ़ाइलों को पढ़ सकता है और DSL को कम-इष्टतम हिस्से को संभालने देने के बजाय निचले स्तर के कोड को इनलाइन कर सकता है। फिर से, "खोजने योग्य" (searchable) PTX दस्तावेज़ को संदर्भ के रूप में देना यहाँ बहुत मददगार है।

सब कुछ एक साथ जोड़ने के लिए आखिरी चीज़ है प्रोफाइलिंग (Profiling)। अगर आपका एजेंट NCU (Nsight Compute Systems) CLI तक पहुँच सकता है, तो आप उसे ऊपर दिए गए सत्यापन योग्य फीडबैक लूप के हिस्से के रूप में आपके कर्नेल के लिए प्रोफाइलिंग करने और एक रिपोर्ट जनरेट करने के लिए कह सकते हैं।

समापन विचार

"तो क्या GPU कर्नेल डेवलपमेंट खत्म हो गया है?"

"हाँ, लेकिन असल में नहीं"

हाँ, क्योंकि लेआउट (layouts), इंडेक्सिंग (indexing), एब्स्ट्रैक्शन और समग्र संरचना का कठिन हिस्सा पर्याप्त संदर्भ वाले एजेंटों द्वारा काफी हद तक हल किया जा सकता है। आप 2-3 हफ्तों के काम को आसानी से 1-2 दिनों तक कम कर सकते हैं। नहीं, क्योंकि असली बाधा कर्नेल से हटकर सत्यापन पर आ गई है। अब काम करने का कोई एक तरीका नहीं है: आपका संदर्भ और हार्नेस (harness) जितना बेहतर होगा, प्रक्रिया उतनी ही तेज़ होगी। विशेष मामलों को और भी अधिक लाभ होगा, और एक अच्छा हार्नेस ही आपको चाहिए होगा।

अंत में, एजेंटों को पूरी तरह स्वायत्त (autonomous) मानने के बजाय, उन्हें वास्तव में स्मार्ट सहायकों के रूप में देखना जिन्हें आप मार्गदर्शन कर सकते हैं, अभी भी आवश्यक है। यहीं पर GPU और कर्नेल की आपकी बुनियादी समझ काम आती है। मानवीय हिस्सा (आप) अभी भी यहाँ आवश्यक है।

यह एक कड़वा-मीठा एहसास है, मुझे पता है :)