Cosmology beyond BAO from the 3D distribution of the Lyman-$\alpha$ forest