Arquitecturas nativas para device: MoE de grano fino y pre-attention router

· Por José Manuel Román

Este post es de la serie sobre rendimiento de inferencia en modelos pequeños. Es la cara arquitectónica de un problema que ya hemos mirado por el lado del régimen de cómputo (el roofline invertido del SLM) y por el lado de la carga de pesos en Del disco a la HBM. Aquí la pregunta es distinta: ¿y si en lugar de adaptar un modelo grande al device, diseñamos el modelo para el device desde el primer commit?

TL;DR

El gesto por defecto para llevar un LLM a un portátil, un móvil o un edge box es coger un denso pensado para cloud y comprimirlo: destilación, poda, cuantización. Es un gesto de reducción: partes de algo grande y le quitas. SmallThinker (arXiv:2507.20984, SJTU IPADS + Zenergize AI) defiende el gesto inverso —diseñar desde cero— y lo articula en tres piezas. Primera: MoE de grano fino, muchos expertos pequeños con muy pocos activados por token, de modo que los parámetros totales N (la capacidad) se desacoplan de los parámetros activados A (el coste de cómputo por token). Segunda: sparse FFN, sparsity de activación tipo ReLU dentro de cada bloque, que añade un segundo nivel de dispersión sobre el primero. Tercera: un pre-attention router que predice qué expertos hará falta antes de ejecutar el bloque de atención y lanza el prefetch de esos pesos desde SSD/flash en paralelo con el cómputo de la atención, ocultando la latencia de almacenamiento —que es el cuello de botella real cuando el modelo no cabe entero en RAM. Los autores reportan SmallThinker-4B-A0.6B y SmallThinker-21B-A3B superando ~20 tok/s en CPU de consumo con Q4_0, consumiendo ~1 GB y ~8 GB de RAM. Los números son interesantes y la dirección es correcta; la metodología de evaluación y el coste de calidad de activar tan poco merecen escepticismo, y a eso dedicamos la última parte.

La analogía: el bibliotecario que se adelanta a tu pedido

Imagina una biblioteca enorme con una sala de lectura pequeña. Tú estás sentado en la sala con un único pupitre: ahí caben pocos libros a la vez (eso es la RAM). El grueso del fondo está en la trastienda, en estanterías largas y lentas de recorrer (eso es el SSD/flash). Y hay un bibliotecario.

El método ingenuo: tú lees, llegas a un punto donde necesitas un libro concreto, lo pides, y entonces el bibliotecario se levanta, va a la trastienda, lo busca y vuelve. Mientras tanto, tú esperas con la página abierta sin avanzar. Cada vez que necesitas un libro nuevo, pagas el viaje completo a la trastienda. La sala de lectura está la mayor parte del tiempo esperando, no leyendo.

El método de SmallThinker: el bibliotecario es listo y se adelanta. Mientras tú todavía estás leyendo el índice del capítulo —averiguando de qué va, relacionando ideas, lo que en el modelo es el bloque de atención—, él ya ha mirado por encima de tu hombro, ha predicho qué tres o cuatro libros vas a pedir y se ha ido a la trastienda a buscarlos. Para cuando terminas el índice y formulas el pedido, los libros ya están sobre tu pupitre. No has esperado: el viaje a la trastienda ocurrió en paralelo con tu lectura del índice.

La analogía se sostiene en cuatro detalles:

  • El pupitre pequeño es la RAM; la trastienda lenta es el SSD/flash.
  • Los libros son los expertos del MoE: solo unos pocos están sobre el pupitre en cada momento.
  • Leer el índice es el bloque de atención; pedir y usar los libros es el bloque FFN/expertos.
  • El bibliotecario que predice y se adelanta es el pre-attention router: la predicción se hace antes, y el viaje a buscar (el prefetch) se solapa con la lectura del índice (la atención).

La pregunta cuantitativa que recorre todo el post es: ¿llega el bibliotecario a tiempo? Solo se oculta la espera si el viaje a la trastienda dura menos que tu lectura del índice. Esa es la condición t_{\text{atención}} \ge t_{\text{prefetch}}, y la haremos con números.

Comprimir un denso vs. diseñar para device

Conviene poner los dos enfoques en frío, porque no son grados de lo mismo: son filosofías distintas.

Enfoque A — comprimir un denso pensado para cloud. Partes de, digamos, un modelo denso de 7B–14B entrenado para correr en una RTX 4090 (24 GB, Ada Lovelace) o en un cluster genérico 4×H100 SXM (320 GB, NVLink, FP8 nativo). Para meterlo en un device aplicas tres palancas, cada una con su post propio: destilación (entrenas un student pequeño que imita al teacher), poda (eliminas pesos o estructuras enteras) y cuantización agresiva (bajas a 4 bits o menos). El modelo resultante sigue siendo denso: todos sus parámetros se activan en cada token. Has reducido el número de parámetros, pero el patrón de cómputo es el del cloud, solo que más pequeño.

Enfoque B — diseñar para device desde cero. Aquí las restricciones del device entran en la arquitectura, no en una fase posterior de compresión. Las restricciones son tres y muy concretas:

  1. Cómputo débil. Una CPU de portátil o un SoC móvil hace órdenes de magnitud menos FLOPs que una GPU de datacenter. Esto empuja a minimizar los parámetros activados por token, no los totales.
  2. Poca RAM. No caben decenas de GB. Esto empuja a tener residente solo lo imprescindible y a streamear el resto.
  3. Almacenamiento lento. El SSD o la flash a la que te ves obligado a streamear tiene un ancho de banda muy inferior al de la HBM de una GPU. Esto convierte la I/O de almacenamiento en el cuello de botella, y empuja a ocultarla.

SmallThinker es el enfoque B llevado al detalle: cada una de esas tres restricciones tiene una respuesta arquitectónica. El cómputo débil se ataca con MoE de grano fino + sparse FFN (minimizar A). La RAM escasa se ataca con streaming desde SSD (residente ≈ A + caché, no N). El almacenamiento lento se ataca con el pre-attention router (ocultar la I/O tras la atención). No es casual que las tres piezas encajen: cada una resuelve una restricción, y juntas se refuerzan.

Un matiz importante, para no caer en el hype: el enfoque B no es gratis ni universalmente superior. Requiere entrenar un modelo nuevo (no reutilizas pesos existentes), y el techo de calidad de un modelo con A muy pequeño está intrínsecamente acotado, como veremos. El argumento no es “B gana siempre”, sino “para el régimen del device, B ataca los cuellos correctos, y A solo los ataca de refilón”.

Dos niveles de sparsity

La idea central de capacidad es vieja y bien entendida en MoE: separar capacidad de coste de cómputo. En un MoE, el modelo tiene N parámetros totales repartidos en expertos, pero para cada token solo se activan A parámetros (los del top-k de expertos que el router elige). El coste de cómputo por token escala con A; la capacidad de conocimiento escala con N. SmallThinker aplica esta idea en dos niveles superpuestos.

Nivel 1 — MoE de grano fino. “Grano fino” significa muchos expertos pequeños en vez de pocos expertos grandes, con muy pocos activados por token. En vez de, digamos, 8 expertos de los que activas 2, tienes decenas de expertos de los que activas un puñado. Con expertos más pequeños, el mismo A se reparte entre más combinaciones posibles, lo que da granularidad fina al router y mantiene A muy bajo respecto a N. El resultado es un cociente N/A agresivo: mucha capacidad, poquísimo cómputo por token.

Nivel 2 — sparse FFN (sparsity de activación tipo ReLU). Este nivel es ortogonal y opera dentro de cada FFN. Con una no-linealidad tipo ReLU, una fracción grande de las neuronas de la capa intermedia produce exactamente cero para un token dado. Una neurona que sale a cero no contribuye nada a la salida: su multiplicación matriz-vector se puede saltar. Esto es sparsity de activación: predecible token a token, y aprovechable para no cargar ni multiplicar las filas/columnas de peso correspondientes a neuronas inactivas. Es el mismo fenómeno que explotan trabajos como Deja Vu o PowerInfer; SmallThinker lo incorpora de fábrica eligiendo activaciones que lo favorecen.

El efecto combinado, en una frase: N grande (capacidad), A minúsculo (coste de cómputo por token ≈ proporcional a A), y además dentro de ese A una fracción de las multiplicaciones se ahorra por la sparsity de activación. Es sparsity sobre sparsity.

MoE clásico (grano grueso)MoE de grano fino + sparse FFN

8 expertos grandes · activa 2A grande por experto · granularidad gruesa

muchos expertos pequeños · activa pocosA total bajo · granularidad fina

+ sparse FFN dentro de cada experto activo:neuronas a 0 (ReLU) →se saltan en el cómputo

Capacidad = N (todos los expertos) · Coste/token ≈ A (activados) · y dentro de A, sparse FFN ahorra másEl truco: subir N sin subir A. La granularidad fina permite un cociente N/A mucho más agresivo.

El pre-attention router: predecir y prefetchar

Aquí está la pieza específica del paper, y la que da nombre al post. El problema que resuelve es de scheduling de I/O, no de calidad.

Cuando el modelo no cabe entero en RAM, los pesos de los expertos viven en SSD/flash y se cargan bajo demanda. El flujo ingenuo de una capa MoE es secuencial: ejecutas la atención, luego el router decide qué expertos tocan, luego cargas esos expertos desde SSD (esperando), luego ejecutas la FFN de esos expertos. El paso de carga es una espera pura: la CPU está bloqueada esperando bytes del SSD. En el régimen del device, donde el SSD es lento, ese tiempo de espera domina el step de decode.

El pre-attention router rompe la secuencialidad invirtiendo el orden de la decisión. La observación es que el router no necesita la salida de la atención de esta misma capa para hacer una predicción razonable de qué expertos harán falta: puede predecirlo a partir del estado que ya tiene antes de ejecutar la atención. Así que:

  1. Antes de ejecutar el bloque de atención de la capa, el router predice los expertos que se necesitarán.
  2. Lanza el prefetch de esos expertos desde SSD/flash de forma asíncrona.
  3. En paralelo, la CPU ejecuta el bloque de atención —que es cómputo puro, no necesita el SSD.
  4. Cuando la atención termina, los expertos prefetchados ya están (idealmente) en RAM, y la FFN procede sin esperar.

El I/O de almacenamiento se ha solapado con el cómputo de atención. Es exactamente el bibliotecario que va a la trastienda mientras tú lees el índice.

Ingenuo (secuencial): la carga desde SSD bloqueaatenciónroutercarga expertos desde SSD(espera bloqueante)FFN expertost_total = t_att + t_load + t_ffn

Pre-attention router: el prefetch se solapa con la atención

hilo de cómputo (CPU)router (pre)atención (t_att)FFN expertos

hilo de I/O (SSD)prefetch expertos (t_prefetch)

El prefetch queda oculto sit_att ≥ t_prefetch: para cuando la atención termina, los expertos ya están en RAM.Si t_prefetch > t_att, asoma una burbuja de espera (t_prefetch − t_att) antes de la FFN. Ese es el caso a evitar.

La condición de ocultamiento es la desigualdad de arriba: el prefetch se oculta completamente si y solo si

$$t_{\text{atención}} ;\ge; t_{\text{prefetch}}.$$

Si la atención tarda más que cargar los expertos, la carga es gratis (ya estaba hecha). Si los expertos son demasiado grandes o el SSD demasiado lento, t_prefetch > t_att y asoma una burbuja de espera igual a t_prefetch − t_att. Por eso el diseño necesita que A sea pequeño (expertos pequeños → menos bytes a prefetchar → t_prefetch bajo) y que el grano sea fino: las dos cosas que hace el nivel 1 de sparsity no son solo para ahorrar FLOPs, son para que el prefetch quepa debajo de la atención.

Las matemáticas que importan

Footprint de memoria: N residente vs. A + caché

El parámetro que decide si el modelo cabe es cuánto tienes que tener residente en RAM a la vez.

  • Todo en RAM. Si exiges que todos los expertos estén cargados, el footprint es \approx N (todos los parámetros, multiplicados por bytes/parámetro según la cuantización). Para un 21B esto es prohibitivo en un device.
  • Streaming desde SSD. Si solo mantienes residentes los expertos activos más una caché de los recientes/probables, el footprint cae a \approx A + \text{caché}. Los pesos que no están en RAM viven en SSD y se prefetchan cuando toca. Aquí está el ahorro real: el residente escala con A, no con N.

La parte no-experta del modelo (embeddings, atención, router, layernorms) sí está siempre residente, pero en un MoE de grano fino el grueso de N está en los expertos, así que la aproximación residente ≈ A + caché + parte_densa es buena.

El cálculo de prefetch, con números

Pongamos los números de la analogía. Supón un SSD de consumo a 5 GB/s de lectura secuencial y un experto cuantizado de tamaño X MB. El tiempo de cargar un experto es

$$t_{\text{1 experto}} = \frac{X \text{ MB}}{5000 \text{ MB/s}} = \frac{X}{5000}\ \text{s} = \frac{X}{5}\ \text{ms}.$$

Concretemos X. En SmallThinker-4B-A0.6B con Q4_0 (~0.5 byte/param efectivo contando overhead de bloques), un experto pequeño de, digamos, 4M parámetros pesa \approx 4\text{M} \times 0.5 = 2 MB. Cargarlo cuesta t_{\text{1 experto}} = 2/5 = 0.4 ms.

Ahora la pregunta de scheduling: si el bloque de atención de la capa toma Y ms, ¿cuántos expertos puedo prefetchar mientras la atención corre? El número es

$$n_{\text{prefetch}} = \left\lfloor \frac{Y}{t_{\text{1 experto}}} \right\rfloor = \left\lfloor \frac{Y \cdot 5}{X} \right\rfloor.$$

Con Y = 2 ms de atención y X = 2 MB por experto: n_{\text{prefetch}} = \lfloor 2 \times 5 / 2 \rfloor = 5 expertos. Es decir, en la ventana de atención de esa capa el SSD alcanza a traer 5 expertos. Si el top-k de la capa activa ≤ 5 expertos, el prefetch los oculta todos y t_prefetch ≤ t_att: latencia de carga cero. Si la capa necesitara 8 expertos, traerías 5 gratis y pagarías la carga de los 3 restantes como burbuja: (8-5) \times 0.4 = 1.2 ms de espera por capa. De ahí que el diseño quiera grano fino con top-k pequeño: para caber debajo de la ventana de atención.

Dos observaciones críticas sobre este cálculo:

  • Los 5 GB/s son lectura secuencial idealizada. Los expertos están dispersos en disco; lecturas aleatorias 4K en un SSD de consumo van mucho más lentas. El ancho de banda efectivo puede ser una fracción del nominal, lo que reduce n_{\text{prefetch}}. La metodología que reporte tok/s debería decir si mide con expertos pre-ordenados en disco o con acceso realista.
  • La ventana Y de atención encoge con el contexto corto y al inicio de la generación. Con prompts cortos, la atención es barata y puede que no cubra el prefetch; la ventaja del solapamiento crece con secuencias más largas. Otro detalle que un benchmark honesto debería desglosar.

Footprint de pesos: por qué reportan ~1 GB para un 4B

Hagamos la cuenta del 4B en Q4_0. Cuantización a 4 bits ≈ 0.5 byte/param, más un pequeño overhead de escalas por bloque (Q4_0 añade un FP16 de escala cada 32 pesos, ~0.56 byte/param efectivos). Entonces:

$$4\text{B} \times 0.5\ \text{B/param} \approx 2\ \text{GB}.$$

Es decir, el modelo completo en Q4_0 ocupa ~2 GB en disco. Pero los autores reportan ~1 GB de RAM. ¿Contradicción? No, y entender por qué es entender el diseño:

  • No todos los expertos están residentes. Solo los activados (A = 0.6B) y una caché caben en RAM; el resto vive en SSD y se streamea. 0.6\text{B} \times 0.5 \approx 0.3 GB de expertos activos, más la parte densa (atención, embeddings, router) y una caché de expertos calientes.
  • La sparse FFN reduce el trabajo y el residente útil. Las neuronas que salen a cero no necesitan estar materializadas para ese token.

Sumando expertos activos + parte densa + caché razonable, ~1 GB es plausible. Pero ojo con el matiz: ~1 GB es el residente en RAM, no el footprint total en almacenamiento, que sigue siendo ~2 GB en SSD. Confundir ambos —reportar “1 GB” a secas— es engañoso si el lector entiende “el modelo ocupa 1 GB”. Ocupa 2 GB; mantiene 1 GB en RAM. La distinción importa para un device con 2 GB de almacenamiento libre: ahí no entra.

Análogamente, SmallThinker-21B-A3B: 21\text{B} \times 0.5 \approx 10.5 GB en disco; 3\text{B} \times 0.5 \approx 1.5 GB de expertos activos, y el ~8 GB de RAM reportado incluye expertos activos + caché generosa + parte densa. La caché grande es lo que sube de 1.5 a ~8 GB: mantienes muchos expertos calientes residentes para no golpear el SSD constantemente.

El coste de calidad: el escepticismo necesario

Toda la maquinaria anterior reduce el cómputo por token a \approx A. Pero A = 0.6B activados es muy poco. Aquí es donde hay que poner el freno al entusiasmo:

  • Capacidad de razonamiento acotada. Un modelo que activa 0.6B de parámetros por token tiene, por token, la potencia de cómputo de un modelo de 0.6B, no de 4B. La capacidad total N=4B ayuda a almacenar más conocimiento (más expertos especializados), pero el procesamiento de cada token sigue limitado por A. Para tareas que requieren composición y razonamiento multi-paso intensivo, esto es un techo real, no un detalle.
  • El router es un punto único de fallo de calidad. Si el router de grano fino elige mal los expertos —y con grano fino hay más decisiones que tomar—, la calidad cae sin que ninguna métrica de velocidad lo refleje. El pre-attention router agrava esto: predice los expertos antes de ver la atención, con menos información que un router post-atención. Los autores deberían reportar cuánta calidad se pierde por predecir antes (mismatch entre experto prefetchado y experto que el router post-atención habría elegido).
  • Los ~20 tok/s necesitan letra pequeña. ¿En qué CPU exactamente? ¿Con qué longitud de contexto y de generación (la ventaja del solapamiento depende de Y)? ¿Cold start incluido o steady state? ¿El SSD estaba con los expertos pre-ordenados secuencialmente? Un “supera 20 tok/s” sin esas condiciones es un número de marketing, no de metodología.
  • Comparación justa. La pregunta correcta no es “¿es rápido?”, sino “¿a igualdad de calidad en un benchmark independiente, es más rápido o más pequeño que un denso comprimido equivalente?”. Eso requiere evals que el lector pueda reproducir, no solo tok/s en la máquina de los autores.

Nada de esto invalida la dirección. Diseñar para device es, conceptualmente, el enfoque correcto: ataca los cuellos reales (cómputo, RAM, I/O) en la arquitectura en vez de paliar­los después. Pero “20 tok/s en ~1 GB” es una afirmación de eficiencia, y la eficiencia solo significa algo anclada a un nivel de calidad medido honestamente. Mientras esa ancla no esté clara, el número correcto de escepticismo es alto.

Implicaciones para inferencia on-premise y edge

  • El SSD pasa a ser parte de la jerarquía de inferencia. En cloud, la jerarquía es HBM → RAM. En device, el SSD/flash entra como un nivel más, y su ancho de banda y latencia de acceso aleatorio se vuelven parámetros de rendimiento de primer orden. Esto conecta con Del disco a la HBM: el cold start y el streaming de pesos dejan de ser solo un problema de arranque y pasan a ser parte del steady state.
  • El edge box hetero­géneo gana sentido. En un patrón de entornos mixtos, un modelo nativo-device como SmallThinker corre en el NUC/edge con CPU y SSD, sirviendo localmente, mientras lo pesado se queda en el cluster central. El pre-attention router es lo que hace viable el edge box sin GPU.
  • El capacity planning cambia de ejes. Como discute Capacity planning de inferencia, en device el recurso a planificar no es VRAM sino la terna RAM-residente / ancho-de-banda-SSD / FLOPs-de-CPU. Un modelo con A pequeño y prefetch solapado mueve el cuello de botella de “¿cabe en RAM?” a “¿el SSD alimenta el prefetch a tiempo?”.

Conclusión

SmallThinker es, sobre todo, un cambio de pregunta. No “¿cómo encojo este modelo cloud para que quepa en el device?” sino “¿cómo sería el modelo si lo diseñara para el device desde el primer parámetro?”. La respuesta —MoE de grano fino para desacoplar N de A, sparse FFN para ahorrar dentro de A, y un pre-attention router que oculta la I/O de almacenamiento bajo la atención— ataca las tres restricciones del device (cómputo, RAM, I/O) en la arquitectura, no en una fase de compresión posterior. La condición clave, t_att ≥ t_prefetch, explica por qué las piezas encajan: el grano fino no solo ahorra FLOPs, hace que el prefetch quepa debajo de la atención. Los números reportados (~20 tok/s, ~1 GB / ~8 GB de RAM) son prometedores y la dirección es sólida; el coste de activar tan poco y la falta de detalle metodológico sobre calidad piden cautela. Diseñar para device es la apuesta correcta; medirlo honestamente es la asignatura pendiente.

Ver también

  • MoE inference: el call center con 256 especialistas — la base conceptual de este post: cómo un router enruta tokens a expertos y por qué N y A se desacoplan; léelo primero si MoE te suena lejano.
  • Del disco a la HBM: cold start y carga de modelo — el streaming de pesos desde almacenamiento lento, que aquí deja de ser problema de arranque y pasa a steady state vía prefetch.
  • Knowledge distillation — la palanca canónica del enfoque “comprimir un denso de cloud”, el contrapunto exacto del enfoque nativo-device.
  • Poda de modelos LLM — la otra palanca de reducción; útil para comparar “quitar a un grande” frente a “diseñar pequeño desde cero”.
  • Entornos mixtos NVIDIA + Intel — dónde encaja un modelo nativo-device: el edge box con CPU y SSD que sirve localmente sin GPU.
  • Capacity planning de inferencia LLM on-premise — en device los ejes a planificar son RAM-residente, ancho de banda de SSD y FLOPs de CPU, no VRAM.
  • Roofline invertido en modelos pequeños (hermano de esta serie, próximamente) — el régimen de rendimiento del SLM que explica por qué A pequeño mantiene el decode memory-bound y dónde está el techo real.
  • Self-speculative decoding con early-exit (hermano de esta serie, próximamente) — self-spec aplicado a MoE on-device: cómo acelerar el decode sin draft externo cuando el modelo ya es pequeño.
  • Cuantización agresiva sub-4-bit y ternaria (hermano de esta serie, próximamente) — Q4_0 y más allá en device: ternario y 2-bit para bajar aún más el footprint de expertos en SSD.

Referencias