یہ مضمون CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs نامی نئی تحقیق کا تعارف کرتا ہے، جس کا مرکزی مقصد ٹرانس فارمر ماڈلز کی تربیت کی کارکردگی کو بہتر بنانا ہے، خاص طور پر ان معمولی لیکن جمع ہو کر زبردست وقت کھانے والے "میموری انٹنسیو" آپریشنز کو حل کرنا۔

مضمون کے مصنف، ذریعہ: ماشین کا دل

22 مئی کو، تری داؤ نے ہان گو کا ایک ٹویٹ سوشل میڈیا پر شیئر کیا۔ اس نے مزید لکھا: "کچھ ری رائٹنگ کے بعد، پتہ چلا کہ ٹرانسفارمر کا ہر حصہ GEMM + epilogue (میٹرکس ضرب + اختتام) کی ایک سیریز ہے۔ کچھ بہترین پریمیٹس دیے جانے پر، LLM (اور نئے) تمام ٹرانسفارمر آپریشنز کے لیے لائٹنگ سپیڈ کرنلز لکھ سکتے ہیں!"

Tri Dao FlashAttention سیریز کے مرکزی مصنفین میں سے ایک ہیں، اور یہ ٹویٹ ان کے اس دن جاری کیے گئے ایک پیپر CODA کی طرف اشارہ کرتی ہے۔

• عنوان مقالہ: CODA: ٹرانسفارمر بلاکس کو GEMM-Epilogue پروگرامس کے طور پر دوبارہ لکھنا
• کاغذ کا لنک: https://arxiv.org/abs/2605.19269
• کوڈ کا پتہ: https://github.com/HanGuo97/coda-kernels

یہ نام، "فن نامہ" جیسا لگتا ہے، اور "CUDA" جیسا بولتا ہے۔ MIT، پرنسٹن، Together AI اور Meta کے محققین، ٹرانسفارمر ٹریننگ میں وہ کم سے کم توجہ دی جانے والے، لیکن مستقل وقت کھا جانے والے "ٹکڑے ٹکڑے کمپوٹیشن" کو نظام مند انداز میں ختم کرنے کے لیے ایک نئی پروگرامنگ ایبسترکشن کا استعمال کر رہے ہیں۔

بڑے ماڈل کی تربیت کا "آرام کا ٹیک"

CODA کے حل کرنے والے مسائل کو سمجھنے کے لیے، سب سے پہلے یہ سمجھنا ضروری ہے کہ بڑے ماڈل تربیت کا وقت کہاں جاتا ہے۔

ایک نوویدیا H100 پر ایک 1B پیرامیٹر LLaMA-3 سٹائل کا ماڈل تربیت دینے پر، زیادہ تر لوگ یہ سمجھتے ہیں کہ وقت میٹرکس ضرب اور توجہ کے حسابات پر صرف ہوتا ہے، کیونکہ یہی "اصل کمپوٹیشن" ہیں۔ یہ سمجھنا تقریباً صحیح ہے: میٹرکس ضرب (GEMM) اور توجہ واقعی اکثر کمپوٹیشن کا بڑا حصہ لیتی ہیں۔

لیکن اگر آپ پرفارمنس پروفلر کھول کر تفصیل سے دیکھیں، تو آپ دیکھیں گے کہ کچھ "چھوٹے آپریٹرز" بھی خاموشی سے وقت کا استعمال کر رہے ہیں: نارملائزیشن (RMSNorm)، ایکٹیویشن فنکشنز (SwiGLU، RoPE)، ریزیڈوئل جمع، اور لیئرز کے درمیان ریڈکشن... ان کا الگ الگ کمپوٹیشن لوڈ چھوٹا ہے، لیکن وہ بار بار بڑے انٹرمیڈیٹ ٹینسرز کو گرافکس میموری میں داخل اور باہر لے جاتے ہیں۔

یہی «میموری بینڈ ویتھ بانٹلیک» کہلاتا ہے: جیسے ایک بہترین شیف جو ہر کھانا بنانے کے لیے سامان کو دور کے اسٹور سے لانے اور استعمال کے بعد واپس بھیجے، بجائے اس کے کہ سامان کو اپنے ہاتھ کے قریب ٹیبل پر رکھ دے۔ شیف کی ہاتھ کی رفتار کتنی ہی تیز کیوں نہ ہو، بھٹکانے کا وقت حقیقی ضائع ہے۔

مزید بر ایں، جبکہ نوڈیا کے FP8، FP4 جیسے کم درجہ کی درستگی والے فارمیٹس میٹرکس کی حساب کتاب کو تیز کر رہے ہیں، ان "نقل" آپریشنز کا نسبتی اخراج بڑھ رہا ہے: میٹرکس کی ضرب کو تیز کیا گیا ہے، لیکن ٹینسر کو اندر اور باہر لانے کا اخراج اسی تناسب سے کم نہیں ہوا۔

ایک ڈیٹا سیٹ تجرباتی طور پر واضح ہے: H100 پر TorchTitan کے استعمال سے 1B پیرامیٹرز کا ماڈل تربیت دی جانے پر، میٹرکس ضرب کے علاوہ کے آپریشنز کا ایک نمایاں حصہ مکمل عمل کے وقت پر مشتمل ہوتا ہے، اور FP8 درجہ بندی کے متعارف کرانے سے یہ تناسب مزید بڑھ جائے گا۔

موجودہ پروگرامنگ فریم ورک اس کے لیے تقریباً بے بس ہیں۔ PyTorch Transformer کے کمپیوٹیشن کو آپریٹرز کی ایک ترتیب کے طور پر ظاہر کرتا ہے، جن کے درمیان واضح حدود ہوتی ہیں۔ یہ حدود خودکار مشتق (autograd) کے لیے بہت دوستانہ ہیں، لیکن وہ بالکل آپریٹرز کے درمیان فیوشن آپٹیمائزیشن کو روک دیتی ہیں: ہر آپریٹر کی حد، اکثر ایک غیر ضروری میموری رائٹ بیک کا باعث بن جاتی ہے۔

CODA: "尾声" میں خزانہ چھپا ہوا ہے

کوڈا کا آغاز ایک سادہ مشاہدے سے ہوا۔

ایک ہائی پرفارمنس میٹرکس ضرب (GEMM) کرنل GPU پر دو حصوں میں تقسیم ہوتا ہے: مین لوپ جو مرکزی میٹرکس بلاک ضرب اور جمع کے حسابات کرتا ہے، اور ایپی لوج جو نتائج کو گرافکس میموری میں لکھنے سے پہلے، جیسے بائس جوڑنا، قسم تبدیل کرنا، اور سادہ سکیلنگ، کے لیے اختتامی کام کرتا ہے۔

اختتام کا مقصد یہ ہے کہ اس وقت میٹرکس ضرب کا نتیجہ ابھی تک چپ پر رجسٹر میں "زندہ" ہے، اور ابھی تک جنرل گلوبل میموری میں نہیں اتارا گیا۔ یہ ایک مختصر سنہری ونڈو ہے: اگر اس لمحے مزید کچھ کمپوٹیشن کی جائے، تو ایک بار میموری میں لکھنے اور دوبارہ پڑھنے کا سفر بالکل ختم ہو سکتا ہے۔

CODA کا مرکزی خیال یہ ہے کہ ٹرانسفارمر میں وہ میموری کے لیے زیادہ بھاری آپریشنز، جن میں سے بہت سے کو جبری طور پر الجبرائی طریقے سے دوبارہ پیرامیٹرائز کیا جا سکتا ہے اور اس «اختتامی» ونڈو میں انجام دیا جا سکتا ہے۔

اس کے لیے کچھ ریاضی کی مہارت کی ضرورت ہے۔ سب سے عام GEMM-RMSNorm-GEMM پیٹرن کو مثال کے طور پر لیں: ایک میٹرکس ضرب کا نتیجہ، جس میں ریزیڈوئل جمع، RMS نارملائزیشن، اور پھر دوسری میٹرکس ضرب شامل ہوتی ہے۔ روایتی طریقہ یہ ہے کہ تین الگ الگ آپریٹرز ترتیب سے انجام دیے جائیں، جس میں درمیانی نتائج دو بار گرافکس میموری میں محفوظ ہوتے ہیں۔

CODA ٹیم نے دریافت کیا کہ RMS نارملائزیشن میں رکھے گئے سطر کے سکیل فیکٹر r، چونکہ یہ ہر سطر کے لیے مشترکہ اسکیلر ہے، اس لیے یہ اگلے میٹرکس ضرب کے ساتھ تبادلہ کا قانون پورا کرتا ہے: r کا застосування "دوسرا GEMM" سے پہلے سے "دوسرا GEMM" کے اختتام تک ملتوی کیا جا سکتا ہے۔ ملتوی کرنے کے بعد، پہلے GEMM کے اختتام پر صرف مقامی "پارشل RMS" (partial RMS) کا حساب لگانا ضروری ہوتا ہے، جسے ایک بہت ہلکا مددگار ریڈکشن کرنل ملائے جاتا ہے، جبکہ مکمل RMSNorm کا حساب ختم ہو جاتا ہے۔

اسی طرح کی دوبارہ پیرامیٹرائزیشن SwiGLU، RoPE (روٹیشنل پوزیشنل اینکوڈنگ)، کراس اینٹروپی لاس جیسے آپریشنز کے لیے بھی لاگو ہوتی ہے، اور یہ بیک ورڈ پروپیگیشن کے لیے بھی درست ہے۔ مقالے میں ایک نظریہ ثابت کیا گیا ہے: جب تک فاروڈ پاس کا آخری مرحلہ "بلوک وائز لوکل" ہو، بیک ورڈ پروپیگیشن خود بخود اسی سٹرکچر کو ورثہ میں پاتی ہے۔ تفصیلات کے لیے اصل مقالہ دیکھیں۔

پانچ "بلوکس" اور ایک "لیگو زبان"

CODA ایک مخصوص فیوژن کرنل نہیں بلکہ ایک پروگرامنگ ایبسترکشن سیٹ ہے۔

یہ ماہرین کے ذریعہ بہتر بنائے گئے GEMM مرکزی لوپ کو فکس کرتا ہے، اور اختتام پر پانچ قسم کے قابل مجموعہ بنیادی پریمیٹو کو ظاہر کرتا ہے:

• عنصر وار تبدیلی (بقایا جمع، ایکٹیویشن فنکشن، RoPE)
• ویکٹر لوڈ اور اسٹور (بروڈکاسٹ RMSNorm وزن)
• میٹرکس کو بلاکس میں لوڈ اور اسٹور کریں (ریورس پروپیگیشن کے لیے درمیانی ایکٹیویشنز محفوظ کریں)
• بلاک ریڈکشن (لوکل آر ایم ایس، بلاک لوگ-سم-ایکسپ)
• حالتی تبدیلیاں (آن لائن نارملائزیشن کے لیے max اور sum-exp احصائیات)

ان پانچ قسم کے بلاکس کے ساتھ، توجہ کے علاوہ ایک معیاری ٹرانسفارمر کے فاروارڈ اور بیک ورڈ پروپیگیشن میں تقریباً تمام آپریشنز کو کور کیا جا سکتا ہے۔

زیادہ دلچسپ بات یہ ہے کہ یہ تصور "کون کوڈ لکھے گا" کے لیے کتنی زیادہ لچک رکھتا ہے۔ تحقیقی مقالے میں دو عملی طریقے کا جائزہ لیا گیا: ایک تو انسانی پروگرامرز کے ذریعہ لکھا گیا، دوسرا Claude Code کے ذریعہ تخلیق کیا گیا — CODA کے ابتدائی تعارف، کچھ مثالوں اور عملی ریکارڈ کے ساتھ، AI نے زیادہ تر کرنل کوڈ تیار کیا اور انسان نے ہلکی سی نگرانی کی۔

دونوں ماڈلز کی کارکردگی اعلیٰ سطح تک پہنچ گئی۔ ٹری داؤ نے ٹویٹ میں کہا کہ "LLM اور نئے صارفین ہی سپیڈی کرنل لکھ سکتے ہیں"، جو تجرباتی نتائج کا عملی طور پر اظہار ہے۔

تجربی نتائج

CODA کا بینچ مارک کے لیے مشکل مقابلہ ہے: cuBLAS کے ساتھ torch.compile، اور LLM کے لیے تیار کیے گئے Liger Kernel اور FlashInfer۔

ہر کور کے لیے دو عمل کی جانچ کی گئی: CODA (LLM) جسے Claude Code نے تخلیق کیا، جس میں محققین نے ابتدائی تفصیلات، کچھ مثالیں اور ایک مستقل اپڈیٹ ہونے والی عمل کے نکات کا ریکارڈ فراہم کیا، جس کے بعد AI نے بنیادی کوڈ تیار کیا اور انسانی نگرانی کم سطح پر کی گئی؛ CODA (Human) جسے انسانی پروگرامرز نے الگ سے لکھا، جس میں اسی بلند سطح کی ری پیرامیٹرائزیشن کا خیال استعمال کیا گیا لیکن CODA ابتدائی مجموعے پر انحصار نہیں کیا گیا۔ دونوں گروہوں کے نتائج کا موازنہ cuBLAS + torch.compile، Liger Kernel، اور FlashInfer جیسے بہترین لائبریریز کے ساتھ کیا گیا۔

ایک منفرد آپریٹر کے سطح پر، GEMM-RMSNorm-GEMM جیسے ایک مثالی پیٹرن کے ساتھ، CODA نے 1B، 7B، اور 70B مدل سائز کے لیے پوشیدہ ڈائیمنشن کے مطابق cuBLAS + PyTorch بیس لائن کو عبور کر لیا۔ SwiGLU، RoPE، اور کراس اینٹروپی جیسے آخری کمبینیشنز بھی اسی طرح کا اظہار کرتے ہیں۔

LLM کے ذریعہ تخلیق کردہ کرنل اکثر بنچ مارکس پر دستی طور پر لکھے گئے ورژن کے برابر ہیں، اور کچھ خاص کنفیگریشنز میں ان سے تھوڑا بہتر بھی ہیں۔ یہ ایک نایاب نتیجہ ہے، خاص طور پر GPU کرنل آپٹیمائزیشن جیسے قدیمی طور پر بہت مشکل شعبے میں۔

واپسی کے عمل میں فائدہ خاص طور پر نمایاں ہے: GEMM-Residual-PartialRMS-GEMM کا واپسی کرنے والا کرنل بنیادی اقدار کے مقابلے میں 1.6 سے 1.8 گنا تیز ہے، جبکہ SwiGLU کے واپسی کے عمل میں بھی تقریباً 1.4 سے 1.6 گنا بہتری آئی ہے۔ اس سمت میں، LLM اور انسانی طور پر تیار کردہ حل کے درمیان فرق بھی بہت کم ہے۔ یہ حیرت کی بات نہیں: واپسی کا عمل خود بخود زیادہ درمیانی ٹینسرز تک رسائی کا تقاضا کرتا ہے، جس کی وجہ سے ٹیل فیوژن کا فائدہ زیادہ ہوتا ہے؛ اور CODA کے ابتدائی ڈیزائن بہت واضح ہیں، جس سے AI ماڈل درست طریقے سے ترکیب کر پاتے ہیں۔

مکمل ٹرانسفارمر لیئر کے اینڈ تو اینڈ بنچ مارک میں، کوڈا کا فارورڈ ایکسلریشن مختلف سائزز پر 5% سے 20% کے درمیان ہے، اور بڑے ماڈل سائزز (70B سائز کے ہائڈن ڈائیمنشن کے مطابق) پر اس کا اثر زیادہ واضح ہوتا ہے۔

عددی درستگی کے لحاظ سے، CODA کی ری پیرامیٹرائزیشن RMSNorm اسکیلنگ فیکٹر کے застосування کے وقت کو تبدیل کرتی ہے، لیکن تجربات ظاہر کرتے ہیں کہ اس کی عددی خطائیں PyTorch ریفرنس ایمپلیمنٹیشن کے برابر ہیں، اور کچھ کنفیگریشنز میں یہ خطائیں کم تر بھی ہوتی ہیں — جس کا سبب GEMM مرکزی لوپ میں اعلی درجہ کی دوبارہ جمع کرنے والی یونٹس ہیں۔

CODA کیا کر سکتی ہے: بڑے منظر کی طرف بڑھنے سے پہلے، CODA کی صلاحیتوں کے دائرہ کار کو واضح کریں۔

• دائرہ کار: اسٹینڈرڈ ٹرانسفارمر (جیسے LLaMA آرکیٹیکچر) کے فاروڈ اور بیک ورڈ پروپیگیشن میں، توجہ اور لفظی ایمبیڈنگ کے علاوہ تقریباً تمام کمپوٹیشن، جس میں RMSNorm، ریزیڈوئل جمع، SwiGLU ایکٹیویشن، RoPE روتیشنل پوزیشنل اینکوڈنگ، کراس اینٹروپی لاس، اور ان آپریشنز کے بیک ورڈ گریڈینٹس شامل ہیں۔
• تیزی کا اثر: 1B سے 70B تک کے پوشیدہ ابعاد کے لیے، اکیلے آپریٹر کے سطح پر cuBLAS + torch.compile بنیادی سطح کے مقابلے میں مختلف درجات کی بہتری دکھائی دیتی ہے، جس میں ریٹروگریڈ پروپیگیشن کا فائدہ سب سے زیادہ واضح ہے (کچھ کرنلز 1.6 گنا سے زیادہ)؛ مکمل ٹرانسفارمر لیئر کا اینڈ تو اینڈ فارورڈ تیزی تقریباً 5% سے 20% ہے، جو بڑے ماڈل سائزز میں زیادہ نمایاں ہوتی ہے۔
• کون یہ کر سکتا ہے: CODA، جو CuTeDSL (NVIDIA CUTLASS کا Python DSL) پر مبنی ہے، انسانی پروگرامرز اور AI ماڈل دونوں کے لیے کرنل لکھنے کی سہولت فراہم کرتا ہے، اور دونوں طریقے اعلیٰ کارکردگی حاصل کرتے ہیں۔
• موجودہ پابندیاں: ابھی تک صرف ایک GPU کے سیناریو کو ہی سپورٹ کیا جاتا ہے، ڈسٹریبیوٹڈ ٹریننگ شامل نہیں؛ ریپیرامیٹرائزیشن بنیادی طور پر معیاری ٹرانسفارمر آرکیٹیکچر کے لیے ہے، دیگر آرکیٹیکچرز کے لیے اطلاق کی تصدیق باقی ہے۔

اختتام

CODA صرف ایک منفرد کام نہیں ہے۔ یہ ایک خیال کا عملی جامہ ہے: GPU پر، حقیقی بہتری کا خلیہ اکثر "کیا کیا جائے" میں نہیں، بلکہ "کیسے منتقل کیا جائے" میں ہوتا ہے۔

FlashAttention نے توجہ کے حسابات کو چپ پر میموری میں "رکھ دیا" ہے، جبکہ CODA نے نارملائزیشن اور ایکٹیویشن فنکشنز کو بھی "اندر" لانے کی کوشش کی ہے۔ Triton نے کسٹم کرنلز لکھنے کی رکاوٹ کم کردی ہے، اور ThunderKittens، TileLang جیسے ٹولز اس خلا کو مختلف سطحوں پر مزید تلاش کر رہے ہیں۔ یہ تمام کام ایک ہی سمت کی طرف اشارہ کرتے ہیں: PyTorch آپریٹر گراف کی آسانی اور ہاتھ سے لکھے گئے CUDA کے قریب عملی کارکردگی کو ایک قابل پروگرام فریم ورک میں اصل میں اکٹھا کرنا۔

Tri Dao کے ٹویٹ کا آخری جملہ دوبارہ سوچنے کے قابل ہے: "LLM اور نئے صارفین تمام ٹرانسفارمر آپریشنز کے لیے لائٹنگ سپیڈ کرنلز لکھ سکتے ہیں۔" اس کے پیچھے ایک گہرا منطق ہے: جب پروگرامنگ ایبسترکشن کافی اچھی طرح ڈیزائن کی جائے، تو AI ماڈل خود اپنی ٹریننگ انفراسٹرکچر کے بہترین بنانے میں شامل ہو سکتا ہے۔ یہ سائکل، CODA کا سب سے زیادہ دلچسپ پہلو ہے۔

اس منظر سے، "CODA" نام کا شاید کوئی اور گہرا مطلب ہے۔ کلاسیکی موسیقی میں، کوڈا موسیقی کے اختتامی حصہ کو کہتے ہیں جو پوری ترانہ کو ختم کرتا ہے۔ یہاں، یہ GEMM کور کا "اختتام" ہے — اور اس اختتام کو اچھی طرح لکھنا، شاید ٹرانسفارمر تربیتی نظام کی کارکردگی میں اضافے کا اگلا اہم مرحلہ ہو۔