torch.compile() throws exception when LigerKernel is used

🐛 Describe the bug

... 
  File "/home/tromero/workspace/seahorse/.venv/lib/python3.11/site-packages/torch/_inductor/async_compile.py", line 173, in triton
    kernel = TritonCodeCache.load(kernel_name, source_code)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tromero/workspace/seahorse/.venv/lib/python3.11/site-packages/torch/_inductor/codecache.py", line 3112, in load
    return _module_to_triton_kernel(PyCodeCache.load(source_code), kernel_name)
                                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tromero/workspace/seahorse/.venv/lib/python3.11/site-packages/torch/_inductor/codecache.py", line 3049, in load
    return cls.load_by_key_path(key, path, linemap, attrs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tromero/workspace/seahorse/.venv/lib/python3.11/site-packages/torch/_inductor/codecache.py", line 3062, in load_by_key_path
    mod = _reload_python_module(key, path)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tromero/workspace/seahorse/.venv/lib/python3.11/site-packages/torch/_inductor/runtime/compile_tasks.py", line 45, in _reload_python_module
    exec(code, mod.__dict__, mod.__dict__)
  File "/tmp/torchinductor_tromero/rg/crgwelnbq5utprhg6blafhwyxbaxibrfeh7n53w5xhpi5jkmp26h.py", line 82, in <module>
    _CASTING_MODE_LLAMA = constexpr[0]
                          ^^^^^^^^^
torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
NameError: name 'constexpr' is not defined

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True

Seems to be related to these constexpr:
https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/rms_norm.py#L25

Reproduce

PR that provides a test that repros the bug: https://github.com/linkedin/Liger-Kernel/pull/173

Versions

Environment Report:

Operating System: Linux-6.5.0-44-generic-x86_64-with-glibc2.35
Python version: 3.10.13
PyTorch version: 2.3.0
CUDA version: 12.1
Triton version: 2.3.0
Transformers version: 4.42.3

Liger-kernal version 0.2.1

Start with the reproduction test from PR #173 and inspect the constexpr usage around src/liger_kernel/ops/rms_norm.py#L25. Re-run it with the reported PyTorch, CUDA, Triton, and Liger-Kernel versions, then verify that torch.compile() completes without the NameError and the regression test passes.