UltraQuant: 4-bit KV Caching for Context-Heavy Agents

2026-06-18

AI Summary

This work addresses the significant KV cache pressure and low GPU utilization caused by long prefix reuse and high concurrency in multi-turn, context-intensive agent tasks, which severely constrain inference latency and throughput. To tackle this challenge, the authors propose a 4-bit KV cache compression scheme tailored for such scenarios, featuring asymmetric K/V quantization, Walsh-Hadamard rotation, and the UltraQuant approximation pathway. The design further integrates FP8 queries, UE8M0 group scaling, and native scaled-MFMA support on CDNA4 architectures to enable an efficient decode-attention kernel. Evaluation on AMD GPUs demonstrates that, compared to an FP8 KV baseline, the approach reduces P50 first-token latency by 3.47× in later turns (2.3× overall) and improves output throughput by 1.63×, effectively balancing task quality, cache efficiency, and serving throughput.

Abstract

Context-heavy agents place unusual pressure on the key-value (KV) cache: long prefixes are reused across many short turns, while concurrency determines whether the serving system can keep GPUs utilized. We study 4-bit KV-cache compression for this setting, using TurboQuant-style rotation and codebook quantization as a quality anchor and vLLM FP8 KV caching as the deployment anchor. We report three contributions. First, we frame 4-bit KV caching around multi-round agent workloads where task quality, cache residency, and serving throughput must be measured jointly. Second, we describe the practical design choices needed to make the 4-bit path robust, including asymmetric K/V treatment, Walsh-Hadamard rotation, QJL removal, and block-scale variants. Third, we present serving optimizations on AMD GPUs, including optimized decode-attention kernels and UltraQuant, an FP4 approximation path that uses FP8 queries, FP4 KV tensors, UE8M0 group scales, and native scaled-MFMA support on CDNA4. On a long-context, multi-turn agentic workload, UltraQuant cuts P50 time-to-first-token by 3.47x in the cache-pressured late rounds (2.3x across all rounds) and raises output throughput by 1.63x over the FP8 KV baseline.

Problem

• KV caching
• context-heavy agents
• 4-bit quantization
• serving throughput
• long-context

Innovation

• 4-bit KV caching
• UltraQuant
• context-heavy agents
• Walsh-Hadamard rotation
• AMD CDNA4