Do not import NVSHMEM in the AoT script unless explicitly requested (#1506)

<!-- .github/pull_request_template.md -->

## 📌 Description

<!-- What does this PR do? Briefly describe the changes and why they’re
needed. -->

## 🔍 Related Issues

<!-- Link any related issues here -->

## 🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.

### ✅ Pre-commit Checks

- [ ] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [ ] I have installed the hooks with `pre-commit install`.
- [ ] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## 🧪 Tests

- [ ] Tests have been added or updated as needed.
- [ ] All tests are passing (`unittest`, etc.).

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->

Author: nandor

flashinfer/aot.py

@@ -11,7 +11,6 @@
 
 from .activation import act_func_def_str, gen_act_and_mul_module
 from .cascade import gen_cascade_module
-from .comm.nvshmem import gen_nvshmem_module
 from .fp4_quantization import gen_fp4_quantization_sm100_module
 from .fused_moe import gen_cutlass_fused_moe_sm100_module
 from .gemm import gen_gemm_module, gen_gemm_sm90_module, gen_gemm_sm100_module
@@ -372,6 +371,9 @@ def gen_all_modules(
 
     if add_comm:
         from .comm import gen_trtllm_comm_module, gen_vllm_comm_module
+        from .comm.nvshmem import gen_nvshmem_module
+
+        jit_specs.append(gen_nvshmem_module())
 
         if has_sm100:
             jit_specs.append(gen_trtllm_comm_module())
@@ -381,7 +383,6 @@ def gen_all_modules(
         jit_specs += [
             gen_cascade_module(),
             gen_norm_module(),
-            gen_nvshmem_module(),
             gen_page_module(),
             gen_quantization_module(),
             gen_rope_module(),