From 6da3cd5089bfbf9c0244b775e322a53912ada4b5 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 19 Aug 2026 20:56:55 -0700 Subject: [PATCH] Load a method whose constant name collides with another method's A program can hold several methods, for example one that fills a cache from a prompt and one that runs a single step. The CUDA delegate can share a constant between them so one copy of a weight serves every method that uses it. Sharing is keyed on the constant's name. That key is not reliable on its own. AOTInductor names a lifted constant by its position within one compiled module, so two methods can each own a `_tensor_constant2` holding unrelated data. Sharing that pair points both methods at the same storage and one of them reads the wrong weights. Until now the first such collision failed the whole load, so a valid program with two same-named constants could not be loaded at all. Compare the cached tensor against the one this method expects, and when they differ leave this method's own constant in place rather than failing. The cached copy stays for whichever method it does match. The comparison happens where the weights blob is loaded, so it covers methods that have at least one constant the cache has not seen. A method whose names are all already cached has nothing to compare against without re-uploading the blob, which costs too much device memory to do on every method, so that case still shares on the name alone. The comment at the cache lookup says so. The comparison is on tensor metadata: dtype, rank, sizes, strides, device type and device index. It does not compare the data and cannot, because both handles point at device memory. Two constants agreeing on all of that while holding different values would still be shared. --- backends/cuda/runtime/cuda_backend.cpp | 37 +++++++++++++++++++------- 1 file changed, 28 insertions(+), 9 deletions(-) diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index a828201845f..349082ad690 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -26,6 +26,7 @@ #include #include #include +#include #include // Include SlimTensor headers for CUDA backend @@ -1061,7 +1062,17 @@ class ET_EXPERIMENTAL CudaBackend final } } + // Names this method must load itself because the cached tensor under the + // same name belongs to a different constant. + std::unordered_set not_shared; + // Phase 2 (locked): pure cache lookup against shared_constant_tensors_. + // A cache hit here is provisional: the name matches, but whether the + // tensors match is only known after extraction, so a hit can still be + // rejected below. Only the path that loads the blob can check this. When + // every name is already cached there is nothing to compare against + // without re-uploading the blob, which costs too much device memory to + // do on every method, so that case still shares on the name alone. { std::lock_guard guard(shared_constants_mutex_); for (const auto& [fqn, _] : fqn_to_name) { @@ -1134,17 +1145,22 @@ class ET_EXPERIMENTAL CudaBackend final if (cached_it == shared_constant_tensors_.end()) { // New constant — add to cache. shared_constant_tensors_[fqn] = extracted_it->second; - } else { - // Same FQN seen before — verify the cached tensor is still - // compatible with what THIS method expects. On mismatch the - // helper logs the offending field and returns an error. - ET_CHECK_OK_OR_RETURN_ERROR( - check_cached_constant_match( - fqn, cached_it->second, extracted_it->second), - "Constant '%s' in method '%s' is incompatible with the " - "cached version from a previous method. Refusing to share.", + } else if ( + check_cached_constant_match( + fqn, cached_it->second, extracted_it->second) != Error::Ok) { + // Same name, different tensor. AOTInductor names a lifted constant + // by position within one compiled module, so two methods each own a + // "_tensor_constant2" holding unrelated data, and a name is only + // evidence of sharing when the tensors agree. Sharing this pair + // would alias mismatched storage, so leave this method's own + // constant in place and keep the cached one for whoever matches it. + ET_LOG( + Info, + "Constant '%s' in method '%s' differs from the cached copy; " + "using this method's own copy instead of sharing", fqn.c_str(), method_name.c_str()); + not_shared.insert(fqn); } } ET_LOG( @@ -1173,6 +1189,9 @@ class ET_EXPERIMENTAL CudaBackend final { std::lock_guard guard(shared_constants_mutex_); for (const auto& [fqn, internal_name] : fqn_to_name) { + if (not_shared.count(fqn) != 0) { + continue; + } auto it = shared_constant_tensors_.find(fqn); if (it != shared_constant_tensors_.end()) { pairs.push_back({internal_name.c_str(), it->second});