Quantized flash attention
AssigneeTesla, Inc.
InventorsRitvik RAWAT, Valeriy ROTAN, Alex SINGH, Srihari SAMPATHKUMAR
Life of this patent
Estimating the term…
Abstract
Embodiments include systems and methods for processing sensor data and generating operational instructions for automated devices. The automated device includes hardware (e.g., INT8 MAC units, FP32 SIMD processors) and software operations executing a machine-learning architecture having transformer architectures implementing mixed-precision quantization for ingesting and analyzing various types of sensor data. The machine-learning architecture implements a transformer that includes dynamic quantization layers that dynamically quantize tensors, per-token, for downstream operations of the transformer and machine-learning architecture.
Claims20
A method comprising:
- obtaining, by a processor, a plurality of input tokens representing input data;
- generating, by the processor, a plurality of feature vectors based upon the plurality of input tokens, the feature vectors having a number of bits according to a first precision;
- executing, by the processor, a dynamic quantization to one or more of the feature vectors based upon a level of precision, to update the number of bits of the feature vectors according to a second precision corresponding to the level of precision for the input data to a subsequent computation operation;
- generating, by the processor, a similarity value using the feature vectors;
- generating, by the processor, one or more attention weights using a normalization function applied to the similarity value;
- generating, by the processor, a weighted sum of one or more of the feature vectors using the attention weights;
- and
- generating, by the processor, an intermediate context representation for the plurality of input tokens by aggregating the weighted sums.
A system comprising:
- a processor configured to:
- obtain a plurality of input tokens representing input data;
- generate a plurality of feature vectors based upon the plurality of input tokens, the feature vectors having a number of bits according to a first precision;
- execute a dynamic quantization to one or more of the feature vectors based upon a level of precision, to update the number of bits of the feature vectors according to a second precision corresponding to the level of precision for the input data to a subsequent computation operation;
- generate a similarity value using the feature vectors;
- generate one or more attention weights using a normalization function applied to the similarity value;
- generate a weighted sum of one or more of the feature vectors using the attention weights;
- and
- generate an intermediate context representation for the plurality of input tokens by aggregating the weighted sums.
Dependencies are inferred from the claim text.