JAXBench: πώς η AI μαθαίνει να βελτιστοποιεί kernels για TPU

Η αυτόματη παραγωγή κώδικα από μοντέλα τεχνητής νοημοσύνης έχει προχωρήσει εντυπωσιακά, όμως η δημιουργία ενός σωστού και πραγματικά γρήγορου kernel για εξειδικευμένο hardware παραμένει διαφορετική κατηγορία προβλήματος. Το JAXBench προτείνει ένα ανοικτό, TPU-native benchmark που εξετάζει ακριβώς αυτή τη δυσκολία: αν ένας AI agent μπορεί όχι μόνο να γράψει κώδικα Pallas που μεταγλωττίζεται, αλλά και να παράγει αριθμητικά σωστά αποτελέσματα και να ξεπερνά σε ταχύτητα το ισχυρό baseline του XLA.

Για επαγγελματίες που σχεδιάζουν AI προϊόντα, inference υπηρεσίες ή υποδομές δεδομένων, η εργασία έχει ιδιαίτερη αξία. Δείχνει ότι η απόδοση δεν εξαρτάται απλώς από το μέγεθος του μοντέλου. Η εξειδικευμένη τεκμηρίωση του στόχου, η σωστή δομή αναζήτησης και η αυστηρή μέτρηση στο ίδιο το accelerator μπορούν να αλλάξουν ουσιαστικά το αποτέλεσμα.

Τι μετρά πραγματικά το JAXBench

Το JAXBench περιλαμβάνει 50 workloads σε JAX, επιλεγμένα ώστε να είναι συναφή με σύγχρονες εφαρμογές και να αφήνουν πραγματικό περιθώριο βελτιστοποίησης. Τα 17 προέρχονται από production operators της δημόσιας βιβλιοθήκης MaxText και καλύπτουν αρχιτεκτονικές όπως Llama-3.1, DeepSeek-V3, Mixtral, Mamba-2 και AlphaFold2. Τα άλλα 33 είναι fused operators του KernelBench Level 2, μεταφρασμένα από PyTorch σε ισοδύναμες υλοποιήσεις JAX.

Η επιλογή αυτή αποφεύγει ένα συνηθισμένο πρόβλημα των benchmarks: μικρά ή τεχνητά workloads στα οποία ο χρόνος εκκίνησης και η μεταφορά μνήμης κρύβουν το πραγματικό κόστος του υπολογισμού. Οι ερευνητές ρύθμισαν ξεχωριστά τις διαστάσεις κάθε workload ώστε τα compute-heavy προβλήματα να αξιοποιούν έντονα τις Matrix Multiply Units της TPU v6e. Έτσι, ένα speedup έχει μεγαλύτερη πιθανότητα να προέρχεται από ουσιαστική επιλογή tiling, pipeline ή layout και όχι από θόρυβο μέτρησης.

Γιατί η βελτιστοποίηση TPU δεν είναι μεταφορά τεχνογνωσίας από GPU

Οι TPU διαφέρουν αρχιτεκτονικά από το μοντέλο μαζικού SIMT παραλληλισμού των GPU. Στην TPU v6e υπάρχουν ευρείς SIMD vector registers και systolic MXUs 256×256, ενώ ο προγραμματισμός χαμηλού επιπέδου περνά από το Pallas και το Mosaic backend. Ο kernel author πρέπει να λάβει υπόψη τις ιεραρχίες VMEM, SMEM και HBM, το software pipelining, τα block shapes και τη λεξικογραφική σειρά διάσχισης του grid.

Αυτό έχει άμεση συνέπεια για τα γλωσσικά μοντέλα. Το Pallas εμφανίζεται πολύ λιγότερο συχνά στα training corpora από το CUDA ή το Triton. Ένα μοντέλο μπορεί να είναι γενικά ικανό στον κώδικα αλλά να επινοεί ανύπαρκτα APIs, να δηλώνει λάθος memory spaces ή να παραβιάζει περιορισμούς tiling. Τα σφάλματα αυτά δεν διορθώνονται πάντα με ένα γενικό μήνυμα compiler, επειδή η κρίσιμη γνώση για το hardware δεν περιλαμβάνεται αναγκαστικά στο feedback.

Τα workloads και τα expert baselines

Τα 17 production workloads καλύπτουν flash, grouped-query, multi-head latent, sparse, flex, paged και ragged paged attention, καθώς και GEMM, SwiGLU MLP, sparse mixture-of-experts, Megablox GMM και ragged dot. Περιλαμβάνονται επίσης RMSNorm, cross-entropy, RetNet retention, Mamba-2 state space duality και triangle multiplication από το AlphaFold2. Οι διαστάσεις προέρχονται από τα αντίστοιχα μοντέλα· για παράδειγμα, το GQA workload χρησιμοποιεί 128 query heads, 8 key-value heads και sequence length 4096.

Για οκτώ priority kernels υπήρχαν hand-optimized Pallas implementations από το Tokamax. Η ομάδα δοκίμασε συνολικά 203 block-size configurations ώστε να δημιουργήσει ένα ισχυρό expert upper bound. Η ρύθμιση αυτή έδωσε έως 2,79× βελτίωση έναντι των προεπιλεγμένων παραμέτρων Pallas για το Megablox GMM. Το benchmark επομένως συγκρίνει τους agents τόσο με το XLA όσο και, σε περιορισμένο υποσύνολο, με ανθρώπινα βελτιστοποιημένο κώδικα.

Πώς εξασφαλίζονται correctness και επαναληψιμότητα

Κάθε workload διαθέτει κοινό interface με configuration, δημιουργία bf16 inputs και συνάρτηση υπολογισμού. Οι μεταφράσεις από PyTorch ελέγχθηκαν στην TPU με αντιστοιχισμένες εισόδους και απαίτηση συμφωνίας μέσω jnp.allclose με atol και rtol 10^-2. Αν μια μετάφραση αποτύγχανε, αναδημιουργούνταν ή διορθωνόταν πριν ενταχθεί στη σουίτα.

Για τον χρόνο εκτέλεσης δεν χρησιμοποιήθηκε απλό wall-clock timing. Το jax.profiler καταγράφει Perfetto-compatible traces, μετά από πέντε warmup iterations και κατά τη διάρκεια 50 timed iterations. Η κύρια μέτρηση είναι η διάμεσος των device-side events. Η επιλογή είναι κρίσιμη για kernels κάτω από ένα millisecond, όπου το Python dispatch, ο προγραμματισμός του runtime και ο συγχρονισμός μπορούν σύμφωνα με την εργασία να αποτελούν 10% έως 20% του φαινομενικού χρόνου.

Οι τέσσερις προσεγγίσεις που αξιολογήθηκαν

Η πρώτη μέθοδος είναι best-of-N: ανεξάρτητα one-shot samples με TPU preamble και τον JAX source code, από τα οποία κρατείται το καλύτερο σωστό αποτέλεσμα. Η δεύτερη είναι iterative refinement, όπου ο agent λαμβάνει ανά γύρο compilation errors, correctness results και profiler summaries. Η τρίτη διατηρεί την ίδια iterative διαδικασία αλλά προσθέτει επιμελημένο context από το Autocomp.

Η τέταρτη είναι το ίδιο το Autocomp, ένα ανοικτό framework που αποστάζει δημόσια τεκμηρίωση JAX Pallas και Cloud TPU σε περίληψη αρχιτεκτονικής, αναφορά API, σχολιασμένα παραδείγματα και κανόνες correctness. Χρησιμοποιεί δύο φάσεις beam search: πρώτα μετάφραση του XLA baseline σε Pallas και μετά βελτιστοποίηση απόδοσης. Όλες οι κύριες δοκιμές έγιναν με Gemini 3 Flash και συγκρίσιμο budget έως 144 samples ανά benchmark.

Το βασικό εύρημα: το εξειδικευμένο context αλλάζει την ορθότητα

Στο πλήρες σύνολο των 50 workloads, το απλό best-of-N παρήγαγε σωστό kernel σε 13 περιπτώσεις και έφτασε γεωμετρικό μέσο speedup 1,01×. Το iterative refinement έλυσε 32 από τα 50 και έφτασε 1,18×. Με την προσθήκη του επιμελημένου TPU context, η iterative μέθοδος έλυσε 48 από τα 50 και ανέβηκε στο 1,28×. Το Autocomp έλυσε 45, αλλά πέτυχε τον υψηλότερο γεωμετρικό μέσο, 1,36×, επειδή αφιέρωσε περισσότερο από το budget στη βελτιστοποίηση και όχι μόνο στο debugging.

Η καθαρότερη ένδειξη αφορά την ορθότητα ανά sample: το curated context την αύξησε από 5,8% σε 37,3% χωρίς αλλαγή στον αλγόριθμο αναζήτησης. Παράλληλα, η κατάχρηση του Pallas API παρέμεινε η κυρίαρχη πηγή αποτυχίας. Στο best-of-N, 99,7% των samples απέτυχαν κατά τη μεταγλώττιση ή την πρώτη εκτέλεση λόγω API/runtime προβλημάτων, έναντι 59,8% στο iterative με context και 55,8% στο Autocomp.

Πόσο κοντά έφτασαν οι agents στους ανθρώπινους kernels

Στα οκτώ workloads με hand-tuned references, το Tokamax έδωσε γεωμετρικό μέσο 2,08× έναντι XLA, με floor στο 1× για τη συγκέντρωση. Το Autocomp έφτασε 1,60×, περίπου το 77% του hand-tuned γεωμετρικού μέσου. Ξεπέρασε το Tokamax σε δύο kernels και βρέθηκε στο 68% έως 91% της επίδοσής του σε τέσσερα ακόμη.

Η απόσταση παρέμεινε μεγάλη στα paged και ragged attention workloads, όπου ο χειροκίνητος scheduling είναι ιδιαίτερα σημαντικός. Το hand-tuned ragged paged attention ήταν 6,91× ταχύτερο από XLA στο συγκεκριμένο table, ενώ ο agent δεν παρήγαγε σωστό kernel. Αντίθετα, στο sparse attention το Autocomp έφτασε 2,81×, πάνω από το 0,86× της διαθέσιμης Tokamax παραλλαγής, η οποία είχε ρυθμιστεί για διαφορετικό sparsity pattern.

Μεγαλύτερο μοντέλο ή καλύτερη γνώση στόχου;

Η εργασία εξετάζει επίσης ένα υποσύνολο πέντε kernels με Gemini 3.1 Pro. Η μεγαλύτερη ικανότητα του μοντέλου βελτίωσε όλες τις μεθόδους: το plain iterative ανέβηκε από 1,07× με Flash σε 2,43× με Pro, το iterative με context από 1,59× σε 3,82× και το Autocomp από 2,35× σε 3,79×. Στο μικρό αυτό υποσύνολο, οι δύο context-aware μέθοδοι παρήγαγαν σωστό Pallas και για τα πέντε workloads.

Το αποτέλεσμα δεν σημαίνει ότι το model scale είναι αδιάφορο. Δείχνει όμως ότι η ικανότητα και το context αλληλεπιδρούν. Στο sparsely documented Pallas, η τεκμηρίωση προσφέρει τις προϋποθέσεις για να γίνει η αναζήτηση παραγωγική. Μόλις ένα ισχυρότερο μοντέλο παράγει σωστά seeds χωρίς εκτεταμένο debugging, περισσότερα samples μπορούν να χρησιμοποιηθούν για πραγματική βελτιστοποίηση.

Τι σημαίνει για επιχειρήσεις που χτίζουν AI υποδομές

Το πρώτο πρακτικό μάθημα είναι ότι ένα coding agent δεν πρέπει να αξιολογείται μόνο από το αν ολοκληρώνει μια συνάρτηση. Σε performance-critical συστήματα χρειάζονται τρεις πύλες: compilability, αριθμητική ορθότητα και απόδοση απέναντι σε ισχυρό baseline. Η παράλειψη της δεύτερης μπορεί να δώσει γρήγορο αλλά λάθος κώδικα· η παράλειψη της τρίτης μπορεί να μετατρέψει μια τεχνικά επιτυχημένη αυτοματοποίηση σε ακριβότερη υποδομή.

Το δεύτερο μάθημα είναι η αξία ενός ελεγχόμενου knowledge pack. Αρχιτεκτονική hardware, έγκυρα API references, μικρά δοκιμασμένα examples και κανόνες correctness είναι πιο χρήσιμα από μια αδιαφοροποίητη συλλογή εγγράφων. Για μια ομάδα platform engineering, αυτό μεταφράζεται σε versioned τεκμηρίωση και reproducible harness γύρω από τον agent, με τις ίδιες αρχές που χρειάζεται ένα σύστημα LLM routing για latency, ακρίβεια και κόστος.

Το τρίτο είναι η σωστή επιλογή προβλημάτων. Οι ερευνητές δείχνουν ότι compute-bound GEMM μπορεί ήδη να κορέσει την MXU μέσω XLA και να μην αφήνει αξιόλογο περιθώριο. Αντίθετα, memory-bound attention kernels μπορούν να ωφεληθούν από tiling που περιορίζει την επαναλαμβανόμενη κίνηση προς HBM. Μια επιχείρηση πρέπει λοιπόν να κάνει profiling πριν επενδύσει σε agentic optimization και να κατευθύνει το budget στα πραγματικά bottlenecks.

Όρια της μελέτης και ασφαλής ανάγνωση των αποτελεσμάτων

Τα αποτελέσματα αφορούν συγκεκριμένη σουίτα, TPU v6e, Pallas/Mosaic και τα μοντέλα που δοκιμάστηκαν. Η σύγκριση Gemini 3.1 Pro έγινε μόνο σε πέντε kernels, λόγω κόστους, και οι συγγραφείς τη χαρακτηρίζουν προκαταρκτική. Επίσης, τα expert references υπάρχουν μόνο για οκτώ από τα 17 priority workloads, άρα δεν αποτελούν καθολικό όριο για όλη τη σουίτα.

Η μέτρηση speedup χρησιμοποιεί floor 1× στην ομαδοποίηση: ένας λανθασμένος ή πιο αργός kernel δεν ρίχνει τον γεωμετρικό μέσο κάτω από το baseline. Αυτό είναι λογικό για το σενάριο όπου το σύστημα κρατά τον XLA kernel αν ο υποψήφιος δεν βοηθά, αλλά πρέπει να λαμβάνεται υπόψη όταν συγκρίνονται τα aggregate νούμερα. Η per-benchmark correctness και το fast@N συμπληρώνουν την εικόνα.

Ένα λειτουργικό πλαίσιο για agentic optimization

Μια υπεύθυνη υλοποίηση μπορεί να ακολουθήσει τη λογική του JAXBench: ξεκινά με production-relevant workloads, ορίζει baseline στο ίδιο hardware, χρησιμοποιεί matched inputs και ανοχές κατάλληλες για τον αριθμητικό τύπο, κάνει warmup και device-side profiling και κρατά ξεχωριστά compilation, correctness και latency outcomes. Οι υποψήφιοι kernels δεν πρέπει να περνούν σε παραγωγή χωρίς επανάληψη των ελέγχων και σύγκριση με το υπάρχον σύστημα. Πρόκειται για την ίδια λογική που απαιτεί η επαλήθευση πολιτικών πριν εμπιστευτούμε έναν AI agent.

Η εργασία δεν αποδεικνύει ότι οι AI agents αντικαθιστούν τους kernel engineers. Τεκμηριώνει όμως ότι μπορούν να ανακτήσουν σημαντικό μέρος της ανθρώπινης βελτιστοποίησης όταν έχουν σωστό context και οργανωμένη αναζήτηση. Το πιο χρήσιμο επιχειρηματικό συμπέρασμα είναι ότι η αξία βρίσκεται στο σύστημα γύρω από το μοντέλο: benchmark, τεκμηρίωση, feedback, selection policy, ασφαλές fallback και σαφής ευθύνη όταν η αυτοματοποίηση με AI agents αποτυγχάνει.