ThAME: 3D Memory-Enabled Heterogeneous Accelerator for LLM Mixture of Experts

Artifact for ThAME: 3D Memory-Enabled Heterogeneous Accelerator for LLM Mixture of Experts.

ThAME is a heterogeneous 2.5D system-in-package: a TPU-v6e host chiplet runs the compute-bound prefill phase, the gating network and the core scheduler; 3D DRAM-PNM chiplets hold the KV cache and compute attention; and 3D FeFET-NAND chiplets store the expert weights and compute the experts on 32 compute cores in their CMOS base die, each core a 64×64 systolic array of PEs. Those base-die cores are wired by a hierarchical NoC — a local crossbar per core group under a tree backbone — obtained by multi-objective optimization (Section IV-C of the paper).