dr: no Apple Silicon, a memória linear de um módulo WebAssembly pode ser compartilhada diretamente com a GPU: sem cópias, sem serialização, sem buffers intermediários. A CPU e a GPU leem e gravam os mesmos bytes físicos. Funciona de ponta a ponta: um convidado Wasm preenche uma matriz em sua memória linear, a GPU lê, calcula, escreve de volta e o convidado vê o resultado através do mesmo ponteiro, mesma memória, zero cópias.
Normalmente Wasm e GPUs são separados por um limite de serialização caro: na maioria dos hardwares, transferir dados de uma sandbox de VM para um acelerador significa copiar através de um barramento. A Arquitetura de Memória Unificada da Apple Silicon apaga essa fronteira (sem barramento, mesma memória física), e o que acontece é um tempo de execução onde Wasm é o plano de controle e o GPU é o plano de computaçãocom sobrecarga quase zero entre eles.
Estou construindo algo chamado Madeira flutuante que explora isso para inferência de IA com estado… e este post é sobre a base (como funciona a cadeia de cópia zero, o que medi, o que ela abre). Ainda cedo, ainda cutucando.
Por que isso normalmente é difícil
Informações rápidas, para quem não vive nesta pilha: o WebAssembly oferece uma sandbox. Seu módulo obtém uma matriz de bytes plana (memória linear) e esse é o universo… tudo fora é mediado por chamadas de função “host”. A questão toda é isolamento, portabilidade, determinismo.
GPU também deseja uma matriz de bytes simples, mas de um tipo específico: alinhado à página, fixado, acessível ao mecanismo DMA. Em uma GPU discreta (pense em NVIDIA ou AMD), essa memória fica em um barramento PCIe da CPU, portanto, obter dados da memória linear de um módulo Wasm para a GPU significa: copiar da sandbox para a memória do host e, em seguida, copiar através do barramento para a memória da GPU. Duas cópias, dois acertos de latência e uma estranha incompatibilidade de impedância entre “VM isolada” e “acelerador de hardware”.
Apple Silicon muda a física. A CPU e a GPU compartilham a mesma memória física (Arquitetura de Memória Unificada da Apple) … sem ônibus! Um ponteiro que a CPU pode ler, a GPU também pode ler, da mesma DRAM. A verdadeira questão: você pode passar esse ponteiro pelas camadas de abstração (o tempo de execução do Wasm, a API da GPU) sem que ninguém faça uma cópia defensiva ao longo do caminho?
Acontece que… você pode!
A cadeia de três elos
Três links. Validei cada um antes de tentar compô-los: é o tipo de coisa que se você pular a etapa de isolamento e todo o pipeline quebrar, você não tem ideia de “qual junta está vazando”.
Link 1: mmap oferece memória alinhada à página. No macOS ARM64, mmap com MAP_ANON | MAP_PRIVATE retorna endereços alinhados com 16 KB. Este não é um acidente de sorte, acontece que é o tamanho da página ARM64e mmap alinha por contrato. O alinhamento é importante porque o Metal requer isto.
Link 2: Metal aceita esse ponteiro sem copiar. MTLDevice.makeBuffer(bytesNoCopy:length:) envolve um ponteiro existente como um buffer Metal. No Apple Silicon, este é o caminho de cópia zeroou seja, a GPU acessa o mesmo memória física que a CPU faz. Eu verifiquei a identidade do ponteiro: o MTLBuffer.contents() ponteiro é igual ao original mmap ponteiro. Eu verifiquei sem cópias ocultas: O delta RSS foi de 0,03 MB (ruído de medição), em comparação com 16,78 MB para o caminho de cópia explícita. E a mesma latência de computação de qualquer maneira.
Link 3: Wasmtime permite que você traga seu próprio alocador. Estava na hora MemoryCreator trait permite controlar como a memória linear é alocada. Em vez de deixar o Wasmtime ligar mmap internamente, você mesmo fornece a memória de apoio. eu implemento MemoryCreator para devolver o nosso ter mmap região e Wasmtime’s memory.data_ptr() retorna exatamente o ponteiro que eu entreguei. O módulo Wasm lê e grava por meio da API de memória do Wasmtime; a GPU lê e grava no buffer Metal; ambos estão operando nos mesmos bytes.
A composição: alocar um mmap região, entregue-o para ambos Wasmtime (como a memória linear do ator) e Metal (como buffer de GPU). O módulo Wasm grava dados em deslocamentos conhecidos, a GPU calcula neles no local e os resultados aparecem na memória linear do módulo com não cópias e não transferência explícita de dados.
Testei a cadeia completa com uma multiplicação de matriz 128×128: o módulo Wasm preenche as matrizes A e B, a GPU executa um shader GEMM, o módulo lê o resultado C de volta. Zero erros em 16.384 elementos. Teste pequeno, mas é o tipo de coisa em que ou tudo se alinha ou você pega lixoentão zero erros é o sinal que eu queria.
O que eu medi
Três coisas que me importavam: identidade do ponteiro (é realmente cópia zero?), sobrecarga de memória (alguma cópia oculta aparecendo?) e correção (a GPU vê o que Wasm escreveu?).
Measurement Zero-copy path Copy path
─────────────────────────────────────────────────────────────
Pointer identity mmap == MTLBuffer different addrs
RSS delta (16 MB region) 0.03 MB 16.78 MB
GEMM latency (128×128) ~6.75 ms ~6.75 ms
Correctness (16K elements) 0 errors 0 errors
A equivalência de latência faz sentido: em UMA, a computação em si é idêntica de qualquer maneira. A imagem da memória é onde ela aparece: o caminho de cópia zero essencialmente não tem sobrecarga para tornar os dados acessíveis pela GPU, e o caminho de cópia dobra o consumo de memória.
Em tamanhos pequenos de tensores, ninguém se importa. Na escala dos caches KV na inferência do transformador (centenas de megabytes por conversação), é a diferença entre colocar quatro ou dois atores na memória. Esse é o regime em que eu realmente quero operar, então a parte da memória é importante.
Da cópia zero à inferência
Então agora tenho um primitivo: o Wasm e a GPU compartilham memória sem sobrecarga. O que você faz com isso?
Conectei a cadeia à estrutura MLX da Apple e executei o Llama 3.2 1B Instruct de um ator Wasm: um decodificador de transformador completo escrito em Rust, compilado para um tempo de execução de host nativo, conduzindo inferência na GPU Apple Silicon por meio de chamadas de função de host. (Eu estava com preguiça de criar um caminho de kernel personalizado do zero e … MLX estava lá)
Latências medidas, executando Llama 3.2 1B (quantizado de 4 bits, 695 MB) em um Macbook Pro 2021 M1 (laptop pessoal antigo, irei reavaliar um Mac Studio adequado algum dia, quando puder colocar as mãos nele 😄):
Operation Latency
──────────────────────────────────────
Model load (safetensors) 229 ms (one-time)
Prefill (5 tokens) 106 ms
Per-token generation ~9 ms
Host function boundary negligible
O limite da função host (o despacho Wasm para GPU) não é mensurável em relação ao custo de inferência. Qualquer pessoa que tenha trabalhado com tempos de execução em área restrita provavelmente estremeceu com a ideia de cruzar esse limite por despacho. Neste hardware, não é uma coisa.
Portabilidade de cache KV
Os transformadores mantêm um cache de valores-chave que acumula contexto ao longo dos turnos de conversação, o que normalmente é efêmero (matar o processo, perder o cache, recomeçar). Se você tentou executar inferência local, conhece a sensação.
Porque o cache reside na memória acessível pela GPU que eu controleposso serializá-lo. Então, eu despejo o cache KV no formato safetensors (serialização de tensor ML padrão, nada exótico) e o restauro mais tarde, na mesma máquina ou em uma máquina diferente, ou potencialmente contra um modelo diferente em uma máquina diferente! Esse último ainda não testei em arquiteturas significativamente diferentes… veremos.
Operation Latency Size
───────────────────────────────────────────────────
Serialize (24 tokens) 1.1 ms 1.58 MB (~66 KB/token)
Restore from disk 1.4 ms
Re-prefill from scratch 67.7 ms (the alternative)
───────────────────────────────────────────────────
Speedup from restore: 5.45×
Round-trip fidelity: bit-identical (10/10 tokens match)
5,45× em 24 tokens, e a proporção melhora com o comprimento do contexto: o tempo de restauração é quase constante, o pré-preenchimento é escalonado linearmente. Com 4.096 tokens, a restauração seria cerca de 100× mais rápida que a recomputação (não na verdade empurrou para 4.096 ainda; isso é a matemática do guardanapo extrapolando a partir da forma constante versus linear).
Esta é a base para mobilidade de atores estatais: congelar uma conversa no meio da conversa, movê-la para outro lugar, descongelá-la com todo o contexto intacto. A memória linear do módulo Wasm captura o estado lógico do ator; o cache KV captura o contexto acumulado do mecanismo de inferência. Junto: um instantâneo portátil de uma conversa de IA em execução (ou, pelo menos, esse é o plano 😅).
O que está sendo construído
Driftwood é um tempo de execução para atores Wasm com estado com inferência de GPU. O cadeia de cópia zero é a base: além disso vou adicionar instantâneos do ator (congelar e retomar qualquer conversa), portabilidade de ponto de verificação (mover o estado de inferência entre máquinas) e suporte multimodelo (o formato do instantâneo é independente do modelo, então, em teoria, o a identidade do ator sobrevive às trocas de modelo … qual poder funcionar, revisitarei assim que testar).
Isso tudo é cedo, ainda estamos costurando as coisas. Mas a “física” funciona: o Wasm e a GPU podem compartilhar memória no Apple Silicon sem sobrecarga, o cache KV é portátil e um transformador completo é executado a partir de um ator em sandbox em velocidade nativa. As próximas coisas que quero abordar: se o instantâneo realmente sobrevive a uma troca de modelo, se a cadeia se mantém em modelos maiores e se estou perdendo algum motivo óbvio para que isso caia em escala. Lento e constante…
Mais sobre o modelo de ator e a arquitetura do instantâneo em uma postagem futura, assim que eu tiver realmente lançado algo além do estágio de “trabalhos de física”.
Fonte: theverge
