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.
replyHow exactly the map-reduce will happen though? Won’t you need to do it host-side, or make a lot of reads and writes?
replyAlso, doesn’t it mean that you forgo batching?
deleted
replyMap-reduce is implemented as a rolling calc, see: online softmax in FlashAttention kernels.
replyRollie 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.
replyHow 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.
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.
reply