perf(moe): prefetch the next layer's experts from a router lookahead during offloaded decode

Abierto
  • ¿Algo no está bien?
#677 1 comentario 0 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 1 día

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
5/5
Tiempo estimado
Más de una semana
Aptitud para principiantes
18/100
Tipo de issue
Nueva funcionalidad
Claridad
Bastante claro
Estado de actividad
Activo
Stack tecnológico
python
Área
ai, performance

Línea de trabajo

Start from the offloaded decode path behind --moe-strategy offload and the expert slot-map ensure logic, then read the linked fork commits and PR #74 to see the proposed router lookahead and prefetch install. Any change must keep every slot-map write on the compute stream, so review the race tests before anything else. Done means the FREETOKEN_MOE_PREFETCH=on path gives token-identical greedy output on cold prompts with the slot-map audit enabled.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

In offloaded decode (--moe-strategy offload, decode on the GPU), each layer's expert misses are copied over PCIe only after that layer's router has run, so the copy sits on the critical path. On an RTX 5080 (16 GB, PCIe) serving Qwen3.8-Flash-Next-NVFP4, the expert H2D copies are about half of each decode step.

Proposal: predict layer L's experts one layer early and copy them while layer L-1 computes.

  • Lookahead: run layer L's router on layer L-1's router input (the residual stream changes little between adjacent layers). Run it on a side stream forked at the start of L-1's MoE block.
  • Install: right after L-1's own ensure, install up to N predicted, non-resident experts as low-priority slots: held, usage 0, first to be evicted. Copy them on a dedicated high-priority stream behind L-1's demand copy. L's demand copy then waits on that stream.
  • Correctness constraint: every slot-map writer (demand ensure and prefetch install) must stay on the compute stream. Only reads and byte copies may cross streams. A first version that installed from the predictor stream behind an event join corrupted 1–6 of 60 cold greedy prompts. Two writers racing on the slot map was enough, even though an nsys trace showed no ordering violation.

Measured (Qwen3.8-Flash-Next-NVFP4, scored eviction with the shared-expert overlap as the baseline, N = 4):

baseline with prefetch
decode, 32K context 43.4 tok/s 53.0 tok/s (+22%)
decode, 255K context 43.0 tok/s 50.8 tok/s
demand misses per token 206 111
VRAM +148 MiB

The output is lossless: greedy outputs were token-identical on 420 cold nonce prompts, with a slot-map audit enabled.

Implementation (on my fork, against an older main):

  • 8c171b6: router lookahead (layer L's router on layer L-1's input).
  • 47d4811: prefetch into low-priority slots behind FREETOKEN_MOE_PREFETCH=on.
  • 3a0ff77: install prefetches on the compute stream, so the slot map has a single writer.
  • The full series, including the race tests and the slot-audit/verify debug modes: https://github.com/gdevenyi/FreeToken/pull/74
Lenguaje dominante
Python
Estrellas
14.2k
Forks
1.4k
Merge medio
1 d 10 h
PR fusionados (30 d)
107

Preparar el entorno

Este proyecto no incluye contenedor de desarrollo, Dockerfile ni guía de contribución, así que la configuración corre por tu cuenta: empieza por su README y consulta nuestra guía para la primera contribución para los pasos generales.

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de FlashML-org/FreeToken

Todos los issues de FlashML-org/FreeToken

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.