pypto.experimental.online_softmax
Applicable Products
- Ascend 950PR/Ascend 950DT: Supported
- Atlas A3 training products/Atlas A3 inference products: Not supported
- Atlas A2 training products/Atlas A2 inference products: Not Supported
Description
This is a custom API with many constraints. Its stability is not guaranteed.
This operator performs block-wise online Softmax computation and computes local statistics on the input scores along dimension 0. The operator first multiplies scores by scale, then computes the maximum value and the exponential sum of each column, and outputs the unnormalized exponential results. This API is typically used in block-wise attention scenarios such as FlashAttention, in conjunction with pypto.experimental.online_softmax_update to update the global maximum, exponential sum, and intermediate outputs block by block.
Prototype
online_softmax(scores: Tensor, scale: float) -> Tuple[Tensor, Tensor, Tensor]Parameters
| Parameter | Input/Output | Description |
|---|---|---|
| scores | Input | Source operand. Supported data type: DT_FP32. Empty tensors are not supported. Two-dimensional tensors are supported. The shape is [k_len, q_len]. |
| scale | Input | float type. Scalar used to scale scores, commonly set to 1.0 / sqrt(head_dim). |
Return Value
Returns three output tensors:
| Return Value | Description |
|---|---|
| exp_scores_bf16 | Exponential results after scaling and subtracting the column maximum, with data type DT_BF16 and the same shape as scores. |
| column_max | Local maximum of each column, with data type DT_FP32 and shape [1, q_len]. |
| column_sum | Local exponential sum of each column, with data type DT_FP32 and shape [1, q_len]. |
Constraints
- This is a custom API. Its stability is not guaranteed.
- The data type of scores supports only DT_FP32.
- In the current version, dimension 0 is not tiled. The constraint
scores.shape[0] <= vec_tile[0]must be satisfied.
Examples
TileShape Setting Example
Before calling this operation API, set the TileShape through set_vec_tile_shapes.
The dimension settings of the TileShape must be consistent with the input scores. In the current version, dimension 0 is not tiled. The constraint scores.shape[0] <= vec_tile[0] must be satisfied.
API Call Example
import pypto
scores = pypto.tensor([128, 128], pypto.DT_FP32)
scale = 1.0 / (128 ** 0.5)
pypto.set_vec_tile_shapes(128, 64)
exp_scores_bf16, column_max, column_sum = pypto.experimental.online_softmax(scores, scale)