TileLens: Efficiently Using Large-Granularity Memory Systems with Transparent Two-Dimensional Memory Layout

2026-07-04

AI Summary

This work addresses severe read amplification in large-granularity memory systems (LGMS) caused by the mismatch between one-dimensional memory layouts and the two-dimensional tiled computation patterns prevalent in large language model (LLM) inference. To align memory access with computational boundaries, the authors propose a tile-major two-dimensional memory layout through a hardware-software co-design approach. They introduce TileLens, comprising a GPU domain-specific language extension (TileLens-SW) and enhancements to tensor memory accelerators (TileLens-HW), augmented with adaptive prefetching. Notably, this framework efficiently supports tiled matrix multiplication without requiring kernel code modifications. Evaluated on an HBM-enhanced GPU, the approach incurs less than 1% inference latency overhead compared to a pure HBM baseline for Qwen-3 30B and Llama-3.1 70B, while achieving speedups of 1.61× to 6.49× over conventional memory layouts.

Abstract

Large Language Model (LLM) inference is bottlenecked by the capacity and bandwidth of GPU High-Bandwidth Memory (HBM). Recent proposals, such as High-Bandwidth Flash (HBF) and RoMe, offer higher capacity or bandwidth than HBM, but require a minimum access granularity of kilobytes. We show that these Large-Granularity Memory Systems (LGMS) can degrade the performance of tiled matrix-multiplication, which is the dominant operation in LLM inference, by up to an order of magnitude. The root cause of the slowdown is read amplification, where memory requests fetch far more data than the tile actually needs. This waste stems from a fundamental mismatch between the two-dimensional nature of compute tiles and the one-dimensional memory layout, leading to each request spilling well beyond the tile boundaries. To mitigate read amplification, we propose to use tile-major layout for LGMS. Rather than storing data as a one-dimensional strip, tile-major layout reshapes each contiguous memory block into a two-dimensional rectangle, aligning memory granularity with tile boundaries. To ease the adoption of tile-major layout on GPUs, we propose TileLens, lightweight software and hardware extensions that collectively cover major classes of GPU kernels. TileLens-SW extends GPU DSLs so that DSL-based kernels can adopt tile-major in global memory by changing only the layout descriptor. TileLens-HW extends the Tensor Memory Accelerator (TMA) for transparent tile-major support in TMA-based kernels without code changes. We evaluate TileLens on a cycle-level simulator using matrix-multiplication kernels from Qwen-3 30B and Llama-3.1 70B. Combining a tile-major layout with an adaptive hardware prefetcher, TileLens achieves near-HBM performance on HBF-augmented GPUs with a 5us HBF NAND read latency, reducing the geomean slowdown from 1.61-6.49x with conventional layouts to within 1% of an HBM-only baseline.

Problem

• Large-Granularity Memory Systems
• LLM inference
• read amplification
• memory layout
• tiled matrix-multiplication

Innovation

• tile-major layout
• Large-Granularity Memory Systems
• read amplification
• Tensor Memory Accelerator
• LLM inference