Dynamic Shapes and Kernel Cache

Dynamic Shapes

Causes other APIs (such as the kernel cache) to treat the graph as a dynamic shape graph.

The API to achieve the above is:

graph.set_dynamic_shape_enabled(true)

Kernel Cache

The kernel cache significantly reduces plan build time by re-using a previously compiled kernel for a given execution plan. Kernel caching is enabled only for dynamic shape graphs.

If a graph’s kernel cache attribute is set, the kernel cache will store the kernel which was compiled for the graph’s execution plan. On future same-topology operation graphs, the kernel cache may bind the previously compiled kernel to the execution plan to avoid recompilation.

The API to create a kernel cache is:

auto kernel_cache = std::make_shared<cudnn_frontend::KernelCache>();

The API to set a dynamic shape graph’s kernel cache is:

graph.set_kernel_cache(kernel_cache)

Kernel cache introspection

Starting with cuDNN 9.27, two accessors report the state of a finalized kernel cache:

int64_t revision = 0;
int64_t size     = 0;
kernel_cache->revision(revision);   // counter that shows if the contents changed
kernel_cache->size(size);           // number of entries

The same two accessors are available in Python:

revision = kernel_cache.revision()
size = kernel_cache.size()

Both accessors need a finalized kernel cache. They return an error before build().

• Use revision(), and not size(), to find if the cache changed. cuDNN can replace data that is already in the cache, so size() is not a reliable measure. An insertion, a replacement, a removal, or an eviction each add one to the revision counter. A lookup does not change it. A change in the counter shows, for example, that a new serialization is worth the cost.
• Two different shapes can use the same kernel cache entry, and size() can stay the same after a build with a new shape.
• The revision counter is local to the process. It is not part of the serialized data, and a from_json() load starts a new count. Two kernel caches have no common count.

Thread-safety contract for KernelCache

The KernelCache API is internally synchronized with a plain std::mutex. Multiple threads may call these methods concurrently on the same KernelCache instance without external synchronization.

build() is idempotent under the lock: the first caller to acquire the mutex initializes and finalizes the descriptor; subsequent callers return OK immediately without re-initializing it. This makes it safe for N graph threads sharing one KernelCache to each call build() concurrently — exactly one backend descriptor is created and finalized.

from_json() is synchronized the same way as build(). Call it before sharing the KernelCache with any thread that will call build().

Override Shape

Override shape allows supplying at execution time tensor shapes that differ from the shapes used when building the graph. A single execution plan can thus support multiple dynamic shapes without rebuilding the graph for each shape.

Typical usage: build the graph and execution plan once with a “cache shape”, then on each execute call pass the actual shapes for that run via override_uids, override_shapes, and override_strides.

API to enable override shape:

graph.set_override_shape_enabled(true)

Call this before building the graph (together with other options such as set_dynamic_shape_enabled or set_kernel_cache). It supports chaining; when the return type is Error, use .is_good() to check success.

Execution API with overrides:

graph->execute(handle, variant_pack, workspace_ptr, override_uids, override_shapes, override_strides)

When using override shape, query the workspace for the actual runtime shape immediately before allocating the workspace:

int64_t workspace_size = 0;
graph->get_workspace_size(handle, workspace_size, override_uids, override_shapes, override_strides);
Surface<int8_t> workspace(workspace_size);
graph->execute(handle, variant_pack, workspace.devPtr, override_uids, override_shapes, override_strides);

This runtime workspace query may return a different size than the workspace queried for the cache shape used to build the graph. Allocate the workspace after this query for each set of override shapes.

Where:

• override_uids: list of tensor UIDs whose shapes are being overridden
• override_shapes: new shape for each tensor in override_uids (each element is a std::vector<int64_t>)
• override_strides: new stride for each tensor in override_uids

The three vectors must have the same length. Tensors not listed in override_uids keep the shapes defined at graph build time.