La notizia è circolata poche ore fa: Helion, il linguaggio dedicato di PyTorch per scrivere kernel ML performanti e portabili, ora ha un backend TPU maturo, frutto del lavoro congiunto di Meta e Google. Su un carico di flash attention — operazione centrale per i transformer moderni — il kernel generato automaticamente da Helion ha toccato 838 TFLOPs su TPU v7, circa il 79% dell'utilizzo massimo di un tensor core. Il dato fa notizia, ma non è la parte più importante.

Come l'autotuner sceglie tra memoria e calcolo

La sfida vera sulle TPU è orchestrare il trasferimento dei dati dalla memoria HBM alla veloce VMEM on-chip, sovrapponendolo ai calcoli sulla matrice MXU e sulle unità vettoriali. Helion, tramite il suo autotuner, decide in automatico tra due schemi di pipeline radicalmente diversi.

La prima strategia, emit_pipeline, carica i blocchi K e V dall'HBM a ogni nuovo blocco di Q. È flessibile e scala su sequenze lunghe, ma introduce bolle di inattività nelle unità di calcolo perché a ogni cambio di Q si attende il primo blocco KV. La seconda, unroll, precarica interamente K e V nella VMEM quando le dimensioni lo consentono, eliminando del tutto i tempi morti: i calcoli procedono senza interruzioni. Il rovescio della medaglia è il consumo di VMEM, che cresce linearmente con la lunghezza della sequenza e rende questa strada impraticabile oltre certe soglie.

I numeri parlano chiaro: su una testa di attenzione con batch 8, 32 teste, dimensione 256, a S=8k la versione con emit_pipeline produce 653 TFLOPs, mentre unroll sale a 892. Con S=32k, unroll esaurisce la memoria e si ripiega su emit_pipeline, che comunque arriva a 695 TFLOPs. L'autotuner sceglie in base alla forma dell'input, senza che lo sviluppatore debba pensare a nulla.

Portabilità hardware: la fine del lock-in?

La vera posta in gioco va oltre i teraflops. Helion è nato come DSL portabile, e lo stesso kernel scritto per GPU viene compilato anche per TPU. In un mercato dove Google spinge Ironwood (TPU v7) con prestazioni paragonabili al NVIDIA B200 ma un TCO potenzialmente inferiore, la possibilità di scrivere il kernel una volta e usarlo su entrambi i silici diventa un'arma per i team ML. Chi ha già cluster GPU on-premise può sperimentare con il cloud TPU senza riscrivere la logica di calcolo, e chi progetta architetture ibride può decidere dove eseguire l'inference o il fine-tuning in base a costi e disponibilità, non alla compatibilità del codice.

Non è un dettaglio. L'astrazione introdotta da Helion (e da iniziative parallele come Triton) sposta il potere contrattuale dai fornitori di chip agli utilizzatori, perché riduce drasticamente i costi di switch. Se il compilatore riesce a estrarre automaticamente il massimo da ogni acceleratore, la scelta dell'hardware diventa una questione di TCO, latenza e posizione geografica, non di competenze specialistiche su Pallas o CUDA. A guadagnarci sono le organizzazioni con stack on-premise che possono ora aprire un canale verso il cloud senza rifare da capo l'ingegneria dei kernel.

I test sulle operazioni comuni mostrano che Helion batte TorchTPU eager (1.55x in media geometrica) e supera lievemente torch.compile via XLA (1.12x). I guadagni maggiori arrivano su kernel con fusioni difficili da scoprire automaticamente, come la flash attention. Su matmul e layer norm, dove XLA lavora già bene, le prestazioni sono comparabili. La differenza la fa la logica di autotuning, che sulle operazioni ad alta intensità di memoria sceglie la pipeline giusta senza intervento manuale.

Per chi sta valutando deployment ibridi, AI-RADAR offre su /llm-onpremise un framework analitico per soppesare questi trade-off. Intanto, il backend TPU di Helion rimane legato alla libreria TorchTPU, non ancora pubblica (è attesa entro fine anno). Ma la direzione è segnata: hardware heterogeneity gestita al livello del framework, con un compilatore abbastanza intelligente da scegliere la strategia di pipeline migliore. Non è solo una notizia tecnica: è un segnale che l'ecosistema software sta diventando il vero ago della bilancia.