I read about the Diffie-Hellman Algorithm in cryptography in which we choose a natural number $g\in\{0,\ldots p-1\}$ for a prime number $p$. Then two other natural numbers $a$ and $b$ are chosen. In the first step the numbers $$A=g^a\bmod p\\ B=g^b\bmod p$$ are computed. Notice that this is not a congruence, but instead the modulus of the division. (I will denote congruences with the usual '$\equiv'$)
Then $$ K_{1}:=B^{a}\bmod p\stackrel{(*)}{=}(g^{b}\bmod p)^{a} \bmod p=(g^{b})^{a}\bmod p=g^{ba}\bmod p$$ and $${\displaystyle K_{2}:=A^{b}\bmod p=(g^{a}\bmod p)^{b}\bmod p=(g^{a})^{b}\bmod p=g^{ab}\bmod p}$$ and therefore $K_1=K_2$. We will not discuss the cryptographical meaning of this, just take it as equations as they stand. In $(*)$ it is used that for a prime $p$ we have elementary multiplication e.g. ${\displaystyle \left((3{\bmod {5}})\cdot (3{\bmod {5}})\right){\bmod {5}}=9{\bmod {5}}}$
As I know this does not hold for non prime integers.
Now here is the thing that confused me: When I use congruence notation we have $$A\equiv g^a\ \pmod p\\B\equiv g^b\pmod p$$
It follows by elemtary properties of the congruence relation that $$A^b\equiv g^{ab}\pmod p$$
and likewise for $B^a$. We then get $K_1\equiv K_2$. This is not equivalent to the statement before where we got equality, but it does not rely on that the modulus has to be prime.
Now in cryptographical terms I don't think this would be of concern since I can choose the smallest modulus or make it unique in a similar way.
Is what I wrote mathematically correct? Is "the price to pay for equality" that we have to consider a prime modulus? Why does the Diffie-Hellman algorithm rely on choosing a prime number as modulus?