Saltar a contenido

MoE top_k override (--moe-top-k)

Reduce el numero de experts activados por token en modelos Mixture of Experts como Qwen3-30B-A3B, intercambiando una pequeña cantidad de calidad por un aumento significativo en el throughput de decodificacion.

Estado: flag opt-in. El comportamiento por defecto no cambia. Los numeros de calidad que se muestran son para Qwen3-30B-A3B-4bit en M4 Max 128 GB. Verifica con tu modelo antes de usarlo en cargas de produccion.

Que hace

Qwen3-30B-A3B se entrena con top_k=8. cada token selecciona 8 de 128 experts. En Apple Silicon con batch=1 durante la decodificacion, la multiplicacion de matrices de experts (SwitchGLU) es la parte más costosa del computo por capa, y ese costo escala de forma aproximadamente lineal con top_k. Reducir top_k en tiempo de inferencia ha demostrado (LExI 2025, Lynx 2024) preservar la mayor parte de la calidad entrenada mientras reduce materialmente el tiempo de decodificacion.

--moe-top-k N itera cada capa del modelo cargado y, en cada capa que tenga .mlp.switch_mlp (es decir, un bloque sparse-MoE), establece top_k = N. Las capas densas y los modelos densos no se modifican: el flag es un no-op para ellos.

Uso

# Server
vllm-mlx serve mlx-community/Qwen3-30B-A3B-4bit \
  --continuous-batching \
  --moe-top-k 4

# Bench
vllm-mlx bench mlx-community/Qwen3-30B-A3B-4bit --moe-top-k 4

El flag se rechaza si N es mayor que el top_k entrenado del modelo (solo tiene sentido reducirlo, nunca aumentarlo).

Impacto medido

Throughput de decodificacion (M4 Max 128 GB, batch=1, greedy)

top_k tok/s vs baseline
8 (baseline) 126.5 -
6 136.1 +7.6%
5 140.3 +10.9%
4 147.3 +16.5%

Calidad (Qwen3-30B-A3B-4bit, lm-evaluation-harness, MLX backend)

top_k MMLU (acc) GSM8K (exact match) Delta vs baseline
8 TBD TBD -
6 TBD TBD TBD
5 TBD TBD TBD
4 TBD TBD TBD

MMLU: 200 muestras seleccionadas aleatoriamente, 0-shot. GSM8K: 100 muestras seleccionadas aleatoriamente, 0-shot, exact-match estricto.

Estos numeros son indicativos: los conjuntos de evaluacion completos son más grandes y desplazarian la precision absoluta, pero no el delta relativo entre configuraciones de forma significativa.

Paridad de salida greedy

Con top_k=4 en el checkpoint de 4 bits observamos los primeros 16 tokens generados identicos al baseline en todos los prompts de prueba que usamos. Esto sugiere que top_k=4 no cambia el argmax en los pasos iniciales de decodificacion: el modelo es internamente robusto a eliminar la mitad de sus experts activados.

Con top_k=3 o menor, la calidad comenzaria a degradarse de forma visible (no medido aquí; inferido del paper LExI), por lo que el flag no permite bajar por debajo de 1 en la capa de validación de configuración. Sin embargo, el piso recomendado para produccion es top_k=4.

Cuando usarlo y cuando no

Usalo cuando: - Ejecutas un Qwen3 MoE (o compatible: Qwen3.5 MoE, Gemma-MoE) y el throughput de decodificacion con un solo usuario es tu cuello de botella. - Tienes una carga de trabajo donde una pequeña perdida de calidad es aceptable a cambio de una mejora visible en latencia. - Despliegas en hardware limitado por ancho de banda de memoria (Apple Silicon serie M) donde el gather de experts domina el tiempo de decodificacion por paso.

No lo uses cuando: - Sirves modelos densos: el flag es un no-op y no aporta nada. - Te importa la precision en el top-1% de suites de evaluacion de leaderboard. - Ejecutas generaciones largas de chain-of-thought o "modo thinking", donde el acantilado de calidad puede ser más pronunciado que lo que sugiere MMLU en 0-shot.

Combinacion con otras optimizaciones

Este flag se compone con la cuantizacion. En Qwen3-30B-A3B-4bit nuestra combinacion medida es:

  • 4-bit + top_k=8: 126.5 tok/s (baseline)
  • 4-bit + top_k=4: 147.3 tok/s (+16.5%)
  • 3-bit + top_k=8: 138.6 tok/s (+9.6%)
  • 3-bit + top_k=6: 147.1 tok/s (+16.3%) . divergencia de calidad medible
  • 3-bit + top_k=4: 157.3 tok/s (+24%) . la calidad de salida se rompe (el modelo respondió una pregunta diferente en nuestra prueba de humo)

3-bit + top_k=4 acumulo el error numérico más alla del punto donde el argmax es estable. Usa a lo sumo un parámetro agresivo: ya sea 4-bit + top_k=4 o 3-bit + top_k=6. Ambos dan aproximadamente el mismo tok/s (~147) con perfiles de calidad muy distintos.

Internos

  • Helper de parcheo: vllm_mlx.scheduler.apply_moe_top_k_override(model, k)
  • Se aplica en Scheduler.__init__ despues de cargar el modelo.
  • Tests: tests/test_moe_top_k.py. cubre modelos densos, arquitecturas mixtas y rutas de validación.

Referencias