LLM Inference Scheduling With Hybrid Batching and Split Memory Access
Find Innovative SolutionsGenerate Solutions
Solution Overview
Problem
Existing large language model (LLM) inference processes face bottlenecks due to low processor utilization and unbalanced computation and memory access latency, particularly in the decoding phase, which limits overall throughput and latency.
Innovation Solution
Implement a hybrid batching mechanism that combines decoding and prefilling tokens into larger batches, and split memory access tasks to be performed concurrently with computations, optimizing throughput and latency by increasing parallelism and balancing computation and memory access.
Engineering Contradictions & Design Principles
Engineering Contradiction Analysis
1Productivity
If traditional sequential processing is used for LLM inference, then computation accuracy is maintained, but processor utilization is low and latency is high
Solution Approach 1:
The patent segments the LLM inference process into distinct phases (prefilling phase and decoding phase) and further divides computations into multiple stages. By segmenting the processing workflow, the system can overlap different phases and stages, allowing processor units to work concurrently on different segments, thereby improving throughput and reducing idle time without compromising computation accuracy.
Solution Approach 2:
The patent implements preliminary action by performing prefilling computations (processing the entire prompt) before decoding begins, and by pre-loading necessary model weights and activations into high-speed memory. This preliminary preparation allows the decoding phase to proceed efficiently with minimal waiting time, reducing overall latency while maintaining processor utilization.
2Productivity
If batch size is increased to improve throughput, then processor utilization improves, but memory access latency increases due to unbalanced computation and memory access
Solution Approach 1:
The patent dynamically adjusts the batching strategy and memory access patterns based on the computation stage. During prefilling, larger batches are processed together, while during decoding, the system dynamically manages batch processing to balance memory access. This dynamic adaptation allows the system to maintain high throughput while optimizing memory access latency for different computational phases.
Solution Approach 2:
The patent ensures continuity of useful action by overlapping memory access operations with computation operations. While processors are computing activations for one batch, memory systems are simultaneously loading weights for the next batch or storing results. This continuous overlapping eliminates idle time in both computation and memory subsystems, maintaining high throughput without increasing memory access latency.
3Productivity
If decoding and prefilling are processed separately, then computation accuracy is maintained, but processor utilization is low
Solution Approach 1:
The patent merges the prefilling and decoding phases into a unified processing framework where both operations can coexist and overlap. Multiple batches can be in different phases simultaneously - some batches prefilling while others are decoding. This merging allows continuous processor utilization across different batch stages while maintaining the computational integrity and accuracy of each phase through proper isolation and ordering of operations.
Data Source
AI summary
A system comprising: a processing circuit; and a memory storing instructions, which, based on being executed by the processing circuit, cause the processing circuit to perform: identifying a first computation performed by a machine learning model; scheduling a first memory access task associated with a first portion of a first data with respect to the first computation; identifying a second computation performed by the machine learning model, wherein the second computation and the first computation are separate computations; and scheduling a second memory access task associated with a second portion of the first data with respect to the second computation, wherein the first portion and the second portion are different portions of the first data.