CuTeDSL στην Perplexity

Στην Perplexity, χρησιμοποιούμε την εσωτερική μηχανή εξυπηρέτησης Runtime-Optimized Serving Engine (ROSE) για την εξυπηρέτηση μοντέλων που κυμαίνονται από embeddings έως LLMs ενός τρισεκατομμυρίου παραμέτρων. Βρίσκεται πίσω από όλα τα API, συμπεριλαμβανομένων των Sonar, Search και Embeddings, και η μηχανή εκτέλεσής μας κατασκευάστηκε για να δίνει τη δυνατότητα στους ερευνητές να εξερευνηθούν τα καταλληλότερα μοντέλα και αρχιτεκτονικές, ενώ παράλληλα αναπτύσσονται γρήγορα οι καλύτεροι υποψήφιοι στην παραγωγή. Για να το πετύχουμε αυτό, πρέπει να είμαστε σε θέση να υλοποιούμε, να αναπτύσσουμε και να βελτιστοποιούμε την εξαγωγή συμπερασμάτων (inference) για αρχιτεκτονικές αιχμής με ταχύ ρυθμό. Σε αυτήν την ανάρτηση ιστολογίου, εξερευνούμε πώς τα kernels CuTeDSL ενσωματωμένα στη μηχανή εξαγωγής συμπερασμάτων μας επιτρέπουν να κατασκευάσουμε τα απαραίτητα kernels GPU για να φέρουμε τα μοντέλα σε μέγιστη απόδοση σε GPU της NVIDIA.

Εξαγωγή συμπερασμάτων στην Perplexity AI

Τα προσαρμοσμένα μοντέλα, κατασκευασμένα από διάφορες ομάδες στην Perplexity, φιλοξενούνται εσωτερικά σε GPU NVIDIA Hopper και Blackwell χρησιμοποιώντας τη μηχανή εξαγωγής συμπερασμάτων ROSE που κατασκευάστηκε από την ομάδα τεχνητής νοημοσύνης για εξαγωγή συμπερασμάτων (AI Inference team). Το ROSE κατασκευάστηκε αρχικά για την καλύτερη εξυπηρέτηση προσαρμοσμένων μοντέλων Llama τόσο για την αποκωδικοποίηση γλώσσας όσο και για την ταξινόμηση, με μια διεπαφή συμβατή με τον διακομιστή εξαγωγής συμπερασμάτων NVIDIA Triton Inference Server. Με την πάροδο του χρόνου, εξελίχθηκε σε μια πλήρη μηχανή εξαγωγής συμπερασμάτων που εξυπηρετεί σημαντικά LLMs, καθώς και μοντέλα κατάταξης, ταξινόμησης, βαθμολόγησης και άλλα μοντέλα που βασίζονται σε transformers για ένα ευρύ φάσμα ομάδων με σκοπό την τροφοδοσία των προϊόντων της Perplexity. Ο σχεδιαστικός στόχος του ROSE είναι να είναι απλό: με την πάροδο του χρόνου, επενδύεται προσπάθεια για την επίτευξη μέγιστης απόδοσης, αλλά η κύρια προτεραιότητα είναι η προσαρμογή στις συνεχώς εξελισσόμενες αρχιτεκτονικές και περιπτώσεις χρήσης για τα μοντέλα μας.

Μεγέθυνση

Σε υψηλό επίπεδο, το ROSE αποτελείται από μια μηχανή η οποία προσαρμόζει ένα μοντέλο σε μια διεπαφή με την οποία μπορούν να αλληλεπιδράσουν οι πελάτες. Οι μηχανές είναι εξαιρετικά επαναχρησιμοποιήσιμες, διαχειριζόμενες τον προγραμματισμό και την ομαδοποίηση αιτημάτων, καθώς και την ανάκτηση αποκωδικοποιημένων tokens, embeddings ή βαθμολογιών από τα υποκείμενα μοντέλα. Η προετοιμασία των συσκευών, τα κανάλια επικοινωνίας μεταξύ διεργασιών και κόμβων, και η φόρτωση βαρών διαχειρίζονται εξ ολοκλήρου από αυτές. Για τα LLMs, η μηχανή είναι υπεύθυνη για την ομαδοποίηση, τον τεμαχισμό και τη δειγματοληψία. Είναι επίσης υπεύθυνη για την κατανομή αποθήκευσης KV, τόσο για πλήρη όσο και για γραμμική προσοχή (attention), καθώς και για την αντιστοίχιση προθεμάτων. Οι μηχανές embedding είναι απλούστερες, καθώς εκτελούν μόνο ομαδοποίηση σε πραγματικό χρόνο (on-line batching). Η εκτέλεση βελτιστοποιείται έτσι ώστε η μηχανή να προετοιμάζει τον επόμενο γύρο εργασίας ενώ το μοντέλο εκτελείται παράλληλα στον επιταχυντή.

Το ROSE διαθέτει μια πλούσια συλλογή από προσαρμοσμένα επίπεδα για τη μοντελοποίηση και τη δημιουργία wrappers γύρω από GPU kernels που έχουν κατασκευαστεί με CuTeDSL, Triton, CUDA, CUTLASS και cuBLAS. Για υλοποιήσεις πολλαπλασιασμού πινάκων και attention, βασιζόμαστε κυρίως σε kernels της NVIDIA από τα CUTLASS και cuBLAS. Το υπόλοιπο pipeline, συμπεριλαμβανομένων των kernels για embedding, norm, MoE και activation, αποτελεί μέρος της kernel library του ROSE. Για να υποστηρίζουμε όσο το δυνατόν καλύτερα το hardware στο οποίο εκτελούμε inference, το ROSE είναι εξοπλισμένο με τα δικά του primitives επικοινωνίας μεταξύ συσκευών για MoE routing, dispatch και combine. Παρότι οι υπόλοιποι kernels είναι σχετικά απλοί, πρέπει να εξειδικεύονται σε βάθος ώστε να επιτυγχάνουν κορυφαία απόδοση σε όλα τα διαφορετικά hidden dimensions και τις συσκευές στις οποίες εκτελούνται. Επιπλέον, βασιζόμαστε σε εξειδίκευση κατά τον χρόνο μεταγλώττισης για να καλύψουμε όλα τα knobs και τις παραμετροποιήσεις των διαφόρων μοντέλων (weight bias έναντι no weight bias για RMS norm, activation functions κ.λπ.). Κατά συνέπεια, αποφασίσαμε να χρησιμοποιήσουμε DSLs για την υλοποίησή τους: αρχικά δημιουργήσαμε το ROSE χρησιμοποιώντας Triton. Ωστόσο, έκτοτε υιοθετήσαμε το CuTeDSL ως το κύριο περιβάλλον για προγραμματισμό GPU.

Γιατί CuTeDSL;

Από μαθηματική άποψη, τα kernels εξαγωγής συμπερασμάτων είναι συνήθως απλοϊκά. Κατά συνέπεια, κατά την ανάπτυξη ενός kernel, η πολυπλοκότητα γενικά δεν έγκειται στην έκφραση του ποιος υπολογισμός πρέπει να γίνει, αλλά στο πώς ακριβώς πρέπει να το κάνει το υλικό. Για να κατασκευάσουμε τα καλύτερα kernels, χρειαζόμαστε μια γλώσσα που προσφέρει τις κατάλληλες αφαιρέσεις πάνω στο υποκείμενο υλικό. Η σωστή γλώσσα μας δίνει τη δυνατότητα να καθορίσουμε τις λειτουργίες υψηλού επιπέδου για την εκτέλεση του υπολογισμού, εξειδικεύοντάς τες και βελτιστοποιώντας τες επιθετικά σε όλα τα σημεία ενός ευρέος χώρου διαμόρφωσης.

Το CuTeDSL είναι ένα DSL βασισμένο στην Python για GPU της NVIDIA, χτισμένο πάνω στην άλγεβρα διάταξης CuTe και το MLIR. Μεταγλωττίζεται just-in-time σε εξαιρετικά βελτιστοποιημένο PTX, διατηρώντας παράλληλα τον έλεγχο των πρωτογόνων υλικού χαμηλού επιπέδου. Στόχος του είναι να προσφέρει τα ίδια χαρακτηριστικά και την ίδια απόδοση με το CUTLASS C++, χωρίς τους δαπανηρούς χρόνους μεταγλώττισης που εμποδίζουν την ταχύτητα ανάπτυξης. Από τις γλώσσες που έχουμε εξετάσει ή με τις οποίες έχουμε εργαστεί (συμπεριλαμβανομένων των CUDA, cuTile, TVM IR και Triton), θεωρούμε ότι είναι η καλύτερη επιλογή για την ανάπτυξη των kernels μας.

Συνολικά, αν και δεν είναι τόσο υψηλού επιπέδου και περιεκτικό όσο το Triton ή το cuTile, διαπιστώσαμε ότι η άλγεβρα διάταξης CuTe (CuTe Layout Algebra) και οι συνθέσεις της (που εκτίθενται μέσω του cute.Tensor) ήταν λογικά υποκατάστατα. Επιπλέον, το CuTeDSL προσφέρει απεριόριστη πρόσβαση σε πρωτόγονα υλικού, παρέχοντάς μας το επίπεδο ελέγχου που απαιτούμε για να πετύχουμε μέγιστη απόδοση σε όλες τις οριακές περιπτώσεις. Με το CUDA, οι ευκαιρίες για εξειδίκευση κατά τον χρόνο μεταγλώττισης ήταν περιορισμένες, καθώς κάθε επιπλέον παράμετρος πρόσθετε αδικαιολόγητο κόστος στους χρόνους μεταγλώττισης ή πολυπλοκότητα στο σύστημα δόμησης. Αντίθετα, με ένα Just-In-Time Compiled DSL, μπορούμε να ειδικευτούμε πιο επιθετικά χωρίς να θυσιάζουμε τον ρυθμό επανάληψης. Σε αντίθεση με το TVM IR, το οποίο χρησιμοποιούσε το CUDA ως ενδιάμεσο βήμα, διαπιστώσαμε ότι ο αγωγός που βασίζεται στο MLIR και μεταγλωττίζει σε PTX είναι πιο ισχυρός και ευκολότερος στον εντοπισμό σφαλμάτων από έναν που εξέπεμπε το CUDA ως ενδιάμεσο βήμα. Επιπλέον, με τη χρήση προσαρμοσμένων λειτουργιών (custom ops) για την περιτύλιξη οδηγιών PTX, διαπιστώσαμε ότι μπορούμε να αξιοποιήσουμε πλήρως όλα τα χαρακτηριστικά του υποκείμενου υλικού.

Εμπειρία προγραμματιστή

Το μεγαλύτερο μέρος του κόστους μεταγλώττισης των CUDA kernels ξοδεύεται συνήθως στη φάση επέκτασης προτύπων (template expansion) των kernels, καθώς η βελτιστοποίηση και η μεταγλώττιση από το στάδιο PTX και έπειτα μπορούν να ολοκληρωθούν μέσα σε λίγα δευτερόλεπτα. Με πολλές παραμέτρους χρόνου μεταγλώττισης, αυτό το στάδιο μεταγλώττισης μπορεί να εκτοξευτεί σε πολλαπλά λεπτά για ένα μη τετριμμένο kernel. Αυτό το σημείο συμφόρησης εντείνεται από το γεγονός ότι οι ευκαιρίες για παράλληλη μεταγλώττιση προκύπτουν μόνο αργότερα, στα επίπεδα IR και PTX, με τον χειρισμό προτύπων C++ να γίνεται διαδοχικά. Ενώ υπάρχουν κόλπα για την αντιμετώπιση αυτού του προβλήματος, μετακινώντας τις στιγμιότυπες προτύπων (template instantiations) σε ξεχωριστές μονάδες μεταγλώττισης, απαιτούν σημαντικές τροποποιήσεις στα συστήματα δόμησης και τη χρήση δυσκίνητων μακροεντολών. Αντίθετα, τα CuTeDSL kernels που μεταγλωττίζονται Just-In-Time (JIT) δεν απαιτούν σύνθετους κανόνες δόμησης και δεν εμποδίζουν τους χρόνους μεταγλώττισης, καθώς ο κύκλος ανάπτυξης μπορεί να επικεντρωθεί σε ένα μόνο σημείο στον χώρο διαμόρφωσης, ενώ προσεκτικοί έλεγχοι μονάδας αργότερα μπορούν να καλύψουν ολόκληρο τον σχεδιαστικό χώρο του kernel.

Παρόλο που η μεταγλώττιση JIT μετατοπίζει το κόστος από τον χρόνο ανάπτυξης στον χρόνο εκτέλεσης, η επιβάρυνση είναι κάτι παραπάνω από αποδεκτή. Όταν ένα μοντέλο αναπτύσσεται, ενεργοποιεί μόνο λίγα από τα σημεία ολόκληρου του χώρου διαμόρφωσης, προκαλώντας επιβαρύνσεις κλίμακας δευτερολέπτων κατά την εκκίνηση/προθέρμανση της μηχανής εξαγωγής συμπερασμάτων. Αντίθετα, η προμεταγλώττιση όλων των διαμορφώσεων ενός kernel θα μπορούσε να διαρκέσει πολλαπλά λεπτά. Επιπλέον, οι επιβαρύνσεις μπορούν να μετριαστούν περαιτέρω με την προσωρινή αποθήκευση (caching) των αντικειμένων μεταγλώττισης JIT.

Το υψηλό κόστος μεταγλώττισης του CUDA εντείνεται από τη δυσκολία εντοπισμού της προέλευσης των σφαλμάτων όταν χρησιμοποιείται αποστολή (dispatch) και εξειδίκευση βασισμένη σε πρότυπα, όπως συμβαίνει με τα CuTe και CUTLASS. Με ένα DSL, τα σφάλματα μπορούν να επισημανθούν στη γραμμή κώδικα που τα παρήγαγε, διευκολύνοντας σημαντικά τον εντοπισμό σφαλμάτων (debugging). Επιπλέον, είναι πολύ απλούστερο να αντιστοιχιστεί ο παραγόμενος κώδικας στη δήλωση DSL από την οποία προήλθε μέσω ενδιάμεσων σταδίων MLIR ή βασιζόμενοι σε πληροφορίες εντοπισμού σφαλμάτων που επισυνάπτονται στο παραγόμενο PTX. Χάρη σε αυτές τις πτυχές του DSL, διαπιστώσαμε ότι είμαστε πολύ πιο παραγωγικοί κατά την ανάπτυξη kernels.

Μεγάλος χώρος διαμόρφωσης

Ενώ τα kernels μοιράζονται πολλές ομοιότητες στα διάφορα μοντέλα, πρέπει εξακολουθητικά να είναι εξοπλισμένα με μια σειρά από μοχλούς διαμόρφωσης που μπορούν να έχουν αντίκτυπο στην απόδοση. Οι σχετικές παράμετροι περιλαμβάνουν:

Κρυφή διάσταση (Hidden Dimension) / Διάσταση κεφαλής (Head Dimension) / Αριθμός κεφαλών προσοχής (Number of Attention Heads)

Τύποι δεδομένων εισόδου/εξόδου

Σχήμα κβάντωσης (Quantization scheme)

Σταθερές χρόνου μεταγλώττισης (biases, συναρτήσεις ενεργοποίησης κ.λπ.)

Μεγέθυνση

Η εξειδίκευση κατά τον χρόνο μεταγλώττισης είναι χρήσιμη τόσο για την απόδοση όσο και για τον εντοπισμό σφαλμάτων. Ιδιαίτερα σε έναν επεξεργαστή pipeline, όπως συμβαίνει με τις GPU, ένας ανύπαρκτος κλάδος (branch) είναι καλύτερος από έναν κλάδο που λαμβάνεται πάντα ή ποτέ. Επιπλέον, η μεταγλώττιση ενός kernel για ένα συγκεκριμένο σύνολο μοχλών διαμόρφωσης έχει ως αποτέλεσμα λιγότερες οδηγίες PTX/SASS, οι οποίες είναι ευκολότερες στην ανάλυση και τον εντοπισμό σφαλμάτων όταν χρησιμοποιείται το cuda-gdb κατά τη διάρκεια μιας έρευνας.

Για ένα δεδομένο μέγεθος παρτίδας (batch size) και διαμόρφωση, η αναμενόμενη κίνηση προς τη DRAM και ο αριθμός των αριθμητικών πράξεων μπορούν να εκτιμηθούν. Στη συνέχεια, μπορούν να ληφθούν βέλτιστες αποφάσεις κατά περίπτωση για τη βέλτιστη κατανομή πλέγματος (grid allocation), τον συγχρονισμό και τη χρήση ορισμένων χαρακτηριστικών υλικού. Το CuTeDSL μας επιτρέπει να ορίζουμε και να επαναχρησιμοποιούμε στοιχεία για βαθιά εξειδίκευση χωρίς να επιφέρουμε απαράδεκτη επιβάρυνση στον χρόνο μεταγλώττισης.

Προπλήρωση και αποκωδικοποίηση εξειδίκευσης

Γενικά, η απόδοση των kernels εξαγωγής συμπερασμάτων πρέπει να βελτιστοποιείται για δύο ξεχωριστές περιπτώσεις χρήσης:

Prefill: μικρό μέγεθος παρτίδας (λίγα chunks, ακολουθίες), πολλά tokens ανά ακολουθία. Η απόδοση (throughput) πρέπει να μεγιστοποιηθεί τόσο για τα embeddings όσο και για τα tokens εισόδου των LLMs. Καθυστέρηση (latency) στις δεκάδες έως εκατοντάδες milliseconds.

Decode: μεγέθη παρτίδων από μικρά έως μεγάλα, ένα (χωρίς speculative decoding) έως μερικά tokens ανά ακολουθία (draft-target speculation, πρόβλεψη πολλαπλών tokens). Καθυστέρηση κοντά στην καθυστέρηση εκκίνησης του kernel (kernel launch latency).

Για το GEMM και την προσοχή, το στάδιο prefill εκτελείται σε μεγάλο αριθμό tokens και η απόδοση περιορίζεται από την υπολογιστική απόδοση της συσκευής. Η αποκωδικοποίηση (decoding) περιλαμβάνει την ίδια λειτουργία σε μικρά μεγέθη παρτίδων, με τους χρόνους εκτέλεσης να κυριαρχούνται από το κόστος φόρτωσης των ενεργοποιήσεων και των βαρών από τη DRAM στους streaming multiprocessors και τα tensor cores, περιοριζόμενοι αντίθετα από το εύρος ζώνης μνήμης της συσκευής. Παρόλο που το μέγεθος της παρτίδας ή ο συνολικός αριθμός των tokens σε μια παρτίδα είναι ένας σχετικά σαφής δείκτης, το κατώφλι εναλλαγής μεταξύ prefill και decode είναι πιο λεπτό επειδή το εύρος ζώνης μνήμης καθορίζεται επίσης από την κρυφή διάσταση του μοντέλου και τον τύπο δεδομένων της ενεργοποίησης. Συνεπώς, τα kernels χωρίζονται σε υλοποιήσεις prefill και decode εάν αλγοριθμικές διαφορές ή αναθέσεις πλέγματος/μπλοκ (grid/block assignments) θα ωφελούσαν κάποιο από τα δύο.

Η παράμετρος που είναι καθολικά διαφορετική μεταξύ των kernels prefill και decode είναι ο αριθμός των warps σε ένα μπλοκ (ο αριθμός των νημάτων, διαιρεμένος με το μέγεθος warp των 32 νημάτων ανά warp). Είναι επιθυμητό να γίνει αυτή η παράμετρος σταθερά χρόνου μεταγλώττισης, καθώς μπορεί να βοηθήσει στην καταχώρηση (register allocation). Στο CUDA, αυτό γίνεται συνήθως καθορίζοντας τον αριθμό των νημάτων μέσω του __launch_bounds__ και παρέχοντας τον αριθμό των warps ως παράμετρο προτύπου:

Ο αριθμός των καταχωρητών (registers) που είναι διαθέσιμοι σε όλα τα νήματα σε ένα μπλοκ αποτελεί περιορισμό υλικού, 64K στα Hopper και Blackwell. Εάν ο αριθμός των νημάτων είναι γνωστός εκ των προτέρων, ο μεταγλωττιστής γνωρίζει ακριβώς πόρους καταχωρητών μπορεί να εκχωρήσει σε κάθε νήμα και μπορεί να πετύχει τη μεταγλώττιση πιο επιθετικά βελτιστοποιημένων kernels. Ωστόσο, αυτό συνοδεύεται από κόστος στον χρόνο μεταγλώττισης, καθώς τώρα το kernel πρέπει να μεταγλωττιστεί εκ νέου για διαφορετικά μεγέθη warp. Με το CuTeDSL, αυτό δεν αποτελεί πλέον πρόβλημα, καθώς δεν χρειάζεται να πληρώσουμε το κόστος μεταγλώττισης για ολόκληρο τον χώρο διαμόρφωσης εκ των προτέρων. Το DSL αναμένει επίσης το μέγεθος του μπλοκ να είναι μια σταθερή έκφραση, εκτελώντας σιωπηρά τη βελτιστοποίηση που υποδηλώνεται από το __latch_bounds__ στο CUDA.

Εκτός από το να γίνει ο αριθμός των warps σταθερά, είναι επιθυμητό να ποικίλλει μεταξύ prefill και decode. Για την πλήρη αξιοποίηση του εύρους ζώνης μνήμης μιας GPU, δεν αρκεί να εκχωρείται πάντα μεγάλος αριθμός warps ανά μπλοκ, καθώς το πλήρες εύρος ζώνης μπορεί να επιτευχθεί μόνο χρησιμοποιώντας όσο το δυνατόν περισσότερους από τους διαθέσιμους Streaming Multiprocessors (SM). Κατά συνέπεια, για έναν δεδομένο αριθμό tokens και γνωστούς τύπους δεδομένων και κρυφές διαστάσεις, επιλέγουμε πρώτα μια ιδανική διάσταση διανύσματος ανά νήμα (8 για bfloat16, 4 για float32) για να διασφαλίσουμε ότι ο παραγόμενος κώδικας χρησιμοποιεί τις οδηγίες φόρτωσης/αποθήκευσης 128-bit (256-bit στο Blackwell). Γνωρίζοντας το μέγεθος του διανύσματος, μπορούμε να υπολογίσουμε τον αριθμό των warps που απαιτούνται από το kernel. Πρώτα προσπαθούμε να μεγιστοποιήσουμε τον αριθμό των SM: εάν υπάρχουν λιγότερα warps από τα SM, επιλέγεται ένα ενιαίο warp ανά μπλοκ. Αυτή είναι συνήθως η παραλλαγή decode του kernel. Εάν υπάρχουν αρκετά warps, αυξάνουμε σταδιακά τον αριθμό των warps σε 8, 16 ή 32, κάνοντας benchmarking σε ένα άνω όριο πέραν του οποίου συνήθως δεν βλέπουμε πλέον βελτίωση στην απόδοση.

Ενώ το Triton εκτελούσε μερικές από αυτές τις βελτιστοποιήσεις στο παρασκήνιο, διαπιστώσαμε ότι μπορούμε να πετύχουμε παρόμοια αποτελέσματα χρησιμοποιώντας προσεκτικά την άλγεβρα διάταξης CuTe, χωρίς να περιοριζόμαστε από τις ελλείψεις ενός μεταγλωττιστή υψηλού επιπέδου. Ειδικά όταν εμπλέκονται οριζόντιες μειώσεις υπο-warp (sub-warp horizontal reductions), διαπιστώσαμε ότι το CuTeDSL μας πρόσφερε τον λεπτομερή έλεγχο που απαιτείται για τη δημιουργία βέλτιστου κώδικα.

Συγχρονισμός πλέγματος (Grid synchronization)

Σε ορισμένες καταστάσεις, για την κατανομή της εργασίας σε όσο το δυνατόν περισσότερα μπλοκ, μπορεί να χρησιμοποιηθεί το cooperative_groups::this_grid().sync() για τον συγχρονισμό όλων των συμμετεχόντων μπλοκ. Ισχύουν περιορισμοί: δεδομένου ότι ένα μπλοκ δεν μπορεί να αποπρογραμματιστεί από έναν SM μόλις αρχίσει η εκτέλεση, για να εγγυηθεί η πρόοδος προς τα εμπρός, ο αριθμός των μπλοκ δεν μπορεί να υπερβαίνει τον αριθμό των SM που είναι διαθέσιμοι στη συσκευή, ώστε να διασφαλιστεί ότι ανά πάσα στιγμή όλα τα μπλοκ μπορούν να ενημερώνουν και να κάνουν polling στις τοποθεσίες καθολικής μνήμης (global memory) που εμπλέκονται στον συγχρονισμό. Αυτός ο περιορισμός επιβάλλεται από το cudaLaunchCooperativeKernel. Το κόστος ενός φραγμού πλέγματος (grid barrier) είναι περίπου 3us στο Hopper, συν το κόστος εκκαθάρισης όλων των προηγούμενων λειτουργιών μνήμης μέσω ενός σιωπηρού φραγμού μνήμης σε ορισμένες καταστάσεις. Συνολικά, αυτό είναι συγκρίσιμο με την επιβάρυνση του διαχωρισμού ενός kernel σε δύο εκκινησεις, οδηγώντας σε μια απόφαση σχεδιασμού ευαίσθητη στην απόδοση. Παρά το γεγονός ότι το Programmatic Dependent Launch (PDL) μπορεί να μειώσει τις επιβαρύνσεις εκκίνησης από το Hopper και έπειτα, για τα kernels αποκωδικοποίησης (decode), η αποφυγή της επιβάρυνσης αξίζει τον κόπο.

Εάν ένας από τους φόρτους εργασίας απαιτεί περισσότερα μπλοκ από τον αριθμό των SM, το kernel πρέπει να διαχωριστεί, βασιζόμενο στον σιωπηλό φραγμό των εκκινήσεων kernel. Αυτό συμβαίνει κυρίως για το prefill. Εάν ο αριθμός των μπλοκ που είναι ιδανικός για όλους τους φόρτους εργασίας χωράει στο πλέγμα, προτιμάται ένας φραγμός πλέγματος, όπως συμβαίνει με τα kernels αποκωδικοποίησης που λειτουργούν σε μικρότερο αριθμό tokens.

Στο CuTeDSL, αυτές οι αποφάσεις μπορούν να ληφθούν στον κώδικα εκκίνησης από την πλευρά του host και τα kernels μπορούν να συντεθούν και να συμπλεχθούν με φραγμούς πλέγματος όπως απαιτείται, όπως φαίνεται στο παράδειγμα παρακάτω. Με μια τέτοια υβριδική στρατηγική, τα kernels αποκωδικοποίησης μπορούν να βελτιστοποιηθούν ώστε να φτάσουν σε ελάχιστες καθυστερήσεις και τα kernels prefill μπορούν να βελτιστοποιηθούν για μέγιστη απόδοση (throughput). Σε αντίθεση με τις λύσεις που βασίζονται σε CUDA, αυτή η λογική μπορεί να οριστεί πιο περιεκτικά και με λιγότερη επιβάρυνση στον χρόνο μεταγλώττισης.

Στην περίπτωση των MoE all-to-all kernels μας, όπως περιγράφεται λεπτομερώς παρακάτω, η ευελιξία να κατασκευάζουμε kernels με ή χωρίς φραγμούς πλέγματος και να τα βελτιστοποιούμε κατά τον χρόνο μεταγλώττισης μας βοήθησε να επιτύχουμε σημαντικά κέρδη απόδοσης μέσω του CuTeDSL.

CuTeDSL Kernels

QK Norm

Το πρώτο kernel που κατασκευάσαμε στο CuTeDSL ήταν ένα εξειδικευμένο kernel RMS norm, βελτιστοποιημένο για μικρές διαστάσεις. Το QK norm, που χρησιμοποιείται στο επίπεδο αυτοπροσοχής (self-attention) των πρόσφατων μοντέλων Qwen, εφαρμόζεται ανά κεφαλή (per-head), σε 64, 128 ή 256 στοιχεία. Σε αντίθεση με το κανονικό RMS norm, το οποίο λειτουργεί σε κρυφές διαστάσεις που υπερβαίνουν τις χιλιάδες, απαιτώντας μια οριζόντια μείωση μέσω τόσο της κοινόχρηστης μνήμης (shared memory) όσο και των καταχωρητών, το QK norm απαιτεί μείωση μόνο μεταξύ των warps, με τη δυνατότητα συμπίεσης πολλαπλών κεφαλών σε ένα ενιαίο warp. Αυτό κατέστησε μια περίπτωση χρήσης όπου μπορούσαμε να βελτιστοποιήσουμε επιθετικά κατά τον χρόνο μεταγλώττισης για να βρούμε τις καλύτερες αναθέσεις token σε warp για όλες τις διαστάσεις κεφαλής και τους τύπους δεδομένων.

Το RMS norm είναι μια αρκετά απλή λειτουργία, αν και απαιτεί ένα οριζόντιο άθροισμα για τον υπολογισμό της τετραγωνικής νόρμας ενός διανύσματος:

Δεδομένου ότι οι κεφαλές προσοχής είναι συνήθως μικρές και οι λειτουργίες ανά νήμα είναι σε μεγάλο βαθμό vectorized, ένα ενιαίο warp μπορεί να υπολογίσει την τετραγωνική νόρμα. Για να μεγιστοποιηθεί το εύρος ζώνης μνήμης, κάθε νήμα θα πρέπει να επεξεργάζεται 128 bits εισόδου, γεγονός που μεταφράζεται σε 8 τιμές bfloat16. Κατά συνέπεια, ένα μισό-warp (half-warp) είναι επαρκές για την επεξεργασία μιας κεφαλής, και η τετραγωνική νόρμα μπορεί να προσδιοριστεί χρησιμοποιώντας μια οριζόντια μείωση που εκτελείται από το shuffle_sync_bfly. Ξεκινώντας από τη διάσταση της κεφαλής και τον τύπο δεδομένων, μπορούμε να προσδιορίσουμε τον αριθμό των κεφαλών ανά warp και να καθορίσουμε τα ορίσματα και τον αριθμό των γύρων μείωσης κατά τη στιγμή της μεταγλώττισης.

Στη συνέχεια, μπορούμε να κατασκευάσουμε ένα εξειδικευμένο kernel με σχετικά λίγες γραμμές κώδικα, βασιζόμενοι σε αυτόματα vectorized φορτώσεις/αποθηκεύσεις και χρησιμοποιώντας λειτουργίες χαμηλότερου επιπέδου ειδικές για την αρχιτεκτονική για τη μείωση warp. Το kernel καθορίζει ποια κεφαλή (Q ή K) θα επεξεργαστεί και εντοπίζει τα αντίστοιχα βάρη που πρέπει να εφαρμοστούν. Φορτώνει ένα ολόκληρο διάνυσμα από την κεφαλή και αθροίζει πρώτα τοπικά. Στη συνέχεια, με βάση τον αριθμό των νημάτων ανά κεφαλή, υπολογίζεται η τετραγωνική νόρμα. Τέλος, το RMS norm, αυτόματα vectorized, προσδιορίζεται και εγγράφεται πίσω στη μνήμη. Οι παράμετροι που είναι ειδικές για το μοντέλο, όπως τα WEIGHT_BIAS και EPS, υπόκεινται σε σταθερή αναδίπλωση (constant-folded).

Σε σύγκριση με ένα κανονικό kernel RMS norm, αυτή η εξειδικευμένη παραλλαγή αποδίδει 2-3x καλύτερη απόδοση (throughput) και καθυστέρηση τόσο για το prefill όσο και για το decode.

Mixture-of-Experts Dispatch/Combine

Το MoE dispatch και combine ανακατεύουν (shuffle) τα tokens εισόδου στους ειδικούς (experts) στους οποίους δρομολογούνται, εκτελώντας μια σταθμισμένη συσσώρευση αφού τα MoE GEMM kernels εκτελέσουν τους σχετικούς πολλαπλασιασμούς πινάκων. Η απόδοση έχει ιδιαίτερο ενδιαφέρον, καθώς διάφορα προϊόντα βασίζονται σε μοντέλα MoE διαφόρων μεγεθών. Με βάση τους τύπους ανάπτυξης, διακρίνουμε διαφορετικά kernels προς βελτιστοποίηση:

TP=1: Τοπικό kernel για μία μόνο GPU. Το Dispatch ανακατεύει τα tokens στους σχετικούς experts και το combine τα συσσωρεύει. Χρησιμοποιείται κυρίως για αναπτύξεις Qwen3 30B σε H100.

DP=1, TP ≠ 1: Tensor-parallel, μονο-κομβικό (singlenode) kernel. Η είσοδος αναπαράγεται (replicated), το dispatch επιλέγει τα tokens για την τρέχουσα βαθμίδα (rank), και το combine τα συσσωρεύει σε όλες τις βαθμίδες. Συσσώρευση μέσω NVLink ή NVLS. Χρησιμοποιείται για μεγάλους προεπεξεργαστές (prefillers).

DP ≠ 1, NVLink: Data-parallel kernel με προαιρετικό tensor parallelism. Κάθε βαθμίδα στέλνει tokens σε κάθε άλλη βαθμίδα που φιλοξενεί συγκεκριμένους experts. Οι πληροφορίες δρομολόγησης συγκεντρώνονται, τα tokens μεταφέρονται μέσω NVIDIA NVLink. Χρησιμοποιείται για αναπτύξεις H200 ή GB200 παράλληλων δεδομένων (data-parallel).

DP ≠ 1, InfiniBand: Peer-to-Peer data-parallel και tensor-parallel kernel που λειτουργεί σε NVLink και NVIDIA Quantum InfiniBand. Χρησιμοποιείται για αναπτύξεις Hopper πολλαπλών κόμβων.

Τοπικό (Local)

Το τοπικό kernel μετρά τον αριθμό των tokens που έχουν ανατεθεί σε κάθε expert με βάση τον πίνακα δρομολόγησης και εκχωρεί ένα μοναδικό ευρετήριο σε κάθε token εντός της λίστα των tokens που ανήκουν σε κάθε expert. Μετά το βήμα δρομολόγησης, τα tokens αντιγράφονται σε προαιρετικά γεμισμένα (padded) buffers για να περάσουν σε ένα MoE GEMM kernel. Χρησιμοποιώντας τις τοποθεσίες των tokens, το combine kernel εκτελεί στη συνέχεια μια συσσώρευση σε καταχωρητές, εγγράφοντας τα τελικά tokens σε ένα buffer εξόδου.

Η φάση αποστολής (dispatch) ποικίλλει μεταξύ prefill και decode: για λίγα tokens, υπάρχουν ελάχιστες πληροφορίες δρομολόγησης για ανάγνωση, οπότε τις επεξεργαζόμαστε σε ένα ενιαίο μπλοκ και ειδοποιούμε το υπόλοιπο πλέγμα, με μεμονωμένα μπλοκ να αντιγράφουν τμήματα tokens. Η ανάθεση token καθορίζεται με ατομική αύξηση μετρητών στην κοινόχρηστη μνήμη. Δεδομένου ότι το prefill πρέπει να επεξεργαστεί περισσότερες πληροφορίες δρομολόγησης, η δρομολόγηση χωρίζεται σε πολλαπλά μπλοκ σε ένα ξεχωριστό kernel. Τα ευρετήρια (indices) εκχωρούνται πρώτα τοπικά χρησιμοποιώντας μετρητές στην κοινόχρηστη μνήμη και στη συνέχεια συγκεντρώνονται σε όλα τα μπλοκ στην καθολική μνήμη για να προσδιοριστεί η θέση κάθε token. Το πλέγμα αυτού του kernel περιορίζεται στον αριθμό των SM προκειμένου να είναι δυνατή η χρήση φραγμών σε επίπεδο πλέγματος (grid-wise barriers). Ένα επόμενο kernel εκτελεί τις αντιγραφές, εκκινούμενο σε πολύ μεγαλύτερο αριθμό μπλοκ και warps για την πλήρη αξιοποίηση του διαθέσιμου εύρους ζώνης μνήμης.

Το combine kernel είναι σχετικά απλό, καθώς εκτελεί vectorized συσσώρευση, βασιζόμενο στις πληροφορίες δρομολόγησης και στις μετατοπίσεις (offsets) που υπολογίστηκαν από το dispatch. Το CuTeDSL αποδείχθηκε πολύτιμο εδώ, καθώς ο υπολογισμός μπορούσε να εκφραστεί μέσω απλών τανυστών CuTe που μεταγλωττίζονται σε αποδοτικό κώδικα για όλους τους τύπους δεδομένων και τις κρυφές διαστάσεις.

Μονο-κομβικό (Singlenode)

Σε μια ανάπτυξη παράλληλων τανυστών (tensor-parallel deployment), οι πληροφορίες δρομολόγησης και τα tokens εισόδου αναπαράγονται. Το τοπικό dispatch kernel μπορεί να επαναχρησιμοποιηθεί πλήρως με μια μικρή αλλαγή για τη δρομολόγηση και την επιλογή μόνο των tokens που έχουν ανατεθεί στους experts που βρίσκονται στην τρέχουσα βαθμίδα. Με ορισμένες παραμέτρους και εξειδίκευση κατά τον χρόνο μεταγλώττισης, τα τοπικά dispatch kernels επαναχρησιμοποιούνται.

Δεδομένου ότι ο αριθμός των βαθμίδων (ranks) είναι συνήθως μικρότερος από τον αριθμό των experts στους οποίους δρομολογείται ένα token, στη μέση περίπτωση μπορεί να υποτεθεί ότι κάθε token έχει τουλάχιστον ένα από τα αντίγραφά του δρομολογημένο σε κάθε βαθμίδα. Σε αυτήν την κατάσταση, το τοπικό combine kernel μπορεί να επαναχρησιμοποιηθεί για τη συγκέντρωση των τοπικών αντιγράφων ενός token, ακολουθούμενο από ένα all-reduce για την άθροιση στις συμμετέχουσες βαθμίδες. Στην περίπτωση του kernel prefill, εκκινούμε ξεχωριστά kernels, ωστόσο το decode συγχωνεύει και τις δύο λειτουργίες σε μια ενιαία εκκίνηση.

Για το all-reduce, χρησιμοποιούμε απλό NVLink εάν συμμετέχουν μόνο δύο βαθμίδες, μεταβαίνοντας σε NVLS όταν εμπλέκονται 4 ή περισσότερες βαθμίδες για αύξηση της απόδοσης κατά 10-20%. Χάρη στο DSL, μπορούμε να μοιραστούμε το μεγαλύτερο μέρος της λογικής μείωσης και των ανταλλαγών peer buffer μεταξύ των MoE και των κανονικών all-reduce kernels. Αυτή είναι μια κατάσταση όπου ο στενός έλεγχος του υλικού είναι ευεργετικός, καθώς απαιτούνται λειτουργίες DSL μέσω των multimem.ld_reduce και multimem.st για το NVLS. Η υλοποίηση της ίδιας της μείωσης ακολουθεί σε μεγάλο βαθμό το παράδειγμα all_reduce από τα NCCL tests.

Παράλληλα δεδομένων (Data Parallel)

Τα kernels Παράλληλων Δεδομένων βασίζονται στο NVLink για τη συγκέντρωση tokens για πολλαπλές συσκευές. Στη φάση αποστολής, πριν από την ανταλλαγή tokens, οι ομότιμοι (peers) δρομολογούν πρώτα τοπικά για να εντοπίσουν την τιμή-στόχο μετατόπισης σε κάθε βαθμίδα για κάθε token που στέλνουν. Στη συνέχεια, αρχίζουν να αντιγράφουν τα tokens σε buffers λήψης που έχουν κατανεμηθεί για κάθε αποστολέα ώστε να αποφευχθούν συγκρούσεις. Χρησιμοποιούνται εγγραφές release-store ανά token για να σηματοδοτηθεί η ολοκλήρωση. Από την πλευρά του δέκτη, οι πληροφορίες δρομολόγησης συγκεντρώνονται ώστε να μπορούν να συμπυκνωθούν όλα τα tokens σε ένα ενιαίο buffer για το MoE GEMM. Οι δέκτες κάνουν poll στη σημαία (flag) που γράφτηκε από τον ομότιμο με ένα acquire-load για να περιμένουν την ολοκλήρωση των εγγραφών.

Στην πλευρά του combine, τα tokens αντιγράφονται μέσω NVLink στις βαθμίδες προέλευσης, με μια σημαία ανά token που υποδεικνύει την παραλαβή. Το combine kernel είναι παρόμοιο με τις τοπικές και μονο-κομβικές παραλλαγές, με μικρές διαφορές στη διευθυνσιοδότηση και κάποια επιπλέον λογική συγχρονισμού.

Κατασκευάσαμε μια βιβλιοθήκη περιτυλιγμάτων PTX (wrappers) για την έκθεση φορτώσεων και αποθηκεύσεων release-acquire σε επίπεδο συστήματος και συσκευής για την υλοποίηση αυτού του kernel. Βελτιώνοντας την κοκκομετρία (granularity) των σημαιών συγχρονισμού και κατανέμοντας περισσότερα μπλοκ για την αποστολή και λήψη tokens μέσω NVLink, πετύχαμε βελτιώσεις 2x σε σχέση με τα προηγούμενα CUDA kernels μας.

Peer-to-Peer

Τα Peer-to-Peer kernels μας πέτυχαν προηγουμένως απόδοση αιχμής σε GPU Hopper συνδυασμένες με προσαρμογείς ConnectX-7 InfiniBand. Έκτοτε έχουμε μεταφέρει τα combine kernels στο CuTeDSL για περαιτέρω κέρδη στην απόδοση. Δεδομένου ότι οι εκδόσεις CUDA βασίζονταν ήδη σε μεγάλο βαθμό σε εξειδίκευση κατά τον χρόνο μεταγλώττισης, αποδείχθηκαν τα σημεία συμφόρησης χρόνου μεταγλώττισης στο αποθετήριό μας, αναγκάζοντάς μας να παραλείψουμε ορισμένες παραμέτρους. Με τη μετάπτωση στο DSL μεταγλώττισης JIT, ειδικευτήκαμε για περισσότερες παραμέτρους και προσαρμόσαμε τις κατανομές πλέγματος, εισάγοντας περαιτέρω διαφορές μεταξύ prefill και decode για περίπου 3% αύξηση στην καθυστέρηση.

Υποστήριξη Blackwell

Από την προοπτική των ROSE kernels, η υποστήριξη Blackwell εμπίπτει σε δύο κατηγορίες: kernels ανεξάρτητα από την αρχιτεκτονική και ειδικά για την αρχιτεκτονική. Τα kernels που είναι ειδικά για την αρχιτεκτονική, βασιζόμενα στα Tensor cores, απαιτούν ρητή εργασία για προσαρμογή. Ωστόσο, χάρη στη χρήση του DSL, τα kernels που είναι ανεξάρτητα από την αρχιτεκτονική προσαρμόζονται ασήμαντα.

Μια GPU Blackwell διαθέτει έως και 160 SM, σε αντίθεση με τους 132 του Hopper. Για τα περισσότερα kernels, αυτό είναι μια σταθερά χρόνου μεταγλώττισης και η λογική στον κώδικα εκκίνησης από την πλευρά του host το λαμβάνει υπόψη για τον προσδιορισμό του βέλτιστου kernel (prefill έναντι decode) προς εκτέλεση. Επιπλέον, το Blackwell εισάγει επίσης οδηγίες φόρτωσης 256-bit, οι οποίες μπορούν εύκολα να εκπέμψουν αυξάνοντας το πλάτος διανύσματος (τελευταία διάσταση των τανυστών CuTe) για τη βελτίωση της αξιοποίησης του εύρους ζώνης μνήμης.

Το Blackwell εισάγει υποστήριξη υλικού για σημαντικό αριθμό σχημάτων κβάντωσης (MXFP8, MXFP4, NVFP4). Η υποστήριξη για αυτά υλοποιήθηκε κυρίως με την τροποποίηση υπαρχόντων kernels για κλιμάκωση ανά μπλοκ FP8 (block-wise scaling). Η υποστήριξη JIT είναι ζωτικής σημασίας εδώ, καθώς νέες λειτουργίες μπορούν να υποστηριχθούν χωρίς να εκρήγνυται ο χώρος διαμόρφωσης χρόνου μεταγλώττισης. Στο παρελθόν, διαπιστώσαμε ότι το Triton ήταν άβολο στην έκφραση οριζόντιων μειώσεων σε λίγα στοιχεία, πάνω από το μισό ή το ένα τέταρτο ενός warp. Ο λεπτομερέστερος έλεγχος επί των μειώσεων ήταν ζωτικής σημασίας για την επίτευξη βέλτιστης απόδοσης κατά τον υπολογισμό μιας απόλυτης μέγιστης τιμής σε 16 ή 32 στοιχεία για NVFP4 και MXFP8/MXFP4 αντίστοιχα.

Δημιουργία Kernel με Υποστήριξη AI

Ως μια σχετικά νέα γλώσσα, τα LLMs είχαν ελάχιστη έκθεση στον κώδικα CuTeDSL κατά την προ-εκπαίδευσή τους. Επιπλέον, το στυλ του τρέχοντος κώδικα μας, με kernels βελτιστοποιημένα για ενεργοποιήσεις και νόρμες (norms), διέφερε ελαφρώς από τα υπάρχοντα παραδείγματα CUTLASS ή Quack, τα οποία ήταν βαριά βελτιστοποιημένα για οδηγίες TMA και πολλαπλασιασμού πινάκων ειδικές για την αρχιτεκτονική. Η επίτευξη ενός σημείου όπου τα LLMs θα μπορούσαν να βοηθήσουν στη συγγραφή kernel ήταν μια σταδιακή διαδικασία.

Κατασκευάσαμε χειροκίνητα ένα αρχικό corpus από kernels και συναρτήσεις βοηθητικού προγράμματος (utility functions), ακολουθώντας συνεπή πρότυπα, προτού μπορέσουμε να βασιστούμε σε agents για τη μετατροπή kernel. Καθώς το corpus προχωρούσε, μπορούσαμε να εκτελέσουμε πρόχειρες μεταφράσεις από CUDA σε CuTeDSL, με σημαντική χειροκίνητη επανεπεξεργασία να ακολουθεί. Ειδικότερα, διαπιστώσαμε ότι τα μοντέλα δυσκολεύονταν με τις μετατροπές δείκτη σε τανυστή (pointer-to-tensor) και το vectorization, κολλώντας προσπαθώντας να διορθώσουν σφάλματα τύπων όταν στόχος ήταν να γίνει ένα τεστ να μεταγλωττιστεί. Έχουν δείξει ελάχιστη κατανόηση της φύσης των kernels που βασίζεται σε πλέγμα, συγχέοντας τον συγχρονισμό νημάτων και πλέγματος. Το πιο απογοητευτικό ήταν ο πειρασμός να δημιουργηθούν μέσω hallutination ανύπαρκτες μεταφράσεις βιβλιοθηκών ειδικών για CUDA αντί να γίνει δόμηση έναντι των δικών μας αντικαταστάσεων CuTeDSL. Με την πάροδο του χρόνου, καθώς το corpus των CuTeDSL kernels μεγάλωνε, διαπιστώσαμε ότι τα LLMs μπορούν να κατασκευάσουν το μεγαλύτερο μέρος του κώδικα σκελετού (skeleton code) και να συμπληρώσουν τα κενά, επιτρέποντάς μας να επικεντρωθούμε στον σχεδιασμό υψηλού επιπέδου.

Συμπέρασμα και μελλοντική εργασία

Σε σύγκριση με το CUDA, διαπιστώνουμε ότι το CuTeDSL προσφέρει πολύ καλύτερη απόδοση στον χρόνο μεταγλώττισης, επιτρέποντάς μας να ειδικευτούμε σε περισσότερες παραμέτρους και να πετύχουμε βελτιώσεις απόδοσης. Ενώ τα kernels είναι ελαφρώς πιο περιλογικά λόγω της έλλειψης ορισμένων αφαιρέσεων σε σύγκριση με το Triton, διαπιστώνουμε ότι η πολυπλοκότητα δικαιολογείται καθώς ο λεπτομερέστερος έλεγχος του υλικού μας επιτρέπει να κατασκευάζουμε ταχύτερα kernels.

Κατά τη στιγμή της συγγραφής, το χάσμα απόδοσης μεταξύ των kernels CUTLASS C++ και CuTeDSL είναι ελάχιστο, ανοίγοντας συναρπαστικές ευκαιρίες για σύνθεση και συγχώνευση GEMM kernels υψηλής απόδοσης με επικοινωνία μεταξύ συσκευών, κανονικοποίηση (normalization) ή ενεργοποιήσεις. Με την προσήλωσή μας σε αυτό το DSL, θα είμαστε σε θέση να εξάγουμε περαιτέρω απόδοση από τις συσκευές μας για να βελτιώσουμε την απόδοση (throughput) σε όλους τους φόρτους εργασίας μας, από embeddings έως LLMs.

Παραπομπές

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/

Παραπομπή

@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}
}