BrIM: A Branching In-Memory Accelerator

While memory-centric (MC) AI accelerators have shown significant promise in AI computing by overcoming memory bottlenecks, these architectures still rely on traditional processors for some of the critical AI-computing tasks. Most notably, Tokenization, which is an integral part of processing Large Language Model (LLM) AI algorithms, involves decomposing the input data into smaller meaningful units (tokens), a task of a highly branching nature. Existing memory-centric systems, which typically do not support branching tasks, offload such workloads to separate general-purpose 'host' processors. However, such task offloading incurs inevitable back-and-forth movement of data between the external processor and the in-memory AI accelerator, leading to expected latency and energy overheads and a loss of performance. To alleviate such issues, we propose a programmable Look-up Table (LUT)-based in-DRAM processing architecture that efficiently processes branching kernels as well as data-parallel AI-oriented workloads so as to support Tokenization, along with the rest of the LLM workloads within the same memory chip. Our proposed solution offers remarkably superior performance optimization of tokenization compared to the traditional CPU+MC systems, with up to 8 × higher utilization of the compute bandwidth.