Working with Runtime Cache

TensorRT-RTX by default compiles GPU kernels during runtime. Runtime caching helps reduce the startup overhead of GPU kernel compilation by storing compiled kernels on disk for reuse.

Overview

When runtime cache is enabled, kernels compiled at runtime can be saved to a local cache file, allowing future runs to load them directly instead of recompiling. Clients can significantly improve performance in workflows with frequent kernel reuse or repeated application runs, resulting in faster startup times, and a smoother user experience. Runtime caching is especially beneficial in production environments or iterative development workflows where minimizing latency is critical.

Compatibility Checks

When using a pre-populated runtime cache, the cache may have been created in a different or outdated environment. Thus, the runtime cache makes the following checks against the runtime environment to ensure the cache is reusable:

• The runtime environment’s GPU SM version should be equivalent to that of the cached version.

• The runtime environment’s TensorRT-RTX version should be greater than or equal to that of the cached version.

• The runtime environment’s CUDA version should be greater than or equal to that of the cached version.

If the above compatibility checks are not met, the pre-compiled runtime cache will not be used and the cached contents will be replaced after the current execution.

APIs

Runtime caching is a feature introduced to TensorRT-RTX’s execution context creation. Ensure the application has completed the necessary steps up to deserializing the TensorRT-RTX engine. The deserialized engine should be available to create an execution context.

Creating the Runtime Cache

During context creation and inference, TRT-RTX will compile the necessary kernels for inference. Each of them will be added to the runtime cache without duplication.

1. Create a runtimeConfig object from the engine. Set the appropriate allocation strategy for the execution context.

   C++

   IRuntimeConfig* runtimeConfig = engine->createRuntimeConfig();
   
   runtimeConfig->setExecutionContextAllocationStrategy(ExecutionContextAllocationStrategy::kSTATIC);
   

   Python

   runtime_config = engine.create_runtime_config()
   
   runtime_config.set_execution_context_allocation_strategy(trt.ExecutionContextAllocationStrategy.STATIC)
   

2. Create the runtimeCache object and set it to the runtimeConfig object.

   C++

   IRuntimeCache* runtimeCache = runtimeConfig->createRuntimeCache();
   
   runtimeConfig->setRuntimeCache(*runtimeCache);
   

   Python

   runtime_cache = runtime_config.create_runtime_cache()
   
   runtime_config.set_runtime_cache(runtime_cache)
   

3. Create the execution context with the configured runtimeConfig object.

   C++

   IExecutionContext *context = engine->createExecutionContext(runtimeConfig);
   

   Python

   context = engine.create_execution_context(runtime_config)
   

Load and Save the Runtime Cache

Your application may need to run inference on the same or similar models repeatedly. In such cases, saving the runtime cache to disk allows previously compiled GPU kernels to be loaded and reused across runs.

TensorRT provides sample utility functions to load the cache file on disk to memory for reuse. The utility function loadTimingCacheFile, initially used for build-time timing cache, can be shared by runtime caches as well. The loaded bytes can be used for deserialization before the runtime cache is placed into the runtimeConfig object.

1. Load the runtime cache and run inference.

   C++

   std::vector<char> loadedCacheBytes
       = samplesCommon::loadTimingCacheFile(sample::gLogger, ".\runtime.cache");
   
   if (!loadedCacheBytes.empty())
   {
       std::vector<uint8_t> runtimeCacheBytes(
                           loadedCacheBytes.begin(), loadedCacheBytes.end());
   
       runtimeCache->deserialize(
                    runtimeCacheBytes.data(), runtimeCacheBytes.size());
   
       runtimeConfig->setRuntimeCache(*runtimeCache);
   }
   

   Python

   // Use TensorRT’s polygraphy library to load and deserialize cache files
   from polygraphy import util
   
   runtime_cache_file = ".\runtime.cache"
   with util.LockFile(runtime_cache_file):
       try:
           loaded_cache_bytes = util.load_file(runtime_cache_file)
           if loaded_cache_bytes:
               runtime_cache.deserialize(loaded_cache_bytes)
       except:
           G_LOGGER.warning(
               f"Did not find runtime cache at: {runtime_cache_file}. ")
   
   runtime_config.set_runtime_cache(runtime_cache)
   

2. After inference is complete, serialize the runtime cache to save it to disk. Save the runtime cache in binary format (for example, std::ios::binary) instead of a text file.

   C++

   // get the runtime config from the execution context
   IRuntimeConfig* runtimeConfig = context->getRuntimeConfig();
   
   // get the runtime cache from the runtime config
   IRuntimeCache* runtimeCache = runtimeConfig->getRuntimeCache();
   
   // serialize the cache into a memory blob
   IHostMemory* hostMemory = runtimeCache->serialize();
   assert(hostMemory != nullptr);
   
   // save the serialized cache to disk
   samplesCommon::saveTimingCacheFile(
                 sample::gLogger, ".\runtime_cache", hostMemory);
   

   Python

   // get the runtime config from the execution context
   runtime_config = context.get_runtime_config()
   
   // get the runtime cache from the runtime config
   runtime_cache = runtime_config.get_runtime_cache()
   
   // serialize the cache into a memory blob and save to disk
   with util.LockFile(runtime_cache_file):
       with runtime_cache.serialize() as buffer:
           util.save_file(buffer, runtime_cache_file,
                          description="runtime cache")
   

tensorrt_rtx Example

Runtime caching is added to tensorrt_rtx with the --runtimeCacheFile flag, which takes in a file path to the runtime cache file on disk. Ensure that the provided file path has the read and write permissions.

# sample command on Windows
tensorrt_rtx --onnx=sample.onnx --runtimeCacheFile=.\runtime.cache

The first tensorrt_rtx run will fill the cache with the compilation information and serialize it to the specified file. The following tensorrt_rtx runs can reuse the cache file to speed up inference. The acceleration is greatest when the runtime cache is used for the same or similarly-structured models.