How do you calculate spectral flatness from an FFT?
If the definition of the flatness dictates that you use a power spectrum, then yes, you should square the magnitudes as the reference from the SciPy documentation indicates. In the equation that you referenced where you didn't see a squaring, I don't think you can read much into it; it says that

$$
S_{flatness} = \frac{\exp\left(\frac{1}{N} \sum_k \log (a_k)\right)}{\frac{1}{N} \sum_k a_k}
$$

but I don't see a definition for $a_k$ anywhere. If you want the spectrum to be proportional to the power in each bin, you need to square.