Exploiting Distributed Shared Memory in GPUs for Next-Gen AI workloads

Tree-structured prefix sharing is now standard in large language model (LLM) serving: a radix-tree KV cache lets multiple requests reuse common context prefixes and reduces memory capacity pressure. Yet the attention kernels used with these caches usually retain a conventional execution pattern in which each thread block independently loads its KV tiles from global memory. As a result, shared prefixes are still fetched repeatedly, and cross-request reuse is lost at the point where long-context decoding spends most of its bandwidth. Recent prefix-aware kernels reduce this redundancy, but their cooperation is limited by the shared-memory capacity of a single Streaming Multiprocessor (SM), forcing intermediate results that cross SM boundaries through global memory.

We present Capillary, an attention-kernel design that exploits Distributed Shared Memory (DSM) on NVIDIA Hopper GPUs to break this single-SM barrier. Capillary co-locates thread blocks that share KV prefixes within a hardware thread-block cluster and transfers partial attention results directly between SMs, avoiding global-memory round trips. A lightweight DSM-aware cost model selects query grouping and cluster configuration at each decoding step, balancing saved HBM traffic against DSM bandwidth and cluster scheduling constraints