Ενεργοποίηση μοντέλων τρισεκατομμυρίων παραμέτρων στο AWS EFA

Στη Perplexity, χρησιμοποιούμε τα καλύτερα μοντέλα για το προϊόν μας, τα API μας και τις ομάδες έρευνάς μας. Τα μεγάλα μοντέλα ανοιχτού κώδικα Mixture-of-Experts, όπως το Kimi-K2, θέτουν ιδιαίτερες προκλήσεις, καθώς οι μεγαλύτεροι κόμβοι εξαγωγής συμπερασμάτων με 8x GPU NVIDIA H200 δεν μπορούν να τα εξυπηρετήσουν αποδοτικά, καθιστώντας αναγκαία την ανάπτυξη σε πολλαπλούς κόμβους. Παρουσιάζουμε ένα σύνολο πυρήνων για παραλληλισμό εμπειρογνωμόνων που επιτυγχάνουν καθυστερήσεις αιχμής σε ConnectX-7, ξεπερνώντας την απόδοση του DeepEP. Οι ίδιοι πυρήνες είναι επίσης οι πρώτοι που επιτυγχάνουν βιώσιμες καθυστερήσεις στο AWS Elastic Fabric Adapter (EFA), επιτρέποντας την ανάπτυξη μοντέλων τρισεκατομμυρίων παραμέτρων.

Δοκιμάστε τους πυρήνες μας στο GitHub και διαβάστε ολόκληρη την ερευνητική εργασία στο arXiv.

Εισαγωγή

Το Mixture-of-Experts (MoE) έχει αναδειχθεί ως η τυπική αρχιτεκτονική για την κλιμάκωση σε εκατοντάδες δισεκατομμύρια ή τρισεκατομμύρια παραμέτρους, προσφέροντας παράλληλα λογική καθυστέρηση κατά τον χρόνο εξαγωγής συμπερασμάτων. Τα μοντέλα MoE αντικαθιστούν το πυκνό επίπεδο ενός transformer με ένα σύνολο εμπειρογνωμόνων και ένα μικρό γραμμικό επίπεδο δρομολόγησης που αποφασίζει με ποιον εμπειρογνώμονα θα πολλαπλασιαστεί ένα token. Ομαδοποιώντας τα tokens που σχετίζονται με τον ίδιο εμπειρογνώμονα, το πρόβλημα γίνεται αραιό και εύκολα παραλληλοποιήσιμο, καθώς διαφορετικοί εμπειρογνώμονες μπορούν πλέον να φιλοξενούνται σε διαφορετικές GPU που κατανέμονται σε πολλαπλούς κόμβους.

Σε αντίθεση με άλλες μορφές παραλληλισμού, όπως ο Tensor Parallelism (TP) και ο Data Parallelism (DP), που μπορούν να υλοποιηθούν εύκολα χρησιμοποιώντας βιβλιοθήκες συλλογικής επικοινωνίας όπως οι torch.distributed και NCCL, η δρομολόγηση MoE περιλαμβάνει αραιή επικοινωνία peer-to-peer. Ένας πυρήνας αποστολής πρέπει να διασπά το σύνολο των εισερχόμενων tokens και να τα αποστέλλει στους rank που φιλοξενούν τους experts, ενώ το αντίστροφό του πρόβλημα, ο combine kernel, πρέπει να παραλαμβάνει τα επεξεργασμένα αντίστοιχά τους και να υπολογίζει τον σταθμισμένο μέσο όρο τους. Αυτές οι μεταφορές δεν αντιστοιχούν σε συλλογικές λειτουργίες και απαιτούν προσαρμοσμένους πυρήνες για την επίτευξη χαμηλής καθυστέρησης.

Στη Perplexity, αναπτύσσουμε τόσο φορητούς πυρήνες μεταξύ κόμβων όσο και εξειδικευμένους πυρήνες εντός του κόμβου, παρουσιάζοντας την εργασία μας με μεγαλύτερη λεπτομέρεια σε ένα προηγούμενο άρθρο σχετικά με την επικοινωνία MoE. Δείξαμε επίσης σε μια άλλη ανάρτηση ιστολογίου ότι η ανάπτυξη σε πολλαπλούς κόμβους μπορεί να προσφέρει χαμηλότερη καθυστέρηση και υψηλότερη απόδοση για μεγάλα μοντέλα MoE.

Σε αυτό το άρθρο, εμβαθύνουμε στην υλοποίηση των νέων μας πυρήνων μεταξύ κόμβων, οι οποίοι βελτιώνουν σημαντικά την απόδοση της προηγούμενης εργασίας μας που βασίζεται στο NVSHMEM. Εκτός από την παρουσίαση αποτελεσμάτων αιχμής στο ConnectX-7, αποδεικνύουμε επίσης ότι εντός του χρόνου εκτέλεσης εξαγωγής συμπερασμάτων, οι πυρήνες μπορούν να χρησιμοποιηθούν για την αποτελεσματική εξυπηρέτηση των μεγαλύτερων μοντέλων ανοιχτού κώδικα.

Εξαγωγή συμπερασμάτων μέσω InfiniBand

Ο παραλληλισμός τανυστών και ο παραλληλισμός εμπειρογνωμόνων μπορούν να κλιμακωθούν αποτελεσματικά εντός ενός μόνο κόμβου, καθώς το NVLink προσφέρει υψηλή απόδοση (900GB/s) με καθυστέρηση λίγων μικροδευτερολέπτων. Παρόλο που τα συστήματα NVIDIA GB200 και GB300 είναι εξοπλισμένα με επαρκή μνήμη για να χωρέσουν τα βάρη και τις κρυφές μνήμες KV μοντέλων τρισεκατομμυρίων παραμέτρων, οι αναπτύξεις σε παρουσία AWS p5en με έως και 8 GPU H200 περιορίζονται στα 1120GB HBM, τα οποία πρέπει να κατανεμηθούν μεταξύ των βαρών του μοντέλου και των κρυφών μνημών KV. Κατά συνέπεια, ορισμένα μοντέλα απαιτούν αναπτύξεις μεταξύ κόμβων μέσω InfiniBand.

Αν και το InfiniBand εξακολουθεί να είναι εξαιρετικά γρήγορο, προσφέρει απόδοση μόνο 400Gbps (50GB/s), προσθέτοντας δεκάδες έως εκατοντάδες μικροδευτερότερα καθυστέρησης για τη δρομολόγηση MoE. Ευτυχώς, οι πρόσθετες καθυστερήσεις μπορούν να κρυφτούν μέσω της χρήσης κοινόχρηστων εμπειρογνωμόνων, micro-batching και επικάλυψης υπολογισμού-επικοινωνίας για την παροχή βιώσιμης απόδοσης σε ανταγωνιστικό κόστος. Το DeepSeek πρωτοπόρησε σε αυτές τις τεχνικές μέσω του DeepSeek-R1, το οποίο αναπτύχθηκε σε συστάδες GPU H800 και προσαρμογείς NVIDIA ConnectX-7 InfiniBand.

Στις συστάδες AWS, η κατάσταση είναι διαφορετική λόγω της χρήσης των προσαρμοσμένων Elastic Fabric Adapters (EFA). Παρόλο που αυτοί προσφέρουν μέγιστη απόδοση 400Gbps στις συλλογικές λειτουργίες που χρησιμοποιούνται συνήθως σε φόρτους εργασίας εκπαίδευσης, υστερούν ελαφρώς έναντι του ConnectX-7 στα μεγέθη μηνυμάτων που ανταλλάσσονται κατά την αποστολή και τον συνδυασμό MoE. Επιπλέον, το EFA δεν υποστηρίζει GPUDirect Async, απαιτώντας ένα νήμα μεσολάβησης CPU (CPU proxy thread) που γεφυρώνει τη GPU με τη NIC προκειμένου να ξεκινήσουν οι μεταφορές. Εκτός από την αυξημένη πολυπλοκότητα, οι πρόσθετες συναλλαγές PCIe προσθέτουν υπερκεφαλίδα αξίας μικροδευτερολέπτων σε όλες τις μεταφορές. Πλαίσια όπως το NVSHMEM επιτρέπουν στους συγγραφείς πυρήνων να παρέχουν υλοποιήσεις ανεξάρτητες από τη συσκευή, ωστόσο τα γενικά API συνεπάγονται υπερκεφαλίδα που καθιστά τις αναπτύξεις αναποτελεσματικές τόσο από άποψη κόστους όσο και από άποψη καθυστέρησης.

Pplx-kernels βασισμένοι σε NVSHMEM

Οι προηγούμενοι πυρήνες μας αφαιρούσαν την υλοποίηση του υποκείμενου transport μέσω του NVSHMEM. Βασίζονταν στο nvshmemx_putmem_signal_nbi_warp για τη μεταφορά tokens ένα προς ένα, χρησιμοποιώντας ατομικούς μετρητές για τον συγχρονισμό των peers. Παρόλο που αυτό πρόσφερε ικανοποιητική απόδοση σε NIC ConnectX-7 με IBGDA, η υλοποίηση IBRC του ConnectX-7 με proxy ήταν σημαντικά πιο αργή, ενώ η υλοποίηση EFA δεν κατάφερε να προσφέρει καθυστερήσεις κάτω από ένα millisecond.

Χρησιμοποιήσαμε τους πυρήνες ως πεδ δοκιμών για το υποκείμενο υλικό, εντοπίζοντας παράλληλα την πιο κατάλληλη διεπαφή για έκθεση στα μοντέλα μας. Υποθέσαμε ότι οι συχνές μεταφορές και αλληλεπιδράσεις μεταξύ των πυρήνων και των NICs ή των νήματος μεσολάβησης (proxy threads) εμπόδιζαν την απόδοση. Επιπλέον, η υλοποίηση μας δεν βασιζόταν σε καμία εγγύηση από το υλικό εκτός από την αξιοπιστία, γεγονός που υποδηλώνει ότι θα μπορούσαμε να επιταχύνουμε τα πράγματα μειώνοντας την υπερκεφαλίδα των πρωτοκόλλων συγχρονισμού.

Με βάση αυτές τις παρατηρήσεις, σχεδιάσαμε τους πυρήνες μας με βάση το EFA, εκθέτοντας την ίδια διεπαφή και βελτιστοποιώντας στο παρασκήνιο τις αλληλεπιδράσεις του device kernel με τον host proxy για μέγιστη απόδοση.

MoE all-to-all over EFA

Μεγέθυνση

Οι νέοι πυρήνες ακολουθούν μια υβριδική αρχιτεκτονική CPU-GPU, όπου οι GPU kernels αλληλεπιδρούν με το μοντέλο στη συσκευή και ένα νήμα proxy που εκτελείται στον host CPU χειρίζεται τις αλληλεπιδράσεις με το NIC. Οι νέοι πυρήνες επαναχρησιμοποιούν το ίδιο TransferEngine που αναπτύξαμε για τις μεταφορές KV cache. 

Η δρομολόγηση MoE εκτελείται από έναν πυρήνα αποστολής (dispatch kernel), ο οποίος λαμβάνει τα tokens και τους εμπειρογνώμονες ανά token για να τα δρομολογήσει από το μοντέλο, αποστέλλοντάς τα μέσω του δικτύου στις τάξεις που φιλοξενούν τους εμπειρογνώμονες. Αφού οι ομαδοποιημένοι πυρήνες GEMM πολλαπλασιάσουν τα tokens με τα βάρη των εμπειρογνωμόνων, ο πυρήνας συνδυασμού μεταφέρει τα tokens στις αρχικές τους τάξεις και υπολογίζει τον σταθμισμένο μέσο όρο τους πριν τα επιστρέψει στο μοντέλο.

Τόσο ο πυρήνας αποστολής όσο και ο πυρήνας συνδυασμού χωρίζονται σε ένα μισό αποστολέα και ένα μισό δέκτη για να διευκολύνουν την επικάλυψη και το micro-batching. Ο αποστολέας ανακατεύει τα tokens σε buffers αποστολής, ώστε να μπορούν να εγγραφούν σε ομότιμους (peers) μέσω μιας ενιαίας λειτουργίας εγγραφής RDMA. Μετά την εκτέλεση του αποστολέα, η GPU μπορεί να εκτελέσει άλλες εργασίες (όπως πολλαπλασιασμό κοινόχρηστου εμπειρογνώμονα) ενώ οι μεταφορές μέσω του δικτύου εκκρεμούν στο παρασκήνιο. Οι πυρήνες δέκτη μπλοκάρουν μέχρι να ληφθούν όλες οι μεταφορές. Στο παρασκήνιο, ένα νήμα εργασίας στην πλευρά του κεντρικού υπολογιστή (host-side worker thread) χειρίζεται τις μεταφορές: ειδοποιείται από τη GPU όταν τα buffers αποστολής έχουν γεμίσει, μετά από το οποίο αποστέλλει τις εγγραφές RDMA. Μόλις η NIC σηματοδοτήσει την ολοκλήρωση, ξεμπλοκάρει τον δέκτη, ο οποίος επιστρέφει τα tokens στο μοντέλο.

Οι πυρήνες είναι στενά συνδεδεμένοι με τον κεντρικό υπολογιστή, επικοινωνώντας μεταξύ τους τόσο μέσω ενοποιημένης μνήμης (unified memory) όσο και μέσω του GDRCopy. Η ενοποιημένη μνήμη CUDA επιτρέπει στη μνήμη της συσκευής ή τη μνήμη του κεντρικού υπολογιστή να χαρτογραφηθούν στους χώρους διευθύνσεων και των δύο, χειριζόμενη διαφανώς την επικοινωνία μέσω PCIe.

Παρότι η ενοποιημένη μνήμη είναι κατάλληλη για μαζικές μεταφορές, η polling χαμηλής καθυστέρησης γίνεται μέσω του GDRCopy. Το GDRCopy χρησιμοποιεί το GPUDirect RDMA για να διαβάζει από και να γράφει στη μνήμη GPU με καθυστέρηση της τάξης των μικροδευτερολέπτων.

Εκτός από το RDMA, εντός ενός κόμβου οι πυρήνες αξιοποιούν επίσης το NVLink για την ανταλλαγή tokens. Αυτό μπορεί να αποφορτίσει ένα σημαντικό κλάσμα της κυκλοφορίας, έως και το ένα τέταρτο για το EP32 και έως και το ένα όγδοο για το EP64, επιτυγχάνοντας συνολικά χαμηλότερες καθυστερήσεις.

Buffers αποστολής και λήψης (Send and Receive Buffers)

Ο σχεδιασμός των buffers αποστολής για τον πυρήνα αποστολής είναι κρίσιμος για την επίτευξη βιώσιμων καθυστερήσεων που περιλαμβάνουν τον μεσολαβητή κεντρικού υπολογιστή (host proxy). Εάν υπάρχουν N τάξεις, καθένας από τους οποίους αποστέλλει T tokens, καθένα δρομολογημένο σε R εμπειρογνώμονες από το σύνολο E, μια τάξη μπορεί να λάβει έως και N * T * max(R, E / N) tokens στη χειρότερη δυνατή περίπτωση, εάν όλα δρομολογηθούν σε αυτήν. Επιπλέον, αυτό το ελάχιστο όριο μπορεί να επιτευχθεί μόνο όταν οι αποστολείς τοποθετούν τις εγγραφές τους διαδοχικά, καθώς τα ιδιωτικά buffers λήψης θα κατανάλωναν πάρα πολύ μνήμη. Για να μπορούν οι αποστολείς να τοποθετούν τις εγγραφές τους διαδοχικά χωρίς συγκρούσεις, η πλήρης γνώση των πληροφοριών δρομολόγησης πρέπει να κατανέμεται σε κάθε τάξη. Κατά συνέπεια, οι πυρήνες αποστολής μας ανταλλάσσουν πρώτα τις μετρήσεις tokens ανά εμπειρογνώμονα, μετά από τις οποίες μπορούν καθένας να προσδιορίσει πού πηγαίνουν οι εγγραφές τους και από ποιο offset λαμβάνουν tokens από κάθε τάξη.

Η απλή ανταλλαγή πληροφοριών δρομολόγησης θα είχε ως αποτέλεσμα αδικαιολόγητα υψηλές καθυστερήσεις, καθώς κάθε τάξη θα έπρεπε να περιμένει να φτάσουν οι εγγραφές πριν αποστείλει οποιεσδήποτε δικές του εγγραφές. Αντίθετα, δεσμεύουμε μια μικρή ποσότητα ιδιωτικού χώρου ανά αποστολέα σε κάθε τάξη δέκτη και αποστέλλουμε μερικά από τα tokens μαζί με τις εγγραφές των πληροφοριών δρομολόγησης. Αυτό εξασφαλίζει πλήρη αξιοποίηση του εύρους ζώνης καθ' όλη τη διάρκεια της διαδικασίας αποστολής. Αφού ανταλλαγούν οι διαδρομές, τα υπόλοιπα tokens μεταφέρονται μέσω μιας ενιαίας εγγραφής ανά απομακρυσμένο ομότιμο (peer). Η ανταλλαγή πληροφοριών δρομολόγησης και τα buffers σταθερού μεγέθους, μαζί με τις μεταφορές που τα γεμίζουν, απεικονίζονται στο παρακάτω σχήμα.

Μεγέθυνση

Οι πυρήνες συνδυασμού είναι ουσιαστικά απλούστεροι, καθώς οι πληροφορίες δρομολόγησης είναι ήδη διαθέσιμες σε αυτούς από το στάδιο της αποστολής. Μπορούν να αντιγράψουν τα tokens με μία μόνο εγγραφή ανά ομότιμο, μετά την οποία διαβάζουν από το buffer λήψης για να υπολογίσουν τον σταθμισμένο μέσο όρο των tokens στους τανυστές εξόδου.

Για το NVLink, τα buffers αποστολής και λήψης εντός ενός κόμβου χαρτογραφούνται στους ομότιμους, ώστε να μπορούν να ωθήσουν ή να τραβήξουν απευθείας tokens από τις γειτονικές τάξεις.

Αποστολή και συνδυασμός (Dispatch and Combine)

Τα παρακάτω σχήματα λεπτομερώς τον χρόνο και τις αλληλεπιδράσεις μεταξύ της GPU, της CPU και της NIC, δείχνοντας μία επανάληψη αποστολής και συνδυασμού σε εναλλαγή (interleaved) με τους πυρήνες του μοντέλου.

Η εργασία ξεκινά στην αποστολή (dispatch send), με τον πυρήνα να συγκεντρώνει τις πληροφορίες δρομολόγησης και να υπολογίζει τις μετρήσεις tokens ανά εμπειρογνώμονα μέσω ατομικής συσσώρευσης σε κοινόχρηστη μνήμη εντός ενός μόνο μπλοκ. Μόλις οι μετρήσεις tokens είναι διαθέσιμες, αντιγράφονται στη μνήμη του κεντρικού υπολογιστή και ο μεσολαβητής (proxy) ειδοποιείται αμέσως να αρχίσει να αποστέλλει τις πληροφορίες δρομολόγησης σε όλες τις άλλες τάξεις. Παράλληλα, η αποστολή συνεχίζει να υπολογίζει το αθροιστικό άθροισμα των μετρήσεων tokens, προσδιορίζοντας τα offsets που απαιτούνται για τη συσκευασία tokens σε συνεχόμενα buffers αποστολής ανά ομότιμο. Μόλις τα offsets είναι διαθέσιμα, όλα τα tokens που πρόκειται να αποσταλούν μέσω EFA συσκευάζονται στα buffers αποστολής. Ο μεσολαβητής ειδοποιείται να δημοσιεύσει τις εγγραφές στα buffers αποστολής των ομότιμων μεταξύ κόμβων. Εντός του κόμβου, οι πυρήνες ωθούν επίσης αμέσως tokens σε ιδιωτικά buffers εντός του ίδιου κόμβου μέσω NVLink.

Μεγέθυνση

Μεγέθυνση

Μόλις μια τάξη λάβει πληροφορίες δρομολόγησης από όλους τους ομότιμους, τις κεντρικοποιεί και υπολογίζει διάφορα offsets που είναι απαραίτητα για τον εντοπισμό tokens για αποστολή και συνδυασμό. Αυτή η εργασία γίνεται στην CPU ενώ οι μεταφορές εκκρεμούν ήδη, επομένως δεν είναι κρίσιμη για την καθυστέρηση. Στη συνέχεια, μόλις γίνουν γνωστά όλα τα offsets, οι εγγραφές για το υπόλοιπο των tokens, τα οποία υπερβαίνουν το μέγεθος του ιδιωτικού buffer, δημοσιεύονται και η φάση αποστολής ολοκληρώνεται.

Στην πλευρά του παραλήπτη, ο πυρήνας πρώτα αντλεί tokens από τον ίδιο κόμβο που υπερβαίνουν τη χωρητικότητα των ιδιωτικών buffers μέσω NVLink. Στη συνέχεια, μπλοκάρει μέχρι το host proxy να σηματοδοτήσει την ολοκλήρωση όλων των μεταφορών. Μόλις ληφθεί το σήμα, τα tokens μετακινούνται από τα buffers λήψης στο output tensor που περνά στο group GEMM. Εφόσον τα buffers επαναχρησιμοποιούνται μεταξύ dispatch και combine, τόσο τα NVLink όσο και τα RDMA barriers διασφαλίζουν ότι τα buffers του παραλήπτη έχουν διαβαστεί πριν επιτραπεί στον επόμενο πυρήνα να τα αντικαταστήσει.

Οι πυρήνες συνδυασμού ακολουθούν ένα περίπου παρόμοιο μοτίβο, ανακατεύοντας tokens, ειδοποιώντας τον μεσολαβητή και περιμένοντας την ολοκλήρωση. Ωστόσο, τα μεταδεδομένα για τη δημοσίευση των εγγραφών τους είναι άμεσα διαθέσιμα και όλα τα απομακρυσμένα offsets είναι γνωστά, επομένως όλα τα tokens που μεταφέρονται εντός του κόμβου ωθούνται μέσω NVLink στους δέκτες και τα tokens μεταξύ κόμβων μεταφέρονται μέσω μιας ενιαίας εγγραφής. Πριν επιστρέψει τα αποτελέσματα, ο πυρήνας συσσωρεύει τα tokens από τα buffers δέκτη, γράφοντάς τα στους τανυστές εξόδου.

TransferEngine για δρομολόγηση MoE

Για να δημιουργήσουμε το πρωτότυπο της υλοποίησής μας, επαναχρησιμοποιήσαμε αρχικά το Transfer Engine που κατασκευάστηκε για αποσυγκεντρωμένο prefill και decode προκειμένου να εκδώσουμε και να παρακολουθήσουμε εγγραφές RDMA. Ωστόσο, σε αντίθεση με την αποσυγκεντρωμένη αποκωδικοποίηση όπου κάθε κόμβος δημοσιεύει μεγάλο αριθμό εγγραφών στον ίδιο ομότιμο, η αποστολή MoE δημοσιεύει μικρές μεταφορές σε όλους τους ομότιμους σε μία λειτουργία. Η απόδοση εκτός κουτιού δεν ήταν η βέλτιστη και η υλοποίηση μας υστερούσε σε σχέση με το DeepEP στο ConnectX-7 κατά περίπου 20us. Κατά συνέπεια, εξειδικεύσαμε το TransferEngine ώστε να υποστηρίζει δύο λειτουργίες, scatter και barrier, για τη βελτιστοποίηση των αλληλεπιδράσεων με μεγάλο αριθμό ομότιμων.

Τόσο οι λειτουργίες scatter όσο και οι λειτουργίες barrier εκθέτουν ομάδες ομότιμων που μπορούν να προ-καταχωριστούν για την προ-συμπλήρωση των δομών δεδομένων που χρησιμοποιούνται από την υποκείμενη υλοποίηση μεταφοράς (πρότυπα Work Request (WR)). Το Scatter αντιγράφει διαφορετικά τμήματα από το ίδιο buffer προέλευσης σε διαφορετικά offsets σε όλα τα μέλη μιας ομάδας ομότιμων. Το Barrier ειδοποιεί όλους τους ομότιμους μέσω ενός άμεσου (immediate). Η τοποθέτηση στην ουρά και ο χειρισμός αυτών των λειτουργιών μαζικά εξοικονομεί πολύτιμα μικροδευτερότερα στην κρίσιμη διαδρομή. Στο EFA, όπου δύο NICs 200Gbps αθροίζονται για να παρέχουν εύρος ζώνης 400Gbps, κατανέμουμε τους ομότιμους της ομάδας στις δύο NICs αντί να κατανέμουμε τα byte κάθε μεταφοράς.

Παρόλο που αρχικά αναπτύξαμε το TransferEngine αποκλειστικά στο EFA χρησιμοποιώντας libfabric, οι πυρήνες MoE λειτουργούν στο ConnectX-7 χωρίς να αλλάξουν ούτε μία γραμμή κώδικα μετά την προσθήκη υποστήριξης ConnectX-7 στο TransferEngine χρησιμοποιώντας libibverbs. Σε σύγκριση με το EFA που βασίζεται σε Scalable Reliable Datagram (SRD), το ConnectX-7 απαιτεί επιπλέον εργασία για τη ρύθμιση σύνδεσης και τη διαχείριση ομότιμων. Χρησιμοποιούμε ένα ζεύγος ουρών Unreliable Datagram (UD) για την εκκίνηση του Reliable Connection (RC). Δημιουργούμε δύο ζεύγη ουρών για κάθε απομακρυσμένο ομότιμο, ένα για SEND/RECV δύο κατευθύνσεων και ένα για WRITE_IMM μίας κατεύθυνσης, επειδή τόσο το RECV όσο και το WRITE_IMM καταναλώνουν λειτουργία RECV με τη σειρά, ενώ οι δύο λειτουργίες απαιτούν διαφορετικά μεγέθη buffer. Επιπλέον, χαλαρώνουμε τη σειρά συναλλαγών PCIe μεταξύ του ConnectX-7 και της μνήμης GPU. Παρόμοια με το EFA, εφαρμόζουμε τη βελτιστοποίηση προτύπων WR. Επιπλέον, κατά την αποστολή πολλαπλών μικρών μηνυμάτων στον ίδιο ομότιμο, χρησιμοποιούμε αλύσωση WR (WR chaining) για τη βελτίωση του ρυθμού μηνυμάτων.

Δεδομένου ότι η εργασία στην πλευρά της GPU είναι ίδια και στους δύο προσαρμογείς, ωθήσαμε την υποκείμενη υλοποίηση μεταφοράς στα όριά της για να ξεπεράσουμε την απόδοση του DeepEP, ώστε να τονίσουμε ότι το χάσμα μεταξύ EFA και δρομολόγησης MoE ConnectX-7 καθορίζεται από τις δυνατότητες του υποκείμενου υλικού.

Αξιολόγηση

Καθυστέρηση αποστολής/συνδυασμού (Dispatch/Combine Latency)

Αξιολογήσαμε τις καθυστερήσεις αποστολής και συνδυασμού σε συστήματα H200 εξοπλισμένα με προσαρμογείς ConnectX-7 και EFA στη τυπική διαμόρφωση DeepSeek-V3 και Kimi-K2 με κρυφή διάσταση 7168, αποστολή fp8 κλιμακούμενη σε επίπεδο μπλοκ και συνδυασμό bf16. Λαμβάνουμε υπόψη μέγεθος παρτίδας (batch size) 128 για την αποστολή και 4096 για τον συνδυασμό. Στο ConnectX-7, συγκρίνουμε με το DeepEP και τους προηγούμενους pplx-kernels βασισμένους σε NVSHMEM, χρησιμοποιώντας τόσο GPUDirect Async (IBGDA) όσο και έναν μεσολαβητή κεντρικού υπολογιστή (IBRC). Επιτυγχάνουμε συνδυασμένες καθυστερήσεις 459us, 582us και 692us σε EP16, EP32 και EP64, αντίστοιχα. Ξεπερνάμε άλλους πυρήνες βασισμένους σε EFA που βασίζονται σε γενικούς μεσολαβητές, όπως το UCCL-EP, οι οποίοι επιτυγχάνουν 519us, 966us και 1159us.

Μεγέθυνση

Στην αποκωδικοποίηση, παρά το γεγονός ότι οι πυρήνες μας είναι μερικά μικροδευτερότερα πιο αργοί από το DeepEP, επιτυγχάνουμε συνολικά χαμηλότερες καθυστερήσεις καθώς οι πυρήνες συνδυασμού μας είναι ουσιαστικά ταχύτεροι. Αν και χρησιμοποιούμε έναν μεσολαβητή κεντρικού υπολογιστή, λόγω των μαζικών μεταφορών μας ξεπερνάμε την απόδοση όλων των υλοποιήσεων χωρίς μεσολαβητή που βασίζονται σε IBGDA. Δείχνουμε επίσης ότι η υλοποίηση του μεσολαβητή έχει σημασία, καθώς ο εξαιρετικά εξειδικευμένος μεσολαβητής μας είναι κατά τάξεις μεγέθους ταχύτερος από τη γενική υλοποίηση NVSHMEM. Παρά το γεγονός ότι το EFA είναι διπλάσιο αργό από το ConnectX-7 κατά τη μεταφορά πακέτων 256KB, τα οποία είναι κοντά στο τυπικό μέγεθος πακέτου που ανταλλάσσεται από τη δρομολόγηση MoE, το χάσμα μεταξύ των δύο προσαρμογέων είναι συνολικά λιγότερο ουσιαστικό.

Μεγέθυνση

Παρόλο που οι βελτιστοποιημένοι για αποκωδικοποίηση πυρήνες μας μπορούν να κλιμακωθούν σε μεγέθη παρτίδων prefill χωρίς καμία επιπλέον ρύθμιση, η απόδοση υστερεί σε σχέση με το DeepEP. Αυτό οφείλεται κυρίως στο γεγονός ότι συσσωρεύουμε tokens μόνο στην πλευρά του δέκτη του συνδυασμού, ενώ το DeepEP συσσωρεύει στον αποστολέα για να μειώσει την ποσότητα των δεδομένων που μεταφέρονται. Ωστόσο, αυτή η βελτιστοποίηση μπορεί να έχει επιπτώσεις στην ακρίβεια. Επιπλέον, δεδομένου ότι οι πυρήνες μεταφέρουν δεδομένα σε μία εγγραφή, αξιοποιούμε περισσότερη μνήμη για τα buffers αποστολής και λήψης.

Απόδοση Kimi-K2 / DeepSeek-V3 (Throughput)

Για να αποδείξουμε τη βιωσιμότητα των πυρήνων μας, κάνουμε benchmark στην απόδοση end-to-end της εσωτερικής μας μηχανής εξαγωγής συμπερασμάτων που εξυπηρετεί τα DeepSeek-V3 (671B παράμετροι) και Kimi-K2 (1T παράμετροι) σε πολλαπλές παρουσίες p5en εξοπλισμένες με Elastic Fabric Adapters (EFA) και GPU H200. Εκτελούμε το μοντέλο σε λειτουργία επαλήθευσης πρόχειρου (draft verification mode) και αναφέρουμε την καθυστέρηση υποθέτοντας ποσοστό αποδοχής 80% χρησιμοποιώντας ένα ενιαίο επίπεδο MTP.

Στο DeepSeek-V3, συγκρίνουμε με ένα baseline εντός κόμβου EP=8 που εκτελεί αποδοτικούς πυρήνες dispatch και combine μόνο μέσω NVLink. Οι αναπτύξεις μεταξύ κόμβων συνολικά είτε ισοφαρίζουν είτε ξεπερνούν την απόδοση της μονοκόμβιας ανάπτυξης, ιδιαίτερα σε μεσαίου μεγέθους batches, δείχνοντας ότι η κλιμάκωση σε πολλαπλούς κόμβους πράγματι προσφέρει βελτιώσεις στη συνολική throughput.

Στο Kimi-K2, οι πυρήνες καθιστούν δυνατή την εξυπηρέτηση. Λόγω του μεγέθους του, το μοντέλο δεν μπορεί να χωρέσει στον ίδιο κόμβο χρησιμοποιώντας Data Parallelism (DP) σε έναν κόμβο H200, επομένως απαιτείται ανάπτυξη μεταξύ κόμβων για την εξυπηρέτηση του μοντέλου. Καθώς το μοντέλο έχει λιγότερες κεφαλές προσοχής (attention heads), πετυχαίνουμε καλύτερες καθυστερήσεις από το DeepSeek-V3 και μπορούμε να εξυπηρετήσουμε το μοντέλο στο EFA με βιώσιμη καθυστέρηση.

Μεγέθυνση

Μεγέθυνση

Έχουμε επίσης εξετάσει το micro-batching, αν και αυτό βελτίωσε τις καθυστερήσεις μόνο σε πολύ υψηλά batch sizes ή όταν χρησιμοποιούνται περισσότερα από ένα draft tokens. Για να επωφεληθεί το splitting των grouped GEMM problems σε δύο, τα token batch sizes πρέπει να είναι επαρκώς πυκνά. Το μεγαλύτερο μέρος του χρόνου εξακολουθεί να δαπανάται στο MoE dispatch και combine, γεγονός που δείχνει ότι οποιεσδήποτε περαιτέρω βελτιώσεις στους πυρήνες θα βελτιώσουν συνεπώς τις καθυστερήσεις και θα διευκολύνουν το micro-batching.

Μελλοντική εργασία

Εργαζόμαστε στενά με τους μηχανικούς της AWS για τη βελτίωση της απόδοσης στο EFA. Από την πλευρά μας, θα ακολουθήσουμε τις ενημερώσεις του libfabric οι οποίες θα μειώσουν την υπερκεφαλίδα του επιπέδου δεδομένων. Επιπλέον, θα πειραματιστούμε με το efa-direct για να μειώσουμε περαιτέρω τις υπερκεφαλίδες του χώρου χρήστη (userspace) και να διασυνδεθούμε απευθείας με το υλικό.

Αναφορές

Τεχνική αναφορά DeepSeek R1

Τεχνική αναφορά Kimi K2

DeepEP

UCCL-EP

Βελτίωση της απόδοσης δικτύου συστημάτων HPC με χρήση των NVIDIA Magnum IO NVSHMEM και GPUDirect Async

Παραπομπή

@online{perplexity2025EnablingTrillionParameterModelsOnAwsEfa,
  author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
  title = {{Enabling trillion-parameter models on AWS EFA}},
  date = {2025-11-04},
  url = {https://www.perplexity.ai/hub/blog/enabling-trillion-parameter-models-on-aws-efa}
}