Επιτάχυνση του Sonar μέσω εικασίας

Η εικαστική αποκωδικοποίηση επιταχύνει τα LLM Sonar μέσω επαλήθευσης πρόχειρου μοντέλου

ΣυγγραφείςPerplexity Research

Η εικαστική αποκωδικοποίηση (speculative decoding) επιταχύνει την ταχύτητα παραγωγής των Μεγάλων Γλωσσικών Μοντέλων (LLM) χρησιμοποιώντας ένα γρήγορο και μικρό μοντέλο πρόχειρης σχεδίασης για την παραγωγή υποψηφίων ολοκλήρωσης που επαληθεύονται από το μεγαλύτερο μοντέλο-στόχο. Σύμφωνα με αυτό το σχήμα, αντί για μία εκτέλεση του δαπανηρού στόχου που παράγει ένα μόνο διακριτικό (token), εκπέμπονται πολλαπλά σε ένα μόνο βήμα. Εδώ παρουσιάζουμε τις λεπτομέρειες υλοποίησης διαφόρων ειδών εικαστικής αποκωδικοποίησης, που εφαρμόζονται στο Perplexity για τη μείωση της καθυστέρησης μεταξύ των διακριτικών στα μοντέλα Sonar.

Εικαστική Αποκωδικοποίηση

Η εικαστική αποκωδικοποίηση αξιοποιεί τη δομή των φυσικών γλωσσών και τη φυλογενετική φύση των μετασχηματιστών (transformers) για την επιτάχυνση της παραγωγής διακριτικών. Ακόμη και αν τα μεγαλύτερα μοντέλα, όπως το Llama-70B, διαθέτουν περισσότερη γνώση από τα μικρότερα, όπως το Llama-1B, σε ορισμένες απλούστερες εργασίες αποδίδουν παρόμοια. Αυτή η επικάλυψη υποδεικνύει ότι ορισμένες ακολουθίες παράγονται καλύτερα από τα λιγότερο δαπανηρά μοντέλα, αφήνοντας τα πολύπλοκα προβλήματα στα μεγαλύτερα. Η πρόκληση έγκειται στον καθορισμό του ποιες ολοκληρώσεις είναι καλύτερες και αν η παραγωγή του μικρότερου μοντέλου είναι της ίδιας ποιότητας με εκείνη του μεγαλύτερου.

Ευτυχώς, τα LLM είναι αυτοπαλίνδρομοι μετασχηματιστές: όταν τους δίνεται μια ακολουθία από tokens, εξάγουν την κατανομή πιθανότητας του επόμενου token. Επιπλέον, τα logits που προκύπτουν από τα ενδιάμεσα χαρακτηριστικά που σχετίζονται με τα tokens στην ακολουθία εισόδου δείχνουν επίσης πόσο πιθανό είναι το μοντέλο να εκδώσει ακριβώς αυτά τα tokens. Αυτή η ιδιότητα επιτρέπει την κερδοσκοπική εκτέλεση: αν μια ακολουθία tokens δημιουργείται από ένα μικρότερο μοντέλο ξεκινώντας από ένα πρόθεμα εισόδου, μπορεί να περαστεί από το μεγαλύτερο μοντέλο ώστε να καθοριστεί πόσο καλά ευθυγραμμίζεται με το μοντέλο-στόχο. Κάθε πρόθεμα των υποψηφίων βαθμολογείται με μια πιθανότητα και επιλέγεται το μακρύτερο που βρίσκεται πάνω από ένα κατώφλι αποδοχής. Ως πρόσθετο όφελος, το μοντέλο-στόχος παρέχει επίσης δωρεάν ένα επόμενο token: αν ένα draft model δημιουργεί n tokens, μπορούν να εκδοθούν έως n + 1 σε ένα βήμα.

Ένα πρόχειρο μοντέλο προτείνει συνεχίσεις· το μοντέλο στόχος δέχεται ένα πρόθεμα και παράγει ένα διακριτικό αντικατάστασης.

Κατά το χρόνο εξαγωγής συμπερασμάτων (inference), η διαδικασία εικαστικής δειγματοληψίας μπορεί να διαιρεθεί σε περίπου 4 στάδια:

Προγεμίσμα: τόσο το μοντέλο-στόχος όσο και το draft model πρέπει να εκτελεστούν πάνω στην ακολουθία εισόδου ώστε να συμπληρωθούν οι εγγραφές της KV cache. Παρότι ορισμένα σχήματα, όπως το Medusa, χρησιμοποιούν απλούστερα πυκνά επίπεδα για πρόβλεψη, σε αυτήν την ανάρτηση εστιάζουμε σε drafts βασισμένα σε μετασχηματιστές που χρειάζονται τις δικές τους KV caches.

Παραγωγή πρόχειρου (Draft generation): το πρόχειρο μοντέλο επαναλαμβάνεται για να παράγει έναν αριθμό σταθερών διακριτικών. Η πρόχειρη ακολουθία μπορεί να είναι γραμμική ή το μοντέλο μπορεί να εξερευνήσει μια δομή που μοιάζει με δέντρο έως ένα δεδομένο βάθος (EAGLE, Medusa). Εδώ, εστιάζουμε σε γραμμικές ακολουθίες.

Αποδοχή: το μοντέλο-στόχος εκτελείται πάνω στη draft ακολουθία, δημιουργώντας logits που αντιστοιχούν σε κάθε draft token. Προσδιορίζεται το μήκος της μακρύτερης αποδεκτής ακολουθίας.

Παραγωγή στόχου (Target generation): δεδομένου ότι ο στόχος παρήγαγε logits, στη θέση ανατιστοιχίας ή στο τέλος της ακολουθίας τα logits αντιστοιχούν σε ένα μεταγενέστερο διακριτικό. Αυτά τα logits μπορούν να υποβληθούν σε δειγματοληψία για να παρέχουν ένα ισχυρό διακριτικό από τον στόχο, ολοκληρώνοντας την ακολουθία.

Υπάρχουν διάφορες μέθοδοι για την υλοποίηση της εικαστικής αποκωδικοποίησης. Σε αυτήν την ανάρτηση, θα εστιάσουμε στα σχήματα που χρησιμοποιήσαμε για να επιταχύνουμε τα μοντέλα Sonar χρησιμοποιώντας ένα εσωτερικό μοντέλο 1B, καθώς και στους μηχανισμούς πρόβλεψης που κατασκευάζουμε για την επιτάχυνση μοντέλων στην κλίμακα του DeepSeek.

Στόχος-Πρόχειρο (Target-Draft)

Η εικαστική αποκωδικοποίηση μπορεί να επιτευχθεί συνδυάζοντας ένα υπάρχον μικρό LLM ως πρόχειρο μοντέλο με ένα μοντέλο-στόχο για τη δημιουργία υποψήφιων ακολουθιών. Στην παραγωγή, έχουμε επιταχύνει το Sonar χρησιμοποιώντας ένα μοντέλο Llama-1B προσαρμοσμένο (fine-tuned) στο ίδιο σύνολο δεδομένων με τον στόχο. Αν και αυτή η προσέγγιση δεν απαιτούσε την εκπαίδευση ενός πρόχειρου από την αρχή, το μικρό μοντέλο εξακολουθεί να χρησιμοποιεί σημαντική χωρητικότητα cache KV και εισάγει μια ελαφρά επιβάρυνση προπλήρωσης, αυξώντας το TTFT.

Σύμφωνα με αυτό το σχήμα, ο αποκωδικοποιητής κάνει εικασίες μόνο σε παρτίδες αποκωδικοποίησης (decode-only), παράγοντας διακριτικά μέσω τυπικής δειγματοληψίας κατά την προπλήρωση ή σε μικτές παρτίδες προπλήρωσης-αποκωδικοποίησης. Στο στάδιο της προπλήρωσης, τα logits του στόχου λαμβάνονται αμέσως σε δειγματοληψία για να προπληρωθεί επίσης το νεοπαραγόμενο διακριτικό στη μνήμη cache KV του πρόχειρου. Το πρόχειρο δεν λαμβάνει ακόμη δειγματοληψία, αλλά τα logits που παράγει μεταφέρονται στο στάδιο αποκωδικοποίησης.

Εικασία αποκωδικοποίησης από διακριτικά, που δείχνει το prefill, τη δημιουργία πρόχειρου, τη δειγματοληψία και την αποδοχή από το μοντέλο στόχο.

Στην αποκωδικοποίηση, το πρόχειρο μοντέλο προχωρά, κάνοντας δειγματοληψία του κορυφαίου διακριτικού σε κάθε στάδιο. Αφού επιτευχθεί το επιθυμητό μήκος πρόχειρου, τα διακριτικά εκτελούνται μέσω του μοντέλου-στόχου για να παραχθούν τα logits βάσει των οποίων ο δειγματολήπτης εντοπίζει το αποδεκτό μήκος ακολουθίας. Η αποδοχή καθορίζεται συγκρίνοντας τις πλήρεις κατανομές πιθανότητας από το πρόχειρο και τον στόχο. Δεδομένου ότι ο στόχος εξάγει πάντα ένα σύνολο logits που ακολουθεί την αποδεκτή πρόχειρη ακολουθία, αυτό υποβάλλεται σε δειγματοληψία για την παραγωγή μιας επιπλέον εξόδου. Επειδή το πρόχειρο μοντέλο δεν έχει δει ακόμη αυτό το αποδεκτό διακριτικό, εκτελείται ξανά για να συμπληρώσει τις αντίστοιχες καταχωρήσεις cache KV σε προετοιμασία για το επόμενο βήμα αποκωδικοποίησης, μεταφέροντας ξανά τα logits.

EAGLE

Το EAGLE είναι ένα σχήμα εικαστικής αποκωδικοποίησης το οποίο εξερευνά πολλαπλές πρόχειρες ακολουθίες, που παράγονται μέσω μιας δενδροειδούς διέλευσης πιθανών πρόχειρων διακριτικών. Ένα σταθερό (EAGLE) ή δυναμικά διαμορφωμένο (EAGLE-2) δέντρο εξερευνάται χρησιμοποιώντας διαδοχικές εκτελέσεις των πρόχειρων διακριτικών, λαμβάνοντας υπόψη τους υποψηφίους Top-K σε κάθε κόμβο αντί να ακολουθεί το διακριτικό με τη μεγαλύτερη βαθμολογία σε μια γραμμική ακολουθία. Στη συνέχεια, οι ακολουθίες βαθμολογούνται και επιλέγεται η μακρύτερη κατάλληλη για συνέχιση, προσθέτοντας επίσης ένα επιπλέον διακριτικό από τον στόχο.

Η επαλήθευση πρόχειρων πολλαπλών ακολουθιών χρησιμοποιεί μια διακλαδιζόμενη δομή διακριτικών και μια αντίστοιχη μάσκα προσοχής.

Για να επιτευχθεί ακριβέστερη πρόβληψη, ένα πρόχειρο μοντέλο EAGLE προβλέπει όχι μόνο βάσει διακριτικών, αλλά χρησιμοποιώντας επίσης τα χαρακτηριστικά στόχου (καταστάσεις κρυφού επιπέδου - hidden states) του μοντέλου-στόχου. Το μειονέκτημα του EAGLE είναι η ανάγκη εκπαίδευσης προσαρμοσμένων, μικρών πρόχειρων μοντέλων που είναι αρκετά ακριβή ώστε να παράγουν κατάλληλους υποψηφίους εντός ενός προϋπολογισμού χαμηλής καθυστέρησης. Συνήθως, ένα πρόχειρο μοντέλο είναι ένα ενιαίο στρώμα μετασχηματιστή πανομοιότυπο με ένα στρώμα αποκωδικοποιητή του αρχικού μοντέλου, το οποίο συνδέεται στενά με τον στόχο μέσω σύνδεσης με τις ενσωματώσεις (embeddings) του και τις προβολές lm_head. Επειδή αυτό απαιτεί μικρότερη χωρητικότητα cache KV, το EAGLE έχει μικρότερο αποτύπωμα μνήμης.

Για την επαλήθευση δενδροειδών ακολουθιών στο μοντέλο-στόχο, πρέπει να χρησιμοποιηθούν προσαρμοσμένες μάσκες προσοχής (attention masks). Δυστυχώς, η χρήση μιας προσαρμοσμένης μάσκας προσοχής για ολόκληρη την ακολουθία επιβραδύνει σημαντικά την προσοχή για ρεαλιστικά μήκη εισόδου (έως και 50%), αναιρώνοντας μέρος της επιτάχυνσης που μπορεί να επιτευχθεί μέσω της εικασίας. Δεν έχουμε αναπτύξει ακόμη πλήρη εξερεύνηση δέντρου στην παραγωγή για αυτόν τον λόγο, εστιάζοντας αντίθετα στην ειδική περίπτωση της πρόβλεψης ενός διακριτικού μέσω σχημάτων τύπου MTP που παρουσιάζονται στην τεχνική αναφορά DeepSeek-V3.

MTP

Αυτό το σχήμα είναι παρόμοιο με την αποκωδικοποίηση πρόχειρου-στόχου, με την εξαίρεση ότι οι κρυφές καταστάσεις χρησιμοποιούνται μαζί με τα διακριτικά για την πρόβλεψη. Πρέπει να γίνει ελαφρώς περισσότερη εργασία τόσο στα στάδια προπλήρωσης όσο και αποκωδικοποίησης σε σύγκριση με την κανονική εικασία πρόχειρου-στόχου. Το πρόχειρο μοντέλο χρησιμοποιεί τόσο διακριτικά όσο και κρυφές καταστάσεις: το διακριτικό t_{i+1} λαμβάνει δειγματοληψία από τα logits L_i που αντιστοιχούν στο διακριτικό t_i, τα οποία με τη σειρά τους προέρχονται από τις κρυφές καταστάσεις H_i. Συνεπώς, τα buffer διακριτικών εισόδου πρέπει να μετατοπιστούν κατά ένα βήμα προς τα αριστερά σε σχέση με τα διανύσματα κρυφών καταστάσεων που εξάγονται από τον στόχο. Το παρακάτω σχήμα επισημαίνει τις αντιστοιχίες που χρησιμοποιούνται για την εκπαίδευση, καθώς και τη μετατόπιση κατά την εξαγωγή συμπερασμάτων.

Εικασία αποκωδικοποίησης από διακριτικά και χαρακτηριστικά, που δείχνει το prefill, τη δημιουργία πρόχειρου και την επαλήθευση από το μοντέλο στόχο.

Η ροή αποκωδικοποίησης είναι αρκετά παρόμοια με την αποκωδικοποίηση πρόχειρου-στόχου, με την εξαίρεση ότι μεταφέρονται τόσο οι κρυφές καταστάσεις όσο και τα logits. Η υλοποίησή μας μοιράζεται όλα τα σχετικά logits επεξεργασίας και δειγματοληψίας, εξειδικεύοντας μόνο τις κλήσεις προς τα εμπρός (forward) του μοντέλου. Όταν προβλέπονται πολλαπλά διακριτικά, το πρόχειρο μοντέλο χρησιμοποιεί πρόχειρες κρυφές καταστάσεις για πρόβλεψη, συμπληρώνοντας επίσης καταχωρήσεις cache KV με βάση τα δικά του χαρακτηριστικά. Μακροπρόθεσμα, αυτό μπορεί να υποβαθμίσει την ακρίβεια. Στη συνέχεια, κατά την εκτέλεση του πρόχειρου μοντέλου για τη συμπλήρωση της καταχώρησης cache KV για την πρόβλεψη του στόχου, το εκτελούμε σε ολόκληρη την ακολουθία λαμβάνοντας ως εισόδους τις ακριβέστερες κρυφές καταστάσεις του στόχου. Δεδομένου ότι αυτά τα πρόχειρα μοντέλα είναι μικρά, το πρόσθετο κόστος επεξεργασίας των επιπλέον διακριτικών είναι αμελητέο.

Εκπαίδευση Κεφαλών MTP (Training MTP Heads)

Για να επωφεληθούμε από το MTP, κατασκευάσαμε την υποδομή που απαιτείται για την εκπαίδευση κεφαλών MTP προσαρτημένων στα προσαρμοσμένα μοντέλα μας στα σύνολα δεδομένων του Perplexity, εκτελώντας σε έναν κόμβο με 8 συσκευές H100. Σε περίπου μία ημέρα, μπορούμε να κατασκευάσουμε κεφάλες για μοντέλα που κυμαίνονται από Llama-1B έως Llama-70B και DeepSeek V2-Lite. Για μεγαλύτερα μοντέλα, βασιζόμαστε σε κεφάλες MTP που κατασκευάστηκαν κατά τη διαδικασία προσαρμογής.

Ο στόχος της εκπαίδευσης MTP είναι να ταιριάξει τις πρόχειρες κρυφές καταστάσεις και τα logits που παρεκτείνονται από τις κρυφές καταστάσεις του στόχου με τα logits του επόμενου διακριτικού και τις κρυφές καταστάσεις του στόχου. Δεδομένου ότι η εξαγωγή συμπερασμάτων για κρυφές καταστάσεις είναι δαπανηρή, τις προ υπολογίζουμε χρησιμοποιώντας την υλοποίηση του μοντέλου-στόχου που είναι βελτιστοποιημένη για εξαγωγή συμπερασμάτων, ώστε να χρησιμοποιηθούν κατά την εκπαίδευση. Ωστόσο, για την επικύρωση της υλοποίησης MTP εξαγωγής συμπερασμάτων και τη διασφάλιση ότι οι αριθμητικές διαφορές λόγω ποσόστωσης ή βελτιστοποιήσεων δεν εμποδίζουν τα αποτελέσματα, για την εκτίμηση της απώλειας επικύρωσης και της ακρίβειας επαναχρησιμοποιούμε πλήρως την υλοποίηση εξαγωγής συμπερασμάτων τόσο του στόχου όσο και των πρόχειρων μοντέλων.

Κατά την κλιμάκωση από το σύνολο δεδομένων ShareGPT που χρησιμοποιήθηκε στην αρχική εργασία σε μεγαλύτερα δείγματα, παρατηρήσαμε ότι η αρχιτεκτονική κεφαλής MTP που περιγράφηκε και υλοποιήθηκε στην εργασία EAGLE απέτυχε να εκπαιδευτεί για μοντέλα μεγέθους 70B. Σε αντίθεση με το ShareGPT που περιείχε μεγαλύτερο αριθμό συντομότερων ακολουθιών, εκπαιδεύουμε σε έναν ελαφρώς μικρότερο αριθμό ουσιωδώς μεγαλύτερων προτροπών (prompts). Επειδή οι αρχικές κεφάλες EAGLE απέκλιναν ελαφρώς στη δομή από έναν τυπικό μετασχηματιστή, επαναφέραμε ορισμένα στρώματα κανονικοποίησης RMS (RMS Normalization) που είχαν αφαιρεθεί. Διαπιστώσαμε ότι αυτό όχι μόνο επέτρεψε στην εκπαίδευση να συγκλίνει, αλλά ενίσχυσε επίσης την ακρίβεια των κεφαλών κατά μερικές ποσοστιαίες μονάδες.

Το μοντέλο στόχος, το Eagle Head και το προτεινόμενο head, με δεσμευμένα embeddings και heads γλωσσικού μοντέλου.

Όχι μόνο οι κανόνες στρωμάτων διευκολύνουν την εκπαίδευση, αλλά η επανεισαγωγή των κανόνων είναι επίσης μαθηματικά διαισθητική. Οι κεφαλές MTP επαναχρησιμοποιούν τις ενσωματώσεις και τις προβολές logits του μοντέλου-στόχου, καθώς μπορεί να είναι σημαντικού μεγέθους (περίπου 2 GB για το Llama 70B). Κατά την εκπαίδευση, αυτά παγώνουν και η προσδοκία είναι ότι το στρώμα MTP μαθαίνει να ενσωματώνει προβλέψεις στον ίδιο διανυσματικό χώρο με αυτόν που έμαθε το στρώμα προβολής του αρχικού μοντέλου κατά την εκπαίδευση. Με την παράλειψη των κανόνων, αναμένεται από ένα ενιαίο MLP να μάθει την ίδια συνάρτηση με ένα MLP που ακολουθείται από έναν κανόνα, γεγονός που εμποδίζει τη σύγκριση μεταξύ των κρυφών καταστάσεων του πρόχειρου και των μοντέλων-στόχων.

Εξαγωγή Συμπερασμάτων με Εικαστική Αποκωδικοποίηση

Στη μηχανή εξαγωγής συμπερασμάτων, προκειμένου να παραχθούν διακριτικά για ακολουθίες εισόδου, πρέπει πρώτα να ομαδοποιηθούν σε παρτίδες λογικού μεγέθους, στη συνέχεια πρέπει να δεσμευτούν σελίδες στη μνήμη cache KV για τα επόμενα διακριτικά. Τα διακριτικά εισόδου και οι πληροφορίες σελίδας KV συσκευάζονται στη συνέχεια σε ένα buffer που μεταδίδεται (broadcast) σε όλες τις παράλληλες βαθμίδες που εκτελούν το μοντέλο. Τέλος, τα μεταδεδομένα αντιγράφονται στη μνήμη της GPU και το μοντέλο εκτελείται για να παραχθούν τα logits από τα οποία γίνεται η δειγματοληψία του επόμενου διακριτικού.

Σε αντίθεση με ορισμένες υλοποιήσεις που συνδέουν χαλαρά έναν inference server draft και στόχου μέσω ενός wrapper που ενορχηστρώνει τα αιτήματα μεταξύ τους, τα ζεύγη draft-στόχου μας είναι στενά συζευγμένα και προχωρούν μαζί στη γεννήτρια. Ο προγραμματισμός batch και η κατανομή KV page μοιράζονται μεταξύ των μοντέλων για όλες τις μορφές speculative decoding: αυτό ενοποιεί τη λογική που γεφυρώνει ένα μοντέλο με τον ευρύτερο inference server, καθώς όλα εκθέτουν το ίδιο interface.

Ο χρόνος εκτέλεσης εξαγωγής συμπερασμάτων στο Perplexity είναι διαμορφωμένος γύρω από το FlashInfer, το οποίο καθορίζει τα μεταδεδομένα που πρέπει να κατασκευαστούν προκειμένου να διαμορφωθεί και να προγραμματιστεί ο πυρήνας προσοχής. Δεδομένων ορισμένων ακολουθιών εισόδου που σχηματίζουν μια παρτίδα, για προπλήρωση, αποκωδικοποίηση ή επαλήθευση, πρέπει να γίνει εργασία στην πλευρά της CPU για τη δέσμευση ενδιάμεσων buffers και τη συμπλήρωση ορισμένων σταθερών buffers που χρησιμοποιούνται στην προσοχή. Αυτή η εργασία επιπροσθέτως του κόστους προγραμματισμού παρτίδων και δέσμευσης σελίδων KV, τα οποία επίσης επιφέρουν καθυστερήσεις που πρέπει να κρυφτούν προκειμένου να μεγιστοποιηθεί η αξιοποίηση της GPU.

Ενώ παραλληλίσαμε πλήρως την εργασία στην πλευρά της CPU και της GPU για εξαγωγή συμπερασμάτων χωρίς εικασία, διαπιστώσαμε ότι η ισορροπία CPU-GPU για την εικαστική αποκωδικοποίηση είναι πιο περίπλοκη. Η κύρια πρόκληση προκύπτει από το γεγονός ότι ο αριθμός των αποδεκτών διακριτικών καθορίζει το μήκος της ακολουθίας για μια μεταγενέστερη εκτέλεση, εισάγοντας ένα σημείο συγχρονισμού GPU-CPU που είναι δύσκολο να αποφευχθεί. Πειραματιστήκαμε με διαφορετικά σχήματα προγραμματισμού προκειμένου να κρύψουμε καλύτερα την καθυστέρηση της εργασίας της CPU.

Πρόγραμμα Πρόχειρου-Στόχου (Draft-Target Schedule)

Παρά το γεγονός ότι είναι μικρότερο από ένα μοντέλο-στόχο, όταν χρησιμοποιείται ολόκληρο LLM ως πρόχειρο, εξακολουθεί να εισάγει σημαντική καθυστέρηση στη GPU, παρέχοντας κάποιο περιθώριο για την απόκρυψη δαπανηρών λειτουργιών CPU. Δεδομένου ότι τα μικρότερα μοντέλα δεν επωφελούνται από τον παραλληλισμό τανυστών (tensor parallelism), υπάρχει ανατιστοιχία μεταξύ του αριθμού των βαθμίδων (ranks) στις οποίες κατατμίζεται ένας στόχος και ένα πρόχειρο. Στην υλοποίησή μας, το πρόχειρο μοντέλο εκτελείται μόνο στη βαθμίδα ηγέτη (leader rank) μιας ομάδας TP.

Η CPU, το gloo και η GPU συνεργάζονται στα βήματα εικασίας πρόχειρου-στόχου, συμπεριλαμβανομένης της δημιουργίας πρόχειρου και της αποδοχής.

Όπως επισημάνθηκε προηγουμένως, ένα βήμα αποκωδικοποίησης μεταφέρει logits στην επόμενη εκτέλεση. Αυτό μας επιτρέπει να επικαλύψουμε μία εκτέλεση του πρόχειρου μοντέλου με την εργασία προγραμματισμού παρτίδων στην πλευρά της CPU. Αφού συγκεντρωθεί η παρτίδα, επαναλαμβανόμενες κλήσεις προς τον δειγματολήπτη και το πρόχειρο παράγουν τα πρόχειρα διακριτικά. Παράλληλα, η παρτίδα για επαλήθευση συγκεντρώνεται για το μοντέλο-στόχο και συγχρονίζεται με τους παράλληλους εργαζόμενους. Τα logits του στόχου επαληθεύονται και υποβάλλονται σε δειγματοληψία για να προσδιοριστούν τα αποδεκτά μήκη ακολουθίας. Σε αυτό το σημείο, ο συγχρονισμός GPU-CPU είναι απαραίτητος προκειμένου να προσδιοριστούν τα μεταγενέστερα μήκη ακολουθίας. Δεδομένου ότι το πρόχειρο μοντέλο εκτελείται μόνο στον κόμβο ηγέτη, η παρτίδα του ρυθμίζεται διαδοχικά και η εκτέλεσή του ξεκινά για να συμπληρώσει τις καταχωρήσεις cache KV του με το επιπλέον διακριτικό που παρήγαγε ο στόχος. Τα logits που παράγονται από αυτήν την πρόχειρη εκτέλεση στην τρέχουσα εκτέλεση θα χρησιμοποιηθούν για τη δειγματοληψία του πρώτου πρόχειρου διακριτικού στην επόμενη εκτέλεση. Το πιο σημαντικό, όσο εκτελείται το πρόχειρο, η επόμενη παρτίδα μπορεί να προγραμματιστεί.

Πρόγραμμα MTP για Ένα Διακριτικό

Παρόλο που ο χρόνος εκτέλεσης δεν παρέχει ακόμη εξερεύνηση πρόχειρου δέντρου τύπου Eagle, υλοποιήσαμε μια ειδική περίπτωση αυτού του σχήματος, λαμβάνοντας υπόψη μια γραμμική ακολουθία πρόχειρων διακριτικών που παράγονται από ένα μοντέλο μεγέθους ενός ενιαίου στρώματος αποκωδικοποιητή μετασχηματιστή. Αυτό το σχήμα μπορεί να χρησιμοποιηθεί για πρόβλεψη πρόχειρου χρησιμοποιώντας τα βάρη ανοιχτού κώδικα του DeepSeek R1. Η υποπερίπτωση της πρόβλεψης ενός ενιαίου διακριτικού είναι ενδιαφέρουσα, καθώς τα μεγάλα στρώματα MTP επιτυγχάνουν επαρκώς υψηλά ποσοστά αποδοχής για να δικαιολογήσουν την επιβάρυνσή τους.

Ο προγραμματισμός MTP είναι κάπως πιο σύνθετος, καθώς το draft model είναι πολύ πιο γρήγορο, αποκρύπτοντας λιγότερη καθυστέρηση από την πλευρά της CPU. Επιπλέον, το draft είναι κατατμημένο μαζί με το μοντέλο-στόχο, απαιτώντας κοινές μεταφορές μνήμης για πληροφορίες batch. Μια εκτέλεση ξεκινά με τη μεταφορά των πληροφοριών batch και τη δειγματοληψία του πρώτου token από τα carry-over logits, παρόμοια με το προηγούμενο σχήμα. Στη συνέχεια, εκτελείται το μοντέλο-στόχος για την επικύρωση των tokens, επεξεργάζοντας 2 * D tokens, όπου D είναι το decode batch size. Αυτό είναι ιδανικό για micro-batching σε μοντέλα Mixture-of-Experts (MoE) πάνω από πιο αργές διασυνδέσεις, όπως το InfiniBand, καθώς η batch χωρίζεται ομοιόμορφα σε δύο μισά. Οι κρυφές καταστάσεις του μοντέλου-στόχου μεταφέρονται στην επόμενη εκτέλεση του draft, ενώ τα logits περνούν στον sampler για επαλήθευση.

Ο προγραμματισμός MTP σε CPU, gloo και GPU για επιλογή top-1, εκτέλεση στόχου, αποδοχή και πρόχειρη δημιουργία.

Εκτελώντας μια περιορισμένη ποσότητα επιπλέον εργασίας στη GPU, αποφεύγουμε τον συγχρονισμό CPU-GPU μετά την αποδοχή της πρόχειρης ακολουθίας. Αφού μετατοπιστούν τα διακριτικά εισόδου των στόχων, ένας πυρήνας εισάγει τα επόμενα διακριτικά του στόχου στις αντίστοιχες θέσεις τους. Το πρόχειρο εκτελείται ξανά με τις ίδιες πληροφορίες παρτίδας με τον στόχο, συμπληρώνοντας καταχωρήσεις cache KV και κατασκευάζοντας τα logits και τις κρυφές καταστάσεις για την επόμενη εκτέλεση, κάνοντας κάποια περιττή εργασία σε διακριτικά που δεν έγιναν αποδεκτά. Σε αυτές τις καταστάσεις, η καθυστέρηση της αχρησιμοποίητης εργασίας είναι ελάχιστα μετρήσιμη λόγω του μικρού μεγέθους του πρόχειρου μοντέλου. Παράλληλα με την πρόχειρη εκτέλεση, τα μήκη ακολουθίας προσδιορίζονται στη CPU και ο προγραμματισμός της επόμενης παρτίδας ξεκινά, χωρίς να χρειάζεται να περιμένετε τον τερματισμό της εργασίας της GPU.

Η επιβάρυνση της πρόσθετης εργασίας στο πρόχειρο στρώμα δεν είναι αισθητή στην προσοχή, ωστόσο τα στρώματα MLP είναι πιο προβληματικά. Δεδομένου ότι οι οδηγίες πολλαπλασιασμού πινάκων γεμίζουν (pad) σε όριο 64 κατά μήκος της διάστασης του αριθμού των διακριτικών, εάν ο διπλασιασμός δεν απαιτεί σημαντικά περισσότερα μπλοκ, η επιβάρυνση κρύβεται. Για μεγαλύτερες πρόχειρες ακολουθίες η επιβάρυνση είναι πιο δαπανηρή και το σχήμα που χρησιμοποιείται για κανονικά μοντέλα πρόχειρου-στόχου λειτουργεί καλύτερα.

Παραπομπές

Γρήγορη Εξαγωγή Συμπερασμάτων από Μετασχηματιστές μέσω Εικαστικής Αποκωδικοποίησης

EAGLE: Η Εικαστική Δειγματοληψία Απαιτεί Επανεξέταση της Αβεβαιότητας Χαρακτηριστικών

EAGLE-2: Ταχύτερη Εξαγωγή Συμπερασμάτων Γλωσσικών Μοντέλων με Δυναμικά Πρόχειρα Δέντρα

EAGLE-3: Κλιμάκωση της Επιτάχυνσης Εξαγωγής Συμπερασμάτων Μεγάλων Γλωσσικών Μοντέλων μέσω Δοκιμής Κατά την Εκπαίδευση

Medusa: Απλό Πλαίσιο Επιτάχυνσης Εξαγωγής Συμπερασμάτων LLM με Πολλαπλές Κεφαλές Αποκωδικοποίησης

Κανονικοποίηση Στρώματος (Layer Normalization)

FlashInfer: Αποτελεσματική και Προσαρμόσιμη Μηχανή Προσοχής για την Εξυπηρέτηση Εξαγωγής Συμπερασμάτων LLM

Τεχνική Αναφορά DeepSeek-V3

Παραπομπή

BibTeX
@online{perplexity2025AcceleratingSonarThroughSpeculation,
  author = {{Perplexity} and Nandor Licker and Erin McKnight},
  title = {{Accelerating Sonar through speculation}},
  date = {2025-06-10},
  url = {https://perplexity.coolkeypoint.com/hub/blog/accelerating-sonar-through-speculation}
}