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?
By inductive extension of the Congruence Sum, Product and Power operation rules it follows the value of a polynomial expression is preserved (stays congruent) if we replace arguments of said constituent operations by congruent arguments (only for the base (vs. expt) in powers). Therefore, in particular, we can replace $\,g^a\,$ by $\,g^a\bmod p\,$ and $\,g^b\,$ by $\,g^b\bmod p\,$ in $\,(g^a)^b = (g^b)^a\,$ to get
$$\begin{align} (g^a\bmod p)^b&\equiv\, (g^b\bmod p)^a\!\!\pmod{\!p},\ \ \text{so taking remainders}\\[.3em] \color{#c00}\Longrightarrow\ \ (g^a\bmod p)^b\bmod p \,&=\, (g^b\bmod p)^a\bmod p \end{align}$$
because $\,\ a\equiv b\pmod{\!p}\,\color{#c00}{\Longrightarrow}\, a\bmod p\, =\, b\bmod p\,\ $ [and conversely]
Remark $\ $ Generally, as above, it is straightforward to prove properties of the mod operator by first converting to congruence form, then applying the simpler arithmetical laws of congruences then, finally, converting back to mod operator form by taking remainders (normal forms) as the last step. While it is possible to perform such proofs without congruences, generally that will be messier, and less conceptual.
See here for background on the distinction between mod as an operator vs. congruence relation.