Machine learning model optimization method and device, computer device, readable storage medium and program product
By performing static operator pre-compilation and dynamic operator combination on large-scale language models, the problem of low resource utilization efficiency in existing technologies is solved, achieving more efficient end-to-end inference speed and throughput.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- SHANGHAI BIREN TECH CO LTD
- Filing Date
- 2026-02-13
- Publication Date
- 2026-04-21
AI Technical Summary
In existing technologies, the decomposition of the computation graph of large-scale language models into fine-grained operators leads to frequent communication and parameter passing between the CPU and GPU, resulting in high end-to-end inference latency and low resource utilization efficiency.
By performing computational logic analysis on the original model, it is divided into static operators and dynamic operators. Adjacent static operators are pre-compiled into static computation units, which are then combined with dynamic computation units to form an optimized model, reducing the number of kernel startups and improving resource utilization.
It significantly reduces kernel startup scheduling overhead, improves end-to-end inference speed and overall throughput, and optimizes online service performance.
Smart Images

Figure CN121707015B_ABST
Abstract
Description
Technical Field
[0001] This application relates to the field of machine learning technology, and in particular to a machine learning model optimization method, apparatus, computer device, computer-readable storage medium, and computer program product. Background Technology
[0002] With the rapid development of machine learning technology, especially large-scale language model technology, its application in many fields such as text generation, dialogue interaction, and code completion is becoming increasingly in-depth, which puts forward the need for high throughput and low latency for online inference services of models.
[0003] In related technologies, to improve the utilization efficiency of hardware computing resources such as GPUs during model computation and enhance service efficiency, the entire computation graph is decomposed into a series of fine-grained basic operators (such as linear layers, normalization layers, activation functions, etc.). Each time a batch of inference is executed, the scheduler needs to initiate a separate kernel startup call for each operator in the computation graph. For a complex Transformer layer, this could mean a dozen or more independent startups. Each kernel startup involves communication between the CPU and GPU, parameter passing, context preparation, and other operations, the cumulative time of which accounts for a considerable proportion of the end-to-end inference latency, becoming a bottleneck limiting further performance improvements.
[0004] Therefore, how to improve the resource utilization efficiency of large language models is a technical problem that needs to be solved. Summary of the Invention
[0005] Therefore, it is necessary to provide a method, apparatus, computer device, computer-readable storage medium, and computer program product that can optimize machine learning models with improved resource utilization efficiency to address the aforementioned technical problems.
[0006] Firstly, this application provides a method for optimizing a machine learning model, the method comprising:
[0007] Obtain the original model;
[0008] The original model is subjected to computational logic analysis to obtain the corresponding original computation graph; wherein, the original computation graph represents the data dependencies between the various operators included in the original model;
[0009] Determine how the computational logic of each operator changes due to the input data, and classify each operator into a static operator or a dynamic operator based on the changes.
[0010] The adjacent static operators in the original computation graph are pre-compiled to obtain at least one static computation unit; wherein the pre-compilation is used to enable the adjacent static operators to be executed by a single scheduling instruction;
[0011] Generate dynamic computation units corresponding to the dynamic operators in the original computation graph to obtain at least one dynamic computation unit;
[0012] Based on the data dependencies, the at least one static computing unit is combined with the dynamic computing unit to obtain the optimized model corresponding to the original model.
[0013] In some embodiments, performing computational logic analysis on the original model to obtain the corresponding original computational graph of the original model includes:
[0014] The computational operations included in the original model are traversed, and each independent computational operation is constructed as a node corresponding to an operator in the original computational graph;
[0015] Based on the data input-output relationship between the operator nodes, a directed edge is established in the original computation graph from the source node of the output data to the destination node of the input data.
[0016] In some embodiments, the change indicates whether the computational logic of the operator changes due to changes in the input data of the current operator;
[0017] The step of determining the computational logic of each operator in response to changes in input data, and classifying each operator into a static operator or a dynamic operator based on these changes, includes:
[0018] For each of the operators, determine whether the computational logic of the current operator includes decision logic that depends on the content of the input tensor; the decision logic includes at least one of conditional branching, dynamic indexing, or sorting and selection operations;
[0019] If the computational logic includes the decision logic, the current operator is determined as the dynamic operator;
[0020] If the computational logic does not include the decision logic, the current operator is determined as the static operator.
[0021] In some embodiments, precompiling adjacent static operators in the original computation graph to obtain at least one static computation unit includes:
[0022] In the original computation graph, all the maximum connected subgraphs consisting of static operator nodes and their direct data dependency edges are identified.
[0023] For each identified maximum connected subgraph, the static operator sequence corresponding to the maximum connected subgraph is executed in the model inference environment to obtain the computational flow record corresponding to the maximum connected subgraph; the computational flow record includes at least one of the following: the start instruction of all operators in the current execution process, the computational logic, and the memory access operation.
[0024] The computation flow record is compiled, and the compilation result is encapsulated into an independent executable object that can be replayed by the single scheduling instruction, thus obtaining a pre-compiled static computation unit.
[0025] In some embodiments, generating dynamic computation units corresponding to the dynamic operators in the original computation graph to obtain at least one dynamic computation unit includes:
[0026] For each of the aforementioned dynamic operators, a unified function call prototype is defined for each dynamic operator. The function call prototype is used to characterize the memory address parameters of the input data and output data of the dynamic operator.
[0027] The native computation kernel of the dynamic operator is encapsulated into a callable function that conforms to the function call prototype, thereby obtaining the dynamic computation unit corresponding to each dynamic operator.
[0028] In some embodiments, combining the at least one static computing unit with the at least one dynamic computing unit according to the data dependency relationship to obtain the optimized model corresponding to the original model includes:
[0029] Based on the data dependencies between the operator nodes in the original computation graph, the connection order between the static computation unit and the dynamic computation unit is determined;
[0030] The static computing unit and the dynamic computing unit are combined according to the connection order to obtain the optimized model.
[0031] In some embodiments, the original model includes a large language model based on the Transformer architecture; the optimized model includes at least one Transformer layer; each Transformer layer includes a first static computation unit, a first dynamic computation unit, and a second static computation unit connected in sequence; the first static computation unit is used for preprocessing of attention mechanism computation; the first dynamic computation unit is used for performing attention mechanism computation; and the second static computation unit is used for postprocessing of attention mechanism computation results.
[0032] In some embodiments, the original model includes a dense feedforward network model; the second static computation unit encapsulates all computation operators in the current Transformer layer after computation by the attention mechanism.
[0033] In some embodiments, the original model includes a sparse expert hybrid network model; the second static computing unit encapsulates the router linear layer operator of the current Transformer layer; the optimized model is followed by a second dynamic computing unit and a third dynamic computing unit in sequence after the second static computing unit; wherein, the second dynamic computing unit is used to perform an expert selection operation based on the router output; the third dynamic computing unit is used to perform sparse matrix multiplication calculations for the selected expert based on the result of the expert selection operation.
[0034] In some embodiments, the data processing using the optimized model includes:
[0035] Receive multiple online inference requests and dynamically combine the multiple online inference requests into an input batch tensor according to a preset capacity constraint;
[0036] The input batch tensor is input into the optimized model to execute the static calculation unit and the dynamic calculation unit in the optimized model in sequence;
[0037] The calculation results of the optimized model are collected and output as the data processing results corresponding to the online inference request.
[0038] Secondly, this application also provides a machine learning model optimization apparatus, the apparatus comprising:
[0039] The acquisition module is used to acquire the original model;
[0040] The analysis module is used to perform computational logic analysis on the original model to obtain the corresponding original computation graph of the original model; wherein, the original computation graph represents the data dependencies between the various operators included in the original model;
[0041] The determination module is used to determine how the computational logic of each operator changes due to the input data, and to classify each operator into a static operator or a dynamic operator based on the changes.
[0042] A pre-compilation module is used to pre-compile adjacent static operators in the original computation graph to obtain at least one static computation unit; wherein, the pre-compilation is used to enable the adjacent static operators to be executed by a single scheduling instruction;
[0043] A generation module is used to generate dynamic computation units corresponding to the dynamic operators in the original computation graph, thereby obtaining at least one dynamic computation unit.
[0044] The combination module is used to combine the at least one static computing unit with the dynamic computing unit according to the data dependency relationship to obtain the optimized model corresponding to the original model.
[0045] Thirdly, this application also provides a computer device, including a memory and a processor, wherein the memory stores a computer program, and the processor executes the computer program to implement the steps included in any of the aforementioned machine learning model optimization methods.
[0046] Fourthly, this application also provides a computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the steps included in any of the aforementioned machine learning model optimization methods.
[0047] Fifthly, this application also provides a computer program product, including a computer program that, when executed by a processor, implements the steps included in any of the aforementioned machine learning model optimization methods.
[0048] The aforementioned machine learning model optimization methods, devices, computer equipment, computer-readable storage media, and computer program products, through the partitioning of dynamic and static operators and the pre-compilation of static operators, integrate a large number of discretely executed static operators in the original model into static computation units that can be executed by a single scheduling instruction. This directly and significantly reduces the kernel startup scheduling overhead caused by the traditional fine-grained operator execution method, comprehensively improves the end-to-end inference speed and overall throughput of large language models, and fully optimizes the performance indicators of online services. Attached Figure Description
[0049] To more clearly illustrate the technical solutions in the embodiments of this application or related technologies, the drawings used in the description of the embodiments of this application or related technologies will be briefly introduced below. Obviously, the drawings described below are only some embodiments of this application. For those skilled in the art, other related drawings can be obtained based on these drawings without creative effort.
[0050] Figure 1 This is a flowchart illustrating a machine learning model optimization method in one embodiment;
[0051] Figure 2 This is a diagram illustrating the application environment of a machine learning model optimization method in one embodiment.
[0052] Figure 3 This is a structural block diagram of a machine learning model optimization device in one embodiment;
[0053] Figure 4 This is an internal structural diagram of a computer device in one embodiment. Detailed Implementation
[0054] To make the objectives, technical solutions, and advantages of this application clearer, the following detailed description is provided in conjunction with the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are merely illustrative and not intended to limit the scope of this application.
[0055] It should be noted that the terms "first," "second," etc., used in this application can be used to describe various elements, but these elements are not limited by these terms. These terms are only used to distinguish the first element from the second element. The terms used in this application include and have, and any variations thereof, are intended to cover non-exclusive inclusion. The term "multiple" in this application refers to two or more. The terms used in this application and / or refer to one of the embodiments, or any combination of multiple embodiments.
[0056] Explanation of related terms:
[0057] Transformer: A deep neural network architecture based on self-attention. Its core feature is its ability to process the entire input sequence in parallel and effectively capture long-term dependencies within the sequence. The model typically consists of stacked encoders and decoders, with each layer containing sub-modules such as self-attention and feed-forward networks. A significant number of large language models (such as GPT, BERT, and LLaMA) are built upon this architecture.
[0058] Attention (self-attention mechanism): A core component of the Transformer model, used to calculate the association weights between different positions in the input sequence.
[0059] FFN (Feed-Forward Network): The module that follows Attention in the Transformer layer. It usually consists of two linear layers and an activation function, and is the standard structure of dense models.
[0060] MoE (Mixture of Experts): A sparse model architecture in which each layer contains multiple expert feedforward networks, and some experts are dynamically selected to participate in the computation through routers.
[0061] n-flight Batching (Dynamic Batching / Online Batching): A scheduling strategy that combines batches in real time as requests arrive to balance latency and throughput.
[0062] max_num_tokens (maximum total number of tokens): The upper limit of the sum of tokens for all requests in the dynamic batch pool.
[0063] max_batch_size (maximum batch size): The upper limit on the number of requests that can be accommodated simultaneously in a dynamic batch pool.
[0064] vLLM, TGI, TensorRT-LLM: Three common large model inference optimization systems / frameworks.
[0065] Graph (Computation Graph): A GPU optimization technique used to capture a series of kernel startup and memory operations into a single replayable unit to reduce startup overhead.
[0066] Kernel launch refers to the operation of initiating a computation kernel function execution on the GPU, and each launch incurs a fixed overhead.
[0067] RMSNorm (Root Mean Square Normalization): A layer normalization technique used to stabilize the training and inference of deep neural networks.
[0068] Linear / ColumnParallel / RowParallel (linear layer / column-parallel linear layer / row-parallel linear layer): basic matrix multiplication operations, which are the main body for model parameter calculation.
[0069] SiLU (Sigmoid Linear Unit), GeLU (Gaussian Error Linear Unit), and SwiGLU (Swish-Gated Linear Unit): These are three activation functions commonly used in feedforward neural networks.
[0070] Residual-Add: An operation that adds the input and output of a module to mitigate gradient vanishing.
[0071] RoPE (Rotary Positional Encoding): A method for injecting absolute and relative positional information into the query (Q) and key (K) vectors in an attention mechanism.
[0072] Softmax (soft maximization function): A mathematical function used in attention mechanisms to normalize attention scores into a probability distribution.
[0073] FlashAttention (an efficient self-attention computation algorithm and kernel implementation that optimizes memory access through I / O awareness): is a representative of dynamic computing kernels.
[0074] Top-K (Top-K selection): The operation of selecting the K largest values and their indices from a set of values. In MoE, it is used to select experts.
[0075] MoE-Expert-GEMM (MoE Expert Matrix Multiplication): refers to the sparse matrix multiplication calculation performed after routing the data to K selected experts.
[0076] O_Linear (Output Linear Layer): After the attention mechanism, the multi-head attention output is projected back to a linear layer in the model dimension.
[0077] Router_Linear: The linear layer in the MoE model used to compute the weights of each token for each expert.
[0078] KV-cache (Key-Value Cache): In autoregressive generation, to accelerate computation, the mechanism caches the key and value tensors calculated in all previous time steps for later use.
[0079] TTFT (Time to First Token): The time required from issuing a request to receiving the first token generated by the model is a key indicator for measuring system response speed.
[0080] like Figure 1As shown, a machine learning model optimization method is provided and applied to an artificial intelligence chip. In this embodiment, the artificial intelligence chip is any one of GPU (Graphics Processing Unit), TPU (Tensor Processing Unit), NPU (Neural Network Processing Unit), DPU (Deep Learning Processing Unit), APU (Accelerated Processing Unit), and GPGPU (General-Purpose Graphics Processing Unit). This embodiment does not specifically limit the specific type of chip, and the following description uses GPGPU as an example.
[0081] Reference Figure 2 The diagram shows a schematic of a GPGPU. A GPGPU is actually an array of Streaming Processor Clusters (SPCs), including, for example,... Figure 2 The diagram shows streaming processor clusters 1, ..., M, where M is a positive integer greater than 1. In a graphics processing unit (GPU), one streaming processor cluster processes one computational task, or multiple streaming processor clusters process one computational task. Multiple streaming processor clusters share data through a global cache or global memory. For example... Figure 2 As shown, taking streaming processor cluster 1 as an example, one streaming processor cluster includes multiple computing units, such as... Figure 2 The system is structured as Compute Unit 1, Compute Unit 2, ..., Compute Unit N, where N is a positive integer. Each Compute Unit (CU) performs arithmetic and logical operations other than matrix calculations such as matrix multiplication and convolution, including operations like accumulation, reduction, and standard addition, subtraction, multiplication, and division. A single Compute Unit contains multiple cores (also called computational kernels). Each core includes an Arithmetic Logic Unit (ALU), a floating-point unit, etc., and is used to execute specific computational tasks. Furthermore, the Compute Unit also includes registers (e.g., ...). Figure 2 The register file and shared cache in a computing unit are used to store source and destination data related to computing tasks in a hierarchical manner. The shared cache in a computing unit is used to share data between the cores of that computing unit.
[0082] In parallel computing, computational tasks are typically executed by multiple threads. These threads are divided into multiple thread blocks before execution in a general-purpose graphics processor (or parallel computing processor), and then dispatched via a thread block distribution module. Figure 2 (Not shown in the image) Multiple thread blocks are distributed to various computation units. All threads in a thread block need to be assigned to the same computation unit for execution. Simultaneously, thread blocks are broken down into minimum execution thread bundles (or simply warps), each containing a fixed number (or less than this fixed number) of threads, for example, 32 threads. Multiple thread blocks can execute in the same computation unit or in different computation units.
[0083] In each computing unit, the thread beam scheduling / distribution module ( Figure 2 (Not shown in the diagram) Thread bundles are scheduled and allocated so that multiple computing cores within the computing unit can run thread bundles. Depending on the number of computing cores in the computing unit, multiple thread bundles within a thread block can execute concurrently or in a time-sharing manner. Multiple threads within each thread bundle execute the same instructions. Memory execution instructions are issued to the shared cache within the computing unit or further issued to intermediate-level caches, global caches, or global memory for read and write operations, etc. Figure 2 As shown, the streaming processor cluster 1 also includes a tensor operation unit, which is used to perform tensor calculations, such as matrix multiplication, convolution operations, etc.
[0084] Reference Figure 1 As shown in the figure, this application provides an optimization method for distributed operators, applied to artificial intelligence chips. Each streaming processor cluster of the artificial intelligence chip is equipped with a global address mapping unit, which records the shared memory address information of the streaming processor clusters of the artificial intelligence chip involved in the distributed operator. The method may include the following steps:
[0085] Step 202: Obtain the original model.
[0086] The original model can be a machine learning model with fixed parameters, trained in previous steps. The original model defines the computational functions from input data to output results. The original model can be stored as a serialized file, such as a model file using ONNX (Open Neural Network Exchange) format, PyTorch's TorchScript format, or TensorFlow's Saved Model format. These formats encapsulate the model's network structure (operator types, connection relationships) and all trained weight parameters. Specifically, the original model can include large language models, such as the GPT, LLaMA, and Mixtral series models based on the Transformer architecture.
[0087] Given the characteristics of large language models, such as massive parameter counts, complex computational graphs, and sensitivity to inference latency, their computational graphs are typically composed of dozens or even hundreds of stacked Transformer layers. Each layer contains complex modules such as self-attention, feedforward networks (FFN), or expert hybrids (MoE). In traditional inference, these modules are decomposed into numerous fine-grained operators (such as multiple linear layers, normalization layers, and activation functions), resulting in a linear multiplication of kernel launches with the number of model layers. This becomes the main bottleneck for end-to-end latency (accounting for 15%-25%). Therefore, optimizing large language models can significantly improve the efficiency and resource utilization of data processing (such as inference computation).
[0088] Optionally, the original model can be a large language model based on the Transformer architecture (such as GPT, LLaMA, BERT, etc.), or it can be various machine learning models that rely on complex computational graphs, such as including but not limited to: text processing models: such as machine translation, text summarization, sentiment analysis, etc.; image and video processing models: such as object detection, image generation, video understanding, etc.; speech recognition and synthesis models: such as speech-to-text, voiceprint recognition, speech synthesis, etc.; multimodal fusion models: such as visual language models, cross-modal retrieval, etc.; scientific computing and simulation models: such as physical simulation, molecular dynamics, climate prediction, etc.
[0089] Step 204: Perform computational logic analysis on the original model to obtain the corresponding original computation graph of the original model; wherein, the original computation graph represents the data dependencies between the various operators included in the original model.
[0090] The computational logic refers to the operational steps and data flow rules followed by the original model during forward inference. The computational logic can include all basic computational operations (i.e., operators) contained in the original model, as well as the execution order and data dependencies between these operators. The computational logic of the original model can be analyzed using a pre-defined model parser (such as the graph optimization tool in ONNX Runtime or the jit.trace mechanism in PyTorch). By traversing the network definition of the original model, the parser identifies each layer and each independent computational unit included in the original model. A computational unit refers to an independently schedulable execution module encapsulated at the software level. A computational unit can include: a static computational unit: formed by fusing a group of adjacent static operators with fixed computational logic through pre-compilation techniques (such as graph capture and replay), which can be triggered by a single scheduling instruction (such as cudaGraphLaunch) to execute all its internal operators consecutively. A dynamic computational unit: a callable function module formed by encapsulating dynamic operators whose computational logic depends on input data through standardized interfaces, retaining the dynamism and high performance of its original kernel. The computation unit in this embodiment serves as a logical execution unit, used to reconstruct the computation graph of the model at the software level. It aggregates fine-grained operators into larger-grained execution modules, thereby reducing kernel startup times, optimizing hardware scheduling efficiency, and improving overall inference performance. The original computation graph can be a directed graph data structure. In the original computation graph, each node represents an operator (such as a matrix multiplication layer or a normalization layer), and each directed edge represents the data dependency between nodes, pointing from the source operator that generates data to the destination operator that consumes data. The original computation graph can be viewed as a visualization and formalization of the computation process of the original model, providing a structured data foundation for subsequent automated analysis, classification, and reconstruction, serving as the basis for subsequent model optimization processes.
[0091] Step 206: Determine the changes in the input data of the computation logic of each operator, and classify each operator into a static operator or a dynamic operator based on the changes.
[0092] Operators, as the basic units constituting the computation graph, are functions that perform specific mathematical operations. For example, operators may include linear transformations (Linear / GEMM), layer normalization (RMSNorm), and activation functions (SiLU, GeLU). Input data may include tensors received by operators during forward inference. The characteristics of the input data may include the shape of the input tensor (e.g., batch size, sequence length) and its specific numerical content (e.g., the embedding vector of a token). Operators refer to indivisible basic computational nodes in the model's computation graph, such as matrix multiplication, normalization, and activation functions; each operator performs a specific mathematical operation.
[0093] Variation is used to characterize whether the internal execution path or control flow of an operator changes due to differences in the shape or content of its input data. For example, for static operators, their execution path remains fixed. For instance, the Linear operator (Y = XW + B) performs the same matrix multiplication and addition operations regardless of the input value; the RMSnorm operator performs a fixed normalization formula. That is, the computation of static operators does not depend on the specific value or length of the input. Correspondingly, for dynamic operators, their execution path includes conditional decisions dependent on the input. For example: the FlashAttention operator: needs to dynamically generate a causal mask based on the length S of the input sequence, and its Softmax computation also depends on the specific numerical value of the input. That is, it is a dynamic operator dependent on the shape and value of the input data. The MoE-Top-K operator: needs to sort the weights output by the router and select the Top-K expert indices. This is a dynamic operator dependent on numerical values (sorting). The MoE-Expert-GEMM operator: needs to gather the corresponding parameters from the expert weight pool based on the dynamically selected Top-K indices. This is an operator dependent on dynamic indices. Therefore, in this embodiment of the invention, if the execution of an operator involves decision logic that depends on input data, such as the above-mentioned conditional branching, dynamic indexing, or sorting selection, it is classified as a dynamic operator; otherwise, it is classified as a static operator.
[0094] Step 208: Precompile adjacent static operators in the original computation graph to obtain at least one static computation unit; wherein the precompile is used to enable the adjacent static operators to be executed by a single scheduling instruction.
[0095] The pre-compilation process involves pre-recording and packaging a sequence of adjacent static operators—which would otherwise require multiple discrete calls at runtime—into a single, efficient executable module. This module only needs to call the kernel once to execute all the encapsulated static operators. Adjacent static operators are consecutive operators with direct data dependencies on the computation graph. Therefore, these logically connected and data-flow-continuous static operators can be packaged together, allowing the pre-compiled static computation unit to function as an independent, self-consistent executable block. It's understood that data dependencies in the computation graph define the execution order and validity of operators. Packaging non-adjacent operators without direct dependencies would disrupt computational logic or prevent the formation of an efficient data flow. For example, the input of operator C depends on the output of operator B, and the input of operator B depends on the output of operator A (A->B->C). This is a valid, adjacent dependency chain. If operators A and C, which are not directly dependent, are forcibly packaged together, skipping the intermediate B, then within the packaged unit, C cannot obtain the correct input data from B during execution, leading to incorrect calculation results. Even if data is passed in through technical means, because there is no inherent, continuous data production and consumption relationship between A and C, the GPU hardware cannot perform deep pipeline optimization, making it difficult to achieve the optimized performance of single-instruction triggered continuous execution.
[0096] Specifically, pre-compilation can be based on runtime computation graph capture techniques (such as CUDA Graph) and may include the following process: In the original computation graph, identify the largest continuous subgraph formed by connecting static operator nodes through direct data dependency edges. In a real GPU execution environment, run this subgraph once with a set of sample data, recording the sequence of all kernel startup commands, memory copies, and other operations initiated during this run, forming a computation flow record corresponding to the subgraph. This computation flow record is then submitted to the runtime system for internal compilation and optimization, resulting in a persistent, independent object (i.e., a static computation unit) that can be repeatedly executed at high speed with a single startup command (such as cudaGraph Launch). It can be understood that a computation unit is an independently schedulable execution module formed by one or more operators through pre-compilation (for static operators) or interface encapsulation (for dynamic operators). Operators are the basic components constituting a computation unit, while the computation unit is the scheduling unit of the optimized model during execution. For example, a single scheduling instruction in CUDA Graph technology could be an API call like cudaGraphLaunch. Before performing the pre-compilation action in this embodiment of the invention, executing three consecutive static operators (e.g., RMSnorm -> Linear -> SiLU) requires the CPU to send three independent kernel startup commands to the GPU sequentially. Each startup involves fixed overhead such as parameter passing and context switching. However, through the pre-compilation action, these three operators are captured and merged into a single executable object named `graphExec`. During execution, the driver only needs to call the `cudaGraphLaunch(graphExec)` instruction once to invoke this single executable object. This instruction instructs the GPU hardware to execute the three kernels—RMSNorm, Linear, and SiLU—continuously in the pre-recorded order, without requiring further CPU intervention or scheduling. This effectively transforms the original discrete three-pass-three-workflow process into a single pass-continuous completion of three tasks, thereby avoiding most of the scheduling overhead. In this embodiment of the invention, by merging the overhead of multiple kernel startups into one, a group of static operators that might otherwise require numerous startup-execution-return loops can be transformed into a single instruction that triggers their continuous, pipelined execution, thereby significantly improving the efficiency of model processing. By fully utilizing the GPU's computation graph capture and replay mechanism, the originally discrete static operator sequence is pre-compiled into a single executable unit, thereby achieving deep optimization of the hardware scheduling mechanism.Specifically, in the static operator pre-compilation stage, the GPU's graph capture capability is utilized to record the kernel startup, memory access, and other operations of multiple adjacent static operators as computational flow records, which are then compiled into executable objects that can be triggered by a single scheduling instruction. This reduces CPU-GPU communication overhead and fully leverages the pipelined parallelism capabilities of the GPU hardware. In dynamic batching (in-flight batching) scheduling, by using preset capacity constraints (such as MaxTokensNum and MaxBatchSize), the system can control the range of shape changes of the input tensor, thereby limiting the amount of pre-compiled graphs stored and avoiding excessive GPU memory usage due to too many graphs. For example, multiple sets of optimized computational graphs can be pre-stored for different total token counts or batch sizes. During actual inference, the matching graph is selected for execution based on the input scale, ensuring both flexibility and memory efficiency. This achieves efficient collaboration between algorithm logic and hardware execution mechanisms. Through the combination of static / dynamic separation, graph pre-compilation, and dynamic batching, the utilization of computing resources and overall inference performance of GPUs and other hardware devices are significantly improved while ensuring the correctness of model functionality.
[0097] Step 210: Generate dynamic computing units corresponding to the dynamic operators in the original computation graph, and obtain at least one dynamic computing unit.
[0098] In this invention, considering that dynamic operators cannot be pre-compiled due to the dynamic nature of their internal logic (such as different conditional branches depending on the input data), in order to ensure that dynamic operators can be seamlessly integrated into a pipeline composed of static computing units without compromising their original high performance and algorithm correctness, this embodiment of the invention adopts only an interface adaptation strategy for dynamic operators. Specifically, a standardized function call prototype is defined for each dynamic operator, specifying its input and output parameters (mainly pointers and descriptors of data memory). Then, the native, highly optimized kernel implementation of the dynamic operator is encapsulated into a callable function conforming to this standard prototype. This encapsulated function can then serve as a dynamic computing unit.
[0099] Step 212: Based on the data dependency relationship, combine the at least one static computing unit with the dynamic computing unit to obtain the optimized model corresponding to the original model.
[0100] In this process, the data dependencies defined in the original computation graph generated in the preceding steps serve as guidance for the reassembly of operators. The static computation units pre-compiled in the preceding steps and the dynamic computation units encapsulated through interfaces are then sequentially assembled according to the order and connection relationships indicated by this guidance, resulting in the optimized model. It is understood that this optimized model is mathematically equivalent to the original model, but after the pre-compilation process described above, its internal execution engine has been replaced by a pipeline composed of fewer, larger-granularity static and dynamic computation units, replacing a large number of fine-grained operators. This improves the processing efficiency and resource utilization of the machine learning model, making it particularly suitable for large language models, such as those containing attention mechanisms.
[0101] Optionally, after obtaining the optimized model, a data processing procedure can be executed through the optimized model. The data processing procedure can be triggered based on an external data processing request, such as a request for inference. Specifically, in the actual inference service, a user request is received, and multiple inference acceleration requests are combined into a batch using a dynamic batching strategy. The batched input tensor is then fed to the optimized model for execution. Thus, the static and dynamic computing units within the optimized model work collaboratively according to a preset pipeline sequence to generate and return the inference result.
[0102] In practical deployments, this application's embodiments, combined with an in-flight batching strategy, can effectively handle dynamic changes in the scale of input data. For example, multiple inference requests can be batched in real-time according to preset capacity constraints (such as the maximum total number of tokens (MaxTokensNum) and the maximum batch size (MaxBatchSize). Optionally, to adapt to inputs of different sizes, multiple sets of static computation units (i.e., pre-compiled graphs) optimized for different input shapes can be pre-stored. For example, when the total number of input tokens falls within a certain preset range, the system directly calls the corresponding pre-compiled graph for execution without recompilation; if the input scale exceeds the range of the pre-stored graph, the system can dynamically trigger a new round of graph capture and compilation, and update the graph cache. This mechanism maintains the flexibility of dynamic batching while avoiding frequent compilation overhead through graph reuse, achieving a balance between memory usage and execution efficiency. It should be noted that the optimized model of this invention is particularly suitable for scenarios with high requirements for latency and throughput in large language model online services. While in-flight batching solutions (vLLM, TGI, etc.) in related technologies can improve throughput, the overhead of kernel-level launches is enormous. In contrast, this invention significantly reduces the computational overhead of each layer through static pre-compilation, resulting in a substantial reduction in end-to-end latency for a single request (especially the generation latency after the first-word delay TTFT), thus improving the user experience. Under dynamic batch scheduling, the overhead of executing a static computation unit in the optimized model is far lower than executing multiple independent operators. This allows the GPU to complete computations more efficiently under the same hardware constraints such as max_num_tokens and max_batch_size, processing more requests per unit time and thereby improving the overall system throughput. Furthermore, the optimized model generated in this invention can simultaneously serve multiple architectures such as Dense (FFN) and Sparse (MoE) using the same binary set. During data processing, regardless of whether the request is for a Chat model (usually Dense) or an expert model (such as Mixtral, or MoE), the system does not need to switch or re-optimize, which greatly simplifies the deployment and maintenance complexity of the production environment and ensures that both models can obtain the same acceleration benefits.
[0103] In summary, the optimized model of this solution inherits the advantages of dynamic batch processing scheduling while fundamentally solving the problem of low computational efficiency. It is particularly suitable for large-scale online language model inference service scenarios that require low latency, high throughput, and compatibility with multiple model architectures.
[0104] In some embodiments, performing computational logic analysis on the original model to obtain the corresponding original computational graph of the original model includes:
[0105] The computational operations included in the original model are traversed, and each independent computational operation is constructed as a node corresponding to an operator in the original computational graph;
[0106] Based on the data input-output relationship between the operator nodes, a directed edge is established in the original computation graph from the source node of the output data to the destination node of the input data.
[0107] The goal of computational logic analysis is to transform and parse the model from high-level abstract layers or modules into a graph representation of low-level atomized operators and the precise data flow between them. Operators are the most basic and indivisible computational units in the original model. The data input-output relationship characterizes how data flows from one operator to the next during computation.
[0108] The process of constructing nodes in the original computation graph can include: first, traversing the entire network structure of the original model. This can be done automatically by a pre-defined model parser. Taking a Transformer model in ONNX format as an example: the parser starts from the model's input nodes and sequentially identifies each computation layer or operation. For each identified independent computation operation, the parser instantiates it as a node in the original computation graph, which represents an operator. For example, when the parser encounters a Linear layer, it creates an operator node representing matrix multiplication and bias addition. When it encounters a LayerNorm or RMSNorm, it creates an operator node representing layer normalization computation. Similarly, activation functions such as SiLU and GELU are also created as corresponding element-wise operation operator nodes. The traversal process continues until all computation operations in the model are covered, thus generating a set containing all operator nodes.
[0109] During or after constructing all operator nodes, it is necessary to determine their execution order. This is achieved by analyzing and establishing data dependency edges. Specifically, this involves analyzing the input and output tensors of each operator. For any two operator nodes A and B, if some input tensor of operator B directly comes from the output tensor of operator A, then a directed edge is established in the graph from node A to node B. The direction of the edge points from the data producer (source node) to the data consumer (destination node). Understandably, the direction of the edge intuitively reflects the order of computation between operators; that is, A should produce data first, and only then can B consume and execute that data.
[0110] For example, suppose we obtain the three operators RMSNorm -> QKV_Linear -> RoPE from the parsing. In the computation graph, this will be represented by three nodes: The RMSNorm node receives the model input and outputs a normalized tensor. The QKV_Linear node receives the output of RMSNorm as input and outputs three projected tensors: query (Q), key (K), and value (V). The RoPE node receives the Q and K tensors output by QKV_Linear as input and outputs the Q and K tensors after injection position encoding.
[0111] Therefore, two corresponding directed edges are established in the original computation graph: one from the RMSNorm node to the QKV_Linear node, and the other from the QKV_Linear node to the RoPE node. These two edges represent the irreversible data flow of normalization, projection, and finally positional encoding.
[0112] Through the steps of node construction and edge establishment described above, the system ultimately generates a complete directed graph data structure, namely the original computation graph. The original computation graph precisely describes which basic operators constitute the original model (i.e., the set of nodes), in what order these operators are executed (through the partial order relation defined by directed edges), and how data flows between operators (the tensors carried by the edges).
[0113] In some embodiments, the change indicates whether the computational logic of the operator changes due to changes in the input data of the current operator;
[0114] The step of determining the computational logic of each operator in response to changes in input data, and classifying each operator into a static operator or a dynamic operator based on these changes, includes:
[0115] For each of the operators, determine whether the computational logic of the current operator includes decision logic that depends on the content of the input tensor; the decision logic includes at least one of conditional branching, dynamic indexing, or sorting and selection operations;
[0116] If the computational logic includes the decision logic, the current operator is determined as the dynamic operator;
[0117] If the computational logic does not include the decision logic, the current operator is determined as the static operator.
[0118] In this context, it's important to consider that different operators in the original model exhibit fundamentally different sensitivity to runtime input data in their computational behavior. Specifically, the computational logic of different operators differs in its dependence on the input data. Input data can include tensors processed during operator execution, whose characteristics include shape (e.g., batch size, sequence length) and specific element values. Decision logic refers to the code mechanisms in the operator implementation that require making choices based on the actual situation of the input data.
[0119] The characteristic of static operators is that their computational logic is deterministic and fixed. Regardless of the value or shape of the input tensor (as long as it meets the basic dimensional requirements), its internal structure follows the exact same sequence of arithmetic or linear algebra operations, without any input-based conditional judgments. For example, the computational expression of a linear transformation operator is Y = XW + B. For any input X, the operator performs the exact same matrix multiplication and vector addition operations. The value of the input X only affects the value of the computational result Y, but does not change the computational path itself, which involves performing one matrix multiplication and one addition. Therefore, linear transformation operators are considered static operators. Another example is layer normalization operators such as RMSnorm, which standardize the input tensor according to a fixed mathematical formula. Regardless of fluctuations in the input value, the operator's steps of calculating the sum of squares and mean, and scaling, remain constant. Finally, activation function operators such as Sigmoid Linear Unit (SiLU) or Gaussian Error Linear Unit (GeLU) apply a pre-defined, branchless nonlinear function to each element of the input tensor. The calculation process also does not depend on the magnitude of the input values.
[0120] Correspondingly, there may also be dynamic operators in the computational logic that depend on decision points based on the input data, and their execution path is dynamically adjusted according to the specific circumstances of each input. The decision logic can include conditional branches: the operator's code contains "if...then..." statements, with the judgment condition based on the input data. For example, in self-attention computation, causal masking is needed to prevent future information leakage. This is usually achieved through a conditional branch: if the currently generated position index i is greater than the historical position index j, then the attention score is set to a minimum value (such as negative infinity). This branch decision depends on the length of the input sequence. Such operators are usually included in optimized attention kernels such as FlashAttention. Another example is sorting and selection operators, which require sorting a set of input values and selecting based on the sorting results. For instance, in expert hybrid networks, the weights output by the linear layer of the router need to undergo a Top-K operation. This operation requires sorting all weight values and then selecting the K largest values and their indices. The sorting and selection process depends entirely on the specific values of the input weights. For example, the memory addresses (indexes) used by dynamic indexing operators for data read and write operations are calculated at runtime based on the input data, rather than being pre-fixed. Similarly, in expert hybrid networks, the expert indices selected based on Top-K need to be dynamically aggregated from the weight tensor storing all expert parameters to determine the weights of those specific experts. This aggregation operation uses dynamic indexes to locate the data, and the index values change with each inference request.
[0121] Each operator in the computation graph is analyzed to check whether its computational implementation includes decision logic such as conditional branches, sorting selection, or dynamic indexing that depend on the shape or value of the input tensor. If it does, the operator's execution path is non-fixed and dynamic, and it is classified as a dynamic operator. If it does not, the operator's execution path is fixed and predictable, and it is classified as a static operator. This classification deepens the description of the original model's computation graph from a structural level to a classification based on computational behavior characteristics, thereby revealing which parts can be aggressively and deterministically optimized in subsequent steps (i.e., for static operators) and which parts need to maintain their flexibility and original state to ensure functional correctness (i.e., for dynamic operators).
[0122] In some embodiments, precompiling adjacent static operators in the original computation graph to obtain at least one static computation unit includes:
[0123] In the original computation graph, all the maximum connected subgraphs consisting of static operator nodes and their direct data dependency edges are identified.
[0124] For each identified maximum connected subgraph, the static operator sequence corresponding to the maximum connected subgraph is executed in the model inference environment to obtain the computational flow record corresponding to the maximum connected subgraph; the computational flow record includes at least one of the following: the start instruction of all operators in the current execution process, the computational logic, and the memory access operation.
[0125] The computation flow record is compiled, and the compilation result is encapsulated into an independent executable object that can be replayed by the single scheduling instruction, thus obtaining a pre-compiled static computation unit.
[0126] To transform the previously fragmented, independent computation tasks into a single, continuously executed computation pipeline, thereby eliminating the overhead of multiple task scheduling and kernel startups, this embodiment of the invention transforms the identified set of static operators into a single, schedulable static computation unit. Specifically, in the original computation graph with its dynamic and static partitioning, a cluster of static operators suitable for packaging is identified. Adjacency refers to static operator nodes connected by directed edges representing direct data dependencies in the computation graph data structure. The maximum connected subgraph is found, consisting of all static operator nodes and their directly dependent edges. Within such a subgraph, any two static operator nodes can be connected by a series of edges (connectivity), and it cannot be expanded outwards to include more static operators with direct dependencies (maximization).
[0127] Example: Suppose that in the first half of a Transformer layer, there exists a static operator sequence of layer normalization -> query / key / value projection -> rotation position encoding, and these operators have a continuous A->B->C data dependency. This sequence will be identified as a maximally statically connected subgraph that meets the following conditions.
[0128] After identifying the static subgraph, it needs to be converted into a repeatable template. Specifically, in a hardware and software environment identical to the actual inference (e.g., a server equipped with a GPU and loaded with the corresponding computational libraries), a set of sample input data capable of triggering the execution of all operators within the subgraph can be prepared. This sample input data can be real or placeholder data with the required structure. The computation process of the subgraph is then initiated. During this process, the underlying runtime system (such as CUDA) does not immediately execute regular, discrete kernel calls, but instead enters a recording mode. It intercepts and records all GPU kernel startup instructions that should have been initiated, related computational parameters (such as matrix size and pointer addresses), and accompanying memory copying or setting operations. All these recorded operation sequences constitute the computation flow record of the static subgraph. It can be understood that this computation flow record is essentially a detailed execution list, precisely describing the order in which kernels are called and which memory operations are performed to complete the computation of this subgraph. Upon receiving the computation flow record, it is internally compiled and optimized. The optimization process may include eliminating redundant operations, optimizing memory access patterns, and merging multiple fine-grained kernels into a more efficient monolithic kernel. The goal of compilation is to generate a low-level execution plan that is highly optimized for the specific hardware. After compilation, the system generates and encapsulates an executable object. This object holds the optimized execution plan and provides a simple startup interface. In CUDA Graph technology, this object is typically referred to as a graph executor of type cudaGraphExec_t.
[0129] To illustrate the pre-compilation example in this embodiment, before pre-compilation, the aforementioned layer normalization -> projection -> rotation position encoding sequence may need to be executed, requiring the CPU to call the GPU driver three times to launch three kernels respectively. Each call involves communication and scheduling latency. After compilation, the execution of the above three operators is merged into a single static computation unit. When execution is needed, the CPU only needs to send one instruction to the GPU (e.g., calling the cudaGraphLaunch() function and passing in the handle of the unit). The GPU will automatically and continuously execute the three kernels sequentially according to the pre-stored, optimized execution plan within the unit, without any CPU intervention. This transformation saves a significant amount of time that was originally consumed in scheduling, allowing the GPU's computing resources to be more concentrated on actual computation, thereby significantly improving the execution efficiency of the static part within the model layer and making a major contribution to reducing overall inference latency.
[0130] In some embodiments, generating dynamic computation units corresponding to the dynamic operators in the original computation graph to obtain at least one dynamic computation unit includes:
[0131] For each of the aforementioned dynamic operators, a unified function call prototype is defined for each dynamic operator. The function call prototype is used to characterize the memory address parameters of the input data and output data of the dynamic operator.
[0132] The native computation kernel of the dynamic operator is encapsulated into a callable function that conforms to the function call prototype, thereby obtaining the dynamic computation unit corresponding to each dynamic operator.
[0133] Dynamic operators, such as attention-based computation kernels or expert-selected routing kernels containing complex control flows, have execution paths dependent on runtime input data. Therefore, they cannot be pre-compiled into a fixed, replayable execution sequence like static operators. Aggressively rewriting or solidifying these operators is not only technically challenging but may also disrupt their carefully optimized internal algorithmic logic, leading to performance degradation or functional errors. Therefore, the principle for handling dynamic operators should be adaptation rather than modification. The goal is to seamlessly integrate their native, efficient kernel implementation into a new execution pipeline composed of static computation units, while maintaining the integrity and performance of their computational logic.
[0134] Specifically, a unified calling convention needs to be established for dynamic operators, i.e., defining function call prototypes. This is equivalent to establishing a clear communication contract between static computation units and dynamic operators. The function call prototype defines all the parameters required to call the dynamic operator, such as the memory address parameter for the data. The memory address parameter for the data can include: Input data pointer: pointing to the device memory address where the input tensor is stored; Output data pointer: pointing to the device memory address where the computation result tensor is stored; Optional tensor descriptor: describing the shape, data type, and memory layout of the tensor; Optional algorithm configuration parameters: such as the number of attention heads, Dropout probability, etc.
[0135] Regardless of the complexity of the internal implementation of dynamic operators, they all follow the same or similar parameter passing pattern externally (e.g., always in the order of input pointer, output pointer, descriptor). This uniformity simplifies the scheduling logic of the execution pipeline, allowing the system to call different types of dynamic operators in a consistent manner. For each dynamic operator, its existing, highly optimized kernel function (e.g., functions in NVIDIA's FlashAttention-2 library, or custom MoETop-K selection kernel functions) is thinly wrapped. This wrapper mainly performs parameter adaptation and format conversion, ensuring that the kernel function can correctly receive parameters from the standardized prototype. The result after wrapping is a dynamic computation unit with a standardized entry point. This unit retains all the algorithmic details and optimizations of the native kernel internally, but presents a predictable calling interface that conforms to the framework conventions externally.
[0136] For example, the native kernel of the FlashAttention operator (such as the flash_attn_v2 function) includes logic such as dynamic sequence mask generation. In this embodiment, a prototype is defined for it, for example, void flash_attn_unit(float* Q, float* K, float* V, float* O, int batch_size, int seq_len, ...). The official flash_attn_v2 kernel call is wrapped in a function conforming to this prototype, forming a dynamic computation unit A. When this unit is called, it performs the native, dynamic FlashAttention computation. Correspondingly, the kernel of MoE-related dynamic operators, such as the Top-K expert selection operator, needs to sort the routing weights and output indices. In this embodiment, a prototype is defined for it, such as void topk_experts(float* router_weights, int* expert_indices, int k, ...), and the corresponding kernel is wrapped to obtain dynamic computation unit B. For another example, the kernel of the expert calculation operator needs to aggregate weights based on dynamic indices and perform calculations. In this embodiment of the invention, a prototype is defined, such as `void moe_expert_gemm(float* input, int* expert_indices, float* expert_weights, float* output,...)`, which is then encapsulated to obtain the dynamic computation unit C. Through interface standardization and kernel encapsulation, this embodiment of the invention successfully modularizes dynamic operators. This allows them to be combined with static computation units later, enabling data transfer and execution triggering through clear interfaces. Thus, without sacrificing the core algorithm performance and dynamic characteristics, they are integrated into the overall optimization framework, forming a highly efficient and unified inference execution pipeline.
[0137] In some embodiments, combining the at least one static computing unit with the at least one dynamic computing unit according to the data dependency relationship to obtain the optimized model corresponding to the original model includes:
[0138] Based on the data dependencies between the operator nodes in the original computation graph, the connection order between the static computation unit and the dynamic computation unit is determined;
[0139] The static computing unit and the dynamic computing unit are combined according to the connection order to obtain the optimized model.
[0140] In this invention, static and dynamic computation units are functionally cohesive modules. To reconnect these modules into a new, efficient whole, the determination of the module connection order depends on the data dependencies recorded in the original computation graph. Each static or dynamic computation unit corresponds to a specific original operator node or subgraph in the original computation graph. Since each unit is encapsulated by one or more original operators, this mapping is explicit. Based on this mapping, directed edges between nodes in the original computation graph (A->B indicates that the output of A is the input of B) are directly converted into connection order requirements between units. If operator node A is encapsulated in unit X, operator node B is encapsulated in unit Y, and there is an edge A->B, then in the optimized model, the execution of unit X must be completed before unit Y, and some output of unit X must be used as the input of unit Y. By traversing and analyzing the entire original computation graph, a global, acyclic execution order list between all units is derived.
[0141] After determining the connection order of all units, the units (static or dynamic) are linked together sequentially according to this order. Linking can include: at the code level, passing the output memory pointer of the previous unit to the input parameter of the next unit; and at the scheduling level, ensuring that the execution engine automatically triggers the execution of the next unit after completing the computation of the previous unit.
[0142] Through the links described above, the discrete units are integrated into a new, directed computation graph structure. This new structure is the execution logic graph of the optimized model. Unlike the original computation graph, the nodes in the new graph are no longer fine-grained basic operators, but rather larger-grained static and dynamic computation units.
[0143] For example, taking the optimization of a typical Transformer layer as an example, the original computation graph may contain the following node sequence: [Norm1, QKV_Proj, RoPE, FlashAttention, O_Proj, Add1, Norm2, FFN_Up, Act, FFN_Down, Add2]. After the dynamic and static partitioning and encapsulation steps mentioned above, we may obtain: Static computation unit 1: encapsulates [Norm1, QKV_Proj, RoPE], Dynamic computation unit 1: encapsulates the native FlashAttention kernel, Static computation unit 2: encapsulates [O_Proj, Add1, Norm2]. (For Dense models) Static computation unit 3: encapsulates [FFN_Up, Act, FFN_Down, Add2].
[0144] Based on the dependencies in the original computation graph, the determined connection order is: Static computation unit 1 -> Dynamic computation unit 1 -> Static computation unit 2 -> Static computation unit 3. Combining them in this order yields the optimized model structure of this Transformer layer. This new structure significantly reduces the number of scheduling operations within the layer while retaining the dynamism and high performance of the core attention computation.
[0145] In some embodiments, the original model includes a large language model based on the Transformer architecture; the optimized model includes at least one Transformer layer; each Transformer layer includes a first static computation unit, a first dynamic computation unit, and a second static computation unit connected in sequence; the first static computation unit is used for preprocessing of attention mechanism computation; the first dynamic computation unit is used for performing attention mechanism computation; and the second static computation unit is used for postprocessing of attention mechanism computation results.
[0146] Among them, the core of the large language model based on the Transformer architecture is composed of multiple structurally similar Transformer layers stacked together. In a standard Transformer layer, the computation usually follows a relatively fixed process: first, the input is preprocessed, such as by normalization; then, self-attention computation is performed to fuse contextual information; finally, the attention output is post-processed, such as by secondary transformation and residual connection, and may be enhanced nonlinearly through a feedforward network. In the embodiments of this invention, considering that the specific operators of different models may vary slightly, this three-stage pattern of preprocessing-core attention computation-postprocessing is common. Therefore, by applying the aforementioned steps of dynamic and static operator partitioning, pre-compilation, and interface standardization to this type of model, a Transformer layer will be reconstructed into three main computational units connected in sequence, forming an efficient execution pipeline: The first static computational unit: This unit is responsible for the preprocessing of the attention mechanism computation. It encapsulates all the necessary static operators with fixed computational logic before entering the core attention computation. For example, it may include layer normalization of the layer input tensor and linear transformations that project it into query, key, and value vectors respectively. Furthermore, the rotational position encoding calculation, which injects positional information into these vectors, is usually included in this unit due to its fixed logic. After this unit completes its execution, it prepares well-formatted Q, K, and V tensors with positional information for attention calculation. First Dynamic Calculation Unit: This unit is dedicated to attention mechanism calculation and is the core of the entire layer's computation and performance bottleneck. It can encapsulate native, highly optimized attention kernels, such as FlashAttention. This kernel receives Q, K, and V tensors from the first static calculation unit, performs self-attention calculations including complex logic such as dynamic mask generation and Softmax, and outputs a new representation after context fusion. Due to its large amount of dynamic control flow, it is reserved as an independent dynamic unit. And the Second Static Calculation Unit: This unit is responsible for post-processing the attention mechanism calculation results. It can include output projection linear transformation to restore the dimension of the attention output, residual connection operations to add the projection result to the initial layer input (through residual connections), and another layer normalization of the added result; these operations are all static operators.
[0147] For a large language model containing dozens or hundreds of Transformer layers, each layer is reconstructed using the aforementioned three-stage structure. Ultimately, the entire forward propagation process of the optimized model manifests as multiple repetitions of these three-layer mini-pipelines. This structural shift brings performance improvements. For example, a typical DenseTransformer layer might have its internal operators divided into: a first static unit (3 operators), a first dynamic unit (1 kernel), and a second static unit (3 operators). Before optimization, executing these 7 computations might require 7 or more independent kernel startups. After optimization, it only requires 3 startups (two for static units and one for dynamic units), significantly reducing scheduling overhead and thus significantly reducing end-to-end latency in model inference.
[0148] In some embodiments, the original model includes a dense feedforward network model; the second static computation unit encapsulates all computation operators in the current Transformer layer after computation by the attention mechanism.
[0149] Dense feedforward (FFN) models can refer to standard Transformer architectures, such as GPT, LLaMA (non-MoE version), and BERT. These standard Transformer architectures connect a fixed, parameter-dense feedforward neural network after the attention mechanism of each Transformer layer. This FFN module typically consists of two linear transformation layers and a non-linear activation function, such as SwiGLU (SiLU-gated linear unit) or standard GeLU activation. Its computational logic is completely static throughout the inference process. Regardless of the input, the FFN performs the exact same computational sequence: up-projecting the input, applying the activation function, and then down-projecting. It does not contain any conditional branches, dynamic indexing, or sorting / selection operations.
[0150] Therefore, this embodiment of the invention extends the scope of this encapsulation for dense feedforward network models. After completing the basic post-processing described above, all static operators contained in the entire feedforward neural network module belonging to that layer are also identified, captured, and integrated into the same second static computation unit. This extended second static computation unit encapsulates a complete and continuous computation chain from the attention output to the final output of that layer.
[0151] Taking an FFN activated by SwiGLU as an example, the integrated second static computation unit may encapsulate the following operator sequence:
[0152] Attention Output Projection: A linear transformation that maps the dimension of the attention output back to the hidden layer dimension of the model. First Residual Connection: Adds the projection result to the initial input of this Transformer layer. Layer Normalization: Standardizes the result of the residual addition. FFN Up Projection: The first linear transformation that expands the dimension (e.g., to 4 times the hidden layer dimension). Activation Function: Performs SwiGLU computation (including SiLU activation and element-wise multiplication). FFN Down Projection: The second linear transformation that reduces the dimension back to the hidden layer dimension. Second Residual Connection: Adds the FFN output to the layer's normalized input (or the value before the first residual connection, depending on the specific architecture) to form the final output of this layer.
[0153] None of the operators mentioned above contain dynamic control flow, and therefore can be safely and continuously captured into the same pre-compiled static computation unit. For dense feedforward network models, this invention achieves maximum staticization and pipelined optimization of the inference process of this type of model by encapsulating all static computations at the end of the layer into a second static computation unit, thereby obtaining considerable end-to-end acceleration.
[0154] In some embodiments, the original model includes a sparse expert hybrid network model; the second static computing unit encapsulates the router linear layer operator of the current Transformer layer; the optimized model is followed by a second dynamic computing unit and a third dynamic computing unit in sequence after the second static computing unit; wherein, the second dynamic computing unit is used to perform an expert selection operation based on the router output; the third dynamic computing unit is used to perform sparse matrix multiplication calculations for the selected expert based on the result of the expert selection operation.
[0155] Among them, sparse expert hybrid network models, such as Mixtral and Grok-1, replace the standard, fixed feedforward network with a MoE module in the Transformer layer. The core of the MoE module is a router and a set of expert networks. The computation process is typically as follows: the input passes through a router's linear layer to generate weights for each expert; then, a dynamic Top-K operation selects the K experts with the highest weights; finally, the input is routed to these selected experts for computation, and the results are weighted and merged. It is evident that the sparse expert hybrid network model introduces dynamic sparsity. Its computation path is not fixed but depends on the router dynamically selecting different subsets of experts based on the weights calculated for each input. The Top-K selection and expert computation based on dynamic indexing are typical dynamic operators.
[0156] Therefore, in this embodiment of the invention, for the MoE model, since its latter part of the computation involves dynamic decisions that cannot be statically rendered, the entire latter part cannot be encapsulated into a second static computation unit as in the case of the Dense model. A precise boundary needs to be set for static optimization, which is defined at the router linear layer operator. The router linear layer is a standard linear transformation with fixed computational logic that does not depend on the input content for decision-making; therefore, it belongs to static operators. Thus, the second static computation unit of each Transformer layer in the MoE model is encapsulated up to the router linear layer operator of that layer. It internally contains all static operators from the attention output until the generation of expert route weights, typically including: attention output projection, first residual connection, layer normalization, and the final router linear layer. After the second static computation unit completes the calculation of route weights, subsequent dynamic routing and computation need to be handled by a dedicated dynamic computation unit. This results in the following clear chain structure: 1. Second dynamic computation unit (expert selection unit): used to perform expert selection operations based on router outputs. It encapsulates the dynamic Top-K selection kernel. The kernel receives the routing weight tensor output from the second static computation unit, sorts it, and selects the K experts with the largest weight values, along with their corresponding indices and weights. This is a typical dynamic operator that includes sorting and selection logic. The third dynamic computation unit (expert computation unit) performs sparse matrix multiplication calculations based on the expert selection results for the selected experts. It encapsulates the dynamic sparse matrix multiplication kernel. This kernel receives expert indices from the second dynamic computation unit, dynamically gathers the parameters of the corresponding K experts from the expert weight library based on these indices, then performs calculations on the input data (or processed data) with the weights of these experts (e.g., performing multiple small, parallel matrix multiplications), and finally sums the K results according to their weights.
[0157] In summary, the computational pipeline of an optimized MoE Transformer layer is as follows: First static computation unit (attention preprocessing) -> First dynamic computation unit (attention computation) -> Second static computation unit (postprocessing to route weight generation) -> Second dynamic computation unit (Top-K expert selection) -> Third dynamic computation unit (sparse expert computation). Compared to the Dense model structure described in the previous embodiments, this invention demonstrates how the same optimization framework can gracefully adapt to a completely different model architecture (MoE) by adjusting the encapsulation boundary of the second static computation unit and connecting different dynamic unit chains thereafter. Whether it's a Dense or MoE model, the core static -> dynamic -> static front-end pattern remains consistent; the difference lies only in the back-end processing chain. This achieves deep optimization while maintaining its core dynamic sparsity characteristics, enabling the same optimization system to efficiently serve two mainstream model architectures.
[0158] In some embodiments, the data processing using the optimized model includes:
[0159] Receive multiple online inference requests and dynamically combine the multiple online inference requests into an input batch tensor according to a preset capacity constraint;
[0160] The input batch tensor is input into the optimized model to execute the static calculation unit and the dynamic calculation unit in the optimized model in sequence;
[0161] The calculation results of the optimized model are collected and output as the data processing results corresponding to the online inference request.
[0162] The optimized model's primary application scenarios can include online inference services for large language models, such as open-domain dialogue, content generation, and code completion. In these scenarios, the server continuously receives real-time requests of varying lengths from different users, requiring consideration of two key metrics: first-word latency (the time from when a user sends a request to when the first word is received) and throughput (the number of requests processed per unit time or the total number of tokens). This embodiment of the invention continuously receives asynchronously arriving online inference requests and maintains a request pool. When a new request arrives, it is judged based on preset capacity constraints. These constraints may include: maximum total tokens: the upper limit of the total number of tokens in all requests in the pool, preventing excessive single-time computation from causing memory overflow or a surge in latency; maximum batch size: the upper limit of the number of requests allowed to be accommodated simultaneously in the pool, controlling concurrency.
[0163] When the request pool meets any constraint (e.g., the total number of tokens reaches its limit, or the waiting time times out), all requests in the pool are combined to trigger a model inference. Their input data is concatenated into an input batch tensor with a uniform batch dimension. This strategy effectively improves GPU utilization while ensuring that the response time of individual requests does not deteriorate indefinitely through constraint control.
[0164] The optimized model executes sequentially according to its internally constructed pipeline (as described in the previous embodiments, a structure consisting of alternating static and dynamic computation units). For batch input, each computation unit in the model (whether static or dynamic) operates on the entire batch of data.
[0165] In dynamic batch processing, the number of operators to be processed in each scheduling step is not reduced. In traditional solutions, a large batch still requires a kernel startup for each operator. However, in this embodiment of the invention, the static computing unit merges the startup of multiple static operators into one. This means that when processing batch data of the same size, the total number of kernel startups required by the optimized model is significantly reduced. This saved scheduling overhead directly translates into shorter GPU computation queuing time and faster batch completion time. Dynamic computing units (such as FlashAttention and MoE routing) have their native kernels optimized for batch processing when processing batch data. This method embodiment ensures that the optimal performance of the core complex computation parts is preserved by seamlessly calling them through standardized interfaces. After the optimized model is completed, the output results are collected, split and returned according to the original order of the request, completing a full online inference service.
[0166] Because the number of kernel startups at the model layer is drastically reduced, the computation execution time for each batch is shortened, thereby reducing the average end-to-end latency of requests. Users can obtain complete responses faster. Under the same hardware and scheduling constraints (max_num_tokens, max_batch_size), because single-batch processing is faster, the GPU can complete more batches per unit time, thus significantly improving the overall system throughput. The hard constraint mechanism of dynamic batch processing is fully preserved. New requests will not be blocked for a long time due to waiting for the execution of an excessively large static graph, the first-word latency metric remains stable and predictable, and the user experience is good. Whether it is the Dense or MoE model, after optimization by this method embodiment, it can be loaded and executed by the same service framework. In actual deployment, there is no need to maintain two different optimization engines or service code for models with different architectures, which greatly simplifies the operation and maintenance complexity and ensures that both models can obtain the same acceleration benefits.
[0167] Therefore, the data processing described in this embodiment fully demonstrates the end-to-end process from online request access to efficient calculation by the optimized model, and finally the return of the result. It verifies that this invention is not only an offline model optimization tool, but also a complete solution that can be directly applied to production environments, significantly improving the performance and efficiency of online services for large language models.
[0168] Secondly, such as Figure 3 As shown, this application also provides a machine learning model optimization device 300, the device comprising:
[0169] Acquisition module 302 is used to acquire the original model; the original model includes a large language model;
[0170] Analysis module 304 is used to perform computational logic analysis on the original model to obtain the original computation graph corresponding to the original model; wherein, the original computation graph represents the data dependency relationship between the various operators included in the original model;
[0171] The determining module 306 is used to determine the changes in the calculation logic of each operator due to the input data, and to classify each operator into a static operator or a dynamic operator based on the changes.
[0172] The pre-compilation module 308 is used to pre-compile adjacent static operators in the original computation graph to obtain at least one static computation unit; wherein, the pre-compilation is used to enable the adjacent static operators to be executed by a single scheduling instruction;
[0173] The generation module 310 is used to generate dynamic computing units corresponding to the dynamic operators in the original computation graph, so as to obtain at least one dynamic computing unit.
[0174] The combination module 312 is used to combine the at least one static computing unit with the dynamic computing unit according to the data dependency relationship to obtain the optimized model corresponding to the original model.
[0175] It should be understood that although the steps in the flowcharts of the embodiments described above are shown sequentially according to the arrows, these steps are not necessarily executed in the order indicated by the arrows. Unless explicitly stated herein, there is no strict order restriction on the execution of these steps, and they can be executed in other orders. Moreover, at least some steps in the flowcharts of the embodiments described above may include multiple steps or multiple stages. These steps or stages are not necessarily completed at the same time, but can be executed at different times. The execution order of these steps or stages is not necessarily sequential, but can be performed alternately or in turn with other steps or at least some of the steps or stages in other steps. It is understood that the steps in different embodiments can be freely combined as needed, and all non-contradictory solutions formed by such combinations are within the scope of protection of this application.
[0176] Based on the same inventive concept, this application also provides a machine learning model optimization apparatus for implementing the methods described above. The solution provided by this apparatus is similar to the implementation scheme described in the methods above. Therefore, the specific limitations in one or more machine learning model optimization apparatus embodiments provided below can be found in the limitations of the machine learning model optimization method described above, and will not be repeated here.
[0177] Each module in the aforementioned machine learning model optimization device can be implemented entirely or partially through software, hardware, or a combination thereof. These modules can be embedded in or independent of the processor in a computer device, or stored in the memory of a computer device as software, so that the processor can call and execute the operations corresponding to each module.
[0178] In one exemplary embodiment, a computer device is provided, which may be a terminal, and its internal structure diagram may be as follows: Figure 4As shown, the computer device includes a processor, memory, input / output interfaces, a communication interface, a display unit, and an input device. The processor, memory, and input / output interfaces are connected via a system bus, and the communication interface, display unit, and input device are also connected to the system bus via the input / output interfaces. The processor provides computational and control capabilities. The memory includes non-volatile storage media and internal memory. The non-volatile storage media stores the operating system and computer programs. The internal memory provides an environment for the operation of the operating system and computer programs stored in the non-volatile storage media. The input / output interfaces are used for exchanging information between the processor and external devices. The communication interface is used for wired or wireless communication with external terminals; wireless communication can be achieved through Wi-Fi, mobile cellular networks, Near Field Communication (NFC), or other technologies. When the computer program is executed by the processor, it implements a machine learning model optimization method. The display unit is used to form a visually visible image and can be a display screen, a projection device, or a virtual reality imaging device. The display screen can be an LCD screen or an e-ink screen. The input device of the computer device can be a touch layer covering the display screen, or buttons, trackballs, or touchpads set on the casing of the computer device, or external keyboards, touchpads, or mice, etc.
[0179] Those skilled in the art will understand that Figure 4 The structure shown is merely a block diagram of a portion of the structure related to the present application and does not constitute a limitation on the computer device to which the present application is applied. Specific computer devices may include more or fewer components than those shown in the figure, or combine certain components, or have different component arrangements.
[0180] In one exemplary embodiment, a computer device is provided, including a memory and a processor, wherein the memory stores a computer program, and the processor executes the computer program to implement the steps included in any of the foregoing machine learning model optimization method embodiments.
[0181] In one embodiment, a computer-readable storage medium is provided having a computer program stored thereon, which, when executed by a processor, implements the steps included in any of the aforementioned machine learning model optimization method embodiments.
[0182] In one embodiment, a computer program product is provided, including a computer program that, when executed by a processor, implements the steps included in any of the foregoing machine learning model optimization method embodiments.
[0183] It should be noted that the user information (including but not limited to user device information, user personal information, etc.) and data (including but not limited to data used for analysis, data stored, data displayed, etc.) involved in this application are all information and data authorized by the user or fully authorized by all parties, and the collection, use and processing of the relevant data must comply with relevant regulations.
[0184] Those skilled in the art will understand that all or part of the processes in the methods of the above embodiments can be implemented by a computer program instructing related hardware. The computer program can be stored in a non-volatile computer-readable storage medium, and when executed, it can include the processes of the embodiments of the above methods. Any references to memory, databases, or other media used in the embodiments provided in this application can include at least one of non-volatile memory and volatile memory. Non-volatile memory can include read-only memory (ROM), magnetic tape, floppy disk, flash memory, optical memory, high-density embedded non-volatile memory, resistive random access memory (ReRAM), magnetic random access memory (MRAM), ferroelectric random access memory (FRAM), phase change memory (PCM), graphene memory, etc. Volatile memory can include random access memory (RAM) or external cache memory, etc. By way of illustration and not limitation, RAM can take many forms, such as Static Random Access Memory (SRAM) or Dynamic Random Access Memory (DRAM). The databases involved in the embodiments provided in this application may include at least one type of relational database and non-relational database. Non-relational databases may include, but are not limited to, blockchain-based distributed databases. The processors involved in the embodiments provided in this application may be general-purpose processors, central processing units, graphics processing units, digital signal processors, programmable logic devices, quantum computing-based data processing logic devices, artificial intelligence (AI) processors, etc., and are not limited to these.
[0185] The technical features of the above embodiments can be combined in any way. For the sake of brevity, not all possible combinations of the technical features in the above embodiments are described. However, as long as there is no contradiction in the combination of these technical features, they should be considered to be within the scope of this application.
[0186] The embodiments described above are merely illustrative of several implementation methods of this application, and while the descriptions are specific and detailed, they should not be construed as limiting the scope of this patent application. It should be noted that those skilled in the art can make various modifications and improvements without departing from the concept of this application, and these all fall within the protection scope of this application. Therefore, the protection scope of this application should be determined by the appended claims.
Claims
1. A method for optimizing a machine learning model, characterized in that, The method includes: Obtain the original model; The original model is subjected to computational logic analysis to obtain the corresponding original computation graph; wherein, the original computation graph represents the data dependencies between the various operators included in the original model; Determine how the computational logic of each operator changes due to the input data, and classify each operator into a static operator or a dynamic operator based on the changes. The adjacent static operators in the original computation graph are pre-compiled to obtain at least one static computation unit; wherein the pre-compilation is used to enable the adjacent static operators to be executed by a single scheduling instruction; Generate dynamic computation units corresponding to the dynamic operators in the original computation graph to obtain at least one dynamic computation unit; Based on the data dependencies, the at least one static computing unit is combined with the dynamic computing unit to obtain the optimized model corresponding to the original model.
2. The method according to claim 1, characterized in that, The step of performing computational logic analysis on the original model to obtain the corresponding original computational graph includes: The computational operations included in the original model are traversed, and each independent computational operation is constructed as a node corresponding to an operator in the original computational graph; Based on the data input-output relationship between the operator nodes, a directed edge is established in the original computation graph from the source node of the output data to the destination node of the input data.
3. The method according to claim 1, characterized in that, The change indicates whether the computational logic of the operator changes due to changes in the input data of the current operator; The step of determining the computational logic of each operator in response to changes in input data, and classifying each operator into a static operator or a dynamic operator based on these changes, includes: For each of the operators, determine whether the computational logic of the current operator includes decision logic that depends on the content of the input tensor; the decision logic includes at least one of conditional branching, dynamic indexing, or sorting and selection operations; If the computational logic includes the decision logic, the current operator is determined as the dynamic operator; If the computational logic does not include the decision logic, the current operator is determined as the static operator.
4. The method according to claim 1, characterized in that, The step of precompiling adjacent static operators in the original computation graph to obtain at least one static computation unit includes: In the original computation graph, all the maximum connected subgraphs consisting of static operator nodes and their direct data dependency edges are identified. For each identified maximum connected subgraph, the static operator sequence corresponding to the maximum connected subgraph is executed in the model inference environment to obtain the computational flow record corresponding to the maximum connected subgraph; the computational flow record includes at least one of the following: the start instruction of all operators in the current execution process, the computational logic, and the memory access operation. The computation flow record is compiled, and the compilation result is encapsulated into an independent executable object that can be replayed by the single scheduling instruction, thus obtaining a pre-compiled static computation unit.
5. The method according to claim 1, characterized in that, The generation of dynamic computation units corresponding to the dynamic operators in the original computation graph yields at least one dynamic computation unit, including: For each of the aforementioned dynamic operators, a unified function call prototype is defined for each dynamic operator. The function call prototype is used to characterize the memory address parameters of the input data and output data of the dynamic operator. The native computation kernel of the dynamic operator is encapsulated into a callable function that conforms to the function call prototype, thereby obtaining the dynamic computation unit corresponding to each dynamic operator.
6. The method according to claim 1, characterized in that, The step of combining the at least one static computing unit with the at least one dynamic computing unit according to the data dependency relationship to obtain the optimized model corresponding to the original model includes: Based on the data dependencies between the operator nodes in the original computation graph, the connection order between the static computation unit and the dynamic computation unit is determined; The static computing unit and the dynamic computing unit are combined according to the connection order to obtain the optimized model.
7. The method according to claim 1, characterized in that, The original model includes a large language model based on the Transformer architecture; the optimized model includes at least one Transformer layer; each Transformer layer includes a first static computation unit, a first dynamic computation unit, and a second static computation unit connected in sequence; the first static computation unit is used for preprocessing of attention mechanism computation; the first dynamic computation unit is used for computation of attention mechanism; and the second static computation unit is used for postprocessing of attention mechanism computation results.
8. The method according to claim 7, characterized in that, The original model includes a dense feedforward network model; the second static computation unit encapsulates all computation operators in the current Transformer layer after the attention mechanism computation.
9. The method according to claim 7, characterized in that, The original model includes a sparse expert hybrid network model; the second static computing unit encapsulates the router linear layer operator of the current Transformer layer; the optimized model is connected to a second dynamic computing unit and a third dynamic computing unit in sequence after the second static computing unit; wherein, the second dynamic computing unit is used to perform an expert selection operation based on the router output; the third dynamic computing unit is used to perform sparse matrix multiplication calculation for the selected expert based on the result of the expert selection operation.
10. The method according to claim 1, characterized in that, The data processing using the optimized model includes: Receive multiple online inference requests and dynamically combine the multiple online inference requests into an input batch tensor according to a preset capacity constraint; The input batch tensor is input into the optimized model to execute the static calculation unit and the dynamic calculation unit in the optimized model in sequence; The calculation results of the optimized model are collected and output as the data processing results corresponding to the online inference request.
11. A machine learning model optimization device, characterized in that, The device includes: The acquisition module is used to acquire the original model; The analysis module is used to perform computational logic analysis on the original model to obtain the corresponding original computation graph of the original model; wherein, the original computation graph represents the data dependencies between the various operators included in the original model; The determination module is used to determine how the computational logic of each operator changes due to the input data, and to classify each operator into a static operator or a dynamic operator based on the changes. A pre-compilation module is used to pre-compile adjacent static operators in the original computation graph to obtain at least one static computation unit; wherein, the pre-compilation is used to enable the adjacent static operators to be executed by a single scheduling instruction; A generation module is used to generate dynamic computation units corresponding to the dynamic operators in the original computation graph, thereby obtaining at least one dynamic computation unit. The combination module is used to combine the at least one static computing unit with the dynamic computing unit according to the data dependency relationship to obtain the optimized model corresponding to the original model.
12. A computer device comprising a memory and a processor, wherein the memory stores a computer program, characterized in that, When the processor executes the computer program, it implements the steps of the method according to any one of claims 1 to 10.
13. A computer-readable storage medium having a computer program stored thereon, characterized in that, When the computer program is executed by a processor, it implements the steps of the method according to any one of claims 1 to 10.
14. A computer program product, comprising a computer program, characterized in that, When the computer program is executed by a processor, it implements the steps of the method according to any one of claims 1 to 10.
Citation Information
Patent Citations
Memory allocation method and device, equipment and storage medium
CN119396568A
Deep learning algorithm compiling method, device, and related product
US20220129289A1