The Pulse

Un team di RadixArk mette in servizio Muse Spark Open di Meta su un solo host TPU

Il server SGLang-TorchTPU di PyTorch esegue Muse Spark Open di Meta, modello da 816 miliardi di parametri, su un solo host TPU v7 con quattro chip. Il post riporta risultati di decodifica più rapidi grazie a un kernel a programma singolo e

AI.info Team · Scritto con l’IA · verificato · 2026-11-10

Un host, otto dispositivi

Un team di RadixArk mette in servizio Muse Spark Open di Meta, un modello mixture-of-experts da 816 miliardi di parametri, su un solo host TPU v7 con quattro chip. Il post del team, pubblicato il 10 novembre sul blog di PyTorch, descrive il sistema che esegue SGLang sul backend PyTorch TorchTPU, con kernel Pallas che gestiscono le operazioni critiche per le prestazioni.

Il post afferma che Muse Spark Open ha circa 42 miliardi di parametri attivi per token. I suoi 62 livelli decoder sono tutti livelli mixture-of-experts e ciascuno instrada un token a otto dei 256 esperti.

Ogni chip TPU v7 ha due chiplet, che il framework presenta come dispositivi separati; l’host dispone quindi di otto dispositivi. Il modello è suddiviso tra questi dispositivi mediante parallelismo tensoriale, assegnando una partizione a ciascuno.

Un solo programma per una fase di serving

La modifica principale è un megakernel: un programma Pallas su ciascun dispositivo gestisce tutti i 62 livelli di una fase di decodifica o verifica. SGLang continua a eseguire, attorno a quel programma, l’embedding, il modello preliminare, la scrittura della cache KV, la testa del modello linguistico e il campionamento, afferma il post.

Per una singola richiesta, il megakernel impiega 1,95 millisecondi per token generato, contro i 3,78 millisecondi del precedente percorso «nativo», con una velocità 1,94 volte superiore, riferisce il team. Per due richieste simultanee, il team riporta una velocità 1,85 volte superiore; il guadagno si riduce all’aumentare delle richieste simultanee.

Il post afferma che la fase di verifica viene eseguita in un tempo compreso tra 1,5 e 1,7 volte quello necessario per il semplice trasferimento dei suoi pesi dalla memoria ad alta larghezza di banda. Per evitare che i trasferimenti interrompano il calcolo, il kernel organizza in pipeline il caricamento dei pesi, la lettura della cache KV e i trasferimenti tra dispositivi.

Modifiche all’attenzione e all’elaborazione degli esperti

Per l’attenzione, il team afferma di aver rimosso il massimo progressivo usato nel calcolo online del softmax. Il post sostiene che la normalizzazione delle query e delle key di Muse Spark, insieme al guadagno fisso delle query del modello, limita i punteggi di attenzione e rende superfluo il massimo; senza tale dipendenza, i blocchi della cache KV possono essere elaborati in pipeline.

Nel sistema di serving, secondo il team, il ciclo di attenzione in pipeline ha aumentato del 18% la velocità di decodifica per richiesta, con due richieste e un contesto di 500.000 token. Un test separato del kernel con instradamento casuale ha mostrato che le fasi di decodifica con molte richieste erano dal 13% al 19% più veloci dopo che il team aveva suddiviso l’elaborazione degli esperti in due cicli.

È stata usata anche la decodifica speculativa DFlash2, in cui un modello preliminare più piccolo propone token che il modello completo verifica insieme. Su GSM8K, con accettazione reale, l’azienda riferisce che questo ha ridotto il tempo del megakernel per token generato per una singola richiesta da 3,93 a 1,95 millisecondi, con una velocità 2,0 volte superiore.

Risultati sui carichi di lavoro degli agenti

Per la versione con contesto da 1 milione del benchmark di programmazione AgentX di SemiAnalysis, il post riporta un intervallo che va da 433 token generati al secondo per richiesta, con 6.101 token al secondo per chip, a 25.906 token al secondo per chip, con 95 token al secondo per richiesta. Secondo il team, tutti i risultati TPU mostrati hanno completato le richieste entro 47 secondi al 90º percentile.

Il post descrive AgentX come basato su 393 sessioni Claude Code fornite su base volontaria: ne riproduce la lunghezza delle richieste, i prefissi condivisi, le pause per gli strumenti e la struttura dei sottoagenti, sostituendo i contenuti originali con token sintetici. Il confronto usa l’insieme di tracce del benchmark con contesto da 1 milione, in cui i prompt raggiungono circa 990.000 token; il throughput comprende sia i token in ingresso sia quelli generati, afferma il team.

Il post indica che la configurazione TPU è composta da un host con quattro chip e usa calcoli FP8 e pesi degli esperti NVFP4. Le cifre sui costi sono i costi esterni di proprietà di InferenceX: 1,21 dollari per ora di utilizzo di un chip TPU v7 e 1,86 dollari per ora di utilizzo di una GPU GB200.

Secondo il team, dopo aver allocato i pesi del modello, su ciascun dispositivo rimangono circa 20 GiB per la cache KV. Il post descrive una disposizione sequence-on-lane per i dati della cache FP8 e L2 HiCache su TPU, con espulsione dei dati che tiene conto delle sessioni, come elementi del suo approccio per aumentare la capacità della cache.