Fix missing symbols in trtllm_utils.so

The `trtllm_utils` library also needs symbols from common files, e.g.
the `llmException`. Extend the dependencies to link `trtllm_utils.so`
with common helpers.

```python
>>> flashinfer.utils.get_trtllm_utils_module()
Traceback (most recent call last):
...
OSError: ...flashinfer/data/aot/trtllm_utils/trtllm_utils.so: undefined symbol: _ZTIN12tensorrt_llm6common13TllmExceptionE
```

Fixes: #1167
Signed-off-by: Christian Heimes <christian@python.org>

Author: tiran

flashinfer/aot.py

@@ -31,7 +31,7 @@
 from .quantization import gen_quantization_module
 from .rope import gen_rope_module
 from .sampling import gen_sampling_module
-from .utils import version_at_least
+from .utils import get_trtllm_utils_spec, version_at_least
 
 
 def gen_fa2(
@@ -536,38 +536,7 @@ def main():
         )
     ]
     if has_sm90:
-        jit_specs.append(
-            gen_jit_spec(
-                "trtllm_utils",
-                [
-                    jit_env.FLASHINFER_CSRC_DIR
-                    / "nv_internal"
-                    / "tensorrt_llm"
-                    / "kernels"
-                    / "delayStream.cu",
-                ],
-                extra_include_paths=[
-                    jit_env.FLASHINFER_CSRC_DIR / "nv_internal",
-                    jit_env.FLASHINFER_CSRC_DIR / "nv_internal" / "include",
-                    jit_env.FLASHINFER_CSRC_DIR
-                    / "nv_internal"
-                    / "tensorrt_llm"
-                    / "cutlass_extensions"
-                    / "include",
-                    jit_env.FLASHINFER_CSRC_DIR
-                    / "nv_internal"
-                    / "tensorrt_llm"
-                    / "kernels"
-                    / "internal_cutlass_kernels"
-                    / "include",
-                    jit_env.FLASHINFER_CSRC_DIR
-                    / "nv_internal"
-                    / "tensorrt_llm"
-                    / "kernels"
-                    / "internal_cutlass_kernels",
-                ],
-            ),
-        )
+        jit_specs.append(get_trtllm_utils_spec())
     jit_specs += gen_all_modules(
         f16_dtype_,
         f8_dtype_,

flashinfer/utils.py

@@ -471,13 +471,16 @@ def set_log_level(lvl_str: str) -> None:
     get_logging_module().set_log_level(log_level_map[lvl_str].value)
 
 
-@functools.cache
-def get_trtllm_utils_module():
+def get_trtllm_utils_spec():
     return gen_jit_spec(
         "trtllm_utils",
         [
             jit_env.FLASHINFER_CSRC_DIR
             / "nv_internal/tensorrt_llm/kernels/delayStream.cu",
+            jit_env.FLASHINFER_CSRC_DIR / "nv_internal/cpp/common/envUtils.cpp",
+            jit_env.FLASHINFER_CSRC_DIR / "nv_internal/cpp/common/logger.cpp",
+            jit_env.FLASHINFER_CSRC_DIR / "nv_internal/cpp/common/stringUtils.cpp",
+            jit_env.FLASHINFER_CSRC_DIR / "nv_internal/cpp/common/tllmException.cpp",
         ],
         extra_include_paths=[
             jit_env.FLASHINFER_CSRC_DIR / "nv_internal",
@@ -499,7 +502,11 @@ def get_trtllm_utils_module():
             / "kernels"
             / "internal_cutlass_kernels",
         ],
-    ).build_and_load()
+    )
+
+@functools.cache
+def get_trtllm_utils_module():
+    return get_trtllm_utils_spec().build_and_load()
 
 
 def delay_kernel(stream_delay_micro_secs):