6 ms·

> You could run MACs directly in RAM Sure, MACs are nice. However, unless there other, PIM-specific/optimal, algorithms, regular matrix multiplication algorith

> You could run MACs directly in RAM Sure, MACs are nice. However, unless there other, PIM-specific/optimal, algorithms, regular matrix multiplication algorithms like tiling-based won’t work here I think — how would the tile be shared? By doing read/write all the time?

whatshisface 19d ago
Attention calculations aren't shared across more than one vector during next token prediction (thinking and writing) which this sounds almost perfect for. Per attention layer, for deepseek at 1M context, you want to broadcast a single 1KB vector to 4GB of dot products, and map reduce a 1KB vector back.
- reliabilityguy 19d ago
  How exactly the map-reduce will happen though? Won’t you need to do it host-side, or make a lot of reads and writes? Also, doesn’t it mean that you forgo batching?
  - deleted 19d ago
    [deleted]
  - ACCount37 19d ago
    Map-reduce is implemented as a rolling calc, see: online softmax in FlashAttention kernels.
    - reliabilityguy 19d ago
      Rollie calculation, like the online softmax in FA, implies a centralized computing unit that does the compute and stores the intermediate results in its registers. With PIM you have no centralized compute unit, you have a bunch of memory, and a bunch of MACs all over the place. How would you do map-reduce across multiple DIMMs w/o extra reads/writes? PIM implies some sort of distributed compute, which can work for some cases, but I am not sure LLMs are one of them.
      - whatshisface 19d ago
        Re-arrange x+(y+(z+w)) into (x+y)+(z+w) to accomplish a distributed sum. Shared bus, interconnect fabric. This is a supercomputer on a chip.
      - ACCount37 19d ago
        Is there a single reason why we can't just "distribute" the online softmax? Each die-attached PIM accelerator computes online softmax for its own KVs. Then the central unit gathers the softmax intermediates, one intermediate per die, and uses those to compute the final softmax. The PIM win is that we crater the memory traffic between the central accelerator and the memory dies for attention ops. Most of the attention bandwidth never leaves the memory. This isn't "run the entire LLM in PIM", no - this is "offload the parts of LLM that benefit from PIM the most to PIM".

whatshisface 19d ago

Attention calculations aren't shared across more than one vector during next token prediction (thinking and writing) which this sounds almost perfect for. Per attention layer, for deepseek at 1M context, you want to broadcast a single 1KB vector to 4GB of dot products, and map reduce a 1KB vector back.

reliabilityguy 19d ago
How exactly the map-reduce will happen though? Won’t you need to do it host-side, or make a lot of reads and writes? Also, doesn’t it mean that you forgo batching?
- deleted 19d ago
  [deleted]
- ACCount37 19d ago
  Map-reduce is implemented as a rolling calc, see: online softmax in FlashAttention kernels.
  - reliabilityguy 19d ago
    Rollie calculation, like the online softmax in FA, implies a centralized computing unit that does the compute and stores the intermediate results in its registers. With PIM you have no centralized compute unit, you have a bunch of memory, and a bunch of MACs all over the place. How would you do map-reduce across multiple DIMMs w/o extra reads/writes? PIM implies some sort of distributed compute, which can work for some cases, but I am not sure LLMs are one of them.
    - whatshisface 19d ago
      Re-arrange x+(y+(z+w)) into (x+y)+(z+w) to accomplish a distributed sum. Shared bus, interconnect fabric. This is a supercomputer on a chip.
    - ACCount37 19d ago
      Is there a single reason why we can't just "distribute" the online softmax? Each die-attached PIM accelerator computes online softmax for its own KVs. Then the central unit gathers the softmax intermediates, one intermediate per die, and uses those to compute the final softmax. The PIM win is that we crater the memory traffic between the central accelerator and the memory dies for attention ops. Most of the attention bandwidth never leaves the memory. This isn't "run the entire LLM in PIM", no - this is "offload the parts of LLM that benefit from PIM the most to PIM".