US 2026/0288899 A1Published application

Fast Attention Calculation

OwnerSergei Chirkunov, Szabolcs CSÉFALVAYinventors

Life of this patent

Estimating the term…

Loading patent…

Abstract

A method of performing an attention calculation in a multithreaded processing system includes a processing unit, coupled with a dedicated cache and a dedicated local memory. The processing unit operates threads grouped into workgroups. A first matrix Q, a second matrix KT, and a third matrix V are inputs. In the cache dedicated to the first processing unit, a portion of the first matrix is stored, as a first matrix subunit. In the local memory dedicated to the first processing unit, a portion of the second matrix is stored, as a second matrix subunit. A portion of the third matrix is stored as a plurality of third matrix subunits. The first matrix subunits are multiplied by the second matrix subunits, to generate a subunit of a first product matrix. A product of (i) a softmax over each row of the first product matrix and (ii) the third matrix are calculated to generate at least a portion of an output of the attention calculation.

Claims20

Dependent claims

A method of performing an attention calculation in a multithreaded processing system, wherein the multithreaded processing system comprises one or more processing units, including a first processing unit, each processing unit coupled with a dedicated cache and a dedicated local memory, and wherein each processing unit is configured to operate a plurality of threads grouped into one or more workgroups, the method comprising:

providing as inputs for the attention calculation a first matrix Q, a second matrix KT, and a third matrix V;
storing, in the cache dedicated to the first processing unit, at least a portion of the first matrix as one or more first matrix subunits;
storing, in the local memory dedicated to the first processing unit, at least a portion of the second matrix as a plurality of second matrix subunits;
multiplying the first matrix by the second matrix, comprising multiplying, using the first processing unit, the first matrix subunits by the second matrix subunits, to generate at least one subunit of a first product matrix;
storing at least a portion of the third matrix as a plurality of third matrix subunits;
and
calculating a product of (i) a softmax over each row of the first product matrix and (ii) the third matrix, comprising processing, using the first processing unit, subunits of the first product matrix together with the third matrix subunits, to generate at least a portion of an output of the attention calculation.

4 more elements

Dependencies are inferred from the claim text.