using python platform

Author: apbose

py/torch_tensorrt/dynamo/utils.py

@@ -875,8 +875,12 @@ def _cache_root() -> Path:
     return Path(tempfile.gettempdir()) / f"torch_tensorrt_{username}"
 
 
-def _extracted_dir_trtllm(platform: str) -> Path:
-    return _cache_root() / "trtllm" / f"{__tensorrt_llm_version__}_{platform}"
+def _extracted_dir_trtllm(platform_system: str, platform_machine: str) -> Path:
+    return (
+        _cache_root()
+        / "trtllm"
+        / f"{__tensorrt_llm_version__}_{platform_system}_{platform_machine}"
+    )
 
 
 def download_and_get_plugin_lib_path() -> Optional[str]:
@@ -889,13 +893,14 @@ def download_and_get_plugin_lib_path() -> Optional[str]:
     Returns:
         Optional[str]: Path to shared library or None if operation fails.
     """
+    platform_system = platform.system().lower()
+    platform_machine = platform.machine().lower()
     wheel_filename = (
         f"tensorrt_llm-{__tensorrt_llm_version__}-{_WHL_CPYTHON_VERSION}-"
-        f"{_WHL_CPYTHON_VERSION}-{platform}.whl"
+        f"{_WHL_CPYTHON_VERSION}-{platform_system}_{platform_machine}.whl"
     )
-    platform_system = platform.system().lower()
     wheel_path = _cache_root() / wheel_filename
-    extract_dir = _extracted_dir_trtllm(platform_system)
+    extract_dir = _extracted_dir_trtllm(platform_system, platform_machine)
     # else will never be met though
     lib_filename = (
         "libnvinfer_plugin_tensorrt_llm.so"