Perplexity-এ CuTeDSL

Perplexity-এ, আমরা এমবেডিং থেকে শুরু করে ট্রিলিয়ন প্যারামিটারের এলএলএম (LLM)-এর মতো বিভিন্ন মডেল পরিবেশনের জন্য ইন-হাউস রানটাইম-অপ্টিমাইজড সার্ভিং ইঞ্জিন (ROSE) ব্যবহার করি। Sonar, Search এবং Embeddings সহ সমস্ত API-এর পিছনে কাজ করে, আমাদের রানটাইম গবেষকদের সবচেয়ে উপযুক্ত মডেল এবং আর্কিটেকচার অন্বেষণ করার ক্ষমতা দেওয়ার জন্য তৈরি করা হয়েছে, পাশাপাশি সেরা প্রার্থীদের দ্রুত উৎপাদনে স্থাপন করা যায়। এটি অর্জনের জন্য, আমাদের দ্রুত গতিতে সর্বাধুনিক আর্কিটেকচারের জন্য অনুমান (inference) বাস্তবায়ন, স্থাপন এবং অপ্টিমাইজ করতে সক্ষম হতে হবে। এই ব্লগ পোস্টে, আমরা অন্বেষণ করব কীভাবে আমাদের ইনফারেন্স ইঞ্জিনের সাথে সংহত CuTeDSL কার্নেলগুলি NVIDIA GPU-তে মডেলগুলিকে শীর্ষ পারফরম্যান্সে নিয়ে যাওয়ার জন্য প্রয়োজনীয় জিপিইউ কার্নেল তৈরি করতে আমাদের সহায়তা করে।

Perplexity-এর বিভিন্ন দল দ্বারা তৈরি কাস্টম মডেলগুলি, AI ইনফারেন্স টিম দ্বারা নির্মিত ROSE ইনফারেন্স ইঞ্জিন ব্যবহার করে NVIDIA Hopper এবং Blackwell GPU-তে ইন-হাউস হোস্ট করা হয়। NVIDIA Triton ইনফারেন্স সার্ভারের সাথে সামঞ্জস্যপূর্ণ একটি ইন্টারফেসের মাধ্যমে ভাষা ডিকোডিং এবং ক্লাসিফিকেশন উভয়ের জন্য কাস্টমাইজড Llama মডেলগুলিকে আরও ভালোভাবে পরিবেশন করার জন্য প্রাথমিকভাবে ROSE তৈরি করা হয়েছিল। সময়ের সাথে সাথে, এটি একটি পূর্ণাঙ্গ ইনফারেন্স ইঞ্জিনে পরিণত হয়েছে যা যথেষ্ট বড় LLM-এর পাশাপাশি Perplexity-এর পণ্যগুলিকে চালিত করার জন্য বিস্তৃত দলের জন্য র‍্যাঙ্কিং, ক্লাসিফিকেশন, স্কোরিং এবং অন্যান্য ট্রান্সফরমার-ভিত্তিক মডেল পরিবেশন করে। ROSE-এর ডিজাইনের লক্ষ্য সহজ হওয়া: সময়ের সাথে সাথে, শীর্ষ পারফরম্যান্সে পৌঁছানোর জন্য প্রচেষ্টা বিনিয়োগ করা হয়, তবে প্রধান অগ্রাধিকার হল আমাদের মডেলগুলির সর্বদা বিবর্তিত আর্কিটেকচার এবং ব্যবহারের ক্ষেত্রে খাপ খাইয়ে নেওয়া।

ডাটা টাইপ, হেড ডাইমেনশন, কম্পাইল-টাইম কনস্ট্যান্ট, এবং লঞ্চ সেটিংস জুড়ে কের্নেল কনফিগারেশন স্পেস।

উচ্চ স্তরে, ROSE একটি ইঞ্জিন নিয়ে গঠিত যা একটি মডেলকে এমন একটি ইন্টারফেসে রূপান্তরিত করে যার সাথে ক্লায়েন্টরা ইন্টারঅ্যাক্ট করতে পারে। ইঞ্জিনগুলি অত্যন্ত পুনঃব্যবহারযোগ্য, অনুরোধ শিডিউলিং এবং ব্যাচিং পরিচালনা করে এবং অন্তর্নিহিত মডেলগুলি থেকে ডিকোড করা টোকেন, এমবেডিং বা স্কোর পুনরুদ্ধার করে। ডিভাইসগুলির ইনিশিয়ালাইজেশন, প্রসেস-মধ্যস্থ এবং নোড-মধ্যস্থ যোগাযোগ চ্যানেল এবং ওয়েট লোডিং সবই এগুলির দ্বারা পরিচালিত হয়। LLM-গুলির জন্য, ইঞ্জিনটি ব্যাচিং, চাংকিং, স্যাম্পলিংয়ের জন্য দায়ী। এটি সম্পূর্ণ এবং লিনিয়ার অ্যাটেনশন উভয়ের জন্য, সেইসাথে প্রিফিক্স ম্যাচিংয়ের জন্য কেভি (KV) স্টোরেজ বরাদ্দের জন্যও দায়ী। এমবেডিং ইঞ্জিনগুলি সহজতর, কারণ তারা শুধুমাত্র অন-লাইন ব্যাচিং করে। এক্সিকিউশন এমনভাবে অপ্টিমাইজ করা হয় যাতে মডেলটি একিলারেটরে সমান্তরালভাবে চলার সময় ইঞ্জিনটি কাজের পরবর্তী রাউন্ড প্রস্তুত করে।

ROSE মডেলিংয়ের জন্য কাস্টম লেয়ারের একটি সমৃদ্ধ সংগ্রহ উন্মুক্ত করে, যা CuTeDSL, Triton, CUDA, CUTLASS এবং cuBLAS দিয়ে তৈরি জিপিইউ কার্নেলগুলিকে ব়্যাপ করে। ম্যাট্রিক্স মাল্টিপ্লিকেশন এবং অ্যাটেনশন ইমপ্লিমেন্টেশনের জন্য, আমরা মূলত CUTLASS এবং cuBLAS-এর NVIDIA কার্নেলগুলির উপর নির্ভর করি। এমবেডিং, নরম, MoE এবং অ্যাক্টিভেশন কার্নেল সহ পাইপলাইনের বাকি অংশটি ROSE-এর কার্নেল লাইব্রেরির অংশ। আমরা যে হার্ডওয়্যারে ইনফারেন্স চালাই তাকে সর্বোত্তমভাবে সমর্থন করার জন্য, ROSE MoE রাউটিং, ডিসপাচ এবং কম্বাইনের জন্য নিজস্ব ইন্টার-ডিভাইস কমিউনিকেশন প্রিমিটিভ দিয়ে সজ্জিত। অন্যান্য কার্নেলগুলি তুলনামূলকভাবে সহজ হলেও, সেগুলিতে চালিত সমস্ত বিভিন্ন হিডেন ডাইমেনশন এবং ডিভাইসে শীর্ষ পারফরম্যান্স অর্জনের জন্য সেগুলিকে পুঙ্খানুপুঙ্খভাবে বিশেষায়িত করতে হবে। অতিরিক্তভাবে, আমরা বিভিন্ন মডেলের সমস্ত নকব এবং টুইকগুলিকে (RMS নর্মের জন্য ওয়েট বায়াস বনাম কোনো ওয়েট বায়াস নেই, অ্যাক্টিভেশন ফাংশন ইত্যাদি) মানিয়ে নিতে কম্পাইল-টাইম স্পেশালাইজেশনের উপর নির্ভর করি। ফলস্বরূপ, আমরা সেগুলি বাস্তবায়নের জন্য DSL ব্যবহার করার সিদ্ধান্ত নিয়েছি: আমরা প্রাথমিকভাবে Triton ব্যবহার করে ROSE তৈরি করেছি। তবে, আমরা তখন থেকেই জিপিইউ প্রোগ্রামিংয়ের প্রাথমিক পরিবেশ হিসেবে CuTeDSL গ্রহণ করেছি।

গাণিতিক দৃষ্টিকোণ থেকে, ইনফারেন্স কার্নেলগুলি সাধারণত জটিল নয়। ফলস্বরূপ, একটি কার্নেল তৈরি করার সময়, জটিলতা সাধারণত কোন গণনাটি করতে হবে তা প্রকাশের মধ্যে নিহিত থাকে না, বরং হার্ডওয়্যারটি কীভাবে এটি ঠিকভাবে করবে তার মধ্যে থাকে। সেরা কার্নেলগুলি তৈরি করার সঠিক ভাষাটির প্রয়োজন হয় যা অন্তর্নিহিত হার্ডওয়্যারের উপর সঠিক অ্যাবস্ট্রাকশন অফার করে। সঠিক ভাষাটি আমাদের গণনা সম্পাদন করার জন্য উচ্চ-স্তরের অপারেশনগুলি নির্দিষ্ট করতে সক্ষম করে, একটি বিস্তৃত কনফিগারেশন স্পেস জুড়ে সমস্ত বিন্দুতে সেগুলিকে আক্রমণাত্মকভাবে বিশেষায়িত এবং অপ্টিমাইজ করে।

CuTeDSL হল CuTe লেআউট অ্যালজেব্রা এবং MLIR-এর উপর নির্মিত NVIDIA জিপিইউ-এর জন্য একটি পাইথন-ভিত্তিক DSL। এটি অত্যন্ত অপ্টিমাইজড PTX-এ জাস্ট-ইন-টাইমে কম্পাইল করে, এবং একই সাথে লো-লেভেল হার্ডওয়্যার প্রিমিটিভগুলির উপর নিয়ন্ত্রণ বজায় রাখে। এর লক্ষ্য হল ব্যয়বহুল কম্পাইলেশন সময় ছাড়াই CUTLASS C++ এর মতো একই বৈশিষ্ট্য এবং পারফরম্যান্স অফার করা যা ডেভেলপমেন্টের গতিকে বাধাগ্রস্ত করে। আমরা যে ভাষাগুলি জরিপ করেছি বা কাজ করেছি (CUDA, cuTile, TVM IR এবং Triton সহ), তার মধ্যে আমরা এটিকে আমাদের কার্নেলগুলির বিকাশের জন্য সবচেয়ে উপযুক্ত বলে মনে করি।

সামগ্রিকভাবে, Triton বা cuTile-এর মতো উচ্চ-স্তরের এবং সংক্ষিপ্ত না হলেও, আমরা দেখেছি যে CuTe লেআউট অ্যালজেব্রা এবং এর কম্পোজিশনগুলি (cute.Tensor-এর মাধ্যমে উন্মুক্ত) যুক্তিসঙ্গত বিকল্প ছিল। অতিরিক্তভাবে, CuTeDSL হার্ডওয়্যার প্রিমিটিভগুলিতে অনিয়ন্ত্রিত অ্যাক্সেস অফার করে, যা আমাদের সমস্ত এজ কেসে শীর্ষ পারফরম্যান্স অর্জনের জন্য প্রয়োজনীয় নিয়ন্ত্রণের মাত্রা প্রদান করে। CUDA-র সাথে, কম্পাইল-টাইম স্পেশালাইজেশনের সুযোগ সীমিত ছিল, কারণ প্রতিটি অতিরিক্ত প্যারামিটার কম্পাইলেশন সময়ে একটি অগ্রহণযোগ্য খরচ বা বিল্ড সিস্টেমে জটিলতা যোগ করে। বিপরীতে, একটি জাস্ট-ইন-টাইম কম্পাইলড DSL-এর সাহায্যে, আমরা ইটারেশনের গতি না কমিয়েই আরও আক্রমণাত্মকভাবে বিশেষায়িত করতে পারি। TVM IR-এর বিপরীতে, যা CUDA কে একটি মধ্যবর্তী ধাপ হিসেবে ব্যবহার করেছিল, আমরা PTX-এ নামানো MLIR-ভিত্তিক পাইপলাইনটিকে CUDA-কে একটি মধ্যবর্তী ধাপ হিসেবে নির্গমনকারী পাইপলাইনের চেয়ে বেশি শক্তিশালী এবং ডিবাগ করা সহজ বলে মনে করেছি। অতিরিক্তভাবে, PTX নির্দেশাবলী ব়্যাপ করার জন্য কাস্টম অপ্স ব্যবহারের মাধ্যমে, আমরা দেখেছি যে আমরা অন্তর্নিহিত হার্ডওয়্যারের সমস্ত বৈশিষ্ট্য সম্পূর্ণরূপে কাজে লাগাতে পারি।

CUDA কার্নেলগুলির কম্পাইলেশন খরচের বেশিরভাগ অংশ সাধারণত কার্নেলগুলির টেমপ্লেট এক্সপানশন ধাপে ব্যয় করা হয়, কারণ PTX পর্যায় থেকে শুরু করে অপ্টিমাইজেশান এবং কম্পাইলেশন কয়েক সেকেন্ডের মধ্যে সম্পন্ন করা যেতে পারে। অনেক কম্পাইল-টাইম প্যারামিটার সহ, একটি অ-তুচ্ছ কার্নেলের জন্য কম্পাইলেশনের এই পর্যায়টি কয়েক মিনিটে বিস্ফোরিত হতে পারে। এই বাধাটি এই সত্য দ্বারা আরও প্রকট হয় যে সমান্তরাল কম্পাইলেশনের সুযোগগুলি কেবল পরে, IR এবং PTX স্তরে উদ্ভূত হয়, যেখানে C++ টেমপ্লেট হ্যান্ডলিং ক্রমানুসারে ঘটে। টেমপ্লেট ইনস্টান্টিয়েশনগুলিকে পৃথক কম্পাইলেশন ইউনিটে সরিয়ে নেওয়ার মাধ্যমে এটি মোকাবেলা করার কৌশল বিদ্যমান থাকলেও, সেগুলির জন্য বিল্ড সিস্টেমে উল্লেখযোগ্য পরিবর্তন এবং অবাস্তব ম্যাক্রো ব্যবহারের প্রয়োজন হয়। বিপরীতে, CuTeDSL কার্নেলগুলি যেগুলি জাস্ট-ইন-টাইমে (JIT) কম্পাইল করা হয় সেগুলির জটিল বিল্ড নিয়মের প্রয়োজন হয় না এবং কম্পাইলেশন সময়কে বাধাগ্রস্ত করে না, কারণ ডেভেলপমেন্ট সাইকেলটি কনফিগারেশন স্পেসের একটি মাত্র বিন্দুতে ফোকাস করা যেতে পারে, এবং পরবর্তী যত্নশীল ইউনিট পরীক্ষাগুলি কার্নেলের পুরো ডিজাইন স্পেস কভার করতে পারে।

যদিও JIT কম্পাইলেশন ডেভেলপমেন্ট-টাইম থেকে রানটাইমে খরচ সরায়, তবে ওভারহেড গ্রহণযোগ্যর চেয়েও বেশি। যখন একটি মডেল স্থাপন করা হয়, তখন এটি পুরো কনফিগারেশন স্পেস থেকে মাত্র কয়েকটি বিন্দু সক্রিয় করে, যা ইনফারেন্স ইঞ্জিনের স্টার্টআপ/ওয়ার্মআপের সময় সেকেন্ড-স্কেল ওভারহেড ঘটায়। বিপরীতে, একটি কার্নেলের সমস্ত কনফিগারেশন প্রি-কম্পাইল করলে তা কয়েক মিনিট পর্যন্ত সময় নিতে পারে। অতিরিক্তভাবে, JIT-কম্পাইল করা আর্টিফ্যাক্ট ক্যাশ করে ওভারহেড আরও কমানো যেতে পারে।

CuTe এবং CUTLASS-এর ক্ষেত্রে যেমন হয়, টেমপ্লেট-ভিত্তিক ডিসপাচ এবং স্পেশালাইজেশন ব্যবহার করা হলে ত্রুটির উৎস ট্রেস করার অসুবিধার কারণে CUDA-র ভারী কম্পাইলেশন খরচ আরও বৃদ্ধি পায়। একটি DSL-এর সাহায্যে, ত্রুটিগুলি যে কোডের লাইনটি তৈরি করেছে তাতে নির্দেশ করা যেতে পারে, যা ডিবাগিংকে উল্লেখযোগ্যভাবে সহজ করে তোলে। অতিরিক্তভাবে, মধ্যবর্তী MLIR ধাপগুলির মাধ্যমে বা জেনারেলিটেড PTX-এর সাথে সংযুক্ত ডিবাগিং তথ্যের উপর নির্ভর করে জেনারেট করা কোডটিকে যে DSL স্টেটমেন্ট থেকে এটি উৎপত্তি হয়েছে তার সাথে ম্যাপ করা অনেক সহজ। DSL-এর এই দিকগুলির জন্য ধন্যবাদ, আমরা কার্নেল ডেভেলপ করার সময় নিজেদের অনেক বেশি উৎপাদনশীল বলে মনে করেছি।

কার্নেলগুলি মডেল জুড়ে অনেক মিল শেয়ার করলেও, সেগুলিকে এখনও কনফিগারেশন নকবগুলির একটি পরিসর দিয়ে সজ্জিত করা আবশ্যক যা পারফরম্যান্সের উপর প্রভাব ফেলতে পারে। প্রাসঙ্গিক প্যারামিটারগুলির মধ্যে রয়েছে:

হিডেন ডাইমেনশন/হেড ডাইমেনশন/অ্যাটেনশন হেডগুলির সংখ্যা

ইনপুট/আউটপুট ডেটা টাইপ

কোয়ান্টাইজেশন স্কিম

কম্পাইল-টাইম ধ্রুবক (বায়াস, অ্যাক্টিভেশন ফাংশন ইত্যাদি)

ROSE আর্কিটেকচার, এর ইঞ্জিন, মডেল, কের্নেল লাইব্রেরি, এবং সার্ভিসসহ।

কম্পাইল-টাইম স্পেশালাইজেশন পারফরম্যান্স এবং ডিবাগিং উভয় ক্ষেত্রেই দরকারী। বিশেষ করে একটি পাইপলাইনড প্রসেসরে, যেমন জিপিইউ-এর ক্ষেত্রে হয়, একটি সর্বদা-গৃহীত বা কখনই-গৃহীত-নয় এমন ব্রাঞ্চের চেয়ে একটি অস্তিত্বহীন ব্রাঞ্চ ভালো। অতিরিক্তভাবে, নির্দিষ্ট সেট কনফিগারেশন নকবগুলির জন্য একটি কার্নেল কম্পাইল করার ফলে কম PTX/SASS নির্দেশাবলী পাওয়া যায়, যা তদন্তের সময় cuda-gdb ব্যবহার করার সময় বিশ্লেষণ এবং ডিবাগ করা সহজ।

একটি প্রদত্ত ব্যাচ সাইজ এবং কনফিগারেশনের জন্য, DRAM-এ প্রত্যাশিত ট্রাফিক এবং পাটিগণিত অপারেশনের সংখ্যা অনুমান করা যেতে পারে। পরবর্তীতে, সর্বোত্তম গ্রিড বরাদ্দ, সিঙ্ক্রোনাইজেশন এবং নির্দিষ্ট হার্ডওয়্যার বৈশিষ্ট্যের ব্যবহারের জন্য কেস-বাই-কেস ভিত্তিতে সর্বোত্তম সিদ্ধান্ত নেওয়া যেতে পারে। CuTeDSL আমাদের অগ্রহণযোগ্য কম্পাইল-টাইম ওভারহেড ছাড়াই গভীর স্পেশালাইজেশনের জন্য উপাদানগুলি সংজ্ঞায়িত করতে এবং পুনঃব্যবহার করতে দেয়।

সাধারণত, দুটি পৃথক ব্যবহারের ক্ষেত্রে ইনফারেন্স কার্নেলের পারফরম্যান্স অপ্টিমাইজ করা আবশ্যক:

প্রিফিল: ছোট ব্যাচ সাইজ (কয়েকটি চাংক, সিকোয়েন্স), প্রতি সিকোয়েন্সে অনেক টোকেন। এমবেডিং এবং LLM-এর ইনপুট টোকেন উভয় জন্যই থ্রুপুট সর্বাধিক করতে হবে। দশ থেকে শত মিলিসেকেন্ডের মধ্যে লেটেন্সি।

ডিকোড: ছোট থেকে বড় ব্যাচ সাইজ, এক (কোনো speculative decoding নয়) থেকে প্রতি সিকোয়েন্সে কয়েকটি টোকেন পর্যন্ত (draft-target speculation, multi-token prediction)। লেটেন্সি kernel launch latency-এর কাছাকাছি।

GEMM এবং অ্যাটেনশনের জন্য, প্রিফিল পর্যায়টি বিপুল সংখ্যক টোকেনে চলে এবং পারফরম্যান্স ডিভাইসের কম্পিউট থ্রুপুট দ্বারা আবদ্ধ থাকে। ডিকোডিংয়ে ছোট ব্যাচ সাইজে একই অপারেশন জড়িত থাকে, যেখানে এক্সিকিউশনের সময় ডিভাইসের মেমরি ব্যান্ডউইথের দ্বারা সীমাবদ্ধ থেকে, স্ট্রিমিং মাল্টিপ্রসেসর এবং টেনসর কোরে DRAM থেকে অ্যাক্টিভেশন এবং ওয়েট লোড করার খরচ দ্বারা প্রভাবিত হয়। ব্যাচ সাইজ বা ব্যাচে টোকেনের সামগ্রিক সংখ্যা তুলনামূলকভাবে স্পষ্ট নির্দেশক হওয়া সত্ত্বেও, প্রিফিল এবং ডিকোডের মধ্যে সুইচওভার থ্রেশহোল্ড আরও সূক্ষ্ম কারণ মেমরি ব্যান্ডউইথ মডেলের হিডেন ডাইমেনশন এবং অ্যাক্টিভেশনের ডেটা টাইপ দ্বারাও নির্ধারিত হয়। অ্যালগরিদমিক পার্থক্য বা গ্রিড/ব্লক অ্যাসাইনমেন্টগুলির মধ্যে যেকোনো একটির সুবিধা হলে কার্নেলগুলিকে ফলস্বরূপ প্রিফিল এবং ডিকোড বাস্তবায়নে বিভক্ত করা হয়।

যে প্যারামিটারটি সর্বজনীনভাবে প্রিফিল এবং ডিকোড কার্নেলের মধ্যে আলাদা তা হল একটি ব্লকের ওয়ার্পের সংখ্যা (থ্রেডের সংখ্যা, প্রতি ওয়ার্পে 32 থ্রেডের ওয়ার্প সাইজ দ্বারা ভাগ করা)। এই প্যারামিটারটিকে একটি কম্পাইল-টাইম ধ্রুবক করা বাঞ্ছনীয়, কারণ এটি রেজিস্টার বরাদ্দে সহায়তা করতে পারে। CUDA-তে, এটি সাধারণত __launch_bounds__-এর মাধ্যমে থ্রেডের সংখ্যা নির্দিষ্ট করে এবং একটি টেমপ্লেট প্যারামিটার হিসাবে ওয়ার্পের সংখ্যা প্রদান করে করা হয়:

cpp

template <size_t NUM_WARPS>
__global__ void 
__launch_bounds__(NUM_WARPS * WARP_SIZE, 1)
some_kernel(...) {

}

একটি ব্লকের সমস্ত থ্রেডের জন্য উপলব্ধ রেজিস্টারের সংখ্যা একটি হার্ডওয়্যার সীমাবদ্ধতা, Hopper এবং Blackwell-এ 64K। যদি থ্রেডের সংখ্যা আগে থেকেই জানা থাকে, তবে কম্পাইলার ঠিক জানে যে এটি প্রতিটি থ্রেডে কতগুলি রেজিস্টার বরাদ্দ করতে পারে এবং আরও আক্রমণাত্মকভাবে অপ্টিমাইজ করা কার্নেল কম্পাইল করতে সফল হতে পারে। তবে, এটি কম্পাইলেশন-সময়ের খরচে আসে, কারণ এখন বিভিন্ন ওয়ার্প সাইজের জন্য কার্নেলটি পুনরায় কম্পাইল করতে হবে। CuTeDSL-এর সাথে, এটি আর কোনো সমস্যা নয়, কারণ পুরো কনফিগারেশন স্পেসের জন্য আমাদের কম্পাইলেশনের খরচ আগেভাগে দিতে হবে না। DSL-টি ব্লক সাইজ একটি ধ্রুবক এক্সপ্রেশন হবে বলেও আশা করে, যা অন্তর্নিহিতভাবে CUDA-তে __latch_bounds__ দ্বারা সূচিত অপ্টিমাইজেশানটি সম্পাদন করে।

ওয়ার্পের সংখ্যা ধ্রুবক করার পাশাপাশি, প্রিফিল এবং ডিকোডের মধ্যে এটি পরিবর্তন করা বাঞ্ছনীয়। একটি জিপিইউ-এর মেমরি ব্যান্ডউইথ সম্পূর্ণরূপে ব্যবহার করার জন্য, প্রতি ব্লকে সর্বদা উচ্চ সংখ্যক ওয়ার্প বরাদ্দ করা যথেষ্ট নয়, কারণ উপলব্ধ স্ট্রিমিং মাল্টিপ্রসেসরগুলির (SM) যতটা সম্ভব ব্যবহার করেই কেবল সম্পূর্ণ ব্যান্ডউইথ অর্জন করা যেতে পারে। ফলস্বরূপ, একটি নির্দিষ্ট সংখ্যক টোকেন এবং পরিচিত ডেটা টাইপ এবং হিডেন ডাইমেনশনের জন্য, জেনারেট করা কোডটি 128-বিট লোড/স্টোর নির্দেশাবলী (Blackwell-এ 256-বিট) ব্যবহার করে তা নিশ্চিত করতে আমরা প্রথমে একটি আদর্শ প্রতি-থ্রেড ভেক্টর ডাইমেনশন বেছে নিই (bfloat16-এর জন্য 8, float32-এর জন্য 4)। ভেক্টরের আকার জানার পরে, আমরা কার্নেলের দ্বারা প্রয়োজনীয় ওয়ার্পের সংখ্যা গণনা করতে পারি। আমরা প্রথমে SM-এর সংখ্যা সর্বাধিক করার চেষ্টা করি: যদি SM-এর চেয়ে কম ওয়ার্প থাকে, তবে প্রতি ব্লকে একটি একক ওয়ার্প বেছে নেওয়া হয়। এটি সাধারণত কার্নেলের ডিকোড ভেরিয়েন্ট হয়। যদি পর্যাপ্ত ওয়ার্প থাকে, তবে আমরা ক্রমান্বয়ে ওয়ার্পের সংখ্যা 8, 16 বা 32-এ উন্নীত করি, যার পরে আমরা সাধারণত পারফরম্যান্সে আর কোনো উন্নতি দেখতে পাই না এমন একটি আপার বাউন্ড বেঞ্চমার্ক করি।

Triton হুডের নিচে এই অপ্টিমাইজেশানগুলির কিছু করলেও, আমরা দেখেছি যে একটি উচ্চ-স্তরের কম্পাইলারের ত্রুটি দ্বারা সীমাবদ্ধ না হয়ে, CuTe লেআউট অ্যালজেব্রা সাবধানে ব্যবহার করে আমরা অনুরূপ ফলাফল অর্জন করতে পারি। বিশেষ করে যখন সাব-ওয়ার্প হরিজন্টাল রিডাকশন জড়িত থাকে, আমরা দেখেছি যে CuTeDSL আমাদের সর্বোত্তম কোড তৈরি করার জন্য প্রয়োজনীয় সূক্ষ্ম-দানাযুক্ত (fine-grained) নিয়ন্ত্রণ প্রদান করে।

কিছু পরিস্থিতিতে, যতটা সম্ভব ব্লকে কাজ বিতরণ করতে, সমস্ত অংশগ্রহণকারী ব্লকের সিঙ্ক্রোনাইজ করার জন্য cooperative_groups::this_grid().sync() ব্যবহার করা যেতে পারে। নিষেধাজ্ঞা প্রযোজ্য: যেহেতু একটি ব্লক একবার এক্সিকিউশন শুরু হলে SM থেকে ডি-শিডিউল করা যায় না, তাই ফরোয়ার্ড অগ্রগতি নিশ্চিত করার জন্য, ব্লকের সংখ্যা ডিভাইসে উপলব্ধ SM-এর সংখ্যা অতিক্রম করতে পারে না, যাতে যেকোনো বিন্দুতে সমস্ত ব্লক সিঙ্ক্রোনাইজেশনের সাথে জড়িত গ্লোবাল মেমরি অবস্থানগুলি আপডেট এবং পোল করতে পারে তা নিশ্চিত করা যায়। এই নিষেধাজ্ঞাটি cudaLaunchCooperativeKernel দ্বারা প্রয়োগ করা হয়। Hopper-এ একটি গ্রিড ব্যারিয়ারের খরচ প্রায় 3us, এছাড়া কিছু পরিস্থিতিতে একটি অন্তর্নিহিত মেমরি ব্যারিয়ারের মাধ্যমে সমস্ত পূর্ববর্তী মেমরি অপারেশন ফ্লাশ করার খরচ রয়েছে। সামগ্রিকভাবে, এটি একটি কার্নেলকে দুটি লঞ্চে বিভক্ত করার ওভারহেডের সাথে তুলনীয়, যা একটি পারফরম্যান্স-সংবেদনশীল ডিজাইন সিদ্ধান্তের দিকে পরিচালিত করে। যদিও প্রোগ্রামমেটিক ডিপেন্ডেন্ট লঞ্চ (PDL) Hopper থেকে লঞ্চ ওভারহেড কমাতে পারে, ডিকোড কার্নেলের জন্য, ওভারহেড এড়ানো সার্থক।

ওয়ার্কলোডগুলির মধ্যে কোনো একটির যদি SM-এর সংখ্যার চেয়ে বেশি ব্লকের প্রয়োজন হয়, তবে কার্নেলটিকে বিভক্ত করা উচিত, যা কার্নেল লঞ্চের অন্তর্নিহিত ব্যারিয়ারের উপর নির্ভর করে। এটি মূলত প্রিফিলের ক্ষেত্রে প্রযোজ্য। সমস্ত ওয়ার্কলোডের জন্য যে ব্লকের সংখ্যা আদর্শ তা যদি গ্রিডে ফিট করে, তবে একটি গ্রিড ব্যারিয়ার পছন্দ করা হয়, যেমনটি ছোট সংখ্যক টোকেনে কাজ করা ডিকোড কার্নেলগুলির ক্ষেত্রে হয়।

CuTeDSL-এ, এই সিদ্ধান্তগুলি হোস্ট-সাইড লঞ্চ কোডে নেওয়া যেতে পারে এবং নীচের উদাহরণ দ্বারা চিত্রিত হিসাবে প্রয়োজন অনুসারে গ্রিড ব্যারিয়ারগুলির সাথে কার্নেলগুলি কম্পোজ এবং ইন্টার্লিভ করা যেতে পারে। এই ধরনের একটি হাইব্রিড কৌশলের সাথে, ন্যূনতম লেটেন্সিতে পৌঁছানোর জন্য ডিকোড কার্নেলগুলি অপ্টিমাইজ করা যেতে পারে এবং সর্বোচ্চ থ্রুপুতের জন্য প্রিফিলগুলি অপ্টিমাইজ করা যেতে পারে। CUDA-ভিত্তিক সমাধানগুলির বিপরীতে, এই লজিকটি আরও সংক্ষিপ্তভাবে এবং কম কম্পাইলেশন-টাইম ওভারহেডের সাথে সংজ্ঞায়িত করা যেতে পারে।

python

@cute.kernel
def dispatch_kernel_decode(...):
	route_tokens(...)
	grid_barrier(grid_barrier_ptr)
	copy_tokens_to_experts(...)

@cute.kernel
def dispatch_kernel_prefill_route(...):
	route_tokens(...)
	
@cute.kernel
def dispatch_kernel_prefill_copy(...):
	copy_tokens_to_experts(...)
	
@cute.jit
def dispatch(...):
	num_route_blocks = ...
	num_copy_blocks = ...
	
	if num_copy_blocks < NUM_SMS:
		dispatch_kernel_decode(...).launch(
			grid=(min(max(num_route_blocks, num_copy_blocks), NUM_SMS), 1, 1)
		)
	else:
		dispach_kernel_prefill_route(...).launch(
			grid=(num_route_blocks, 1, 1)
		)
		dispatch_kernel_prefill_copy(...).launch(
			grid=(num_copy_blocks, 1, 1)
		)

আমাদের MoE all-to-all কার্নেলের ক্ষেত্রে, নীচে বিশদভাবে বর্ণিত হিসাবে, গ্রিড ব্যারিয়ার সহ বা ছাড়া কার্নেল তৈরি করার এবং কম্পাইল করার সময় সেগুলিকে অপ্টিমাইজ করার নমনীয়তা থাকা আমাদের CuTeDSL-এর মাধ্যমে উল্লেখযোগ্য পারফরম্যান্স লাভ করতে সাহায্য করেছে।

আমরা CuTeDSL-এ যে প্রথম কার্নেলটি তৈরি করেছি তা ছিল একটি বিশেষায়িত RMS নর্ম কার্নেল, যা ছোট ডাইমেনশনের জন্য অপ্টিমাইজ করা হয়েছে। সাম্প্রতিক Qwen মডেলগুলির সেলফ-অ্যাটেনশন স্তরে ব্যবহৃত QK নর্ম, প্রতি-হেড ভিত্তিতে, 64, 128 বা 256 উপাদানে প্রয়োগ করা হয়। হাজার হাজার ছাড়িয়ে যাওয়া হিডেন ডাইমেনশনে কাজ করে এমন নিয়মিত RMS নর্মের বিপরীতে, যার জন্য শেয়ার্ড মেমরি এবং রেজিস্টার উভয় মাধ্যমেই একটি হরিজন্টাল রিডাকশনের প্রয়োজন হয়, QK নর্মের শুধুমাত্র ওয়ার্প জুড়ে রিডাকশনের প্রয়োজন হয়, যেখানে একটি একক ওয়ার্পে একাধিক হেড স্কুইজ করার সম্ভাবনা থাকে। এটি এমন একটি ব্যবহারের কেস হয়ে ওঠে যেখানে আমরা সমস্ত হেড ডাইমেনশন এবং ডেটা টাইপের জন্য সেরা টোকেন-টু-ওয়ার্প অ্যাসাইনমেন্ট খুঁজে পেতে কম্পাইল করার সময় আক্রমণাত্মকভাবে অপ্টিমাইজ করতে পারি।

RMS নর্ম একটি বেশ সোজা অপারেশন, যদিও একটি ভেক্টরের স্কোয়ার্ড নর্ম গণনা করার জন্য এতে একটি হরিজন্টাল সামের প্রয়োজন হয়:

cpp

r = w * (x * ((x * x).sum(dim=-1) / x.shape[-1] + EPS))

যেহেতু অ্যাটেনশন হেডগুলি সাধারণত ছোট হয় এবং প্রতি-থ্রেড অপারেশনগুলি ভারীভাবে ভেক্টরাইজ করা হয়, তাই একটি একক ওয়ার্প স্কোয়ার্ড নর্ম গণনা করতে পারে। মেমরি ব্যান্ডউইথ সর্বাধিক করার জন্য, প্রতিটি থ্রেডের 128 বিট ইনপুট প্রসেস করা উচিত, যা 8টি bfloat16 মানের সমতুল্য। ফলস্বরূপ, একটি হেড প্রসেস করার জন্য একটি হাফ-ওয়ার্প যথেষ্ট, এবং shuffle_sync_bfly দ্বারা সম্পাদিত একটি হরিজন্টাল রিডাকশন ব্যবহার করে স্কোয়ার্ড নর্ম নির্ধারণ করা যেতে পারে। হেড ডাইমেনশন এবং ডেটা টাইপ থেকে শুরু করে, আমরা প্রতি ওয়ার্পে হেডের সংখ্যা নির্ধারণ করতে পারি এবং কম্পাইল করার সময় আর্গুমেন্ট এবং রিডাকশন রাউন্ডের সংখ্যা নির্দিষ্ট করতে পারি।

আমরা তখন তুলনামূলকভাবে কম লাইনের কোড সহ একটি বিশেষায়িত কার্নেল তৈরি করতে পারি, যা স্বয়ংক্রিয়ভাবে ভেক্টরাইজড লোড/স্টোরের উপর নির্ভর করে এবং ওয়ার্প রিডাকশনের জন্য লোয়ার-লেভেল আর্কিটেকচার-নির্দিষ্ট অপারেশন ব্যবহার করে। কার্নেলটি নির্ধারণ করে কোন হেডে (Q বা K) কাজ করতে হবে এবং প্রয়োগ করার জন্য সংশ্লিষ্ট ওয়েটগুলি চিহ্নিত করে। এটি হেড থেকে একটি সম্পূর্ণ ভেক্টর লোড করে এবং প্রথমে স্থানীয়ভাবে যোগ করে। পরবর্তীতে, প্রতি হেডের থ্রেডের সংখ্যার উপর ভিত্তি করে, স্কোয়ার্ড নর্ম গণনা করা হয়। পরিশেষে, স্বয়ংক্রিয়ভাবে ভেক্টরাইজড হওয়া RMS নর্ম নির্ধারিত হয় এবং মেমরিতে ব্যাক রাইট করা হয়। মডেল-নির্দিষ্ট প্যারামিটার যেমন WEIGHT_BIAS এবং EPS, কনস্ট্যান্ট-ফোল্ড করা হয়।

python

@cute.kernel
def _qk_norm_kernel(
    q: cute.Tensor,
    k: cute.Tensor,
    q_weight: cute.Tensor,
    k_weight: cute.Tensor,
    WEIGHT_BIAS: cutlass.Constexpr[float],
    EPS: cutlass.Constexpr[float],
) -> None:
    # Find the block and thread indices.
    bidx, _, _ = cute.arch.block_idx()
    bdim, _, _ = cute.arch.block_dim()
    tidx, _, _ = cute.arch.thread_idx()

    _, num_q_heads, (threads_per_head, vecs_per_thread, vec_size) = q.shape
    _, num_k_heads, _ = k.shape
    head_dim = threads_per_head * vecs_per_thread * vec_size

    # Determine the thread and the vector to work on.
    vec_id = (bidx * bdim + tidx) // threads_per_head
    thread_id = tidx % threads_per_head

    # Determine the token and head indices.
    token = vec_id // (num_q_heads + num_k_heads)
    head_offset = vec_id % (num_q_heads + num_k_heads)
    is_q = head_offset < num_q_heads

    # Determine which tensor to work with.
    thread_coord = (thread_id, None, None)
    head_coord = (token, head_offset if is_q else head_offset - num_q_heads, thread_coord)

    # Load the head into registers.
    x_ref = q[head_coord] if is_q else k[head_coord]
    w = (q_weight if is_q else k_weight)[thread_coord].load()

    # Compute the squared sum locally.
    x = x_ref.load()
    x_f32 = x.to(cute.Float32)
    x_square = x_f32 * x_f32
    x_square_sum = x_square.reduce(cute.ReductionOp.ADD, init_val=0.0)

    # Sum up across the warp.
    mask = activemask()
    start = cute.arch.WARP_SIZE // (cute.arch.WARP_SIZE // threads_per_head * 2)
    for offset in [start >> i for i in range(5) if (start >> i) > 0]:
        x_square_sum += cute.arch.shuffle_sync_bfly(
            x_square_sum,
            offset=offset,
            mask=mask,
        )

    # RMS norm.
    rms = cute.math.rsqrt(x_square_sum / head_dim + EPS)
    x_norm = ((w + WEIGHT_BIAS) * (x * rms)).to(x.element_type)
    x_ref.store(x_norm)

একটি নিয়মিত RMS নর্ম কার্নেলের তুলনায়, এই বিশেষায়িত ভেরিয়েন্টটি প্রিফিল এবং ডিকোড উভয় ক্ষেত্রেই 2-3 গুণ ভালো থ্রুপুট এবং লেটেন্সি প্রদান করে।

MoE ডিসপাচ এবং কম্বাইন ইনপুট টোকেনগুলিকে সেই এক্সপার্টদের কাছে সাফেল করে যেগুলিতে সেগুলি রাউট করা হয়, MoE GEMM কার্নেলগুলি প্রাসঙ্গিক ম্যাট্রিক্স মাল্টিপ্লিকেশন সম্পাদন করার পরে একটি ওয়েটেড অ্যাক্যুমুলেশন সম্পাদন করে। পারফরম্যান্স বিশেষ আগ্রহের বিষয়, কারণ বিভিন্ন পণ্য বিভিন্ন আকারের MoE মডেলের উপর নির্ভর করে। স্থাপনের প্রকারের উপর ভিত্তি করে, আমরা অপ্টিমাইজ করার জন্য বিভিন্ন কার্নেল আলাদা করি:

TP=1: একটি একক জিপিইউ-এর জন্য স্থানীয় কার্নেল। ডিসপাচ টোকেনগুলিকে সংশ্লিষ্ট এক্সপার্টদের কাছে সাফেল করে এবং কম্বাইন সেগুলিকে অ্যাক্যুমুলেট করে। H100-এ Qwen3 30B স্থাপনের জন্য মূলত ব্যবহৃত হয়।

DP=1, TP ≠ 1: টেনসর-প্যারালাল, সিঙ্গেলনোড কার্নেল। ইনপুটটি প্রতিলিপি করা হয়, ডিসপাচ বর্তমান র‍্যাঙ্কের জন্য টোকেনগুলি নির্বাচন করে এবং কম্বাইন সমস্ত র‍্যাঙ্ক জুড়ে সেগুলিকে অ্যাক্যুমুলেট করে। NVLink বা NVLS জুড়ে অ্যাক্যুমুলেশন। বড় প্রিফিলারগুলির জন্য ব্যবহৃত হয়।

DP ≠ 1, NVLink: ঐচ্ছিক টেনসর প্যারালালিজম সহ ডেটা-প্যারালাল কার্নেল। প্রতিটি র‍্যাঙ্ক নির্দিষ্ট এক্সপার্ট হোস্ট করা প্রতিটি অন্যান্য র‍্যাঙ্কে টোকেন পাঠায়। রাউটিং তথ্য একত্রিত করা হয়, NVIDIA NVLink জুড়ে টোকেনগুলি স্থানান্তরিত হয়। ডেটা-প্যারালাল H200 বা GB200 স্থাপনের জন্য ব্যবহৃত হয়।

DP ≠ 1, InfiniBand: NVLink এবং NVIDIA Quantum InfiniBand উভয় জুড়ে পরিচালিত পিয়ার-টু-পিয়ার ডেটা-প্যারালাল এবং টেনসর-প্যারালাল কার্নেল। মাল্টি-নোড Hopper স্থাপনের জন্য ব্যবহৃত হয়।

স্থানীয় কার্নেল রাউটিং টেবিলের উপর ভিত্তি করে প্রতিটি এক্সপার্টকে নির্ধারিত টোকেনের সংখ্যা গণনা করে এবং প্রতিটি এক্সপার্টের অন্তর্গত টোকেনের তালিকার মধ্যে প্রতিটি টোকেনকে একটি অনন্য সূচক (index) নির্ধারণ করে। রাউটিং ধাপের পরে, MoE GEMM কার্নেলে পাস করার জন্য টোকেনগুলিকে ঐচ্ছিকভাবে প্যাড করা বাফারে অনুলিপি করা হয়। টোকেন অবস্থানগুলি ব্যবহার করে, কম্বাইন কার্নেল তখন রেজিস্টারে একটি অ্যাক্যুমুলেশন সম্পাদন করে, চূড়ান্ত টোকেনগুলিকে একটি আউটপুট বাফারে লেখে।

প্রিফিল এবং ডিকোডের মধ্যে ডিসপাচ পর্যায়টি পরিবর্তিত হয়: অল্প সংখ্যক টোকেনের জন্য, পড়ার মতো সামান্য রাউটিং তথ্য থাকে, তাই আমরা এটি একটি একক ব্লকে প্রসেস করি এবং গ্রিডের বাকি অংশে সংকেত দিই, যেখানে স্বতন্ত্র ব্লকগুলি টোকেনের চাংক অনুলিপি করে। শেয়ার্ড মেমরিতে পরমাণুকভাবে (atomically) কাউন্টার বৃদ্ধি করে টোকেন অ্যাসাইনমেন্ট নির্ধারিত হয়। যেহেতু প্রিফিলকে আরও রাউটিং তথ্য প্রসেস করতে হবে, তাই রাউটিং একটি পৃথক কার্নেলে একাধিক ব্লকে বিভক্ত করা হয়। শেয়ার্ড মেমরিতে কাউন্টার ব্যবহার করে প্রথমে স্থানীয়ভাবে সূচকগুলি নির্ধারিত হয় এবং তারপরে প্রতিটি টোকেনের অবস্থান নির্ধারণ করতে গ্লোবাল মেমরিতে সমস্ত ব্লক জুড়ে একত্রিত করা হয়। গ্রিড-ওয়াইজ ব্যারিয়ার ব্যবহার করতে সক্ষম হওয়ার জন্য এই কার্নেলের গ্রিডটি SM-এর সংখ্যায় বাউন্ডিং করা হয়। একটি পরবর্তী কার্নেল অনুলিপিগুলি সম্পাদন করে, যা উপলব্ধ মেমরি ব্যান্ডউইথ সম্পূর্ণরূপে ব্যবহার করতে অনেক বেশি সংখ্যক ব্লক এবং ওয়ার্প জুড়ে লঞ্চ করা হয়।

কম্বাইন কার্নেলটি তুলনামূলকভাবে সহজ, কারণ এটি রাউটিং তথ্য এবং ডিসপাচ দ্বারা গণিত অফসেটের উপর নির্ভর করে ভেক্টরাইজড অ্যাক্যুমুলেশন সম্পাদন করে। CuTeDSL এখানে মূল্যবান প্রমাণিত হয়েছে, কারণ গণনাটি সাধারণ CuTe টেনসরের মাধ্যমে প্রকাশ করা যেতে পারে যা সমস্ত ডেটা টাইপ এবং হিডেন ডাইমেনশনের জন্য দক্ষ কোডে কম্পাইল হয়।

একটি টেনসর-প্যারালাল স্থাপনে, রাউটিং তথ্য এবং ইনপুট টোকেনগুলির প্রতিলিপি তৈরি করা হয়। বর্তমান র‍্যাঙ্কে অবস্থিত এক্সপার্টদের কাছে নির্ধারিত টোকেনগুলি মাত্র রাউট এবং নির্বাচন করার জন্য সামান্য পরিবর্তনের সাথে স্থানীয় ডিসপাচ কার্নেলটি সম্পূর্ণরূপে পুনরায় ব্যবহার করা যেতে পারে। কিছু প্যারামিটার এবং কম্পাইল-টাইম স্পেশালাইজেশনের সাথে, স্থানীয় ডিসপাচ কার্নেলগুলি পুনরায় ব্যবহার করা হয়।

যেহেতু র‍্যাঙ্কের সংখ্যা সাধারণত একটি টোকেন যে এক্সপার্টদের কাছে রাউট করা হয় তার সংখ্যার চেয়ে কম হয়, তাই গড় ক্ষেত্রে এটি ধরে নেওয়া যেতে পারে যে প্রতিটি টোকেনের অন্তত একটি প্রতিলিপি প্রতিটি র‍্যাঙ্কে রাউট করা আছে। এই পরিস্থিতিতে, অংশগ্রহণকারী র‍্যাঙ্কগুলি জুড়ে যোগ করার জন্য একটি অল-রিডিউস অনুসরণ করে, একটি টোকেনের স্থানীয় অনুলিপিগুলি একত্রিত করতে স্থানীয় কম্বাইন কার্নেলটি পুনরায় ব্যবহার করা যেতে পারে। প্রিফিল কার্নেলের ক্ষেত্রে, আমরা পৃথক কার্নেল লঞ্চ করি, তবে ডিকোড উভয় অপারেশনকে একটি একক লঞ্চে ফিউজ করে।

অল-রিডিউস-এর জন্য, আমরা সাধারণ NVLink ব্যবহার করি যদি মাত্র দুটি র‍্যাঙ্ক অংশগ্রহণ করে, এবং থ্রুপউটে 10-20% বৃদ্ধির জন্য যখন 4 বা তার বেশি র‍্যাঙ্ক জড়িত থাকে তখন NVLS-এ স্যুইচ ওভার করি। DSL-এর জন্য ধন্যবাদ, আমরা MoE এবং নিয়মিত অল-রিডিউস কার্নেলগুলির মধ্যে বেশিরভাগ রিডাকশন লজিক এবং পিয়ার বাফার এক্সচেঞ্জ শেয়ার করতে পারি। এটি এমন একটি পরিস্থিতি যেখানে হার্ডওয়্যারের উপর নিবিড় নিয়ন্ত্রণ উপকারী, কারণ NVLS-এর জন্য multimem.ld_reduce এবং multimem.st এর উপর DSL অপারেশন প্রয়োজন। রিডাকশনের বাস্তবায়ন নিজেই মূলত NCCL পরীক্ষা থেকে all_reduce উদাহরণ অনুসরণ করে।

ডেটা প্যারালাল কার্নেলগুলি একাধিক ডিভাইসের জন্য টোকেন একত্রিত করতে NVLink-এর উপর নির্ভর করে। ডিসপাচ পর্যায়ে, টোকেন বিনিময় করার আগে, পিয়ারগুলি প্রথমে স্থানীয়ভাবে রাউট করে প্রতিটি টোকেনের জন্য প্রতিটি র‍্যাঙ্কে লক্ষ্য অফসেট সনাক্ত করে যা তারা পাঠাচ্ছে। পরবর্তীতে, তারা দ্বন্দ্ব এড়াতে প্রতিটি প্রেরকের জন্য বরাদ্দ করা রিসিভ বাফারে টোকেনগুলি অনুলিপি করা শুরু করে। সমাপ্তির সংকেত দিতে প্রতি-টোকেন রিলিজ-স্টোর রাইট ব্যবহার করা হয়। রিসিভার প্রান্তে, MoE GEMM-এর জন্য একটি একক বাফারে সমস্ত টোকেন ঘনীভূত করতে সক্ষম হতে রাউটিং তথ্য একত্রিত করা হয়। রাইটগুলি সম্পন্ন হওয়ার জন্য অপেক্ষা করতে রিসিভাররা একটি অ্যাকুইর-লোড সহ পিয়ার-লিখিত পতাকা পোল করে।

কম্বাইন সাইডে, রসিদ নির্দেশকারী প্রতি-টোকেন পতাকা সহ, উৎস র‍্যাঙ্কে NVLink জুড়ে টোকেনগুলি অনুলিপি করা হয়। কম্বাইন কার্নেলটি স্থানীয় এবং সিঙ্গেলনোড ভেরিয়েন্টের মতো, যেখানে অ্যাড্রেসিংয়ে সামান্য পার্থক্য এবং কিছু অতিরিক্ত সিঙ্ক্রোনাইজেশন লজিক রয়েছে।

এই কার্নেলটি বাস্তবায়নের জন্য আমরা সিস্টেম এবং ডিভাইস উভয় স্তরেই রিলিজ-অ্যাকুইর লোড এবং স্টোর উন্মুক্ত করতে PTX র‍্যাপারের একটি লাইব্রেরি তৈরি করেছি। সিঙ্ক্রোনাইজেশন পতাকার গ্রানুলারিটি উন্নত করে এবং NVLink জুড়ে টোকেন পাঠাতে এবং গ্রহণ করতে আরও ব্লক বরাদ্দ করে, আমরা আমাদের পূর্ববর্তী CUDA কার্নেলের তুলনায় 2 গুণ উন্নতি অর্জন করেছি।

ConnectX-7 InfiniBand অ্যাডাপ্টারের সাথে যুক্ত Hopper GPU-তে আমাদের পিয়ার-টু-পিয়ার কার্নেলগুলি পূর্বে সর্বাধুনিক পারফরম্যান্স অর্জন করেছিল। পারফরম্যান্স আরও বৃদ্ধির জন্য আমরা তখন থেকে CuTeDSL-এ কম্বাইন কার্নেলগুলি পোর্ট করেছি। যেহেতু CUDA সংস্করণগুলি ইতিমধ্যেই কম্পাইল-টাইম স্পেশালাইজেশনের উপর ব্যাপকভাবে নির্ভর করেছিল, তাই সেগুলি আমাদের রিপোজিটরিতে কম্পাইলেশন-সময়ের বাধা প্রমাণিত হয়েছে, যা আমাদের কিছু প্যারামিটার বাদ দিতে বাধ্য করেছে। JIT-কম্পাইলড DSL-এ মাইগ্রেট করার মাধ্যমে, আমরা আরও প্যারামিটারের জন্য বিশেষায়িত করেছি এবং গ্রিড বরাদ্দ টুইক করেছি, লেটেন্সিতে প্রায় 3% বৃদ্ধির জন্য প্রিফিল এবং ডিকোডের মধ্যে আরও পার্থক্য প্রবর্তন করেছি।

ROSE কার্নেলের দৃষ্টিভঙ্গি থেকে, Blackwell সমর্থন দুটি বিভাগে পড়ে: আর্কিটেকচার-অ্যাগনস্টিক এবং আর্কিটেকচার-নির্দিষ্ট কার্নেল। টেনসর কোরের উপর নির্ভর করে আর্কিটেকচার-নির্দিষ্ট কার্নেলগুলির মানিয়ে নেওয়ার জন্য স্পষ্ট কাজের প্রয়োজন হয়। তবে, DSL ব্যবহারের জন্য ধন্যবাদ, আর্কিটেকচার-অ্যাগনস্টিক কার্নেলগুলি সহজেই মানিয়ে নেয়।

Hopper-এর 132-এর বিপরীতে একটি Blackwell GPU-তে 160টি পর্যন্ত SM রয়েছে। বেশিরভাগ কার্নেলের কাছে, এটি একটি কম্পাইল-টাইম ধ্রুবক এবং হোস্ট-সাইড লঞ্চার কোডের লজিক চালানোর জন্য সর্বোত্তম (প্রিফিল বনাম ডিকোড) কার্নেল নির্ধারণের জন্য এটি বিবেচনা করে। অতিরিক্তভাবে, Blackwell 256-বিট লোড নির্দেশাবলীও প্রবর্তন করে, যা মেমরি ব্যান্ডউইথ ব্যবহার উন্নত করতে ভেক্টরের প্রস্থ (CuTe টেনসরের শেষ ডাইমেনশন) বৃদ্ধি করে সহজেই নির্গমন করা যেতে পারে।

Blackwell উল্লেখযোগ্য সংখ্যক কোয়ান্টাইজেশন স্কিমের জন্য হার্ডওয়্যার সমর্থন প্রবর্তন করে (MXFP8, MXFP4, NVFP4)। এগুলির জন্য সমর্থন মূলত FP8 ব্লক-ওয়াইজ স্কেলিংয়ের জন্য বিদ্যমান কার্নেলগুলিকে টুইক করে বাস্তবায়িত করা হয়েছিল। JIT সমর্থন এখানে অত্যন্ত গুরুত্বপূর্ণ, কারণ কম্পাইল-টাইম কনফিগারেশন স্পেস বিস্ফোরিত না করে নতুন মোডগুলিকে সমর্থন করা যেতে পারে। অতীতে, আমরা দেখেছি যে Triton অল্প সংখ্যক উপাদান জুড়ে, একটি ওয়ার্পের অর্ধেক বা চতুর্থাংশের উপর হরিজন্টাল রিডাকশন প্রকাশ করতে অস্বস্তিকর। যথাক্রমে NVFP4 এবং MXFP8/MXFP4-এর জন্য 16 বা 32 উপাদান জুড়ে একটি পরম সর্বোচ্চ মান গণনা করার সময় সর্বোত্তম পারফরম্যান্স অর্জনের জন্য রিডাকশনের উপর সূক্ষ্ম-দানাযুক্ত নিয়ন্ত্রণ অত্যন্ত গুরুত্বপূর্ণ ছিল।

তুলনামূলকভাবে একটি নতুন ভাষা হওয়ায়, প্রি-ট্রেনিংয়ের সময় LLM-গুলির CuTeDSL কোডের সাথে খুব কম পরিচয় ছিল। অতিরিক্তভাবে, অ্যাক্টিভেশন এবং নর্মের জন্য অপ্টিমাইজ করা কার্নেল সহ আমাদের বর্তমান কোডের স্টাইলটি বিদ্যমান CUTLASS বা Quack উদাহরণের থেকে কিছুটা আলাদা ছিল, যা TMA এবং আর্কিটেকচার-নির্দিষ্ট ম্যাট্রিক্স মাল্টিপ্লিকেশন নির্দেশাবলীর জন্য ভারীভাবে অপ্টিমাইজ করা হয়েছিল। LLM-গুলি কার্নেল লেখায় সহায়তা করতে পারে এমন একটি বিন্দুতে পৌঁছানো একটি ক্রমান্বয়িক প্রক্রিয়া ছিল।

কার্নেল রূপান্তরের জন্য আমরা এজেন্টদের উপর নির্ভর করার আগে, সামঞ্জস্যপূর্ণ প্যাটার্ন অনুসরণ করে আমরা ম্যানুয়ালি কার্নেল এবং ইউটিলিটি ফাংশনের একটি প্রাথমিক কর্পাস তৈরি করেছি। কর্পাসটি এগিয়ে যাওয়ার সাথে সাথে, আমরা উল্লেখযোগ্য ম্যানুয়াল পুনর্কাজের পরে CUDA থেকে CuTeDSL-এ মোটামুটি অনুবাদ করতে পেরেছি। বিশেষ করে, আমরা দেখেছি যে মডেলগুলি পয়েন্টার-টু-টেনসর রূপান্তর এবং ভেক্টরাইজেশনের সাথে লড়াই করেছে, একটি পরীক্ষা কম্পাইল করার লক্ষ্য রাখার সময় টাইপ ত্রুটিগুলি ঠিক করার চেষ্টা করে আটকে গেছে। তারা কার্নেলের গ্রিড-ভিত্তিক প্রকৃতির খুব কম বোঝাপড়া দেখিয়েছে, থ্রেড এবং গ্রিড সিঙ্ক্রোনাইজেশন গুলিয়ে ফেলেছে। সবচেয়ে হতাশাজনকভাবে, আমাদের নিজস্ব CuTeDSL প্রতিস্থাপনের বিরুদ্ধে তৈরি করার পরিবর্তে CUDA-নির্দিষ্ট লাইব্রেরির অবাস্তব অনুবাদ হ্যালুসিনেশন করার একটি প্রলোভন ছিল। সময়ের সাথে সাথে, CuTeDSL কার্নেলের কর্পাস বৃদ্ধি পাওয়ার সাথে সাথে, আমরা দেখতে পেয়েছি যে LLM-গুলি বেশিরভাগ স্কেলেটন কোড তৈরি করতে এবং ফাঁকা জায়গাগুলি পূরণ করতে পারে, যা আমাদের উচ্চ-স্তরের ডিজাইনে ফোকাস করার অনুমতি দেয়।

CUDA-র তুলনায়, আমরা দেখতে পেয়েছি যে CuTeDSL অনেক ভালো কম্পাইল-টাইম পারফরম্যান্স সরবরাহ করে, যা আমাদের আরও প্যারামিটারের জন্য বিশেষায়িত করতে এবং পারফরম্যান্সের উন্নতি অর্জন করতে দেয়। Triton-এর তুলনায় নির্দিষ্ট কিছু অ্যাবস্ট্রাকশনের অনুপস্থিতির কারণে কার্নেলগুলি কিছুটা বেশি ভারবোস হলেও, আমরা দেখতে পেয়েছি যে জটিলতাটি ন্যায়সঙ্গত কারণ হার্ডওয়্যারের উপর সূক্ষ্ম-দানাযুক্ত নিয়ন্ত্রণ আমাদের দ্রুত কার্নেল তৈরি করতে দেয়।

লেখার সময়, CUTLASS C++ এবং CuTeDSL কার্নেলের মধ্যে পারফরম্যান্সের ব্যবধান ন্যূনতম, যা ডিভাইস-মধ্যস্থ যোগাযোগ, নরমালাইজেশন বা অ্যাক্টিভেশনের সাথে উচ্চ-পারফরম্যান্স GEMM কার্নেলগুলি কম্পোজ এবং ফিউজ করার রোমাঞ্চকর সুযোগ উন্মুক্ত করে। এই DSL-এর প্রতি প্রতিশ্রুতিবদ্ধ হওয়ার মাধ্যমে, আমরা এমবেডিং থেকে শুরু করে LLM পর্যন্ত আমাদের ওয়ার্কলোড জুড়ে থ্রুপুট উন্নত করতে আমাদের ডিভাইসগুলি থেকে আরও পারফরম্যান্স বের করতে সক্ষম হব।

https://github.com/Dao-AILab/quack

https://github.com/NVIDIA/cutlass

https://github.com/NVIDIA/nccl-tests

https://docs.nvidia.com/cuda/parallel-thread-execution/

BibTeX

@online{perplexity2026CutedslAtPerplexity,
  author = {{Perplexity} and Nandor Licker and Bex Brooker and Jesse Dwyer and Phi Hoang},
  title = {{CuTeDSL at Perplexity}},
  date = {2026-05-06},
  url = {https://www.perplexity.ai/hub/blog/cutedsl-at-perplexity}
}

Perplexity AI-এ ইনফারেন্স

কেন CuTeDSL?

ডেভেলপার অভিজ্ঞতা

বড় কনফিগারেশন স্পেস

প্রিফিল এবং ডিকোড স্পেশালাইজেশন

গ্রিড সিঙ্ক্রোনাইজেশন

CuTeDSL কার্নেল

QK নর্ম

মিক্সচার-অফ-এক্সপার্টস ডিসপাচ/কম্বাইন

স্থানীয়

সিঙ্গেলনোড

ডেটা প্যারালাল

পিয়ার-টু-পিয়ার

Blackwell সমর্থন

AI-সহায়তাপ্রাপ্ত কার্নেল তৈরি

উপসংহার এবং ভবিষ্যতের কাজ

রেফারেন্স

উদ্ধৃতি