PIM-SHERPA: Software Method for On-device LLM Inference by Resolving PIM Memory Attribute and Layout Inconsistencies

O artigo apresenta o PIM-SHERPA, um método puramente de software que otimiza a inferência de LLMs em dispositivos ao resolver inconsistências de atributos e layout de memória em sistemas com Processamento em Memória (PIM), alcançando economias significativas de capacidade sem comprometer o desempenho.

O essencial

Imagine que você tem um chef de cozinha genial (o seu celular) que precisa cozinhar um prato complexo usando uma receita gigante (o Modelo de Linguagem, ou LLM).

O problema é que a cozinha tem duas fases muito diferentes:

1. A Fase de Preparo (Prefill): Você lê a receita inteira de uma vez para entender o contexto. É rápido e exige muita força mental (computação).
2. A Fase de Servir (Decode): Você serve o prato prato por prato, um pedaço de cada vez. Isso exige muita velocidade para pegar os ingredientes da despensa (memória) sem atrasar o serviço.

Aqui entra a tecnologia PIM (Processamento em Memória). Pense nela como uma despensa inteligente que não só guarda os ingredientes, mas também os corta e mistura enquanto você os pega. Isso é incrível para a "Fase de Servir", pois acelera tudo.

O Grande Problema: A "Batalha das Regras"

O artigo explica que, para usar essa despensa inteligente (PIM), o sistema precisa seguir regras estritas que conflitam com as regras da "Fase de Preparo". É como se a despensa inteligente exigisse que os ingredientes estivessem em caixas de metal (não-cacheáveis) para funcionar, mas o chef, para preparar a receita, precisasse que eles estivessem em prateleiras de vidro acessíveis (cacheáveis) para pegar rápido.

Além disso, a organização dos ingredientes na despensa inteligente é diferente da organização na cozinha principal.

• Na Cozinha (Host): Os ingredientes estão organizados em fileiras e colunas para o chef pegar rápido.
• Na Despensa Inteligente (PIM): Eles precisam estar empilhados de um jeito específico para a máquina cortar rápido.

Soluções antigas (e ruins):
Antes, a única solução era ter duas cópias de tudo: uma organizada para o chef e outra organizada para a máquina.

• Analogia: É como ter que comprar duas geladeiras cheias de comida só para não ter que reorganizar os ingredientes. Isso ocupa o dobro do espaço na sua cozinha (memória do celular), o que é impossível em celulares pequenos.

A Solução: PIM-SHERPA

Os autores criaram o PIM-SHERPA, um "guia de cozinha" (software) que resolve esse conflito sem precisar de duas geladeiras. Eles usam duas estratégias criativas:

1. O "Duplo Balde" (DRAM Double Buffering - DDB)

Imagine que você tem dois baldes pequenos na pia.

• Enquanto o chef está cozinhando o prato atual usando o Balde A, um ajudante (software) já está correndo para a despensa inteligente, pegando os ingredientes do próximo prato, reorganizando-os e colocando no Balde B.
• Assim que o prato atual termina, o chef troca para o Balde B e o ajudante começa a encher o Balde A para o prato seguinte.
• O Truque: O tempo de correr até a despensa e reorganizar os ingredientes acontece enquanto o chef já está cozinhando. O tempo de espera é "escondido" dentro do tempo de trabalho.

2. A "Reorganização Sob Demanda" (OWR)

Para receitas muito longas (quando o usuário pede um texto enorme), o tempo de cozinhar é tão grande que o ajudante pode simplesmente reorganizar os ingredientes do prato atual antes de começar, sem precisar de dois baldes.

• Como o prato é gigante, o tempo que o ajudante gasta reorganizando é tão pequeno comparado ao tempo total de cozinhado que ninguém percebe o atraso.

Por que isso é um marco?

1. Economia de Espaço: Em vez de ocupar o dobro da memória (com duas cópias), o PIM-SHERPA usa apenas uma cópia dos ingredientes e um pequeno espaço extra (os baldes). Isso economiza cerca de 48% da memória do celular.
2. Sem Mudanças de Hardware: Eles não precisaram trocar a despensa ou a geladeira. Tudo foi resolvido com um novo "manual de instruções" (software) que o celular já consegue ler.
3. Velocidade: O sistema fica tão rápido quanto o ideal teórico, permitindo que celulares rodem modelos de IA gigantes (como o Llama 3.2) sem travar e sem precisar de uma internet super rápida.

Resumo da Ópera

O PIM-SHERPA é como um maestro genial que organiza a orquestra (o celular) para que a música (a IA) toque perfeitamente. Ele resolve o conflito entre "como a música é escrita" e "como os músicos tocam", permitindo que o celular use a tecnologia de memória mais rápida (PIM) sem precisar de mais memória física, mantendo tudo rápido e eficiente. É a chave para ter assistentes de IA poderosos direto no seu bolso, sem precisar de servidores na nuvem.

Resumo técnico

1. O Problema: Inconsistências em Sistemas PIM para LLMs

A implantação de Grandes Modelos de Linguagem (LLMs) em dispositivos de borda (como smartphones) enfrenta desafios críticos ao utilizar arquiteturas de Processamento em Memória (PIM), especificamente LPDDR-PIM. O inferência de LLMs ocorre em duas fases distintas com requisitos conflitantes:

1. Fase de Prefill (Preenchimento): Processa sequências longas de entrada. É intensiva em computação (GEMM - Multiplicação Matriz-Matriz) e beneficia-se de uma região de memória cacheável para maximizar a reutilização de dados.
2. Fase de Decode (Decodificação): Processa um token por vez. É intensiva em largura de banda de memória (GEMV - Multiplicação Matriz-Vetor) e utiliza PIM para acelerar o acesso. Para que o PIM seja ativado, os dados devem estar em uma região de memória não cacheável; caso contrário, os hits no cache impedem que as requisições cheguem ao controlador de memória, bloqueando a execução PIM.

As duas inconsistências principais identificadas são:

• Inconsistência de Atributo de Memória: O prefill exige pesos cacheáveis, enquanto o decode exige pesos não cacheáveis para ativar o PIM.
• Inconsistência de Layout de Pesos:
  • Host-friendly: Layout contíguo otimizado para intercalação de canais e bancos (ideal para o processador hospedeiro).
  • PIM-aware: Layout onde linhas inteiras da matriz residem dentro de um único banco de DRAM para maximizar a utilização SIMD local (ideal para PIM).

Soluções Anteriores e Limitações:

• Duplicação de Pesos: Manter duas cópias (uma cacheável/host-friendly e outra não cacheável/PIM-aware) resolve o problema, mas dobra o uso de DRAM, tornando inviável a execução de modelos grandes em smartphones com memória limitada (ex: 8-12 GB).
• Soluções de Hardware (ex: FACIL): Modificam o controlador de memória para mapear endereços dinamicamente, mas exigem alterações de hardware e não resolvem a inconsistência de atributo (cache vs. não cache).

---

2. Metodologia: PIM-SHERPA

O PIM-SHERPA é um método puramente de software que resolve essas inconsistências sem modificar o hardware, utilizando buffers de memória cacheáveis pequenos e cópias de memória "embaralhadas" (swizzled). O sistema opera em dois modos principais:

A. DRAM Double Buffering (DDB)

• Conceito: Utiliza dois buffers pequenos na região cacheável para esconder a latência de rearranjo de pesos durante a computação.
• Funcionamento:
  1. Enquanto a camada atual executa GEMM no buffer 0 (cacheável), os pesos da próxima camada são pré-busca (prefetch) da região não cacheável (PIM-aware) para o buffer 1.
  2. Durante a cópia, ocorre uma Swizzled Memory Copy (SMC): os pesos são convertidos do layout PIM-aware para o layout host-friendly durante a transferência.
  3. Os buffers alternam a cada camada do decodificador.
• Vantagem: A latência de rearranjo é sobreposta (overlapped) com o tempo de computação, eliminando a penalidade de latência para sequências de entrada suficientemente longas.

B. Online Weight Rearrangement with Swizzled Memory Copy (OWR)

• Conceito: Realiza a cópia e rearranjo dos pesos imediatamente antes da computação de cada camada, de forma serial.
• Funcionamento: Utiliza um único buffer cacheável. Os pesos são copiados e rearranjados (SMC) e, em seguida, o GEMM é executado.
• Aplicabilidade: Ideal para cenários com sequências de entrada muito longas, onde o tempo de computação (GEMM) domina o tempo total, tornando a latência estática da cópia insignificante em termos percentuais.
• Vantagem: Implementação mais simples, sem necessidade de sincronização complexa entre threads de cópia e computação.

Técnica Chave: Swizzled Memory Copy (SMC)
É o processo de mover dados da região não cacheável (PIM-aware) para a cacheável (host-friendly) enquanto reorganiza os dados. O sistema calcula o endereço físico na DRAM baseado no layout PIM-aware, lê os dados e os escreve no buffer cacheável no formato contíguo esperado pelo kernel GEMM padrão.

---

3. Contribuições Principais

1. Identificação do Problema: É o primeiro trabalho a identificar e formalizar a "inconsistência de atributo de memória" e a "inconsistência de layout de pesos" entre as fases de prefill e decode em sistemas PIM habilitados para LLMs.
2. Soluções Puramente de Software: Propõe duas técnicas (DDB e OWR) que resolvem ambos os problemas sem exigir duplicação de pesos (que consome muita memória) ou alterações no hardware (controlador de memória).
3. Eficiência de Capacidade: Demonstra que é possível manter o desempenho teórico máximo do PIM reduzindo drasticamente o uso de DRAM.
4. Validação em Plataforma Real: O método foi implementado e testado em um dispositivo comercial (Samsung Galaxy S24+) usando o runtime ExecuTorch, simulando o comportamento do LPDDR5X-PIM.

---

4. Resultados Experimentais

Os testes foram realizados com os modelos Llama 3.2 (1B e 3B) em um Samsung Galaxy S24+ (Exynos 2400).

• Economia de Capacidade de DRAM:

  • Comparado à duplicação de pesos (WD), o PIM-SHERPA economiza 47,8% a 49,7% de capacidade de DRAM.
  • Para o Llama 3.2 3B, a duplicação exigiria ~12 GB (inviável para smartphones), enquanto o PIM-SHERPA permite a execução com apenas ~6 GB.
  • O desempenho é comparável ao da solução teórica FACIL-O (que assume um controlador de memória modificado).

• Desempenho (TTFT - Time To First Token):

  • Para sequências de entrada longas (SL ≥ 128), o S-DDB atinge um TTFT quase idêntico ao do FACIL-O, escondendo eficazmente a latência de rearranjo.
  • O S-OWR mostra um desempenho ligeiramente inferior em sequências curtas devido à latência serial, mas aproxima-se dos outros métodos conforme o SL aumenta (ex: 16,7 TPS em SL 192 para o modelo 3B).

• Aceleração End-to-End:

  • O uso de PIM-SHERPA com LPDDR5X-PIM fornece um speedup de até 3,3x em comparação com uma base sem PIM (apenas CPU/DRAM).
  • Em cenários de memória limitada (onde apenas uma cópia dos pesos cabe), o PIM-SHERPA é a única solução viável para manter a aceleração PIM, enquanto métodos sem PIM (NC-GEMM) sofrem de gargalos de largura de banda severos.

---

5. Significado e Impacto

O PIM-SHERPA é um marco para a comercialização de PIM em dispositivos móveis e de borda.

• Viabilidade Prática: Remove a barreira de memória que impedia o uso de PIM em smartphones (onde a duplicação de pesos tornaria a execução de modelos grandes impossível).
• Flexibilidade: Oferece duas abordagens (DDB para latência otimizada e OWR para simplicidade de implementação) adaptáveis a diferentes cenários de carga de trabalho.
• Universalidade: A abordagem pode ser adaptada para GPUs e NPUs que possuem hierarquias de memória semelhantes (cache + memória principal), sugerindo que a resolução de inconsistências de layout e atributo é um passo crítico para o futuro da aceleração de IA em hardware.

Em resumo, o PIM-SHERPA permite que dispositivos móveis executem LLMs de última geração com a eficiência de largura de banda do PIM, resolvendo conflitos fundamentais de software que anteriormente exigiam soluções de hardware caras ou inviáveis.