In this talk, the speaker from Cog discusses their efforts to build a real-time inference stack for large language models on AMD Instinct GPUs. They highlight the challenge that while AI is fast for humans, it remains slow for machines, especially in agentic workflows like software engineering where inference can take hours. Their goal is to achieve the fastest large language model inference on GPUs, specifically AMD GPUs, due to their advantageous hardware features such as large SRAM caches and register files. To maximize speed, they avoid third-party overhead by writing custom kernels in low-level assembly-like languages, focusing on zero-cost abstractions and hardware-specific optimizations.
A key insight shared is that the CPU often becomes a bottleneck in GPU inference workflows because it is slower at scheduling and queuing GPU kernels. To address this, Cog performs all operations, including token sampling, directly on the GPU to avoid costly CPU-GPU communication overhead. They also emphasize the importance of minimizing kernel launches by consolidating multiple operations into a single “mono kernel” that can generate thousands of tokens continuously without interruption. This approach helps overcome overheads related to kernel launch, memory streaming interruptions, and inter-GPU communication, which otherwise limit throughput despite the GPU’s raw computational power.
The team also developed a new transformer variant called the “Lane Former” to optimize inter-GPU communication by overlapping communication and computation, thereby reducing latency. They use advanced profiling tools to analyze memory streaming and instruction usage, identifying inefficiencies such as memory contention and unnecessary compiler-generated instructions. By hand-optimizing assembly code and carefully managing memory access patterns, they improve performance significantly. Their demo showcases a chat application running at nearly 3,000 tokens per second on AMD GPUs, which is notably faster than typical GPU-based inference speeds.
Augusta, a teammate, then delves into a specific technical optimization related to grid synchronization on the AMD MI300X GPU. Grid synchronization is critical for coordinating computation across GPU compute units but can be a major runtime bottleneck, consuming up to 40% of inference time. Traditional synchronization methods rely on kernel boundaries and hardware barriers, which are too slow or unavailable on the MI300X due to its single execution context design. Instead, they implement a custom barrier using memory polling without atomic counters, initially reducing synchronization overhead from 6,600 nanoseconds to around 750 nanoseconds.
Further optimization was achieved by reverse-engineering the MI300X’s memory architecture, which consists of multiple compute dice (XCDs) and IO dice connected to HBM memory stacks. By understanding how physical memory addresses map to these hardware units, they designed a synchronization scheme that ensures each compute unit accesses memory local to its IO die, minimizing costly cross-die memory hops. This approach not only stabilizes synchronization latency across compute units but also improves bandwidth by up to 50%. The talk concludes by noting that this hardware-aware optimization strategy can be extended to other memory mappings and that duplicating some computations can sometimes be more efficient than synchronizing across the entire device.
