Enhancing LLM Inference Performance with Managed Tiered KV Cache and Intelligent Routing in Amazon SageMaker HyperPod

Optimizing LLM Inference with Managed Tiered KV Cache and Intelligent Routing

Sample Flow for Inference Requests with KV Caching and Intelligent Routing

Managed Tiered KV Cache

Intelligent Routing

Deploying the Managed Tiered KV Cache and Intelligent Routing Solution

Prerequisites

Preparing Your Model Deployment Manifest Files

Observability

Benchmarking

Conclusion

About the Authors

Boosting LLM Inference: Unlocking Speed and Efficiency with Managed Tiered KV Cache and Intelligent Routing

Modern AI applications, particularly those leveraging large language models (LLMs), require swift and cost-effective responses, especially in scenarios involving lengthy documents or extensive conversations. However, as the context length grows, LLM inference can quickly become sluggish and expensive, leading to exponential increases in latency and costs with each interaction.

Understanding the Bottleneck

The primary challenge with LLM inference lies in the need to recalculate attention mechanisms for previous tokens every time a new token is generated. This results in considerable computational overhead and high latency for longer sequences. To tackle this issue, key-value (KV) caching has emerged as a promising solution. It alleviates the bottleneck by storing key-value vectors from previous computations and reusing them, effectively reducing the inference latency and time-to-first-token (TTFT).

Intelligent Routing further enhances the efficiency of LLMs by directing requests with shared prompts to the same inference instance, allowing for the reuse of cached KV data. This not only accelerates processing but also diminishes latency. Despite its potential, many users express struggles with setting up and configuring KV caching and intelligent routing frameworks at production scale, often requiring significant experimental cycles.

Introducing Amazon SageMaker HyperPod Enhancements

We are thrilled to announce that Amazon SageMaker HyperPod now incorporates Managed Tiered KV Cache and Intelligent Routing capabilities via the HyperPod Inference Operator. These innovations deliver remarkable performance enhancements for LLM inference workloads, reducing TTFT by up to 40%, boosting throughput, and lowering compute costs by up to 25% for long context prompts and multi-turn chat dialogues.

Key Features

1. Managed Tiered KV Cache:

   • Automates management of attention states across CPU memory (L1) and distributed tiered storage (L2).
   • Offers configurable cache sizes and eviction policies, optimizing for resource utilization and cost-efficiency.

2. Intelligent Routing:

   • Implements configurable request routing to maximize cache hits using strategies like prefix-aware, KV-aware, and round-robin routing.

3. Observability:

   • Built-in integration with Amazon Managed Grafana for observability of metrics and logs related to the Managed Tiered KV Cache and Intelligent Routing.

How It Works: Sample Inference Flow

Here’s a simplified breakdown of how inference requests are handled with these features:

1. A user sends an inference request to the HyperPod Load Balancer.
2. The Load Balancer forwards the request to the Intelligent Router, which dynamically directs it to the most suitable model pod based on the routing strategy.
3. The model pod checks the L1 cache for frequently used key-value pairs; if not found, it queries the shared L2 cache.
4. If data is still unavailable, full computation is performed, and results are cached for future use.

Benefits Across Industries

These optimizations have transformative impacts across several sectors:

• Legal: Legal teams can analyze 200-page contracts and receive instant answers to follow-up questions, drastically reducing wait times.
• Healthcare: Chatbots facilitate seamless, natural conversations across more than 20 patient dialogue turns, enhancing patient experience.
• Customer Service: High-performance systems manage millions of queries daily while significantly reducing infrastructure costs.

Deployment and Configuration

Prerequisites

• Create a HyperPod cluster with Amazon EKS as orchestrator.
• Confirm the status of the HyperPod cluster and inference operator is running.

Model Deployment Manifest

You can enable the new features by adding configurations to your InferenceEndpointConfig custom CRD file. Here’s how to prepare your deployment:

apiVersion: inference.sagemaker.aws.amazon.com/v1
kind: InferenceEndpointConfig
metadata:
  name: demo
  namespace: default
spec:
  modelName: "Llama-3.1-8B-Instruct"
  instanceType: "ml.g5.24xlarge"
  kvCacheSpec:
    enableL1Cache: true
    enableL2Cache: true
    l2CacheSpec:
      l2CacheBackend: "tieredstorage"
  intelligentRoutingSpec:
    enabled: true
    routingStrategy: prefixaware

Observability and Metrics

With these new features, monitoring through SageMaker HyperPod’s observability capabilities allows you to track KV Cache metrics effectively through the inference dashboard.

Benchmarking Results

Comprehensive benchmarking confirms substantial improvements in real-world performance. In scenarios leveraging the Llama-3.1-70B-Instruct model, significant reductions in TTFT, throughput increases, and cost reductions were achieved, particularly benefiting long context workloads.

Conclusion

Amazon SageMaker HyperPod’s Managed Tiered KV Cache and Intelligent Routing offer powerful solutions to optimize LLM inference performance and costs via efficient memory management and intelligent request routing. Starting today, you can capitalize on these configurations in your HyperPod model deployments.

For more detailed guidance, visit the Amazon SageMaker HyperPod documentation and explore the model deployment getting started guide.

About the Authors

This post has been collaboratively written by a talented team of engineers and product managers at AWS, each bringing a unique set of experiences and skills to drive innovation in AI and machine learning infrastructure. Their collective efforts are focused on making advanced AI capabilities more accessible and efficient for enterprises looking to harness the power of large language models.