Update rms_norm.premake to better support AOT compilation

[voltjia]

pytest output:

============================= test session starts ==============================
platform linux -- Python 3.12.3, pytest-8.3.5, pluggy-1.5.0
rootdir: /home/voltjia/ntops
configfile: pyproject.toml
plugins: jaxtyping-0.3.2, anyio-4.9.0, cov-6.0.0
collected 64 items

tests/test_rms_norm.py ................................................. [ 76%]
...............                                                          [100%]

============================= 64 passed in 35.95s ==============================