A high-performance operator automatic generation system and method based on multi-agent cooperation

Abstract

The application discloses a kind of high-performance operator automatic generation system and method based on multi-agent cooperation, the system is based on the architecture of multi-agent cooperation, by the collaborative work of analysis chain, generation chain, verification chain and arrangement chain agent, realize the automatic conversion of PyTorch operator to high-performance operator;The system integrates hardware information module, knowledge base module, operation type inference module, performance benchmark test module and iterative optimization module, using progressive iterative optimization mechanism and intelligent error analysis repair mechanism, under the premise of guaranteeing correctness, the generated high-performance operator is continuously optimized in performance.The application can face different hardware platforms and multiple operator types, automatically generate operator implementation with high-performance characteristics, significantly reduce the development threshold and engineering cost of GPU operator, improve the automation level and reusability of deep learning operator optimization, have higher engineering application value.

Description

Technical Field

[0001] This invention relates to the fields of artificial intelligence, high-performance computing, and operator optimization technology, specifically to a high-performance operator automatic generation system and method based on multi-agent collaboration, applicable to operator optimization of deep learning models, GPU programming automation, and high-performance computing acceleration. Background Technology

[0002] With the exponential growth in the scale of deep learning models and the widespread adoption of AI applications, GPU resources have become a key bottleneck limiting model deployment and inference performance. To improve throughput and reduce latency, the industry generally adopts targeted GPU operator optimization techniques, replacing operators in high-level frameworks (such as the PyTorch framework) with high-performance operators written manually in CUDA / Triton and other code. However, existing operator optimization methods mainly rely on manual development by a small number of experts with extensive GPU programming experience. This approach has the following common problems: First, the technical threshold is extremely high. Developers need to understand parallel computing theory, GPU architecture and storage hierarchy, the characteristics of different hardware platforms, and optimization strategies for various operator modes. Developers with such capabilities are extremely scarce, resulting in persistently high costs for operator development and maintenance. Second, manually optimized code is usually deeply customized for specific devices and scales, resulting in poor generalization and portability. Operators optimized for new architectures such as A100 often fail to maintain the same performance on V100 or other GPUs. Third, with the rapid evolution of model structures and operator types, older operators struggle to quickly adapt to new computing modes and data layouts, increasing the complexity of system upgrades and feature expansions.

[0003] To reduce labor costs and improve automation, existing technologies have proposed various operator generation schemes based on compilers and automatic tuning. For example, some methods describe the operator computation process using tensor expressions and scheduling templates, and then the compiler performs scheduling search and automatic code generation within a given search space. Another type of method is based on operator graph analysis and pattern matching, sinking some operators to specific libraries or merging them into new kernels to achieve better cache hit rates and parallel efficiency. However, these schemes generally rely on predefined intermediate representations and scheduling rules, which limit their optimization capabilities for complex structures, irregular memory access, or fusion scenarios involving multiple operators. At the same time, the development and maintenance of compilers and automatic tuning systems themselves require significant expert input, making it difficult to incorporate first-hand GPU engineering experience in a timely manner. Furthermore, in real engineering environments, it often requires a long search time to converge to a better performance configuration, resulting in high computational costs.

[0004] In recent years, Large Language Models (LLMs) have demonstrated strong capabilities in general code generation. Some studies have attempted to use LLMs to directly generate GPU kernel code such as CUDA / Triton to accelerate typical operators like attention and matrix multiplication. However, directly applying LLMs to GPU operator optimization still faces several challenges: First, LLMs have significant deficiencies in their knowledge of GPU architecture, storage hierarchy, and scheduling strategies, making it difficult to systematically design high-quality parallel algorithms that meet hardware characteristics and performance constraints. Second, these methods typically employ single-round or limited-round generation processes, lacking systematic analysis and automatic repair mechanisms for compilation errors, runtime errors, and numerical instability, often requiring manual intervention for debugging. Furthermore, existing solutions generally lack a rigorous performance evaluation loop, making it difficult to promptly obtain the actual performance of the generated operators on the target hardware, and also lacking targeted iterative optimization capabilities based on performance feedback.

[0005] In terms of context understanding, existing technologies mostly optimize individual operators or local computational fragments in isolation, lacking a global grasp of the complete PyTorch model structure and its data flow. This makes it difficult to make reasonable operator partitioning and fusion decisions based on operator combination patterns, tensor shape evolution, and hardware execution characteristics. Furthermore, existing methods generally lack a unified mechanism for managing hardware information and domain knowledge. They cannot systematically integrate information such as the number of streaming multiprocessors (SMs) on different GPU architectures, memory bandwidth, and tensor core support into the generation process. They also cannot structurally precipitate and reuse mature optimization experience and common error patterns of typical operators such as matrix multiplication and convolution during the generation stage.

[0006] Among the currently available technical solutions, there is no automated operator generation method that can form a closed loop in the following aspects: First, it lacks a hierarchical architecture based on multi-agent collaboration, which cannot decouple operator analysis, code generation, verification and performance evaluation, process orchestration, etc., and have different roles collaborate to complete them; Second, it lacks a unified framework that organically combines hardware awareness modules, domain knowledge bases, and type inference mechanisms, making it impossible to establish effective linkage between high-level semantics (i.e., PyTorch code), hardware characteristics, and expert experience; Third, it lacks iterative optimization and error analysis mechanisms oriented towards actual engineering scenarios, which cannot automatically record historical best results, perform rollbacks and local corrections during multiple rounds of generation and verification, and is also difficult to output structured diagnostic information for common error types to guide subsequent repair and optimization.

[0007] In summary, existing technologies for automatic operator generation and optimization still suffer from several shortcomings, including heavy reliance on human experience, insufficient adaptability to hardware and operator scenarios, lack of systematic error handling and performance feedback mechanisms, and limited automation. There is an urgent need for a new technical solution that, within a unified framework, deeply integrates the general code generation capabilities of large language models with GPU domain knowledge, hardware context information, structured knowledge bases, and multi-round verification and optimization mechanisms. This solution should achieve an end-to-end automatic conversion and optimization pipeline from PyTorch operators to CUDA / Triton high-performance operators through multi-agent collaboration, thereby reducing development barriers and maintenance costs while improving the performance, robustness, and scalability of generated operators. Summary of the Invention

[0008] The purpose of this invention is to address the problems in existing technologies, such as heavy reliance on human experience in operator optimization, poor adaptability to different hardware and operator scenarios, lack of systematic error handling and performance feedback mechanisms, and low automation. This invention provides a high-performance operator automatic generation system and method based on multi-agent collaboration. Focusing on PyTorch reference operators, this invention designs a hierarchical multi-agent collaborative architecture and its closed-loop optimization mechanism for operator generation tasks, realizing an end-to-end automated pipeline from operator problem modeling, analysis and decision-making, code generation, correctness and performance verification to iterative optimization.

[0009] The objective of this invention is achieved through the following technical solution: The first aspect of this invention provides a high-performance operator automatic generation system based on multi-agent cooperation, comprising: The analysis chain agent is used to receive PyTorch operator code, infer the operation type through a large language model, and generate an operator architecture design scheme by combining hardware information and knowledge base. Generate chain intelligent agents, which are used to call large language models to generate high-performance operator code that is equivalent to PyTorch operator code based on the operator architecture design scheme; The verification chain agent is used to perform static verification, correctness verification, and performance benchmark testing on high-performance operator code, and output the verification results. The orchestration chain agent is used to coordinate the execution order of each agent and continuously improve the operator performance through an iterative optimization mechanism based on the verification results until the preset iteration termination condition is reached. The hardware information module is used to detect and provide the hardware information of the current device, and format it as a hardware context string; The knowledge base module is used to store and retrieve best practice code and optimization knowledge for operators of different operation types to generate knowledge context strings, while also supporting dynamic expansion and updating of the knowledge base; The operation type inference module is used to analyze PyTorch operator code through a large language model to identify its primary and secondary operation types and provide operation type information for the analysis chain agent. The performance benchmarking module is used to provide performance benchmarking services for the verification chain agent, including dynamically determining the correctness verification threshold based on numerical characteristics, and obtaining stable performance data through warm-up runs and multiple benchmark tests. The iterative optimization module records the results of each iteration and supports intelligent error analysis and adjustment of repair strategies based on historical iteration results.

[0010] Furthermore, the hardware information includes device model, computing power, memory size, number of streaming multiprocessors, thread bundle width, and whether tensor cores are supported; The optimization knowledge includes recommended indexing and partitioning strategies, optimization techniques, common problems, and error avoidance solutions; The performance data includes average execution time, speedup ratio, video memory usage, operator throughput, memory bandwidth utilization, and computational intensity. The iteration results include generated code and its corresponding error messages, correctness results, and performance data.

[0011] Furthermore, the analytical chain agent specifically includes: The large language model performs structured parsing and operation type inference on the input PyTorch operator code; it obtains hardware information from the hardware information module and best practice code and optimization knowledge from the knowledge base module to generate an operator architecture design scheme that includes parallelization strategies, memory access modes and thread block configurations; and it provides repair suggestions by combining historical error information during the iterative optimization process.

[0012] Furthermore, the generating chain agent specifically includes: A structured prompt word template is constructed as input to the large language model. The prompt word template includes a problem description, a hardware context string, a knowledge context string, an operator architecture design scheme, historical best code snippets, error summaries, and repair targets. In the first iteration, initial high-performance operator code is generated based on the operator architecture design scheme. In subsequent iterations, error information or performance data from the verification results fed back by the verification chain agent are used to locally repair or optimize the previously generated high-performance operator code while keeping the overall code architecture unchanged.

[0013] Furthermore, the verification chain agent specifically includes: The high-performance operator code undergoes static syntax and constraint checks to eliminate code that clearly violates constraints. The statically validated high-performance operator code is compiled and executed, along with the original PyTorch operator code. Using the output of the original PyTorch operator code as a benchmark, the correctness of the high-performance operator code is verified through dynamic numerical comparison. The performance benchmarking module is then invoked to perform performance benchmark tests on the correctness-verified high-performance operator code, obtaining performance data. Finally, the verification results, including static verification results, correctness verification results, and performance data, are output.

[0014] Furthermore, the orchestration chain agent specifically includes: The order in which the analysis chain agent, generation chain agent, and verification chain agent are called is determined according to the preset iteration strategy, and the current iteration state is maintained in the form of a state machine. The code of the high-performance operator and its performance indicators of the best historical version are maintained. When the verification result of the current iteration version is a failure or its performance indicators are lower than the preset threshold, the rollback mechanism is triggered to roll back to the best historical version and use the best historical version as the starting point of the next round of iteration. When the number of iterations reaches the upper limit, or the performance indicators meet the preset target or are triggered by external conditions, the iteration is terminated and comprehensive summary information including high-performance operators, performance indicators, correctness verification status, and key optimization points is output.

[0015] Furthermore, it also includes: The problem loading and runtime environment initialization module is used to read PyTorch operator code and its description information from the operator problem library according to a unified interface specification, and to build a standardized problem description structure for information exchange between multiple agents, while initializing the test environment.

[0016] Furthermore, the interface specification includes: The first interface is used to obtain the set of initialization parameters required to initialize the model, the set of initialization parameters including weight dimensions and hyperparameter configuration; The second interface is used to obtain a set of input tensor specifications for testing the model. The set of input tensor specifications includes tensor name, tensor shape, data type, and step size information.

[0017] Furthermore, the standardized problem description structure is obtained by encapsulating the initialization parameter set, the input tensor specification set, and the hardware context string from the hardware information module together; The standardized problem description structure includes at least the following fields: operation name field, operation type field, input tensor specification field, output tensor specification field, initialization parameter field, hardware context field, and target performance index field.

[0018] A second aspect of this invention provides an automatic generation method based on the above-described high-performance operator automatic generation system based on multi-agent cooperation, specifically including: The hardware information of the current device is obtained through the hardware information module, and relevant optimization knowledge is loaded through the knowledge base module. The analysis chain agent receives PyTorch operator code, calls the operation type inference module to determine its operation type, and combines hardware information and optimization knowledge to generate an operator architecture design scheme. Based on the operator architecture design scheme, the generated chain agent generates high-performance operator code that is functionally equivalent to PyTorch operator code using a large language model; The verification chain agent performs static verification, correctness verification, and performance benchmark testing on the high-performance operator code, and outputs the verification results. By orchestrating chain agents to coordinate the execution order of each agent, and judging whether the preset iteration termination condition is met based on the verification result, if it is met, the current high-performance operator code is output as the final result; if it is not met, the iterative optimization process is triggered, and the verification result is fed back to the analysis chain agent or the generation chain agent, and the operator performance is continuously improved through the iterative optimization module.

[0019] Compared with the prior art, the beneficial effects of the present invention are as follows: (1) This invention proposes a hierarchical multi-agent collaborative architecture for GPU high-performance operator generation tasks. Through standardized problem description structure, explicit role division and state-driven iterative mechanism, the operator analysis, code generation, verification and optimization links are decoupled and a closed-loop optimization process is formed.

[0020] (2) This invention introduces a unified injection mechanism of hardware perception and domain knowledge base in multi-agent collaboration, and significantly improves the availability and transferability of the generation operator by constraining the generation behavior of the large language model through hardware context and knowledge context.

[0021] (3) This invention proposes a progressive repair strategy based on error type classification and the principle of minimum modification, which maps compilation errors, runtime errors and performance degradation into structured error information, and achieves controllable automatic repair and performance optimization by limiting the scope of modification and rollback mechanism.

[0022] (4) This invention designs a dynamic threshold correctness verification and performance evaluation method suitable for operator generation scenarios. Based on multiple inputs and multiple runs, it provides comprehensive indicators such as speedup ratio, stability and resource consumption, providing reliable feedback for automated iteration.

[0023] (5) Compared with existing schemes that use multiple agents for general task decomposition or code collaboration, this invention defines a dedicated intermediate representation, error ontology and iteration strategy for the specific task of “from PyTorch to high-performance operators”, and realizes a task-specific multi-agent collaboration mechanism. This significantly reduces the development threshold and maintenance cost of GPU operators while ensuring automation and robustness, improves the performance and scalability of generated operators, and enhances the automation level and reusability of deep learning operator optimization. It has high engineering application value. Attached Figure Description

[0024] Figure 1 This is an architecture diagram of a high-performance operator automatic generation system based on multi-agent cooperation provided in an embodiment of the present invention. Figure 2 This is an example flowchart of a multi-agent cooperative generation of high-performance operators provided in an embodiment of the present invention; Figure 3 A flowchart of multi-agent collaboration provided in an embodiment of the present invention; Figure 4 A flowchart illustrating the hardware information module and knowledge base module provided in an embodiment of the present invention; Figure 5 A flowchart of an iterative optimization module provided in an embodiment of the present invention. Detailed Implementation

[0025] Exemplary embodiments will now be described in detail, examples of which are illustrated in the accompanying drawings. When the following description relates to the drawings, unless otherwise indicated, the same numerals in different drawings denote the same or similar elements. The embodiments described in the following exemplary embodiments do not represent all embodiments consistent with the present invention. Rather, they are merely examples of apparatuses and methods consistent with some aspects of the invention as detailed in the appended claims.

[0026] The terminology used herein is for the purpose of describing particular embodiments only and is not intended to be limiting of the invention. The singular forms “a,” “the,” and “the” used in this invention and the appended claims are also intended to include the plural forms unless the context clearly indicates otherwise. It should also be understood that the term “and / or” as used herein refers to and includes any or all possible combinations of one or more of the associated listed items.

[0027] It should be understood that although the terms first, second, third, etc., may be used in this invention to describe various information, this information should not be limited to these terms. These terms are only used to distinguish information of the same type from one another. For example, first information may also be referred to as second information without departing from the scope of this invention, and similarly, second information may also be referred to as first information. Depending on the context, the word "if" as used herein may be interpreted as "when," "when," or "in response to a determination."

[0028] The present invention will now be described in detail with reference to the accompanying drawings. Unless otherwise specified, the features of the following embodiments and implementations can be combined with each other.

[0029] See Figure 1 and Figure 2The high-performance operator automatic generation system based on multi-agent collaboration of the present invention can be deployed on a local GPU server or other computing platform with GPU resources. The core of this high-performance operator automatic generation system includes an AnalysisChain agent, a GenerationChain agent, a ValidationChain agent, an OrchestrationChain agent, a hardware information module, a knowledge base module, an operation type inference module, a performance benchmarking module, and an iterative optimization module. Specifically, the AnalysisChain agent receives PyTorch operator code, infers the operation type through a large language model, and generates an operator architecture design scheme by combining hardware information and the knowledge base. The GenerationChain agent, based on the operator architecture design scheme, calls the large language model to generate high-performance operator code functionally equivalent to the PyTorch operator code. The ValidationChain agent performs static verification, correctness verification, and performance benchmarking on the high-performance operator code and outputs the verification results. The orchestration chain agent coordinates the execution order of various agents and continuously improves operator performance through an iterative optimization mechanism based on verification results until a preset iteration termination condition is reached. The hardware information module detects and provides the hardware information of the current device, formatting it into a hardware context string. This hardware information includes device model, computing power, memory size, number of SMs, warp width, and whether Tensor Cores are supported. The knowledge base module stores and retrieves best practice code and optimization knowledge for operators of different operation types, generating knowledge context strings. It also supports dynamic expansion and updates of the knowledge base, facilitating the integration of new operation types and optimization methods. Optimization knowledge includes recommended indexing and tiling strategies, optimization techniques, common problems, and error avoidance solutions. The operation type inference module analyzes PyTorch operator code using a large language model to identify its primary and secondary operation types, achieving coarse classification to guide subsequent thought processes and optimization, and providing operation type information to the analysis chain agent. The performance benchmarking module provides performance benchmarking services for the verification chain agent, including dynamically determining the correctness verification threshold based on numerical characteristics, and obtaining stable performance data through warm-up runs and multiple benchmark tests. Performance data includes metrics such as average execution time, speedup ratio, memory usage, operator throughput, memory bandwidth utilization, and computational intensity. The iterative optimization module records the iteration results of each round and supports intelligent error analysis and correction strategy adjustments based on historical iteration results, thereby maintaining historically optimal operator code. Iteration results include generated code and its corresponding error information, correctness results, and performance data.The system described in this invention, through the coordinated work of the above modules, can achieve automated generation and progressive optimization of high-performance operators such as PyTorch operators and CUDA / Triton, significantly reducing manual development costs.

[0030] It should be noted that this invention constructs a hierarchical multi-agent collaborative architecture, comprising at least an analysis chain agent, a production chain agent, a verification chain agent, and an orchestration chain agent. Each agent has a clearly defined role and collaborates with others to complete the entire process from operator analysis to code generation, verification, and orchestration optimization. The aforementioned analysis chain agent, production chain agent, verification chain agent, and orchestration chain agent all refer to functional execution units built based on a pre-trained large language model. They perform reasoning processing on contextual information through pre-defined role constraints and task objectives, and can call external tools or service interfaces to complete specific tasks. These agents differ from dialogue models used only for natural language dialogue.

[0031] In addition, this high-performance operator automatic generation system includes several supporting modules, including at least a hardware information module, a knowledge base module, an operation type inference module, a performance benchmarking module, and an iterative optimization module. These modules are tightly coupled with the multi-agent collaborative architecture to jointly realize hardware perception, knowledge injection, and iterative optimization.

[0032] In other embodiments, the high-performance operator automatic generation system may also include a problem loading and runtime environment initialization module, which reads PyTorch operator code and its description information from an operator problem library (such as the KernelBench dataset) according to a unified interface specification, and constructs a standardized problem description structure (ProblemSpec) for information exchange between multiple agents, so as to facilitate the exchange of information between multiple agents in a consistent format, while initializing the test environment.

[0033] Furthermore, in the operator problem library or user-provided reference implementation, the Model class is required to conform to the following interface specification: ① The `get_init_inputs()` function: As the first interface, it is used to obtain the set of initialization parameters required to initialize the model. The initialization parameter set includes weight dimensions, hyperparameter configurations (such as hidden dimensions, convolutional kernel size, etc.), etc. For optional initialization parameters, if the caller does not explicitly provide them, it is assumed that these initialization parameters can be omitted, and the system default value is used. The source of this default value is explicitly recorded as "system default" in the constructed standardized problem description structure.

[0034] ②The get_inputs() function: As the second interface, it is used to obtain the set of input tensor specifications for testing the model. The set of input tensor specifications includes tensor name, tensor shape, data type and stride information.

[0035] Furthermore, when the original PyTorch operator code does not provide the get_inputs() function (i.e., the second interface), the system automatically constructs the input tensor specification set according to the following rules: infers the number of input tensors and basic shape parameters based on the parameter list of the forward function; uses the default test size for dimensions not explicitly given, such as the batch size dimension being 32 by default, and the feature dimension being inferred based on the size of the weight matrix; automatically constructs the input tensor specification set, and marks the source as "auto_infer" in the standardized problem description structure.

[0036] Furthermore, the system abstracts the operator problem into a unified standardized problem description structure. This standardized problem description structure is obtained by encapsulating the initialization parameter set, the input tensor specification set, and the hardware context string from the hardware information module. This standardized problem description structure includes at least the following fields: ① Operator_name field: used to identify the name or number of the operator; ② Operation_type field: used to identify the operation type label output by the operation type inference module; ③ Input_specs field: the set of input tensor specifications used to describe the shape, data type, and stride information of the input tensors; ④ Output_specs field: used to describe the expected shape and data type of the output tensors; ⑤ Initialization_params field: the set of model initialization parameters, including parameter names, values, and source labels (user-specified or system default); ⑥ Hardware_context field: used to record the hardware context string provided by the hardware information module; ⑦ Target_metrics field: such as expected speedup, maximum acceptable memory usage, etc. The above structure is stored internally in JSON or equivalent key-value pair format and serves as the unified input data carrier for the analysis chain agent, the generation chain agent, and the verification chain agent.

[0037] It should be understood that, through the above implementation method, the problem loading and runtime environment initialization module can provide the system with complete input information and runtime context, ensuring that the subsequent analysis, generation and verification steps have a consistent basic environment.

[0038] In this embodiment, see Figure 1 and Figure 2The analysis chain agent receives PyTorch operator code and, based on a large language model, performs structured parsing of the input PyTorch operator code and runtime environment. It infers operation types, extracts operator computation logic and parameter dependencies, and, combined with hardware information and a knowledge base, generates an operator architecture design scheme to guide subsequent code generation. Specifically, the analysis chain agent performs structured parsing and operation type inference of the input PyTorch operator code using a large language model; acquires hardware information from the hardware information module and best practice code and optimization knowledge from the knowledge base module to generate an operator architecture design scheme that includes parallelization strategies, memory access modes, and thread block configurations; and, during iterative optimization, provides repair suggestions based on historical error information.

[0039] Furthermore, the input PyTorch operator code is subjected to structured parsing and operation type inference through a large language model, specifically including: (a1) Based on the large language model, the input PyTorch operator code and runtime environment are structured and parsed to obtain the structural information of the PyTorch operator code: the model initial definition (Model.init) in the PyTorch operator code is subjected to syntactic and semantic analysis to identify the layer type, parameter name and inter-layer dependency relationship defined in the model; the model forward function (Model.forward) in the PyTorch operator code is parsed line by line to construct the operator operation sequence (i.e. operator computation logic).

[0040] Preferably, the analysis chain agent can verify the correctness of the static analysis results by running several forward inferences and combining dynamic shape information.

[0041] (a2) Inferring the operation type of the input PyTorch operator code through a large language model: Call the operation type inference module, send the PyTorch operator code text, operation name and structural information to the large language model for one or more operation type inferences; Based on the operation type returned by the large language model, determine the operation type of the current operator, identify its primary and secondary operation types, such as matrix multiplication (matmul), convolution (conv), two-dimensional convolution (conv2d), element-wise operation, reduction operation, or combinations thereof; The large language model provides a confidence score after analysis based on the operation type of the operator, and the operation type of the operator can be quantified based on the confidence score; When the confidence score is lower than the preset confidence score threshold, the backup rule engine is activated, whereby the backup rule engine classifies the operation type by analyzing the actual computation logic of the operator, that is, the backup inference path based on rules and keywords, and gives the final operation type judgment by combining the results of both.

[0042] Furthermore, by acquiring hardware information from the hardware information module and best practice code and optimization knowledge from the knowledge base module, an operator architecture design scheme is generated, including parallelization strategies, memory access modes, and thread block configurations. Specifically, this includes: (b1) Hardware and Knowledge Base Context Injection: Obtain the hardware information of the current GPU device from the hardware information module, i.e., the hardware context string, including the recommended block size range, mesh configuration suggestions, number of SMs, thread beam width, memory access precautions, etc.; retrieve best practice code and related optimization knowledge from the knowledge base service module according to the operation type of the operator, obtain best practice code, common optimization techniques and common mistakes, error avoidance solutions, etc., and structure them into a knowledge context string; combine the above hardware context string, knowledge context string and the structured information obtained from the structured parsing to form the analysis prompt input, and send it to the large language model.

[0043] (b2) Operator architecture design scheme: Based on the analysis prompts, the large language model performs a comprehensive analysis of the operator's computation mode, memory access mode and parallelization strategy, designs the GPU grid configuration, block size suggestions and memory access mode of the high-performance operator, and gives the overall architecture design of the Triton operator (Triton kernel), that is, the operator architecture design scheme.

[0044] (b3) Structured Output: The analysis chain agent organizes the output of the large language model into a structured form, ultimately outputting a strictly defined analysis result structure. This structure includes optional information such as the PyTorch model analysis field (pytorch_model_analysis), the operator design scheme field (kernel_design), and the block size recommendations field (block_size_recommendations), which drive subsequent generation and validation chain agents. Specifically, the pytorch_model_analysis field includes a list of model layers, a description of the forward steps, key parameter names, and operation sequence descriptions; the kernel_design field includes a list of kernel parameters, tiling strategies, and batch dimension processing methods; and block_size_recommendations includes recommended block sizes (BLOCK_* parameters) and their design rationale. Other optional information includes suggestions for memory layer usage and thread organization methods.

[0045] (b4) Persistently store the above structured analysis results and pass them to the generating chain agent as input.

[0046] It should be understood that, through the above implementation methods, the analysis chain agent can automatically complete the abstraction and decomposition of complex PyTorch operators by human-imposed constraints, providing an accurate design blueprint for subsequent Triton operator code generation.

[0047] In this embodiment, see Figure 1 and Figure 2 The generative chain agent, based on the analysis chain agent output and hardware and knowledge base context, is responsible for automatically generating Triton high-performance operator implementations equivalent to PyTorch operators. Specifically, based on the analysis results and historical version information, the generative chain agent uses structured prompt templates to call the large language model to generate high-performance operator code. The generative chain agent specifically includes: constructing structured prompt templates as input to the large language model, where the prompt templates include fields such as problem description, hardware context string, knowledge context string, operator architecture design scheme, historical best code snippets, error summaries, and repair targets; generating initial high-performance operator code based on the operator architecture design scheme in the first iteration; and in subsequent iterations, combining error information or performance data from the verification chain agent's feedback verification results to perform local repairs or optimizations on the previously generated high-performance operator code, while maintaining the overall code architecture unchanged, and verifying the syntactic correctness and structural integrity of the generated high-performance operator code.

[0048] Furthermore, constructing structured prompt word templates specifically includes: (c1) Collect and analyze the operator architecture design schemes output by the chain agent, including kernel parameters, indexing methods, block dimensions, etc.

[0049] (c2) Collect information such as block size candidates, number of SMs and memory bandwidth provided by the hardware information module, as well as best practices at the operation level provided by the knowledge base module.

[0050] (c3) Based on the error information or performance feedback of the current iteration, construct a repair guide or optimization guide field.

[0051] (c4) Organize the above content into a unified structured prompt word template and use it as input for the large language model.

[0052] Furthermore, during the first iteration, initial high-performance operator code is generated based on the operator architecture design. Specifically, when generating the initial high-performance operator code, the large language model is invoked, requiring its output to follow a preset JSON structure and code block format, including at least: ① a complete Triton kernel function (e.g., named kernel); ② a wrapper function compatible with the PyTorch interface (e.g., named kernel_wrapper); ③ recommended grid configuration and block size (BLOCK_SIZE) setting methods. Subsequently, the generating chain agent performs static analysis on the returned code to confirm the correctness of basic elements such as import statements, function signatures, and the number of parameters.

[0053] Furthermore, in subsequent iterations, a gradual repair strategy based on the "minimum modification principle" is adopted: only local modifications or optimizations are allowed within specified code scopes (e.g., index calculation, block size (BLOCK_SIZE) macros, mask constructors), while maintaining the function interfaces and overall architecture unchanged. Specifically, when the verification results returned by the verification chain agent contain syntax errors, shape mismatches, or mask usage errors, the generation chain agent organizes the error information, i.e., the error type and key information, into a structured error description. Subsequently, in subsequent iterations, the generation chain agent inputs the current version code, the structured error description, and the overall operator architecture constraints provided by the analysis chain agent into the large language model, requesting local repairs or optimizations while maintaining the overall design architecture unchanged.

[0054] Furthermore, the generation chain agent also performs multi-version management: the generation chain agent assigns a unique version identifier to the code generated or repaired in each round, records its source analysis results, hardware context snapshots and corresponding verification results; outputs a structured result containing complete kernel functions, wrapper functions, recommended grid configurations and version identifiers, and submits the structured result corresponding to each version of the code to the verification chain agent for unified testing, so that the orchestration chain agent can select the best implementation.

[0055] In this embodiment, see Figure 2 and Figure 5The verification chain agent is responsible for performing static verification, correctness verification, and performance benchmarking on the generated Triton high-performance operators. Specifically, the verification chain agent includes: performing static syntax and constraint checks on the high-performance operator code, eliminating high-performance operator code that obviously violates constraints; compiling and executing the statically verified high-performance operator code and the original PyTorch operator code, using the output of the original PyTorch operator code as a benchmark, and verifying the correctness of the high-performance operator code through dynamic numerical comparison; calling the performance benchmarking module to perform performance benchmarking on the high-performance operator code that has passed correctness verification, obtaining performance data; and outputting verification results containing static verification results, correctness verification results, and performance data.

[0056] Furthermore, static verification specifically includes checking whether the code satisfies the Triton syntax and runtime constraints without actually executing the Triton kernel. Specifically, it can detect code that "obviously violates the constraints" through static analysis rules, such as program ID (program_id) dimensions exceeding the grid definition, pointer access out of bounds risk, mask shape inconsistent with the access range, unused pointers or illegal type conversions, etc. For code that obviously violates the constraints, it is directly marked as static verification failure, and the corresponding error type and location are returned. The detected violations are classified into a predefined set of error types.

[0057] In determining the error type, errors can be categorized into types such as mesh configuration error, mask usage error, shape mismatch, operator syntax error, and performance problem by matching typical error messages, and key information can be extracted to form a structured description.

[0058] It should be understood that these structured error type descriptions can be placed as error summaries in the prompt word template. The generating chain agent can then use this structured error information, combined with hardware context and knowledge base content, to make targeted corrections to key parts such as grid dimensions, index calculation, mask usage errors, and memory access methods, thereby improving the success rate of automatic repair and reducing the risk of introducing new problems.

[0059] Furthermore, the correctness verification specifically includes: based on the input tensor specifications in the standardized problem description structure, executing multiple sets of test inputs on the Triton / CUDA high-performance operator code and the original PyTorch operator code, and comparing errors using a dynamic threshold strategy based on numerical characteristics. For example, during the matrix multiplication operator generation process, the system generates matrices of different sizes as inputs for both the original PyTorch operator and the generated Triton / CUDA high-performance operator, and compares the calculated matrix results. Only when the error of each set of results is within the threshold is the high-performance operator considered qualified. Specifically, at least one set of test input data is constructed and executed on both the original PyTorch operator code and the generated high-performance operator code; the relative error tolerance and absolute error tolerance are dynamically determined based on the output numerical range of the original PyTorch operator code; it is then determined whether the two sets of execution results are consistent within the error threshold, and if so, the correctness verification is considered passed.

[0060] Specifically, the input set is first constructed using the test inputs provided by `get_inputs()` in `ProblemSpec`. This set must contain at least three inputs with different shapes or numerical distributions. For numerical distributions, a normal distribution with a mean of 0 and a standard deviation of 1 is used for random initialization by default. If a specific distribution is specified in `ProblemSpec`, then sampling is performed according to that distribution. Then, the error threshold is calculated: for each input set, the PyTorch reference implementation and the Triton / CUDA high-performance operator implementation are executed respectively to obtain the output. and ,in The reference values ​​obtained from executing the PyTorch reference implementation. The operator computation values ​​obtained from executing the Triton / CUDA high-performance operator implementation; based on The numerical scale is adaptively set to rtol / atol for calculation. Where rtol is the relative error tolerance and atol is the absolute error tolerance. This is the error threshold. If... Therefore, relative error is used as the primary factor, and rtol = 1e is set. -3 atol = 1e -4 ;like Therefore, absolute error is used as the primary factor, and rtol = 1e is set. -4 atol = 1e -6 For reduction operators or scenarios with large cumulative errors, a more lenient or stricter default threshold can be read from the knowledge base based on the operation type (op_type). For example, for large-scale matrix multiplication, rtol = 5e can be set. -3 atol = 5e -4Then, multiple comparisons and pass / fail conditions are set: for each input group, `torch.allclose(` is executed. , The `rtol` and `atol` functions determine whether the input passes. The version is marked as "correctness verification passed" only if at least three sets of inputs pass. If any set fails, the failed sample and its error statistics (maximum relative error, maximum absolute error) are recorded, and the error type is marked as numeric mismatch (NUMERICAL_MISMATCH) and returned.

[0061] Furthermore, assuming correctness is achieved, performance benchmarking is conducted by running the code multiple times repeatedly, and performance metrics such as average execution time, speedup (PyTorch operator execution time / operator generation time), and GPU memory usage are calculated. Specifically, for high-performance operators that pass correctness, the verification chain agent further measures their average execution time over multiple repeated runs and compares it with the execution time of the PyTorch operator code to calculate performance metrics such as speedup. In addition, the verification chain agent can also record performance metrics such as GPU memory usage and operator throughput, providing a basis for subsequent performance evaluation and optimization.

[0062] Furthermore, the verification chain agent outputs the final verification results. Specifically, the verification chain agent feeds back the static verification results, error types, correctness verification results, and performance data of each version of the high-performance operator code to the orchestration chain agent and the generation chain agent in a unified structure. When outputting verification results, for failures, the verification chain agent focuses on returning the error type and location information; for successes but with insufficient performance, it returns optimization suggestions.

[0063] In this embodiment, see Figure 2 and Figure 3The orchestration chain agent, acting as the scheduling and control hub for multi-agent collaboration, is responsible for managing the iteration process (managing the calling order among multiple agents), maintaining the iteration strategy, and preserving the global optimal solution. Specifically, the orchestration chain agent includes: determining the calling order of the analysis chain agent, generation chain agent, and verification chain agent according to a preset iteration strategy; maintaining the current iteration state (e.g., in analysis, generation, verification, repair, termination) through a state machine; maintaining the historical best version of high-performance operator code and its performance metrics; triggering a rollback mechanism when the verification result of the current iteration version fails or its performance metrics fall below a preset threshold, rolling back to the historical best version, and using the historical best version as the starting point for the next iteration; terminating the iteration when the number of iterations reaches the upper limit, the performance metrics meet the preset target, or external conditions are triggered, and outputting a comprehensive summary information including high-performance operators, performance metrics, correctness verification status, and key optimization points. The external conditions refer to the preset iteration time; for example, a single operator generation practice could be set to 600 seconds, allowing the system to exit after a timeout and avoid getting stuck in an infinite loop and consuming computational resources.

[0064] Furthermore, the iterative process control specifically includes the following steps: After system startup, the orchestration chain agent first calls the analysis chain agent to perform operator analysis and obtain the initial architecture design. Then, it calls the generation chain agent to generate the initial version of the Triton implementation and calls the verification chain agent for verification. Finally, based on a preset iteration strategy, it decides whether to continue iterating, trigger a repair mode, or terminate the process according to the verification results. The iteration strategy refers to the pre-set number of iterations, such as 20 iterations (if the user can set it to 20 iterations) or exiting the iteration when a preset target performance is achieved.

[0065] Furthermore, the orchestration chain agent is used to maintain a record of historical best implementations internally (i.e., the high-performance operator code and performance metrics of the historical best version), including the corresponding Triton code, correctness state and performance metrics. The best record is updated when the performance improvement of the latest candidate implementation exceeds a preset gain threshold and the correctness is passed. That is, when the implementation of the current iteration version passes the correctness test and the performance is better than the historical best version, the best implementation record is updated.

[0066] Furthermore, the orchestration chain agent can trigger a rollback mechanism when the verification result of the current iteration fails or its performance metrics fall below a preset threshold, intelligently selecting subordinate agents to determine the next assigned operation. For example... Figure 4As shown, when the current iteration version has compilation errors, runtime errors (i.e., its verification result is a failure) or its performance is significantly lower than a certain percentage threshold of the historical best value (i.e. performance degradation), the rollback mechanism is triggered to roll back to the historical best version. The orchestration chain agent takes the previous round's best implementation as the starting point for the next round of generation and instructs the generation chain agent to perform local repairs or optimizations based on this version.

[0067] Preferably, by setting a performance degradation threshold, invalid iterations in the wrong direction can be effectively avoided.

[0068] In this embodiment, see Figure 2 and Figure 4 The hardware information module is used to perceive the hardware information of the current GPU hardware device and, based on this, provide hardware optimization suggestions for the analysis chain agent and the generation chain agent, guiding code generation and tuning. The hardware information module is implemented through the following steps: (d1) Hardware information collection: During the system initialization phase, the hardware information of the current GPU device is collected through the underlying runtime interface or deep learning framework API, including but not limited to: device name, computing power, video memory capacity, number of SMs, etc.

[0069] (d2) Hardware Context String Construction: In the example implementation, the hardware information module can return a set of candidate block sizes and memory optimization strategies based on different operator architectures (such as A100), which are used to constrain and guide the large language model to generate high-performance operator code that better suits the hardware characteristics. Based on the collected hardware information, it generates a hardware context string describing the current device capabilities and limitations. This hardware context string is then injected into each agent during the prompt construction phase. Examples include recommended block size ranges and whether tensor cores are suitable. This hardware context can be directly injected into the prompt templates of the analysis and generation chains to guide the large language model in selecting more appropriate parallelization and memory access strategies.

[0070] In this embodiment, see Figure 2 and Figure 4 The knowledge base module provides operator-level optimization experience and common problem references for large language models. Specifically, the knowledge base module includes the following: (e1) Knowledge Base Organization: Knowledge items are organized using operator type as the primary key. For example, separate entries are maintained for matrix multiplication (matmul), 2D convolution (conv2d), and layer normalization (layernorm). Each entry includes at least best practice code, recommended tiling strategies, optimization techniques, common problems, and error avoidance solutions. The knowledge base module is used to introduce domain knowledge and store best practice code, optimization techniques, and common problems for different operation types in a structured form. For example, suggestions on matrix multiplication operations, recordable block algorithms, shared memory usage, block size (BLOCK_SIZE) selection, use of tl.dot (i.e., the Triton dot product function API), and typical performance pitfalls.

[0071] (e2) Knowledge retrieval: The analysis chain agent initiates a query request to the knowledge base module according to the operation type, retrieves the corresponding relevant entries, namely best practice code and related optimization knowledge, forms a knowledge context and injects it into the prompt, guides the large language model to follow mature optimization strategies and avoids repeating the same pitfalls.

[0072] (e3) Knowledge Injection: The knowledge base retrieval results are returned in the form of structured text and injected into the prompts of the analysis chain agent and the generation chain agent as background information when the large language model generates analysis results and code.

[0073] In this embodiment, the operation type inference module is used to perform multi-level type identification on the input PyTorch operator code and operation type name. This module first calls a large language model for intelligent inference, providing the operation type and its confidence level. When the judgment result is unclear or the confidence level is lower than a preset threshold (usually 0.95), it can revert to a backup inference mechanism based on keyword matching and pattern recognition. The inference result is used both to select knowledge base entries and to constrain subsequent operator architecture design and verification strategies, thereby achieving unified processing of operators with different complexities.

[0074] In this embodiment, the performance benchmarking module and the iterative optimization module are mainly used to evaluate and iteratively optimize the performance of different versions of the Triton high-performance operator code and maintain the final output results. The performance benchmarking module is used to verify the correctness of each version of the high-performance operator code; for each version of the high-performance operator code that passes the correctness verification, its average execution time, speedup ratio, and memory usage are recorded. The iterative optimization module is used to gradually improve the high-performance operator code during multiple rounds of generation and verification. In each iteration, the correctness result and performance indicators of the currently generated code are recorded, and the historical best results are maintained; when generation fails or performance degrades significantly, the orchestration chain agent can automatically revert to the best version of the previous round and use it as the starting point for the new round of generation, avoiding invalid iterations and performance regression; this process embodies the "minimum modification principle," gradually approaching a high-quality implementation based on local corrections.

[0075] It should be noted that in each iteration, the best version is selected from all candidate implementations based on a preset optimization objective (such as maximum speedup or minimum resource consumption within a certain threshold). The code path and performance metrics of the best version are then used as the final output of the orchestration chain agent. Furthermore, historical results need to be managed, and intermediate results and key performance data from each iteration need to be saved for subsequent analysis and reproduction.

[0076] As can be seen from the above embodiments, the high-performance operator automatic generation and optimization method based on multi-agent cooperation proposed in this invention can realize the automatic conversion and incremental optimization of PyTorch operators to Triton high-performance operators while ensuring correctness, and has good engineering practical value.

[0077] It is worth mentioning that this invention also provides a method for automatically generating high-performance operators, implemented based on the multi-agent cooperative high-performance operator automatic generation system described in the above embodiments. This method specifically includes the following steps: Step S1: Obtain the hardware information of the current device through the hardware information module, and load relevant optimization knowledge through the knowledge base module.

[0078] Step S2: The analysis chain agent receives the PyTorch operator code, calls the operation type inference module to determine its operation type, and combines hardware information and optimization knowledge to generate an operator architecture design scheme.

[0079] Step S3: Generate chain intelligent agent based on operator architecture design scheme and large language model to generate high-performance operator code that is functionally equivalent to PyTorch operator code.

[0080] Step S4: The verification chain agent performs static verification, correctness verification, and performance benchmark testing on the high-performance operator code, and outputs the verification results.

[0081] Step S5: Coordinate the execution order of each agent by orchestrating the chain agent, and determine whether the preset iteration termination condition is met based on the verification result. If it is met, output the current high-performance operator code as the final result; if it is not met, trigger the iterative optimization process, feed back the verification result to the analysis chain agent or the generation chain agent, and return to step S2 or step S3 for the next round of iteration. The operator performance is continuously improved through the iterative optimization module.

[0082] The above embodiments are only used to illustrate the technical solutions of the present invention, and are not intended to limit it. Although the present invention has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that modifications can still be made to the technical solutions described in the foregoing embodiments, or equivalent substitutions can be made to some of the technical features. Such modifications or substitutions do not cause the essence of the corresponding technical solutions to deviate from the spirit and scope of the technical solutions of the embodiments of the present invention.

Claims

1. A high-performance operator automatic generation system based on multi-agent cooperation, characterized in that, include: The analysis chain agent is used to receive PyTorch operator code, infer the operation type through a large language model, and generate an operator architecture design scheme by combining hardware information and knowledge base. Generate chain intelligent agents, which are used to call large language models to generate high-performance operator code that is equivalent to PyTorch operator code based on the operator architecture design scheme; The verification chain agent is used to perform static verification, correctness verification, and performance benchmark testing on high-performance operator code, and output the verification results. The orchestration chain agent is used to coordinate the execution order of each agent and continuously improve the operator performance through an iterative optimization mechanism based on the verification results until the preset iteration termination condition is reached. The hardware information module is used to detect and provide the hardware information of the current device, and format it as a hardware context string; The knowledge base module is used to store and retrieve best practice code and optimization knowledge for operators of different operation types to generate knowledge context strings, while also supporting dynamic expansion and updating of the knowledge base; The operation type inference module is used to analyze PyTorch operator code through a large language model to identify its primary and secondary operation types and provide operation type information for the analysis chain agent. The performance benchmarking module is used to provide performance benchmarking services for the verification chain agent, including dynamically determining the correctness verification threshold based on numerical characteristics, and obtaining stable performance data through warm-up runs and multiple benchmark tests. The iterative optimization module records the results of each iteration and supports intelligent error analysis and adjustment of repair strategies based on historical iteration results.

2. The high-performance operator automatic generation system based on multi-agent cooperation according to claim 1, characterized in that, The hardware information includes device model, computing power, memory size, number of streaming multiprocessors, thread bundle width, and whether tensor cores are supported. The optimization knowledge includes recommended indexing and partitioning strategies, optimization techniques, common problems, and error avoidance solutions; The performance data includes average execution time, speedup ratio, video memory usage, operator throughput, memory bandwidth utilization, and computational intensity. The iteration results include generated code and its corresponding error messages, correctness results, and performance data.

3. The high-performance operator automatic generation system based on multi-agent cooperation according to claim 1, characterized in that, The analysis chain agent specifically includes: The large language model performs structured parsing and operation type inference on the input PyTorch operator code; it obtains hardware information from the hardware information module and best practice code and optimization knowledge from the knowledge base module to generate an operator architecture design scheme that includes parallelization strategies, memory access modes and thread block configurations; and it provides repair suggestions by combining historical error information during the iterative optimization process.

4. The high-performance operator automatic generation system based on multi-agent cooperation according to claim 1, characterized in that, The generated chain agent specifically includes: A structured prompt word template is constructed as input to the large language model. The prompt word template includes a problem description, a hardware context string, a knowledge context string, an operator architecture design scheme, historical best code snippets, error summaries, and repair targets. In the first iteration, initial high-performance operator code is generated based on the operator architecture design scheme. In subsequent iterations, error information or performance data from the verification results fed back by the verification chain agent are used to locally repair or optimize the previously generated high-performance operator code while keeping the overall code architecture unchanged.

5. The high-performance operator automatic generation system based on multi-agent cooperation according to claim 1, characterized in that, The verification chain agent specifically includes: The high-performance operator code undergoes static syntax and constraint checks to eliminate code that clearly violates constraints. The statically validated high-performance operator code is compiled and executed, along with the original PyTorch operator code. Using the output of the original PyTorch operator code as a benchmark, the correctness of the high-performance operator code is verified through dynamic numerical comparison. The performance benchmarking module is then invoked to perform performance benchmark tests on the correctness-verified high-performance operator code, obtaining performance data. Finally, the verification results, including static verification results, correctness verification results, and performance data, are output.

6. The high-performance operator automatic generation system based on multi-agent cooperation according to claim 1, characterized in that, The orchestration chain agent specifically includes: The order in which the analysis chain agent, generation chain agent, and verification chain agent are called is determined according to the preset iteration strategy, and the current iteration state is maintained in the form of a state machine. The code of the high-performance operator and its performance indicators of the best historical version are maintained. When the verification result of the current iteration version is a failure or its performance indicators are lower than the preset threshold, the rollback mechanism is triggered to roll back to the best historical version and use the best historical version as the starting point of the next round of iteration. When the number of iterations reaches the upper limit, or the performance indicators meet the preset target or are triggered by external conditions, the iteration is terminated and comprehensive summary information including high-performance operators, performance indicators, correctness verification status, and key optimization points is output.

7. The high-performance operator automatic generation system based on multi-agent cooperation according to claim 1, characterized in that, Also includes: The problem loading and runtime environment initialization module is used to read PyTorch operator code and its description information from the operator problem library according to a unified interface specification, and to build a standardized problem description structure for information exchange between multiple agents, while initializing the test environment.

8. The high-performance operator automatic generation system based on multi-agent cooperation according to claim 7, characterized in that, The interface specification includes: The first interface is used to obtain the set of initialization parameters required to initialize the model, the set of initialization parameters including weight dimensions and hyperparameter configuration; The second interface is used to obtain a set of input tensor specifications for testing the model. The set of input tensor specifications includes tensor name, tensor shape, data type, and step size information.

9. The high-performance operator automatic generation system based on multi-agent cooperation according to claim 8, characterized in that, The standardized problem description structure is obtained by encapsulating the initialization parameter set, the input tensor specification set, and the hardware context string from the hardware information module together; The standardized problem description structure includes at least the following fields: operation name field, operation type field, input tensor specification field, output tensor specification field, initialization parameter field, hardware context field, and target performance index field.

10. An automatic generation method for a high-performance operator automatic generation system based on multi-agent cooperation as described in any one of claims 1-9, characterized in that, Specifically, it includes: The hardware information of the current device is obtained through the hardware information module, and relevant optimization knowledge is loaded through the knowledge base module. The analysis chain agent receives PyTorch operator code, calls the operation type inference module to determine its operation type, and combines hardware information and optimization knowledge to generate an operator architecture design scheme. Based on the operator architecture design scheme, the generated chain agent generates high-performance operator code that is functionally equivalent to PyTorch operator code using a large language model; The verification chain agent performs static verification, correctness verification, and performance benchmark testing on the high-performance operator code, and outputs the verification results. The execution order of each agent is coordinated by orchestrating the chain agent, and the preset iteration termination condition is determined based on the verification result. If the condition is met, the current high-performance operator code is output as the final result. If the requirements are not met, an iterative optimization process is triggered, and the verification results are fed back to the analysis chain agent or the generation chain agent to continuously improve the operator performance through the iterative optimization module.

Images