High performance softmax for large models
The softmax operation is pipelined to evenly balanced operations with 2N latency per sharded M dimension of a tensor shaped M*N, resulting in ˜1.8× performance gain. As each operation is fragmented, the pipeline does not assume a fixed-cost fill that can performance-wise hurt significantly for small Tensor dimensions. In addition, this innovative design is Place-and-Route (PNR) friendly as well as resource efficient. Moreover, it is easily parallelized without requiring additional support at the subnet level.
1. A computer-implemented method for conducting softmax computations in a reconfigurable dataflow system, the method comprising:
computing a maximum value and a difference between the maximum value and input values in a first compute stage;
performing exponent calculations in a second compute stage; and
computing a sum and a reciprocal in a third compute stage.
2. The computer-implemented method of claim 1 , further including:
storing the maximum value in a compute unit's temporary storage of the first compute stage and using the stored maximum value in a next compute cycle of the first compute stage.
3. The computer-implemented method of claim 1 , wherein:
the first, second, and third compute stages require first, second, and third compute units in the reconfigurable dataflow system.
4. The computer-implemented method of claim 3 , wherein:
absence of multiple compute units each of the computing, performing, and computing steps with shared fan-in obviates predetermined placement of the first, second, and third compute units when conducting the softmax computations in the reconfigurable dataflow system.
5. The computer-implemented method of claim 1 , wherein:
each of the first, second, and third compute stages requires 2N latency per sharded M dimension of an M*N shape tensor.
6. A non-transitory computer-readable storage medium storing computer program instructions that, when executed on a processor, perform operations comprising:
computing a maximum value and a difference between the maximum value and input values in a first compute stage;
performing exponent calculations in a second compute stage; and
computing a sum and a reciprocal in a third compute stage.
7. The non-transitory computer-readable storage medium of claim 6 , further comprising:
storing the maximum value in a compute unit's temporary storage of the first compute stage and using the stored maximum value in a next compute cycle of the first compute stage.
8. The non-transitory computer-readable storage medium of claim 6 , wherein:
the first, second, and third compute stages require first, second, and third compute units in the reconfigurable dataflow system.
9. The non-transitory computer-readable storage medium of claim 8 , wherein:
absence of multiple compute units each of the computing, performing, and computing steps with shared fan-in obviates predetermined placement of the first, second, and third compute units when conducting the softmax computations in the reconfigurable dataflow system.
10. The non-transitory computer-readable storage medium of claim 6 , wherein:
each of the first, second, and third compute stages requires 2N latency per sharded M dimension of an M*N shape tensor.
11. A system comprising one or more processors coupled to a memory device, the memory device to store computer program instructions that are executable by the one or more processors to perform operations comprising:
computing a maximum value and a difference between the maximum value and input values in a first compute stage;
performing exponent calculations in a second compute stage; and
computing a sum and a reciprocal in a third compute stage.
12. The system of claim 11 , further comprising:
storing the maximum value in a compute unit's temporary storage of the first compute stage and using the stored maximum value in a next compute cycle of the first compute stage.
13. The system of claim 11 , wherein:
the first, second, and third compute stages require first, second, and third compute units in the reconfigurable dataflow system.
14. The system of claim 13 , wherein:
absence of multiple compute units each of the computing, performing, and computing steps with shared fan-in obviates predetermined placement of the first, second, and third compute units when conducting the softmax computations in the reconfigurable dataflow system.
15. The system of claim 11 , wherein:
each of the first, second, and third compute stages requires 2N latency per sharded M dimension of an M*N shape tensor.