Integrating SSD-like high bandwidth flash (HBF) memory directly into GPUs offers large capacity but falls short in speed and write performance compared to existing high bandwidth memory (HBM) and multi-GPU setups, limiting its effectiveness for training and inference of large language models. Consequently, GPU manufacturers prioritize system-level solutions over embedding flash memory on-chip, as the memory bottleneck requires a more comprehensive approach beyond simply adding cheaper flash storage.
The idea of integrating solid-state drives (SSDs) directly into the GPU fabric as a cheap memory option seems promising, especially with recent efforts by SanDisk and SK Hynix to establish a high bandwidth flash (HBF) standard. This new tier of memory aims to provide large capacities, such as 2 TB, directly accessible by the GPU with very high throughput (up to 3 TB/s using the UCIe standard), potentially enabling large language models to stream their weights to the GPU core in less than a second. However, despite the apparent advantages, GPU manufacturers like Nvidia and AMD have been slow to adopt this technology.
One major challenge is that current high-end GPUs already use high bandwidth memory (HBM) with speeds around 4.8 TB/s, which is significantly faster than the proposed HBF. Large models are typically split across multiple GPUs, allowing each GPU to access a smaller portion of the model at very high speeds. This multi-GPU approach currently outperforms the single-GPU HBF solution in terms of speed and efficiency, making HBF less attractive for training large models.
When it comes to inference, using HBF to store KV cache (key-value cache) is problematic because KV cache grows dynamically and requires frequent write operations. NAND flash memory, which underpins SSDs and HBF, is much slower at writes compared to reads due to the need to erase data before rewriting. This inherent limitation means that HBF provides minimal latency improvements during inference, as demonstrated by research showing only about a 1% improvement even with theoretical speed increases.
Training large language models also sees limited benefit from HBF. Training workloads involve complex read-write operations, including generating gradients and updating model weights, which do not align well with the strengths of flash memory. While inference is more memory capacity and bandwidth-bound, training benefits more from scaling the number of nodes and faster interconnect bandwidth rather than increased memory capacity alone. Thus, HBF does not significantly improve training performance.
Ultimately, the memory bottleneck in GPUs remains a critical issue, but the solution may not lie in integrating cheap flash memory directly into the GPU. Nvidia, for example, is focusing on infrastructure-level solutions like their Context Memory Extension (CMX) and unified system pods combining storage and networking to handle large KV caches more effectively. Meanwhile, companies like SanDisk and SK Hynix, who do not control the AI software stack, face challenges in driving adoption of HBF. The problem of expanding GPU memory capacity is important but requires a more holistic approach beyond just adding flash memory at the chip level.
