def finalize(self, final_scale: cute.Float32 = 1.0) -> cute.Tensor:
        """Finalize the online softmax by computing the scale and logsumexp.
        """
        # quad reduction for row_sum as we didn't do it during each iteration of online softmax
        self.row_sum.store(utils.warp_reduce(self.row_sum.load(), operator.add, width=4))
        row_scale = cute.make_fragment_like(self.row_max, cutlass.Float32)
        for r in range(cute.size(self.row_sum)):
            # if row_sum is zero or nan, set acc_O_mn_row to 1.0
            acc_O_mn_row_is_zero_or_nan = self.row_sum[r] == 0.0 or self.row_sum[r] != self.row_sum[r]
            row_scale[r] = (
                cute.arch.rcp_approx(self.row_sum[r] if not acc_O_mn_row_is_zero_or_nan else 1.0)
            ) * final_scale
            row_sum_cur = self.row_sum[r]
            LN2 = math.log(2.0)
            self.row_sum[r] = (
                (self.row_max[r] * self.scale_log2 + utils.log2f(row_sum_cur)) * LN2
                if not acc_O_mn_row_is_zero_or_nan else -cutlass.Float32.inf
            )
        return row_scale
Finalize the online softmax by computing the scale and logsumexp.
        
flash_attn/cute/softmax.py