An efficient automatic tuning method, system and application for dynamic shape model inference

Abstract

The application discloses a kind of high-efficiency automatic tuning methods for dynamic shape model reasoning, the tuning method includes: step one, one or more to be tuned tensor program is normalized and is processed, is converted into the standard form conforming to preset scheduling template;Step two, including standardization tensor program, program feature information and target hardware specification information is passed to tuning engine and optimization is started;Step three, tuning engine receives including the bucketing factor obtained by bucketing strategy and the TopK value indicating the optimal performance of the first K configuration parameters, allocates scheduling template and configuration parameter transmitter, cooperates, and optimizes each barrel corresponding subproblem, obtains the corresponding optimization kernel in each barrel range to complete automatic tuning.The application also discloses a tuning system for implementing the above method, which has a wide range of application scenarios.

Description

Technical Field

[0001] This invention relates to the technical fields of model inference acceleration, deep learning compilers, and code generation optimization, and in particular to an efficient automatic tuning method, system, and application for dynamic shape model inference. Background Technology

[0002] In recent years, deep learning technologies, exemplified by the ChatGPT large model, have developed rapidly, and their innovative breakthroughs have had a profound impact on multiple disciplines and service industries. In deep learning fields such as large model applications, the model inference process is the crucial stage where the model generates predictions or decisions from input data, deploying the trained model to perform tasks in real-world scenarios. The performance of model inference directly affects the system's real-time performance, response speed, scalability, and practical application effectiveness. Typically, the inference process can be represented by a computation graph, where nodes represent operators (such as addition operators, matrix multiplication operators, etc.), and edges represent the transfer of tensor data between operators.

[0003] Accelerating model inference is one of the keys to improving the productivity of AI applications. Current mainstream model inference frameworks typically support calling high-performance operator libraries provided by hardware vendors (such as cuDNN and MKL) to achieve efficient model inference. However, on the one hand, these high-performance operator libraries need to be specifically optimized for particular hardware, thus requiring significant manpower to maintain and update them, especially when supporting new hardware or models. On the other hand, the operators in high-performance operator libraries are usually predefined and may not meet the performance requirements of specific scenarios, limiting flexibility. In contrast, deep learning compilers (such as TVM and XLA) can convert computation graphs into efficient code specific to hardware architectures and further accelerate model inference through operator fusion, memory optimization, and scheduling strategies, reducing the workload for developers. A tensor program is a program representation of a computation subgraph containing multiple operators, obtained by a deep learning compiler through operator fusion optimization. Tensor programs are the core object of low-level optimization in deep learning compilers. Scheduling is a combination of various optimization transformations applied to tensor programs (such as loop splitting, rearrangement, expansion, vectorization, tensor quantization, and software pipelines). Deep learning compilers typically explore different scheduling strategies to optimize tensor programs, ultimately accelerating model inference.

[0004] Dynamic shape models refer to models where the shape of the input tensor changes during inference. Dynamic input is particularly common in fields such as natural language processing and computer vision, especially in applications like chatbots and image generation models. For example, the batch size, sequence length, and image resolution of the input can change dynamically during inference. However, accelerating the inference of dynamic shape models, such as generative models like ChatGPT, remains a challenge. Compared to traditional methods that call operator libraries, methods based on deep learning compilers, while capable of handling dynamic input shapes, still have some shortcomings and limitations:

[0005] 1. Difficulty in effectively handling large-scale dynamic shapes: Existing methods (such as HAOTuner) only consider a small range of dynamic shape changes during inference optimization and attempt to use a scheduling strategy to handle all ranges of dynamic shapes. When faced with a large range of dynamic shape changes, such as for an M×N×K matrix multiplication where M varies from 1 to 2048 and N and K are fixed values, a scheduling strategy may perform well for a small range of M, such as 512-768, but performance will significantly decrease in other ranges, resulting in poor overall optimization results.

[0006] 2. Excessive Tuning Time: Some methods use machine learning algorithms (such as simulated annealing and genetic algorithms) to find optimized implementations of the target tensor program within a given search space. However, to find a better solution, deep learning compilers typically require thousands of searches, and each step needs to be executed on real hardware, making the entire model tuning process extremely time-consuming. Even with current state-of-the-art deep learning compilers, end-to-end model tuning can take hours or even days to complete.

[0007] 3. Lack of Interpretability: Tuning methods based on machine learning algorithms lack interpretability. Because the selection process for scheduling strategies is theoretically uninterpretable, solutions obtained through such methods offer almost no valuable information regarding performance explanation, making it difficult to understand and analyze the reasons for performance issues. When faced with new hardware or models, it is difficult to promptly improve the optimization algorithm based on the tuning results. Summary of the Invention

[0008] To overcome the limitations of existing technologies, this invention provides an efficient automatic tuning method for dynamic shape model inference, as well as a system and application for implementing the method. The method achieves efficient dynamic shape model inference optimization based on the collaborative work of scheduling templates and configuration parameter transmitters.

[0009] The first technical problem this invention aims to solve is how to effectively handle dynamic shape programs with large range variations. Experiments have shown that dynamic shapes within a certain range can achieve high performance using the same scheduling strategy. To address this, this invention proposes an innovative adaptive dynamic bucketing method. Bucketing is adaptively performed based on the type of tensor program and the problem size, with different buckets corresponding to different dynamic shape ranges. For the input dynamic shape, this invention can automatically generate bucketing options through the built-in bucketing algorithm, without user intervention, or the user can specify the bucketing strategy. For each bucket's range, the high-performance kernel for that range is optimized. Finally, all buckets are integrated to achieve excellent overall performance across a large dynamic range.

[0010] The second technical problem addressed by this invention is how to achieve optimized tensor programs within a short tuning time. To this end, this invention proposes a scheduling template and a configuration parameter emitter, as well as a method for their collaborative operation. The scheduling template is a combination of specific scheduling strategies, containing multiple adjustable configuration parameters, such as block shape, number of threads per thread block, number of started thread blocks, and vectorization size. The configuration parameter emitter is used to efficiently provide optimal configuration parameters. Based on a built-in bucketing algorithm, the configuration parameter emitter quickly generates high-performance configuration parameters according to the target tensor program, bucket range, and hardware specifications, and applies them to the scheduling template. The scheduling template and the configuration parameter emitter work together to achieve efficient automatic tuning.

[0011] The third technical problem addressed by this invention is how to design a tuning algorithm with good interpretability. To this end, this invention introduces a micro-performance model based on a user-transparent scheduling strategy. This micro-performance model, derived from optimization theory, clearly quantifies the specific impact of each configuration parameter on performance, helping users understand the basis of tuning decisions. The scheduling template and configuration parameter emitter used in the tuning process follow clear logic and rules, enabling users to track each step of the decision-making process. Therefore, this invention, through pre-experimentation, designs a micro-performance model containing multiple prior heuristic rules, ensuring the interpretability of the entire tuning algorithm.

[0012] Specifically, in this invention, the micro-performance model is essentially a combination of prior heuristic rules summarized from a large amount of experimental data. These rules explain which configuration parameters perform well and why. For example, the configuration parameters for matrix multiplication mentioned below have a high hit rate in the L2 cache, resulting in high equivalent bandwidth, less computation time per block, and thus better performance. Therefore, the micro-performance model improves the interpretability of the tuning algorithm described in this invention.

[0013] This invention provides an efficient automatic tuning method for dynamic shape model reasoning, the tuning method comprising:

[0014] Step 1: Normalize one or more tensor programs to be tuned, converting them into a standard form that conforms to the preset scheduling template;

[0015] Step 2: Transfer the standardized tensor program, program feature information, and target hardware specification information to the tuning engine to start optimization;

[0016] Step 3: The tuning engine receives the bucketing factor obtained from the bucketing strategy and the TopK value representing the top K configuration parameters with the best performance. It allocates scheduling templates and configuration parameter transmitters, and works together to optimize the sub-problems corresponding to each bucket, thereby obtaining the corresponding optimized kernel within each bucket range to complete automatic tuning.

[0017] In one specific implementation, a preprocessing stage is included before step one. The preprocessing stage includes the following steps: converting the computation graph of the model into an intermediate representation, Relax IR, and performing global optimization on the Relax IR through one or more graph-level optimization passes to generate one or more tensor programs to be tuned.

[0018] And / or,

[0019] Step three is followed by a post-processing stage, which includes the following steps: uniformly processing the optimized kernel generated by the tuning engine to generate host-side and device-side code, integrating the code and compiling it into a shared library; the shared library is loaded at runtime, and the optimal kernel is selected for inference based on the dynamic shape to achieve efficient dynamic scheduling.

[0020] In step one, when the tensor program contains only element-wise operations, all nested loops are merged into a single loop;

[0021] And / or,

[0022] Collect feature information from the tensor program, including: the data type of the input and / or output tensors, the static shape dimension, and whether the reduction axis involves dynamic shape, etc.

[0023] And / or,

[0024] Set the upper bound of the tensor procedure to 1024 or 2048;

[0025] And / or,

[0026] When the upper bound of a tensor procedure is set to 1024, the tensor procedure has a dynamic shape and an upper bound marker.

[0027] In step two, the target hardware specifications include: the maximum number of resident threads per streaming multiprocessor, the maximum number of resident thread blocks per streaming multiprocessor, the maximum number of registers that each thread block can use, the maximum shared memory capacity that each thread block can use, the maximum number of threads that each thread block can have, the maximum number of registers that each thread can use, the number of streaming multiprocessors, the number of thread bundle schedulers in the streaming multiprocessors, the size of the thread bundle, the capacity of the L2 cache, the bandwidth of the global memory, the bandwidth of the L2 cache, the bandwidth of the shared memory, the computing power of Float32, the computing power of Float16, the computing power of Int64, the computing power of the tensor core, and the computing power of the target GPU, etc.

[0028] In step three, a bucketing strategy is constructed to break down each dynamic shape tensor program into subproblems of different sizes; subproblems that are applicable to the same optimal scheduling combination are grouped into one bucket.

[0029] And / or,

[0030] Based on the input standardized tensor program, program feature information, and target hardware specification information, a scheduling template and configuration parameter transmitter are allocated.

[0031] The scheduling template contains one or more scheduling strategies with adjustable configuration parameters. The scheduling template includes: a general element-wise operation scheduling template, a transpose operation scheduling template, a general matrix multiplication scheduling template, a WMMA scheduling template, a general matrix-vector multiplication scheduling template, and a general reduction operation scheduling template.

[0032] The general element-wise operation scheduling template is optimized using a scheduling strategy of loop fusion, expansion and vectorization, and is suitable for tensor program scheduling that only involves spatial axes;

[0033] The transpose operation scheduling template, for transpose operations, loads data into shared memory to achieve merged read and write of global memory, thereby gaining additional speedup;

[0034] The general matrix multiplication scheduling template uses the CUDA kernel to schedule tensor programs for matrix multiplication behavior. It sets up a two-level block structure from global memory to shared memory and from shared memory to register files through a combination of scheduling methods such as loop fusion, splitting, rearranging, and vectorization. It also optimizes software pipeline data prefetching and cooperative acquisition of shared memory data.

[0035] The WMMA scheduling template optimizes matrix multiplication scheduling by using a higher-performance tensor core through tensor-based scheduling primitives, based on the general matrix multiplication scheduling template.

[0036] The general matrix-vector multiplication scheduling template further accelerates the computation of the dot product of each element through a combination of scheduling combinations of loop fusion, splitting, and rearrangement, as well as the __shfl_down_sync primitive, making it suitable for optimized scheduling of matrix-vector multiplication.

[0037] The general reduction operation scheduling template uses the __shfl_down_sync primitive to parallelize and accelerate reduction operations, and is suitable for optimized scheduling including reduction summation, reduction maximum value and softmax activation function reduction operations.

[0038] In the preprocessing stage, the graph-level optimization pass includes common subexpression elimination, dead code elimination, constant folding, memory planning, and operator fusion; the tensor program to be tuned represents the optimized computational subgraph, which is composed of multiple operator operations fused into a tensor program;

[0039] And / or,

[0040] In the post-processing stage, the host-side code is responsible for calling the device-side optimization kernel for the corresponding bucket range based on the actual value of the dynamic shape at runtime, so as to achieve efficient dynamic scheduling.

[0041] The present invention also provides an optimization system for implementing the above-described automatic optimization method, the optimization system comprising:

[0042] 1) Preprocessing module

[0043] It is responsible for receiving the input model computation graph and converting it into an intermediate representation (Relax IR) to prepare for the subsequent optimization process.

[0044] The preprocessing module converts the model computation graph into a Relax IR representation and performs global optimization on the computation graph through a graph-level optimization pass, including common subexpression elimination, dead code elimination, constant folding, memory planning, and operator fusion. The output generates several tensor programs to be optimized, each tensor program representing an optimized computation subgraph.

[0045] 2) Feature extraction module

[0046] Extract the feature information of the tensor program to be optimized in order to provide sufficient input information for the tuning engine.

[0047] The feature extraction module normalizes the tensor program, converting it into a standardized form (e.g., matrix multiplication, element-wise operations, etc.), and collects program feature information, including the data type of the input and / or output tensors, the static shape dimension, and whether dynamic shape reduction axes are involved. If the tensor program has a dynamic shape, a corresponding upper bound is set (e.g., {"m":1024}), and a standardized tensor program containing feature information is output for use by the tuning engine.

[0048] 3) Bucketing Module

[0049] The dynamic shape tensor procedure is divided into subproblems of different sizes to reduce the number of problems to be optimized.

[0050] The bucketing module divides the tensor program into multiple buckets based on the dynamic shape change range and bucketing factor. Each bucket represents a group of subproblems of similar size. The module incorporates an adaptive bucketing algorithm. The default bucketing factor for transpose and element-wise operations is 512, while the default bucketing factor for reduction, matrix-vector multiplication, and matrix multiplication is 256. Furthermore, if the user enables advanced optimization options, the module can further divide the first bucket (range 1-256) of the matrix multiplication program into buckets 1-32, 33-128, and 129-256. It instructs the system to use a general matrix-vector multiplication scheduling template for buckets 1-32, a general matrix multiplication scheduling template for buckets 33-128, and a WMMA scheduling template for buckets 129-256, thereby pursuing higher performance for matrix multiplication with small-scale inputs. The bucketing module also supports custom bucketing strategies. By sharing scheduling combinations within buckets, it reduces the need for individual tuning, improves optimization efficiency, and outputs a set of bucketed tensor programs, with each bucket containing multiple subproblems.

[0051] 4) Optimize the engine module

[0052] Based on program characteristics and hardware specifications, the configuration parameters of the tensor program in each bucket are tuned to generate an optimized kernel.

[0053] The tuning engine module receives a standardized tensor program, feature information, target hardware specifications, and TopK values. It uses a scheduling template to generate candidate configuration parameters and selects the TopK configurations with the best performance through actual performance testing. It generates corresponding optimized kernels for sub-problems within each bucket and outputs the optimized kernel code for each bucket.

[0054] 5) Code integration module

[0055] The optimized kernel generated by the integration and tuning engine is used to generate host-side code and device-side code for dynamic scheduling at runtime.

[0056] The code integration module generates host-side code, which is used to select the corresponding bucket and optimized kernel based on the runtime dynamic shape, generate device-side code, which includes optimized kernels for a specific bucket range, and output the integrated host-side code and device-side kernel code, ready for compilation.

[0057] 6) Shared library generation module

[0058] The integrated code is compiled into a shared library for loading and calling at runtime.

[0059] The shared library generation module compiles and integrates the host-side and device-side code into a single shared library file. This shared library file contains optimized kernels for all buckets, ensuring that the model can flexibly adapt to dynamic shape inputs during inference. It outputs the generated shared library file, which is then dynamically loaded during model inference.

[0060] 7) Runtime scheduling module

[0061] During inference runtime, shared libraries are loaded, and the optimal device-side kernel is selected based on the dynamic shape of the input for efficient inference computation.

[0062] The runtime scheduling module determines the corresponding bucket range based on the actual shape of the input tensor, calls the corresponding optimized kernel in the shared library, completes the model inference calculation, and outputs efficient inference of the dynamic shape model.

[0063] This invention also provides applications of the above-mentioned automatic tuning method or automatic tuning system in generating high-performance kernels on hardware platforms, quickly calling dynamic shape models to obtain shared libraries for runtime inference, integrating deep learning frameworks to improve inference speed, and remote tuning in the cloud.

[0064] Compared with the prior art, the present invention has the following advantages and beneficial effects:

[0065] This invention employs a bucketing method, breaking down dynamic shape tensor programs within a certain range into sub-problems of different ranges, enabling efficient handling of large-scale dynamic shapes. By applying multiple efficient kernels to each small range, the tensor program achieves excellent performance even when dealing with the entire large range.

[0066] This invention provides a comprehensive operator scheduling template, covering various commonly used operators in the industry, and utilizes efficient algorithms such as software pipelining, tensor quantization, and Split-K to optimize tensor programs. Furthermore, the scheduling template supports scheduling multiple fusion operators, demonstrating stronger optimization capabilities. The accompanying configuration parameter emitter combines static program information with the target platform's hardware specifications to recommend theoretically superior configuration parameters for the dynamic shape input range. With the support of the configuration parameter emitter, the scheduling template can quickly acquire and apply high-performance configuration parameters and perform performance measurements on actual hardware to ensure the high accuracy of the final selected configuration parameters. This invention enables optimization tasks that previously required hours or even days to be completed in minutes or even seconds, shortening optimization time and ensuring that the optimized dynamic shape tensor program exhibits excellent performance for any shape within the dynamic range.

[0067] The scheduling templates and configuration parameter emitters during the tuning process are based on explicit logic and rules, ensuring that users can track each step of the decision-making process, and the entire tuning algorithm is interpretable. Furthermore, the high-performance configuration parameters selected by the built-in micro-performance model have clear explanatory basis, enabling users to understand why these configuration parameters and scheduling strategies achieve excellent performance within a specific range. When encountering new hardware or new models, users can promptly adjust the scheduling templates and optimization algorithms based on actual performance feedback.

[0068] This invention provides a code generator and runtime environment, allowing users to perform end-to-end inference directly using the runtime environment after model tuning. Furthermore, users can generate multiple versions of high-performance kernel functions based on this invention, encapsulating them into a high-performance operator library to address different dynamic shape problems.

[0069] This invention exhibits high scalability, allowing users to design specific scheduling templates and customize the transmitter configuration parameters to integrate into the proposed framework. Simultaneously, this invention offers high flexibility, allowing users to customize bucketing strategies, tensor parameter generation quantities, TopK values, and the selection of scheduling templates and configuration parameters. For example, for tensor programs exhibiting general matrix multiplication behavior, users can optimize using a general matrix-vector multiplication scheduling template when M(1-64) is small, a general matrix multiplication scheduling template in the range of 65-256, and a WMMA scheduling template in the range of 257-1024.

[0070] The method of this invention provides a 1.1x to 1.7x performance improvement over TVM for end-to-end prefill inference of Llama on RTX3090 after optimization by this invention, with input sequence lengths ranging from 256 to 2048. Attached Figure Description

[0071] To more clearly illustrate the technical solutions in the embodiments of the present invention or the prior art, the drawings used in the description of the embodiments or the prior art will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0072] Figure 1 This is a complete flowchart of the optimization method of the present invention.

[0073] Figure 2 This is a flowchart illustrating the implementation of the optimization engine during the automatic tuning process of this invention.

[0074] Figure 3 This is a flowchart of the scheduling template and configuration parameter allocation process of the present invention.

[0075] Figure 4 This is a schematic diagram of the implementation process of Embodiment 1 of the present invention.

[0076] Figure 5 This is a schematic diagram of the implementation process of Embodiment 2 of the present invention.

[0077] Figure 6 This is a schematic diagram of the implementation process of Embodiment 3 of the present invention.

[0078] Figure 7 This is a schematic diagram illustrating the implementation of the optimized tensor program of this invention.

[0079] Figure 8 This is a schematic diagram of the implementation process of Embodiment 4 of the present invention.

[0080] Figure 9 This is a schematic diagram of the implementation process of Embodiment 5 of the present invention. Detailed Implementation

[0081] The present invention will be further described in detail below with reference to the specific embodiments and accompanying drawings. Except for the contents specifically mentioned below, the processes, conditions, and experimental methods for implementing the present invention are all common knowledge and general knowledge in the art, and the present invention does not have any particular limitations.

[0082] This invention proposes an efficient automatic tuning method for dynamic shape model reasoning, as well as a system and application implemented through the method.

[0083] The flow of the automatic tuning method of the present invention is as follows: Figure 1 As shown, the overall process is divided into three main stages: preprocessing, automatic optimization, and postprocessing.

[0084] 1) Preprocessing stage

[0085] In this phase, the computation graph of the dynamic shape model is first converted into an intermediate representation, Relax IR, and then globally optimized through multiple graph-level optimization passes. These optimizations include common subexpression elimination, dead code elimination, constant folding, memory planning, and operator fusion to generate several tensor programs to be tuned. Each tensor program represents an optimized computational subgraph, and deep learning compilers typically fuse multiple operator operations (such as element-wise addition, matrix multiplication, activation functions, etc.) into a single univariate tensor program. The scheduling template provided by this invention is based on TVM TIR Schedule. TIR scheduling provides a variety of high-level scheduling primitives, allowing developers to flexibly control the computation order, parallelism, and memory layout of computation units. Specifically, through scheduling primitives such as comput_at, comput_inline, and reverse_compute_inline, multiple operators can be executed in parallel and share resources, thereby improving computational performance, reducing memory consumption, and ensuring efficient utilization of hardware resources. This supports the simultaneous scheduling of multiple fusion operators. For example, for the computation subgraph of matrix multiplication → element-wise addition → Gaussian activation function, this invention can directly use the general matrix multiplication scheduling template for alignment and scheduling. Compared to scheduling each unit separately, simultaneous scheduling can effectively reduce the number of memory reads, thereby improving performance.

[0086] 2) Automatic optimization stage

[0087] In this stage, the system optimizes each tensor program individually. First, the tensor program is normalized, converting it into a standard form conforming to a specific preset scheduling template. For example, for a tensor program with matrix multiplication behavior, this invention normalizes it to the form C[bsz,M,N]=A[bsz,M,K]×B[bsz,N,K]; for tensor programs with only element-wise operations, the system merges all nested loops into a single loop. Simultaneously, the system collects feature information from the tensor program, including the data type of the input and / or output tensors, the static shape dimension, and whether dynamic shapes are involved, such as reduction axes. If the tensor program has dynamic shapes and a corresponding upper bound marker (e.g., {"m":1024}), the upper bound of m will be set to 1024; otherwise, the upper bound defaults to 2048. Subsequently, this feature information, along with the program to be tuned, is passed to the tuning engine to initiate the optimization process.

[0088] To effectively handle large-scale dynamic shapes, this invention employs a bucketing method. Based on the range of dynamic shape changes and a bucketing factor, each dynamic shape tensor program is broken down into sub-problems of different sizes. For example, for matrix multiplication of size M×N×K, when m dynamically varies from 1 to 2048, and N and K are fixed values, m=1 and M=2048 represent completely different optimization requirements. However, the optimal scheduling combination for matrix multiplication with M values ​​within a certain range is the same. Therefore, buckets can be used to represent a certain range, dividing a certain range of M values ​​into buckets, with sub-problems within the same bucket enjoying the same scheduling combination. In this way, optimizing the sub-problem for each individual value becomes optimizing the set of sub-problems represented by each bucket, thereby appropriately reducing the number of problems to be optimized and reducing tuning time. By bucketing, sub-problems within a bucket share the same optimal scheduling combination, ensuring consistent performance of dynamic shapes within a certain range and avoiding the problem of poor performance of a single scheduling strategy for large-scale dynamic shapes.

[0089] In one specific implementation, the default bucketing algorithm built into this invention uses 256 as the range for each bucket (i.e., the default value for the bucketing factor below is 256). Experiments have shown that such a division is sufficient to provide good performance over the overall range. Of course, for matrix multiplication, using a bucket range of 128 can further accelerate the performance of a small range, but the performance of 256 is already ideal. This default setting not only simplifies the optimization process but also achieves a good balance between performance and complexity. This invention also supports custom bucketing strategies, which can flexibly divide the dynamic range and further improve the overall optimization effect.

[0090] The tuning engine receives two key parameters: the bucketing factor and the TopK value. The bucketing factor, as the core metric of the partitioning strategy, defines the range of each bucket after dynamic shape partitioning. For example, when the bucketing factor is 256 and the dynamic shape variation range is 1-1024, the range will be divided into four buckets: 1-256, 257-512, 513-768, and 769-1024. Furthermore, if the user provides a custom bucketing strategy, that strategy will be used. The TopK value is used to select the top K configuration parameters with the best performance during the tuning process. The default TopK value is 1. Although increasing the TopK value may improve optimization, a TopK value of 10 is usually sufficient; adding configuration parameter candidates with values ​​exceeding 10 typically has limited performance improvement. In addition, the tuning engine receives target hardware specification information pre-collected from the CUDA programming manual and white paper, including the maximum number of resident threads per stream multiprocessor (SM), the maximum number of resident thread blocks per SM, the maximum number of registers that can be used per thread block, the maximum shared memory capacity that can be used per thread block, the maximum number of threads per thread block, the maximum number of registers that can be used per thread, the number of SMs, the number of thread bundle schedulers in the SM, the size of the thread bundle, the capacity of the L2 cache, the bandwidth of global memory, the bandwidth of L2 cache, the bandwidth of shared memory, the computing power of Float32, the computing power of Float16, the computing power of Int64, the computing power of tensor cores, and the computing power of the target GPU. Subsequently, each bucket represents a sub-problem, and the tuning engine optimizes the sub-problems corresponding to each bucket one by one, obtaining the corresponding optimized kernels within each bucket range and finally completing the automatic tuning.

[0091] 3) Post-processing stage

[0092] In the post-processing stage, the system uniformly processes the results obtained from the tuning engine. For each tensor program, the system integrates the code of each bucket and its optimization kernel, generating host-side code and device-side code respectively. The host-side code is responsible for calling the device-side optimization kernel within the corresponding bucket range based on the actual value of the dynamic shape at runtime, to achieve efficient dynamic scheduling. In addition, the system compiles all the integrated code into a shared library for runtime loading to implement model inference.

[0093] The automatic tuning method of this invention includes the following process and implementation steps:

[0094] Since the efficient automatic tuning method is the most significant innovation of this invention, it will be described in detail below. As mentioned above, in the automatic tuning phase, the system uses the standardized program, program feature information, and target hardware specification information as inputs to the tuning engine. The workflow of the tuning engine is as follows: Figure 2As shown. After initializing with bucketing factors and TopK values, the tuning engine allocates suitable scheduling templates and configuration parameter transmitters based on the input standardized procedures, program characteristic information, and target hardware specifications.

[0095] 1) Scheduling template and configuration parameters for transmitter allocation

[0096] This invention innovatively combines scheduling strategies such as loop splitting, rearrangement, and vectorization, proposing six scheduling templates with specific combinations and their corresponding configuration parameter emitters. Each scheduling template contains different scheduling strategies, applicable to different types of tensor programs, specifically including: a general element-wise operation scheduling template, a transpose operation scheduling template, a general matrix multiplication scheduling template, a WMMA scheduling template, a general matrix-vector multiplication scheduling template, and a general reduction operation scheduling template. Among them, the general element-wise operation scheduling template is optimized using loop fusion, expansion, and vectorization scheduling strategies, suitable for scheduling tensor programs involving only spatial axes, such as transpose, reshape, split, and add operations. For transpose operations, the transpose operation scheduling template is used to achieve merged read and write of global memory by loading data into shared memory, thereby obtaining additional speedup. The general matrix multiplication scheduling template uses the CUDA core to schedule tensor programs performing matrix multiplication. It employs a two-level block structure—from global memory to shared memory and from shared memory to register files—through a combination of loop fusion, splitting, rearrangement, and vectorization scheduling. Optimizations include software pipeline data prefetching and cooperative data acquisition from shared memory. The WMMA scheduling template, building upon the general matrix multiplication template, utilizes tensor-based scheduling primitives to optimize matrix multiplication scheduling using a higher-performance tensor core. Notably, matrix-vector multiplication, as a special subset of matrix multiplication, has significantly different optimization requirements compared to large-scale matrix multiplication. Therefore, a specific general matrix-vector multiplication scheduling template is proposed to meet its specific tuning needs. This template further accelerates the process by parallelizing the dot product calculation for each element through a combination of loop fusion, splitting, and rearrangement scheduling, as well as the `__shfl_down_sync` primitive. The general reduction operation scheduling template supports reduction operations other than matrix multiplication, including reduction summation, reduction maximization, and softmax activation function, and utilizes the __shfl_down_sync primitive to parallelize and accelerate reduction operations. These six scheduling templates cover a variety of operators commonly used in mainstream large language models, image generation models, and temporal prediction models, ensuring broad applicability and flexibility to meet diverse deep learning needs.

[0097] The allocation process for scheduling templates and configuration parameters of transmitters is as follows: Figure 3As shown, the process first determines whether the tensor program involves a reduction operation. If no reduction operation is involved, it further determines whether it involves a transpose operation. If it is a transpose operation, a transpose operation scheduling template and a corresponding configuration parameter transmitter are allocated; otherwise, an element-wise operation scheduling template and a corresponding configuration parameter transmitter are allocated. If a reduction operation is involved and the tensor program contains matrix multiplication, it further determines whether it is a matrix-vector multiplication mode. If so, a general matrix-vector multiplication scheduling template and a corresponding configuration parameter transmitter are allocated. Otherwise, further determination is made based on the target hardware computing power. The target hardware computing power consists of a major revision number X and a minor revision number Y, denoted as XY. Devices with the same major revision number have the same core architecture; for example, the major revision number of the hardware computing power of the NVIDIA Ampere GPU architecture is 8. The minor revision number corresponds to incremental improvements in the core architecture, which may include some new features. If the target hardware computing power is greater than or equal to 7.0, it indicates that the target hardware supports the tensor core, so a WMMA scheduling template and a corresponding configuration parameter transmitter are allocated to it; otherwise, a general matrix multiplication scheduling template and a corresponding configuration parameter transmitter are allocated. If there are no matrix multiplication operations but there are reduction operations, then allocate a general reduction operation scheduling template and a corresponding configuration parameter transmitter.

[0098] 2) Configure parameters for transmission

[0099] After the scheduling template and configuration parameter emitter are allocated, the tuning engine divides the dynamically shaped tensor program into different buckets based on the bucketing factor or the user-specified bucketing strategy. For each bucket, the configuration parameter emitter generates multiple valid configuration parameter candidates based on hardware specifications and program characteristics. For example, for matrix multiplication operators, the choice of thread block size must ensure that the shared memory usage is less than its maximum available capacity; otherwise, a segmentation fault will occur during the runtime of the generated kernel code. The thread-level block size must ensure that the number of registers used by each thread is less than the maximum number of registers that each thread can use; otherwise, registers will overflow into local memory, which, since local memory is much slower than registers, will cause severe performance degradation.

[0100] Next, the configuration parameter emitter evaluates the theoretical performance of each configuration parameter candidate using a built-in micro-performance model. Specifically, the micro-performance model reasonably evaluates the performance of configuration parameters within the corresponding bucket range by combining the characteristics of the target tensor program and the target hardware specifications, accelerating the tuning process while ensuring high performance. For example, the micro-performance model of the general matrix multiplication configuration parameter emitter first determines whether the configuration parameter candidate is aligned with the hardware specifications and the shape of the input tensor, and then estimates the L2 cache hit rate achievable by the block partitioning strategy of the corresponding configuration parameter candidate, thereby obtaining the equivalent memory bandwidth and estimating the computation time required for each block, thus estimating the computation time of the entire task, which serves as the second basis for evaluating configuration parameters. If the estimated times are the same, the theoretical computing power achievable by the block and the computation-to-memory ratio of the thread-level block are further compared. Higher-ranking configuration parameters have higher priority and better theoretical performance. Finally, the configuration emitter selects the top K configuration parameters with the highest theoretical performance, passes them to the scheduling template to generate the corresponding compilation results, and performs actual performance testing on the hardware. The element-wise operation configuration parameter emitter prioritizes configuration parameters with vectorized lengths aligned with hardware specifications and the shape of the input tensor program. Under this condition, larger vectorized lengths have higher priority and theoretically better performance. For the general reduction operation configuration parameter emitter, it first ensures that the number of elements computed by each thread is less than a certain experimental value (this experimental value is hardware-dependent). Under this condition, configuration parameter candidates with larger maximum active thread blocks per streaming multiprocessor have higher priority (representing less synchronization overhead). If the maximum active thread blocks per streaming multiprocessor are the same, then the more threads in a thread block, the higher the priority (representing higher parallelism).

[0101] For matrix-vector multiplication, transpose, and other operations, there are also corresponding micro-performance models to evaluate configuration parameters.

[0102] 3) Actual performance testing

[0103] During actual performance testing, the tuning engine generates tensor data for analysis based on the compilation results and bucket ranges. By default, 10 sets of data are generated for each bucket range to balance tuning time and accuracy. These data are used to compare the actual hardware performance of different configuration parameters within that bucket range. If higher accuracy is required, users can choose to generate more sets of data, but this will increase tuning time.

[0104] Finally, the tuning engine outputs a dictionary-structured result, where each bucket corresponds to an optimized kernel with the best compilation result, thus completing the entire automatic tuning process.

[0105] Example 1: Using the default bucketing strategy, multiple high-performance kernel versions are generated for the matrix multiplication operator on a specific hardware platform. (The following example illustrates the optimization of M×N×K matrix multiplication on an NVIDIA RTX 3090, where N and K are fixed values, and M changes dynamically with an upper bound of 1024.)

[0106] like Figure 4 As shown, the computation graph of the model only contains matrix multiplication operators. First, the computation graph is used as input, preprocessed, and normalized to obtain a normalized tensor program and program feature information. The normalized tensor program and program feature information, along with the target hardware specifications (RTX3090 hardware specifications), are then input into the tuning engine for optimization. Since the target hardware RTX3090 has a computing power of 8.6 > 7.0, the tuning engine allocates a WMMA scheduling template and configuration parameter emitter according to the default built-in algorithm. Based on the default bucketing strategy, the range 1–1024 is divided into four buckets: 1–256, 257–512, 513–768, and 769–1024. The scheduling template and configuration parameter emitter then work together to optimize the tensor program for each bucket range to obtain a high-performance core, resulting in a total of four high-performance cores.

[0107] Example 2: Using a user-provided bucketing strategy, multiple high-performance kernel versions are generated for the matrix multiplication operator on a specific hardware platform. (The following example illustrates the optimization of M×N×K matrix multiplication on an NVIDIA RTX 3090, where N and K are fixed values, and M changes dynamically with an upper bound of 1024.)

[0108] like Figure 5As shown, the computation graph of the model only contains matrix multiplication operators. First, the computation graph is used as input, preprocessed, and normalized to obtain a normalized tensor program and program feature information. The normalized tensor program and program feature information, along with the target hardware specifications (RTX3090 hardware specifications), are then input into the tuning engine for tuning. Since the target hardware RTX3090 has a computing power of 8.6 > 7.0, the tuning engine will allocate a WMMA scheduling template and configuration parameter emitter according to the default built-in algorithm. The user additionally provided a bucketing strategy, instructing the system to divide 1–64 into one bucket and apply the general matrix-vector multiplication scheduling template and configuration parameter emitter, and 65–256 into another bucket and apply the same general matrix multiplication scheduling template and configuration parameter emitter. The remaining range uses the system's default algorithm, dividing 257–512, 513–768, and 769–1024 into buckets respectively, and applying the WMMA scheduling template and configuration parameter emitter. Subsequently, the scheduling template and the configuration parameter transmitter will work together to optimize the tensor program for each bucket range to obtain a high-performance kernel, resulting in a total of five high-performance kernels.

[0109] Example 3: Optimize the dynamic shape model and directly obtain the shared library for runtime inference.

[0110] like Figure 6 As shown, the computational graph of the dynamic shape model is first preprocessed to obtain several tensor programs to be tuned. Then, based on the input target hardware specifications, each of the tensor programs to be tuned is automatically tuned. For example... Figure 7 As shown, each optimized tensor program obtained through tuning has several optimized kernels, which can dynamically call the optimized kernels corresponding to the bucket range based on the actual value of the dynamic shape at runtime. After obtaining each optimized tensor program, this invention integrates its code and generates a shared library, and finally uses the runtime environment provided by this invention for inference. For the Llama2-7B model, in the end-to-end prefill stage of inference on an NVIDIA RTX3090, the model implementation obtained after tuning by this invention can achieve a performance improvement of 1.1 to 1.7 times compared to TVM for input sequence lengths in the range of 256-2048.

[0111] Example 4: Integration with Deep Learning Frameworks

[0112] like Figure 8As shown, this invention can be seamlessly integrated into deep learning frameworks (such as PyTorch, TensorFlow, etc.). After training a dynamic shape model in a deep learning framework, users can directly extract its computational graph and call the tuning engine of this invention to perform automatic optimization. During the tuning process, this invention generates efficient optimization kernels for different dynamic shape ranges by analyzing the tensor procedures in the model. After tuning is complete, the tuning engine will return an optimized shared library, which users can load into the runtime environment of the deep learning framework for efficient inference of the dynamic shape model.

[0113] Example 5: Deploying this automatic tuning system in a cloud environment

[0114] like Figure 9 As shown, this invention can be deployed as a cloud service, providing remote automatic optimization functionality. The client sends a computational graph of the dynamic shape model and the target hardware selection to the server. Upon receiving the optimization request, the server automatically selects a suitable hardware resource environment to execute the optimization process. After optimization is complete, the server returns the optimized shared library to the client for local runtime loading and efficient inference.

[0115] Those skilled in the art will understand that embodiments of this application can be provided as methods, systems, or computer program products. Therefore, this application can take the form of a completely hardware embodiment, a completely software embodiment, or an embodiment combining software and hardware aspects. Furthermore, this application can take the form of a computer program product implemented on one or more computer-usable storage media (including but not limited to disk storage, CD-ROM, optical storage, etc.) containing computer-usable program code. The solutions in the embodiments of this application can be implemented in various computer languages, such as the object-oriented programming language Java and the interpreted scripting language JavaScript.

[0116] This application is described with reference to flowchart illustrations and / or block diagrams of methods, apparatus (systems), and computer program products according to embodiments of this application. It will be understood that each block of the flowchart illustrations and / or block diagrams, and combinations of blocks in the flowchart illustrations and / or block diagrams, can be implemented by computer program instructions. These computer program instructions can be provided to a processor of a general-purpose computer, special-purpose computer, embedded processor, or other programmable data processing apparatus to produce a machine, such that the instructions, which execute via the processor of the computer or other programmable data processing apparatus, generate instructions for implementing the flowchart... Figure 1 One or more processes and / or boxes Figure 1 A device that provides the functions specified in one or more boxes.

[0117] These computer program instructions may also be stored in a computer-readable storage medium that can direct a computer or other programmable data processing device to function in a particular manner, such that the instructions stored in the computer-readable storage medium produce an article of manufacture including instruction means, which are implemented in a process Figure 1 One or more processes and / or boxes Figure 1 The function specified in one or more boxes.

[0118] These computer program instructions may also be loaded onto a computer or other programmable data processing equipment to cause a series of operational steps to be performed on the computer or other programmable equipment to produce a computer-implemented process, thereby providing instructions that execute on the computer or other programmable equipment for implementing the process. Figure 1 One or more processes and / or boxes Figure 1 The steps of the function specified in one or more boxes.

[0119] Although preferred embodiments of this application have been described, those skilled in the art, upon learning the basic inventive concept, can make other changes and modifications to these embodiments. Therefore, the appended claims are intended to be interpreted as including the preferred embodiments as well as all changes and modifications falling within the scope of this application.

[0120] Obviously, those skilled in the art can make various modifications and variations to this application without departing from the spirit and scope of this application. Therefore, if such modifications and variations fall within the scope of the claims of this application and their equivalents, this application also intends to include such modifications and variations.

[0121] The scope of protection of this invention is not limited to the above embodiments. Any variations and advantages that can be conceived by those skilled in the art without departing from the spirit and scope of this invention are included in this invention and are protected by the appended claims.

Claims

1. An efficient automatic tuning method for dynamic shape model reasoning, characterized in that, The optimization method includes: Step 1: Normalize one or more tensor programs to be tuned, converting them into a standard form that conforms to the preset scheduling template; Step 2: Transfer the standardized tensor program, program feature information, and target hardware specification information to the tuning engine to start optimization; Step 3: The tuning engine receives the bucketing factor obtained from the bucketing strategy and the TopK value representing the top K configuration parameters with the best performance. It allocates scheduling templates and configuration parameter transmitters, and works together to optimize the sub-problems corresponding to each bucket, thereby obtaining the corresponding optimized kernel within each bucket range to complete automatic tuning.

2. The optimization method as described in claim 1, characterized in that, In step one, when the tensor program contains only element-wise operations, all nested loops are merged into a single loop; And / or, Collect feature information from the tensor program, including: the data type of the input and / or output tensors, the static shape dimension, and whether the reduction axis involves dynamic shape; And / or, Set the upper bound of the tensor procedure to 1024 or 2048; And / or, When the upper bound of a tensor procedure is set to 1024, the tensor procedure has a dynamic shape and an upper bound marker.

3. The optimization method as described in claim 1, characterized in that, In step two, the target hardware specifications include: the maximum number of resident threads per streaming multiprocessor, the maximum number of resident thread blocks per streaming multiprocessor, the maximum number of registers that each thread block can use, the maximum shared memory capacity that each thread block can use, the maximum number of threads that each thread block can have, the maximum number of registers that each thread can use, the number of streaming multiprocessors, the number of thread bundle schedulers in the streaming multiprocessors, the size of the thread bundle, the capacity of the L2 cache, the bandwidth of the global memory, the bandwidth of the L2 cache, the bandwidth of the shared memory, the computing power of Float32, the computing power of Float16, the computing power of Int64, the computing power of the tensor core, and the computing power of the target GPU.

4. The optimization method as described in claim 1, characterized in that, In step three, a bucketing strategy is constructed to break down each dynamic shape tensor program into subproblems of different sizes; subproblems that are applicable to the same optimal scheduling combination are grouped into one bucket. And / or, Based on the input standardized tensor program, program feature information, and target hardware specification information, a scheduling template and configuration parameter transmitter are allocated.

5. The optimization method as described in claim 4, characterized in that, The scheduling template contains one or more scheduling strategies with adjustable configuration parameters. The scheduling template includes: a general element-wise operation scheduling template, a transpose operation scheduling template, a general matrix multiplication scheduling template, a WMMA scheduling template, a general matrix-vector multiplication scheduling template, and a general reduction operation scheduling template. The general element-wise operation scheduling template is optimized using a scheduling strategy of loop fusion, expansion and vectorization, and is suitable for tensor program scheduling that only involves spatial axes; The transpose operation scheduling template, for transpose operations, loads data into shared memory to achieve merged read and write of global memory, thereby gaining additional speedup; The general matrix multiplication scheduling template uses the CUDA kernel to schedule tensor programs for matrix multiplication behavior. It sets up a two-level block structure from global memory to shared memory and from shared memory to register files through a combination of scheduling methods such as loop fusion, splitting, rearranging, and vectorization. It also optimizes software pipeline data prefetching and cooperative acquisition of shared memory data. The WMMA scheduling template optimizes matrix multiplication scheduling by using tensor cores through tensor-quantized scheduling primitives, based on the general matrix multiplication scheduling template. The general matrix-vector multiplication scheduling template further accelerates the computation of the dot product of each element through a combination of scheduling combinations of loop fusion, splitting, and rearrangement, as well as the __shfl_down_sync primitive, making it suitable for optimized scheduling of matrix-vector multiplication. The general reduction operation scheduling template uses the __shfl_down_sync primitive to parallelize and accelerate reduction operations, and is suitable for optimized scheduling of reduction operations including reduction summation, reduction maximum value, and softmax activation function reduction operations.

6. The optimization method as described in claim 1, characterized in that, Before step one, there is also a preprocessing stage, which includes the following steps: converting the computation graph of the model into an intermediate representation, Relax IR, and performing global optimization on the Relax IR through one or more graph-level optimization passes to generate one or more tensor programs to be tuned; And / or, Step three is followed by a post-processing stage, which includes the following steps: uniformly processing the optimized kernel generated by the tuning engine to generate host-side and device-side code, integrating the code and compiling it into a shared library; the shared library is loaded at runtime, and the optimal kernel is selected for inference based on the dynamic shape to achieve efficient dynamic scheduling.

7. The optimization method as described in claim 6, characterized in that, In the preprocessing stage, the graph-level optimization pass includes common subexpression elimination, dead code elimination, constant folding, memory planning, and operator fusion; the tensor program to be tuned represents the optimized computational subgraph, which is composed of multiple operator operations fused into a tensor program; And / or, In the post-processing stage, the host-side code is responsible for calling the device-side optimization kernel for the corresponding bucket range based on the actual value of the dynamic shape at runtime, so as to achieve efficient dynamic scheduling.

8. A tuning system that implements the tuning method as described in any one of claims 1-7, characterized in that, The optimization system includes: a preprocessing module, a feature extraction module, a bucketing module, an optimization engine module, a code integration module, a shared library generation module, and a runtime scheduling module; The preprocessing module is responsible for receiving the input model computation graph, converting it into an intermediate representation, and performing global optimization on the intermediate representation; The feature extraction module extracts the feature information of the tensor program to be optimized, which is used as the input information of the tuning engine. The bucketing module divides the dynamic shape tensor program into sub-problems of different sizes, reducing the number of problems to be optimized. The tuning engine module tunes the configuration parameters of the tensor program in each bucket based on program characteristics and hardware specifications, generating an optimized kernel. The code integration module integrates the optimized kernel generated by the tuning engine to generate host-side code and device-side code for dynamic runtime scheduling. The shared library generation module compiles the integrated code into a shared library for runtime loading and calling. The runtime scheduling module loads the shared library during inference runtime, selects the optimal device-side kernel based on the dynamic shape of the input, and performs inference calculations.

9. The automatic tuning method as described in any one of claims 1-7, or the automatic tuning system as described in claim 8, is applied in generating high-performance kernels on hardware platforms, quickly calling dynamic shape models to obtain shared libraries for runtime inference, integrating deep learning frameworks to improve inference speed, and cloud-based remote tuning.

Images