Method and apparatus for inference of diffusion transformer networks

Abstract

A method for inference of a diffusion transformer network is disclosed. The method comprises generating a sequence of tokens based on input media data; processing a token of the sequence of tokens by a multi-head self-attention (MHA) residual block of the diffusion transformer network; and processing an output of the MHA residual block by a feed-forward network (FFN) residual block of the diffusion transformer network, wherein the method further comprises performing operations in mix-precision and / or performing two or more operations in a fusion manner, by at least one of the MHA residual block and the FFN residual block.

Description

METHOD AND APPARATUS FOR INFERENCE OF DIFFUSION TRANSFORMER NETWORKSFIELD

[0001] The present disclosure relates generally to artificial intelligence technology, and more particularly, to inference of diffusion transformer networks.BACKGROUND

[0002] In recent years, diffusion models (DMs) have gained increasing popularity in numerous generative tasks due to their impressive performance. However, the computational cost and latency associated with the multistep sampling of DMs pose significant challenges on the practical applications, especially as model size grows. Network quantization has emerged as a promising approach to accelerate the sampling process of DMs. Some studies have managed to quantize DMs into low-bit representations while preserving the quality of generation.

[0003] However, the potential of low-precision networks to expedite DMs sampling has not been fully exploited. Current researches have concentrated exclusively on the computation-intensive operations, like matrix multiplications, which compromise most Multiply-Accumulate Operations (MACs) within the network inference. They typically overlook the optimization of memory-bounded operations such as normalization and activation functions. Given that global memory access (IO) on modern GPUs are expensive, neglecting these aspects could lead to sub-optimal performance.SUMMARY

[0004] The following presents a simplified summary of one or more aspects according to the present disclosure in order to provide a basic understanding of such aspects. This summary is not an extensive overview of all contemplated aspects, and is intended to neither identify key or critical elements of all aspects nor delineate the scope of any or all aspects. Its sole purpose is to present some concepts of one or more aspects in a simplified form as a prelude to the more detailed description that is presented later.

[0005] In an aspect of the disclosure, a method for inference of a diffusion transformer network is provided. The method comprises generating a sequence of tokens based on input media data; processing a token of the sequence of tokens by a multi-head self- attention (MHA) residual block of the diffusion transformer network; and processing an output of the MHA residual block by a feed-forward network (FFN) residual block of the diffusion transformer network, wherein the method further comprises performing operations in mix-precision and / or performing two or more operations in a fusion manner, by at least one of the MHA residual block and the FFN residual block, for example, on an on-chip memory.

[0006] In another aspect of the disclosure, an apparatus for inference of a diffusion transformer network is provided. The apparatus comprises a multi-head self-attention (MHA) residual block for processing a token generated based on input media data; and a feed-forward network (FFN) residual block for processing an output of the MHA residual block, wherein at least one of the MHA residual block and the FFN residual block is configured to perform operations in mix-precision and / or perform two or more operations in a fusion manner, for example, on an on-chip memory.

[0007] In another aspect of the disclosure, an apparatus for inference of a diffusion transformer network is provided. The apparatus may comprise a memory and at least one processor. The at least one processor may be configured to generate a sequence of tokens based on input media data; process a token of the sequence of tokens by a multi-head self-attention (MHA) residual block of the diffusion transformer network; and process an output of the MHA residual block by a feed-forward network (FFN) residual block of the diffusion transformer network, wherein the at least one processor may be further configured to perform operations in mix-precision and / or perform two or more operations in a fusion manner, by at least one of the MHA residual block and the FFN residual block. In some embodiments, the memory is a global memory, the at least one processor is coupled to the global memory, the at least one processor comprises an on chip memory and is configured to perform operations in mix-precision and / or perform two or more operations in a fusion manner, by at least one of the MHA residual block and the FFN residual block on the on-chip memory.

[0008] In another aspect of the disclosure, a computer readable medium storing computer program codes for inference of a diffusion transformer network is provided. The computer program codes, when executed by a processor, may cause the processor to generate a sequence of tokens based on input media data; process a token of the sequence of tokens by a multi-head self-attention (MHA) residual block of the diffusion transformer network; and process an output of the MHA residual block by a feed-forward network (FFN) residual block of the diffusion transformer network, wherein the computer program codes, when executed by a processor, may further cause the  processor to perform operations in mix-precision and / or perform two or more operations in a fusion manner, by at least one of the MHA residual block and the FFN residual block, for example, on an on-chip memory.

[0009] In another aspect of the disclosure, a computer program product for inference of a diffusion transformer network is provided. The computer program product may comprise processor executable computer program codes for generating a sequence of tokens based on input media data; processing a token of the sequence of tokens by a multi-head self-attention (MHA) residual block of the diffusion transformer network; and processing an output of the MHA residual block by a feed-forward network (FFN) residual block of the diffusion transformer network, wherein the computer program product may further comprise processor executable computer program codes for performing operations in mix-precision and / or performing two or more operations in a fusion manner, by at least one of the MHA residual block and the FFN residual block, for example, on an on-chip memory.

[0010] Other aspects or variations of the disclosure will become apparent by consideration of the following detailed description and accompanying drawings.BRIEF DESCRIPTION OF THE DRAWINGS

[0011] The following figures depict various embodiments of the present disclosure for purposes of illustration only. One skilled in the art will readily recognize from the following description that alternative embodiments of the methods and structures disclosed herein may be implemented without departing from the spirit and principles of the disclosure described herein.

[0012] FIG. 1 illustrates a diffusion process of diffusion models (DMs) in accordance with one aspect of the present disclosure.

[0013] FIG. 2 illustrates an inference workflow of a single diffusion transformer (DiT) block with float point 16 precision in accordance with one aspect of the present disclosure.

[0014] FIG. 3 illustrates an optimized inference workflow of a single diffusion transformer (DiT) block with mix precision in accordance with one aspect of the present disclosure.

[0015] FIG. 4 illustrates a block diagram of an apparatus for inference of a diffusion transformer network in accordance with one aspect of the present disclosure.

[0016] FIG. 5 illustrates a flow chart of a method for inference of a diffusion transformer network in accordance with one aspect of the present disclosure.

[0017] FIG. 6 illustrates a block diagram of an apparatus for inference of a diffusion transformer network in accordance with one aspect of the present disclosure.DETAILED DESCRIPTION

[0018] Before any embodiments of the present disclosure are explained in detail, it is to be understood that the disclosure is not limited in its application to the details of construction and the arrangement of features set forth in the following description. The disclosure is capable of other embodiments and of being practiced or of being carried out in various ways.

[0019] Diffusion models are a type of generative model that may generate data by progressively adding noise to a dataset and then learning to reverse this process for generating data by de-noising. Diffusion models may be trained to generate various media data in real world. For example, diffusion models may generate image, video (each frame of video may also be construed as an image) , audio (including speech) , text, 3-D objects, and other types of media data that can be recognized by humans. All these kinds of data that can be generated by diffusion models are collectively referred to as media data herein.

[0020] FIG. 1 illustrates a diffusion process of diffusion models (DMs) in accordance with one aspect of the present disclosure. DMs gradually inject noise into clean media data x0 (such as, an image as shown in the figure) by a fixed forward process:

[0021] where t denotes the timestep and are hyperparameters pre-defined ensuring that  when t is large. Typically, DMs employ the reparameterization trick to sample where is the noise distribution. DMs then define the reverse process through a Markov chain:

[0022] DMs typically parameterize the μθ and Σθ with neural networks in various manners. Using ∈-prediction, the reverse process is trained with the variational lower bound, simplifying the training objective to the mean square error (MSE) between the output of noise prediction network and the ground truth noise:  Then, an image may be generated by removing from a noise  image (such as, xT) the noise predicted by the trained noise prediction network step by step.

[0023] Transformers are a kind of prevalent network architecture due to its excellent scalability, and some studies have been working on transformers for DMs. This work focuses on optimization of diffusion transformers (DiT) , which is a particular type of diffusion model. A brief introduction of DiT architecture is provided here. DiT execute the diffusion process within the latent space. Initially, DiT may convert the spatial representations of media data into a sequence of T tokens z. For the text or speech audio media data of natural language, each token may be a character, a word, or a syllable, etc. For the image media data or video media data consisting of multiple images, DiT may divide an image into a plurality of blocks, convert the spatial representation of each image block into a token, and incorporate position embeddings similar to those in vision transformer. In a controlled generation task, the condition embedding c will also be derived from the timestep embedding t and class label embedding y, e.g., in order to generate an image of a specific kind of animal. The condition embedding c is not necessary. DiT may generate media data randomly without the condition embedding c.

[0024] Subsequently, the sequence of tokens z may be fed into a sequence of N DiT blocks respectively and the condition embedding c may be fed into each of the sequence of N DiT blocks, in an inference process in order to generate corresponding media data, such as images, videos, audios, texts, etc. FIG. 2 illustrates an inference workflow of a single DiT block with float point 16 precision in accordance with one aspect of the present disclosure. Each DiT block primarily consists of three components: (1) a condition multilayer perceptron (MLP) block 210, (2) a multi-head self-attention (MHA) residual block 220, and (3) a feed-forward network (FFN) residual block 230.

[0025] The condition MLP block 210 may regress adaptive dimension-wise scale and shift parameters from condition embedding c, injecting the timestep and class label information into subsequent layers. As illustrated in FIG. 2, in the condition MLP block 210, an activation function (Act2) , such as, GELU, may read the condition embedding c from global memory, causing Ro IO cost. After processing the condition embedding c, the activation function may write the output into the global memory, and then a linear layer, such as, a full connection network (FC5) , may read the output of the activation function from the global memory, causing WoRo IO cost. The full connection network may write its output, such as, scale and shift parameters, into the global memory for MHA residual block 220 and FFN residual block 230, causing W6o IO cost.

[0026] In FIG. 2, considering a tensor transferred from layer A to its subsequent layer B, where N is the context length and C is the hidden dimension, and assuming batch size is 1, Wo / Ro indicates that the IO cost is O (C) and can be ignored as it is relatively small compared to N; Wx / Rx indicates writing / reading xNC bytes to / from global memory; and Wxa / Rxa in MHA indicates writing / reading xN2 bytes to / from global memory because the attention maps are of size N×N.

[0027] In the MHA residual block 220, the residual path begins with an adaptive layernorm (Norm1) , which reads input tokens z. After normalizing the input tokens, the normalization layer writes its output to global memory. These tokens are then transformed by scale and shift layer (SS1) with scale and shift parameters γ1, β1 from the condition MLP block. The MHA module then processes the features as detailed below. The outputs are subsequently scaled by the dimension-wise scale parameters α1 from the condition MLP block in the scale layer (S1) , and added to the signal on the shortcut path by the addition operation (Add1) . The scaling residual method is termed as adaptive residual addition.

[0028] As shown in FIG. 2, in the MHA module, a full connection network (FC1) may perform linear process on the input read from the global memory and write the output embedding into the global memory. The output embedding is respectively multiplied by a query matrix and key matrix, and then the two multiplication results are multiplied through a matrix multiplication operation in the QKT layer, to generate an attention score. After computing a softmax function of the attention score in the SFM layer, the output attention score will be multiplied with a multiplication result of the output embedding of the full connection network (FC1) and a value matrix in the SV matrix multiplication layer. The output self-attention value will then be processed by a full connection network (FC2) , to output to the subsequent layers. The IO cost caused by the operations in the MHA residual block are shown in FIG. 2.

[0029] The FFN residual block 230 is similar to the MHA residual block 220, but replaces the MHA module with feed forward network (FFN) , such as, a two-layer fully connected network. In the FFN module, a full connection network (FC3) may perform linear process on the input read from the global memory and write the output into the global memory. An activation function (Act1) of the output of the full connection network (FC3) is computed and then output to a full connection network (FC4) via global memory. The IO cost caused by the operations in the MHA residual block and the FFN residual block are shown in FIG. 2. Since the data generated by each layer / network / module in the MHA residual block and the FFN residual block is high  precision 16 bit float point data, and is transmitted between various layer / network / module via global memory, the IO cost for inference of the diffusion transformer is huge. A general algorithm of a DiT block is shown below.

[0030] Network quantization has emerged as a promising approach to accelerate the sampling process (including multiple times of inference) of diffusion models. In one aspect, quantization approximates a high precision matrix with a low precision matrix and a scale factor. In this disclosure, INT8 (i.e., 8 bit integer) quantization is studied, and the quantization process is performed by a quantizer: xI=Q (xF; sx) where is the original high precision matrix (such as, FP 16, i.e., 16 bit float point) and  is the quantized matrix in INT8 precision. In the following, we use xF, xI, xI32 to represent a tensor x in FP16, INT8 and INT32 respectively. This disclosure uses static quantization, fixing the FP16-precision step-size parameter sxduring the inference. Then, the dequantizer may be defined as xF=DQ (xI; sx) which converts the integer tensor (can be either INT32 or INT8 depending on where it's derived from) back to full-precision. For example, the following matrix multiplication (MM) may be considered:

[0031] y=xwT, where

[0032] In conventional INT8 inference solutions, the following method is usually adopted to utilize the INT8 TensorCore to achieve 2x speedup:

[0033] 1. Quantize input and weight to INT8: xI=Q (xF; sx) , wI=Q (wF; sw)

[0034] 2. Perform MM in INT8 precision: yI32=xIwI

[0035] 3. Dequantize the output to FP16: yF=DQ (yI32, sxsw)

[0036] However, existing solutions have focused solely on Multiply-Accumulate operations (MACs) , neglecting global memory access (IO) cost, which significantly contribute to the total execution time on modern GPUs. Moreover, the frequent quantization and dequantization operations between linear and non-linear layers impose a substantial burden on memory load-write operations, further exacerbating the issue. For example, EfficientDM and Q-Diffusion are two existing diffusion model quantization methods. EfficientDM quantizes the linear layers while keeping matrix  multiplication in the attention mechanism at FP16. Q-Diffusion go a step further by employing 8-bit precision for matrix multiplication within the attention mechanism. However, when applying EfficientDM and Q-Diffusion to diffusion transformers, neither of them achieves a significant throughput improvement over the FP16 baseline. This may be because of the oversight in optimizing global memory access.

[0037] To address the aforementioned problems, the quantization strategy for the diffusion transformer network's computational workflow is meticulously design in this disclosure. The disclosed approach involves performing almost all operations at low precision and leveraging kernel fusion techniques to minimize IO cost. Executing the memory-constrained non-linear operations at INT8 precision theoretically enables a 2×acceleration compared to FP16 precision. Operator fusion techniques could further eliminate the additional IO associated with quantization-dequantization (Q-DQ) operations and activation functions. We further analyze the sensitivity of different layers to quantization, and develop a mixed-precision inference workflow. Utilizing these techniques, the optimized method (termed as I-DiT herein) may effectively reduce the IO cost of e.g. DiT-XL-2 512×512 model from 6.39 GB of FP16 model to 3.53 GB.

[0038] This computational workflow design may implement highly efficient operators with Triton (which is an intermediate language and compiler for tiled neural network computations) and achieve practical inference speedups on GPUs. We validate I-DiT across multiple image generation benchmarks, evaluating both image generation quality and sampling latency. The experimental results show that I-DiT delivers remarkable speedup compared to FP16 models, while preserving high visual fidelity. In class-conditional image generation on ImageNet dataset at 512×512 resolutions, this method achieves an inference speedup of 2.13× compared to FP16 inference on a single RTX 4090 GPU. Furthermore, this method achieves 22.1%of the peak INT8 TOps (Tera Operations per second) , greatly outperforming existing quantization methods for DMs.

[0039] Considering a tensor transferred from layer A to its subsequent layer B, where N is the context length and C is the hidden dimension and assuming batch size is 1, typically, after computations in A finish, the GPU writes the computation results from on-chip memory (such as, SRAM or cache) to global memory (i.e., external memory) and then reads from global memory before executing computations in B, incurring 4NC bytes data movement when executing in FP16 precision (2 bytes / dimension, 1 read + 1 write) . Representing the tensor in INT8 will halve the amount of IO cost to 2NC bytes. Therefore, low-precision data  representations provide significant inference speedup for memory-bounded non-linear operators. However, the quantizers and dequantizers inserted for precision conversion will introduce additional IO cost. A quantizer reads 2NC bytes in FP16 from and write 1NC bytes in INT8 to global memory. A dequantizer reads 1NC bytes in INT8 from and writes 2NC bytes in FP16 to global memory. To mitigate this overhead, the disclosed method utilizes fusion techniques. Specifically, by fusing the operators of layers A and B, computations in B will directly act on the on-chip memory which holds A's computation results, eliminating one write and one load operation. Quantization and dequantization, being element-wise operations, are typically fusible with their predecessors or successors. In I-DiT, it is carefully designed to absorb quantization and dequantization into other operations, aiming to minimize global memory read / load operations in FP16, thereby significantly reduce the IO cost.

[0040] FIG. 3 illustrates an optimized inference workflow of a single diffusion transformer (DiT) block with mix precision in accordance with one aspect of the present disclosure. Similarly with the original DiT block shown in FIG. 2, the optimized DiT block also comprises a condition MLP block 310, a MHA residual block 320, and a FFN residual block 330.

[0041] Adaptive layernorm in DiT block is an element-wise operator. As shown in FIG. 2, the adaptive layernorm in an original DiT block comprises a Normalize operator Norm (·) (such as, Norm1 in MHA residual block 220 and Norm2 in FFN residual block 230) and a "Scale and Shift" operator SS (·) (such as, SS1 in MHA residual block 220 and SS2 in FFN residual block 230) . Algorithm 2 formalizes the computations and global memory read / write operations within the adaptive layernorm in original DiT block. Both Norm (·) and SS (·) operations require 4NC bytes (2 bytes  / dimension, 1 write + 1 read) , resulting in 8NC bytes of global memory access in total. Here we omit IO for scale and shift parameters as they only incur O (C) IO cost.

[0042] As depicted in FIG. 3 and Algorithm 3, the Norm (·) (such as, Norm1 in MHA  residual block 320) and SS (·) (such as, SS1 in MHA residual block 320) are fused together, reducing data movement by 4NC bytes. The output of the operation SS (·) are quantized to INT8 by the quantizer Q (·) (such as, Q1 in MHA residual block 320) , and the quantizer Q (·) is absorbed into SS (·) . This further reduces the memory writing requirements from 2NC to 1NC bytes. Since the subsequent layer is always a linear layer, the quantized output can be directly fed into the matrix multiplication without additional data conversion. Similarly, in the FFN residual block 330, the operations of Norm2 and SS2 are fused together, the output of the operation SS2 are quantized to INT8 in Q6, and the quantization operation of Q6 is absorbed into Norm2 and SS2. In contrast, conventional quantization methods typically introduce a quantizer Q (·) between the adaptive layernorm and the subsequent linear layer. They add 3NC bytes extra data movement: loading 2NC in FP16 and writing 1NC in INT8, resulting in a total amount of 11NC bytes IO cost, which is even higher than FP16 inference and not efficient. Note that the input of adaptive layernorm (such as, Norm1 and Norm2) is maintained in FP16 precision to preserve model performance.

[0043] FIG. 3 further illustrates the detailed process of the FFN module in the FFN residual block 330. The standard FFN in DiT block consists of (1) a linear layer FC3, (2) an activation function Act1 (such as, GELU) , (3) another linear layer FC4, as shown in FIG. 2. The dropout layers are ignored since we focus solely on the inference phase. As shown in FIG. 3, in the optimized FFN module, the input of FC3 is already in INT8 precision (from the predecessor adaptive layernorm including the quantizer Q6) . Note that FC3 increases the hidden dimension from C to 4C. Algorithm 4 and Algorithm 5 compare the computational workflows of FFN in original DiT and I-DiT.

[0044] Matrix multiplication (MM) on GPU typically require multiple times of reading the input matrices from global memory. We introduce a coefficient κ to denote it.  Consider the multiplication of two matrices and denote the CUDA block size as BM×BN, the total number of blocks is Each block has  inner loops, and within each loop, it needs to load BMBK+BKBN elements from global memory. Each block also needs to store BMBN elements to global memory. To summarize, it needs to perform times read operation and times write operation. For matrix A, it performs read operations, thus we define indicating it needs to perform κ read operations over the whole matrix. For simplicity, we only focus on IO cost regarding the input and output matrices in the algorithm's formalisation, while omit those regarding the weight matrices which is halved in INT8 compared to FP16.

[0045] To minimize the IO cost, the element-wise operator Act1 (such as, GELU) may be fused into FC3, eliminating the corresponding 16NC bytes of IO cost (1 write + 1 read, 2 bytes  / dimension) . Then, a quantizer Q (·) (Q7) is inserted and fused into Act1, reducing IO cost in writing from 8NC of FP 16 to 4NC. The output of FC4 is further quantized by a quantizer Q8 into INT8 and fused with Q8, reducing 1NC bytes.

[0046] FIG. 3 further illustrates the detailed process of the MHA module in the MHA residual block 320. In the MHA module, all the operators'IOs are compressed to INT8, except the input of Softmax (SFM) operator. The output of the full connection network (FC1) is quantized to INT8 by the quantizer Q (·) (Q2) , and the operations of FC1 and Q2 are fused together without global memory access. The input of the QKV linear layer (i.e., QKT layer and SV layer) is naturally in INT8 precision as it's the output of the adaptive layernorm (i.e., FC1 and Q2) . The quantizer Q (·) (Q2) is fused with the linear layer (FC1) to minimize the memory access during the writing operations, which also eliminating the need for data conversion in the QKT and SV matrix multiplications. Similarly, the outputs of Softmax operator, SV and linear layer (FC2) are also quantized respectively by quantizers Q3, Q4 and Q5 to INT8, and all the quantizers are fused into their predecessors. Consequently, nearly all the global memory read / write operations are executed in INT8, and the matrix multiplications are naturally performed in INT8 arithmetic without further data conversion. Note that the input to softmax operation is maintained in FP16 for the consideration of accuracy. Algorithm 6 of Triton implementation of the optimized MHA module is provided below.

[0047] The adaptive residual addition in the MHA residual block 320 and the FFN residual block 330 is also memory-bounded. On the residual branch, the input of the scale S (·) (such as, S1 and S2) appears as INT8 since it is the output of the FFN or Attention module. Fusing the dequantization operator DQ (·) into S (·) may reduce the IO cost by NC bytes. Additionally, the addition operator Add (·, ·) (such as, Add1 or Add2) may be fused with S (·) , further reducing IO cost by 4NC bytes. Since the output of Add (·, ·) serves as the input of the subsequent adaptive layernorm, it remains in FP16 precision. Data movement on the shortcut branch also remains in FP16 for the consideration of accuracy. As a result, the total data movement can be reduced from 10 NC bytes in FP16 to 5 NC bytes.

[0048] FIG. 3 further illustrates an optimized condition MLP block 310. In the condition MLP block 310, the operations of activation function (Act2) and liner layer (FC5) are fused together without global memory access. Moreover, a dequantization operator DQ (·) (DQ3) converting the input condition embedding c into FP16 precision is inserted before Act2, and quantizers Q9 and Q10 are inserted after Act2 and FC5 respectively to convert the outputs of Act2 and FC5 into INT8 precision. The dequantization operation DQ3 and quantization operations Q9 and Q10 are also fused together with the Act2 and FC5 layers without incurring additional global memory access.

[0049] In FIG. 3, the solid line with arrow represents data transmission in FP16 precision, and the dashed line with arrow represents data transmission in INT8 precision. In other words, a mix-precision computational workflow is carefully designed in the I-DiT block. Although in the embodiments of the disclosure the mix-precision is described by an example of a mix of FP16 precision and INT8 precision, the mix-precision may also comprise a mix of FP32 and INT8, or a mix of FP32 and INT16, etc., i.e. a mix of a relatively high precision and a relatively low precision. Only certain layers of I-DiT block are quantized without affecting the generation quality. A per-channel quantization is applied to the weight parameters, and a per-tensor static quantization may be applied to the activations, fixing the step-size parameters in the activation quantizers. In one embodiment, the layers of matrix multiplications MM (including both those within linear layers and those involved in the attention mechanism, such as, FC1-FC5, QKT, SV) , non-linear activation functions Act (such as, Act1 and Act2) , and adaptive residual addition Add (such as, Add1 and Add2) can be quantized, while the layers of adaptive layernorm (such as, Norm1 and Norm2) and softmax operation (SFM) are not quantized. In other embodiments, only the layers of MM and Act can be quantized, or only the layers of MM can be quantized.

[0050] FIG. 4 illustrates a block diagram of an apparatus 400 for inference of a diffusion transformer (DiT) network in accordance with one aspect of the present disclosure. The apparatus 400 may be configured to implement a single DiT block of a sequence of DiT blocks in a DiT network. During an inference phase of the DiT network, it may generate media data, such as images, videos, audios, texts, 3-D objects, etc.

[0051] In one embodiment, each DiT block of the network may comprise an MHA residual block 420 for processing a token generated based on input media data. The input media data may be randomly generated noise media data, such as, a noise image, or a noise audio. Each DiT block of the network may also comprise an FFN residual block 430 for processing an output of the MHA residual block 420 for the following layers of the DiT network to generate random media data, such as, a random image, or a random speech audio.

[0052] The MHA residual block 420 and / or the FFN residual block 430 may be configured to perform operations in mix-precision and perform two or more operations in a fusion manner on an on-chip memory. For example, the computation workflow of MHA residual block 420 may be the same as the MHA residual block 320, and the computation workflow of FFN residual block 430 may be the same as the FFN residual block 330, as described in connection with FIG. 3.

[0053] Specifically, the MHA residual block 420 may be configured to perform, on the token of the sequence of tokens, a first normalization operation, a first scale and shift operation, and a first quantization operation sequentially in a fusion manner on the on-chip memory; perform, on an output of the first quantization operation, an MHA process; and perform, on an output of the MHA process, a first dequantization operation, a first scale operation, and a first addition operation with the token sequentially in a fusion manner on the on-chip memory. The MHA process may be implemented by an MHA module configured to perform, on the output of the first quantization operation, a first full connection network (FC) operation and a second quantization operation sequentially in a fusion manner on the on-chip memory; perform, on an output of the second quantization operation, a matrix multiplication operation with a query matrix and a key matrix to output an attention score; perform, on the output attention score, a normalized exponential (i.e., Softmax) function operation and a third quantization operation sequentially in a fusion manner on the on-chip memory; perform, on the output of the third quantization operation, a matrix multiplication operation with a value matrix and the output of the second quantization operation, and a fourth quantization operation sequentially in a fusion manner on the on-chip memory; and perform, on the output of the fourth quantization operation, a second FC operation and a fifth quantization operation sequentially in a fusion manner on the on-chip memory.

[0054] The FFN residual block 430 may be configured to perform, on the output of the MHA residual block, a second normalization operation, a second scale and shift operation, and a sixth quantization operation sequentially in a fusion manner on the on-chip memory; perform, on an output of the sixth quantization operation, an FFN process; and perform, on an output of the FFN process, a second dequantization operation, a second scale operation, and a second addition operation with the output of the MHA residual block sequentially in a fusion manner on the on-chip memory. The FFN process may be implemented by an FFN module configured to perform, on the output of the sixth quantization operation, a third full connection network (FC) operation, a first activation function operation, and a seventh quantization operation sequentially in a fusion manner on the on-chip memory; and perform, on the output of the seventh quantization operation, a fourth FC operation and an eighth quantization operation sequentially in a fusion manner on the on-chip memory.

[0055] In another embodiment of controlled or manipulated generation task, as shown in FIG. 4, each DiT block of the network may further comprise a condition multilayer perceptron (MLP) block 410 for processing a condition embedding, to generate scale  and shift parameters for at least one of the MHA residual block 420 and the FFN residual block 430. The condition embedding may be generated based on an input manipulation condition, such as, a text or speech audio of a desired image. Then, under the control of the condition MLP block 410, the MHA residual block 420 and the FFN residual block 430 may generated the desired media data, such as, an image with a specific class of animal. The computation workflow of the condition MLP block 410 may be the same as the condition MLP block 310, as described in connection with FIG. 3. For example, the condition MLP block 410 may be configured to perform, on the condition embedding, a third dequantization operation, a second activation function operation, a ninth quantization operation, a fifth full connection network (FC) operation and a tenth quantization operation sequentially in a fusion manner on the on-chip memory.

[0056] FIG. 5 illustrates a flow chart of a method 500 for inference of a diffusion transformer network in accordance with one aspect of the present disclosure. The inference of the diffusion transformer network may be performed multiple times repeatedly in a sampling process to implement a generation task. The method 500 may be performed by each DiT block of the diffusion transformer network as described above in connection with FIG. 4. The method 500 may also be performed by an apparatus for inference of the diffusion transformer network comprising at least on processor and a memory coupled to the at least on processor as described below in connection with FIG. 6.

[0057] In one embodiment, in a random generation task, the method 500 may generate a sequence of tokens based on input media data in step 510. The input media data may be randomly generated noise media data. The diffusion transformer network may generate corresponding media data by removing the noise. The media data comprises at least one of text data, image data, audio data, and video data. In step 520, the method 500 may process a token of the sequence of tokens by a multi-head self-attention (MHA) residual block of the diffusion transformer network. In step 530, the method 500 may processing an output of the MHA residual block by a feed-forward network (FFN) residual block of the diffusion transformer network. The MHA residual block and the FFN residual block may perform operations in mix-precision and performs two or more operations in a fusion manner on an on-chip memory, in order to reduce the global memory access and speedup the inference procedure.

[0058] In another embodiment, in a controlled or manipulated generation task, the method 500 may further comprise processing a condition embedding by a condition  multilayer perceptron (MLP) block to generate scale and shift parameters for at least one of the MHA residual block and the FFN residual block. The condition embedding may be generated based on an input manipulation condition, such as, a text or speech audio of a desired image.

[0059] The condition MLP block, the MHA residual block, and the FFN residual block may be the condition MLP block 310 / 410, the MHA residual block 320 / 420, and the FFN residual block 330 / 430, as described above in connection with FIGs. 3 and 4.

[0060] For example, in the condition MLP block, the method 500 may perform, on the condition embedding, a third dequantization operation (DQ3) , a second activation function operation (Act2) , a ninth quantization operation (Q9) , a fifth full connection network operation (FC5) and a tenth quantization operation (Q10) sequentially in a fusion manner on the on-chip memory.

[0061] In the MHA residual block , the method 500 may perform, on the token of the sequence of tokens, a first normalization operation (Norm1) , a first scale and shift operation (SS1) , and a first quantization operation (Q1) sequentially in a fusion manner on the on-chip memory; perform, on an output of the first quantization operation, an MHA process; and perform, on an output of the MHA process, a first dequantization operation (DQ1) , a first scale operation (S1) , and a first addition operation (Add1) with the token sequentially in a fusion manner on the on-chip memory. The MHA process may further comprise performing, on the output of the first quantization operation, a first full connection network operation (FC1) and a second quantization operation (Q2) sequentially in a fusion manner on the on-chip memory; performing, on an output of the second quantization operation, a matrix multiplication operation (QKT) with a query matrix and a key matrix to output an attention score; performing, on the output attention score, a normalized exponential function (Softmax) operation and a third quantization operation (Q3) sequentially in a fusion manner on the on-chip memory; performing, on the output of the third quantization operation, a matrix multiplication operation (SV) with a value matrix and the output of the second quantization operation (Q2) , and a fourth quantization operation (Q4) sequentially in a fusion manner on the on-chip memory; and performing, on the output of the fourth quantization operation, a second FC operation (FC2) and a fifth quantization operation (Q5) sequentially in a fusion manner on the on-chip memory.

[0062] In the FFN residual block, the method 500 may perform, on the output of the MHA residual block, a second normalization operation (Norm2) , a second scale and shift operation (SS2) , and a sixth quantization operation (Q6) sequentially in a fusion  manner on the on-chip memory; perform, on an output of the sixth quantization operation, an FFN process; and perform, on an output of the FFN process, a second dequantization operation (DQ2) , a second scale operation (S2) , and a second addition operation (Add2) with the output of the MHA residual block sequentially in a fusion manner on the on-chip memory. The FFN process may further comprise performing, on the output of the sixth quantization operation, a third full connection network operation (FC3) , a first activation function operation (Act1) , and a seventh quantization operation (Q7) sequentially in a fusion manner on the on-chip memory; and performing, on the output of the seventh quantization operation, a fourth FC operation (FC4) and an eighth quantization operation (Q8) sequentially in a fusion manner on the on-chip memory.

[0063] According to method 500, in the MHA residual block, the operations of Norm1, SS1 and Q1 are fused together; the operations of DQ1, S1 and Add1 are fused together, the operations of FC1 and Q2 are fused together, the operations of SFM and Q3 are fused together; the operations of SV and Q4 are fused together; the operations of FC2 and Q5 are fused together. In the FFN residual block, the operations of Norm2, SS2 and Q6 are fused together; the operations of DQ2, S2 and Add2 are fused together; the operations of FC3, Act1 and Q7 are fused together; the operations of FC4 and Q8 are fused together. In the condition MLP block, the operations of DQ3, Act2, Q9, FC5 and Q10 are fused together. In the fusion manner, the method 500 may perform the fused operations on the on-chip memory (such as, SRAM) of the processor (such as, GPU) without accessing external memory (i.e., global memory) .

[0064] Moreover, the method 500 may perform the operations of a DiT block during the inference phase in mixed-precision. For example, the method 500 may perform the operations of liner layers (such as, FC1-FC4) and matrix multiplications (such as, QKT and SV, etc. ) in a low precision (such as, INT8) ; and the method 500 may perform the operations of Softmax and adaptive residual addition (such as, Add1 and Add2) in a high precision (such as, FP16) . The method 500 may speedup the inference procedure while maintaining the generation quality of the diffusion transformer network.

[0065] FIG. 6 illustrates a block diagram of an apparatus 600 for inference of a diffusion transformer network in accordance with one aspect of the present disclosure. The apparatus 600 may comprise a memory 610 and at least one processor 620. The memory 610 may be external memory (i.e., global memory) for processor 620. The processor 620 may be coupled to the memory 610 and may further comprise on-chip memory 625, such as, cache. The processor 620 may be configured to perform the  method 500 described above with reference to FIG. 5. The processor 620 may be a general-purpose processor, a graphic processor (GPU) , a neural processor, or may also be implemented as a combination of computing devices, e.g., a combination of a DSP and a microprocessor, multiple microprocessors, one or more microprocessors in conjunction with a DSP core, or any other such configuration. The memory 610 may store the input data (such as, images) , output data, data generated by processor 620, and / or instructions executed by processor 620. The on-chip memory 625 may temporarily cache the intermedia data generated by the processor 620 when performing the operations for inference of a diffusion transformer network.

[0066] The various operations, blocks, modules, models, layers and networks described in connection with the disclosure herein may be implemented in hardware, software executed by a processor, firmware, or any combination thereof. According an embodiment of the disclosure, a computer program product for inference of a diffusion transformer network may comprise processor executable computer program codes for performing the method 500 described above with reference to FIG. 5. According to another embodiment of the disclosure, a computer readable medium may store computer program codes for inference of a diffusion transformer network. The computer program codes when executed by a processor may cause the processor to perform the method 500 described above with reference to FIG. 5. The computer readable media includes both non-transitory computer storage media and communication media including any medium that facilitates transfer of a computer program from one place to another. Any connection may be properly termed as a computer-readable medium. Other embodiments and implementations are within the scope of the disclosure.

[0067] The preceding description of the disclosed embodiments is provided to enable any person skilled in the art to make or use the various embodiments. Various modifications to these embodiments will be readily apparent to those skilled in the art, and the generic principles defined herein may be applied to other embodiments without departing from the scope of the various embodiments. Thus, the claims are not intended to be limited to the embodiments shown herein but is to be accorded the widest scope consistent with the following claims and the principles and novel features disclosed herein.

Claims

1.A method for inference of a diffusion transformer network, comprising:generating a sequence of tokens based on input media data;processing a token of the sequence of tokens by a multi-head self-attention (MHA) residual block of the diffusion transformer network; andprocessing an output of the MHA residual block by a feed-forward network (FFN) residual block of the diffusion transformer network,wherein the method further comprises performing operations in mix-precision and / or performing two or more operations in a fusion manner, by at least one of the MHA residual block and the FFN residual block.2.The method of claim 1, wherein the media data comprises at least one of text data, image data, audio data, and video data.3.The method of claim 1, wherein the MHA residual block comprises:performing, on the token of the sequence of tokens, a first normalization operation, a first scale and shift operation, and a first quantization operation sequentially in a fusion manner on an on-chip memory;performing, on an output of the first quantization operation, an MHA process; andperforming, on an output of the MHA process, a first dequantization operation, a first scale operation, and a first addition operation with the token sequentially in a fusion manner on the on-chip memory.4.The method of claim 3, wherein the MHA process comprises:performing, on the output of the first quantization operation, a first full connection network (FC) operation and a second quantization operation sequentially in a fusion manner on the on-chip memory;performing, on an output of the second quantization operation, a matrix multiplication operation with a query matrix and a key matrix to output an attention score;performing, on the output attention score, a normalized exponential function operation and a third quantization operation sequentially in a fusion manner on the on-chip memory;performing, on the output of the third quantization operation, a matrix multiplication operation with a value matrix and the output of the second quantization operation, and a fourth quantization operation sequentially in a fusion manner on the on-chip memory; andperforming, on the output of the fourth quantization operation, a second FC operation and a fifth quantization operation sequentially in a fusion manner on the on-chip memory.5.The method of claim 1, wherein the FFN residual block comprises:performing, on the output of the MHA residual block, a second normalization operation, a second scale and shift operation, and a sixth quantization operation sequentially in a fusion manner on a on-chip memory;performing, on an output of the sixth quantization operation, an FFN process; andperforming, on an output of the FFN process, a second dequantization operation, a second scale operation, and a second addition operation with the output of the MHA residual block sequentially in a fusion manner on the on-chip memory.6.The method of claim 5, wherein the FFN process comprises:performing, on the output of the sixth quantization operation, a third full connection network (FC) operation, a first activation function operation, and a seventh quantization operation sequentially in a fusion manner on the on-chip memory; andperforming, on the output of the seventh quantization operation, a fourth FC operation and an eighth quantization operation sequentially in a fusion manner on the on-chip memory.7.The method of claim 1, further comprising:processing a condition embedding by a condition multilayer perceptron (MLP) block to generate scale and shift parameters for at least one of the MHA residual block and the FFN residual block.8.The method of claim 7, wherein the condition MLP block comprises:performing, on the condition embedding, a third dequantization operation, a second activation function operation, a ninth quantization operation, a fifth full connection network (FC) operation and a tenth quantization operation sequentially in a fusion manner on an on-chip memory.9.The method of claim 1, wherein the inference of the diffusion transformer network is performed multiple times repeatedly in a sampling process.10.An apparatus for inference of a diffusion transformer network, comprising:a multi-head self-attention (MHA) residual block for processing a token generated based on input media data; anda feed-forward network (FFN) residual block for processing an output of the MHA residual block,wherein at least one of the MHA residual block and the FFN residual block is configured to perform operations in mix-precision and / or perform two or more operations in a fusion manner.11.The apparatus of claim 10, further comprising:a condition multilayer perceptron block for processing a condition embedding, to generate scale and shift parameters for at least one of the MHA residual block and the FFN residual block.12.An apparatus for inference of a diffusion transformer network, comprising:a memory; andat least one processor, wherein the least one processor is configured to perform the method of one of claims 1-9.13.The apparatus of claim 12, wherein the memory is a global memory, the at  least one processor is coupled to the global memory, the at least one processor comprises an on chip memory and is configured to perform operations in mix-precision and / or perform two or more operations in a fusion manner, by at least one of the MHA residual block and the FFN residual block on the on chip memory.14.A computer readable medium, storing computer program codes for inference of a diffusion transformer network, the computer program codes when executed by a processor, causing the processor to perform the method of one of claims 1-9.15.A computer program product for inference of a diffusion transformer network, comprising: processor executable computer program codes for performing the method of one of claims 1-9.

Images