diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 1f989ab1b5b..921c15ae7af 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -288,6 +288,11 @@ class ET_EXPERIMENTAL CudaBackend final get_constant_dtype, AOTInductorModelContainerGetConstantDtype); LOAD_OPTIONAL_SYMBOL( get_constant_data_size, AOTInductorModelContainerGetConstantDataSize); + LOAD_OPTIONAL_SYMBOL( + get_constant_type, AOTInductorModelContainerGetConstantType); + LOAD_OPTIONAL_SYMBOL( + get_constant_from_folded, + AOTInductorModelContainerGetConstantFromFolded); LOAD_OPTIONAL_SYMBOL( extract_constants_map, AOTInductorModelContainerExtractConstantsMap); LOAD_OPTIONAL_SYMBOL( diff --git a/backends/cuda/runtime/cuda_delegate_handle.h b/backends/cuda/runtime/cuda_delegate_handle.h index 93c11eaeb52..f42f0f2382f 100644 --- a/backends/cuda/runtime/cuda_delegate_handle.h +++ b/backends/cuda/runtime/cuda_delegate_handle.h @@ -56,6 +56,15 @@ using AOTInductorModelContainerGetConstantDataSizeFunc = aoti::AOTInductorModelContainerHandle container_handle, size_t idx, size_t* data_size); +using AOTInductorModelContainerGetConstantTypeFunc = aoti::AOTIRuntimeError (*)( + aoti::AOTInductorModelContainerHandle container_handle, + size_t idx, + int32_t* type); +using AOTInductorModelContainerGetConstantFromFoldedFunc = + aoti::AOTIRuntimeError (*)( + aoti::AOTInductorModelContainerHandle container_handle, + size_t idx, + bool* from_folded); struct CudaWeightStorage { void* data{nullptr}; size_t nbytes{0}; @@ -286,6 +295,11 @@ struct CudaDelegateHandle : public aoti::AOTIDelegateHandle { // storage against it before binding. AOTInductorModelContainerGetConstantDataSizeFunc get_constant_data_size{ nullptr}; + // With these, FQN weight binding checks completeness itself instead of + // AOTI's full update, which warns once per off-graph KV buffer on every load. + AOTInductorModelContainerGetConstantTypeFunc get_constant_type{nullptr}; + AOTInductorModelContainerGetConstantFromFoldedFunc get_constant_from_folded{ + nullptr}; // The per-thread stream. Nothing owns it: the value is a fixed sentinel the // driver resolves to a different stream on each host thread, so releasing the diff --git a/backends/cuda/runtime/cuda_weight_cache.cpp b/backends/cuda/runtime/cuda_weight_cache.cpp index 017629f4760..104b4461318 100644 --- a/backends/cuda/runtime/cuda_weight_cache.cpp +++ b/backends/cuda/runtime/cuda_weight_cache.cpp @@ -140,6 +140,59 @@ bool is_offgraph_kv_fqn(const std::string& fqn) { return fqn.rfind(kOffGraphKVFqnPrefix, 0) == 0; } +// torch::aot_inductor::ConstantType. +constexpr int32_t kAotiParameter = 1; +constexpr int32_t kAotiBuffer = 2; +constexpr int32_t kAotiTensorConstant = 3; + +bool can_check_constants_bound(const CudaDelegateHandle& handle) { + return handle.get_constant_type && handle.get_constant_from_folded && + handle.get_constant_data_size; +} + +// The same rule as AOTI's validate_full_update, without its warning for each +// absent buffer: folded constants, buffers, tensor constants and empty +// parameters may be absent, every other constant must be bound. +Error check_constants_bound( + const CudaDelegateHandle& handle, + size_t num_constants, + const std::unordered_set& bound) { + for (size_t index = 0; index < num_constants; ++index) { + bool from_folded = false; + int32_t type = 0; + size_t data_size = 0; + ET_CHECK_OK_OR_RETURN_ERROR( + handle.get_constant_from_folded( + handle.container_handle, index, &from_folded), + "Failed to read CUDA AOTI constant folding at index %zu", + index); + ET_CHECK_OK_OR_RETURN_ERROR( + handle.get_constant_type(handle.container_handle, index, &type), + "Failed to read CUDA AOTI constant type at index %zu", + index); + ET_CHECK_OK_OR_RETURN_ERROR( + handle.get_constant_data_size( + handle.container_handle, index, &data_size), + "Failed to read CUDA AOTI constant size at index %zu", + index); + if (from_folded || type == kAotiBuffer || type == kAotiTensorConstant || + (type == kAotiParameter && data_size == 0)) { + continue; + } + const char* name = nullptr; + ET_CHECK_OK_OR_RETURN_ERROR( + handle.get_constant_name(handle.container_handle, index, &name), + "Failed to read CUDA AOTI constant name at index %zu", + index); + ET_CHECK_OR_RETURN_ERROR( + name != nullptr && bound.count(name) != 0, + InvalidProgram, + "CUDA AOTI constant '%s' has no bound FQN weight", + name == nullptr ? "" : name); + } + return Error::Ok; +} + } // namespace bool CudaWeightCache::is_serialized(const void* data, size_t size) { @@ -650,16 +703,25 @@ Error CudaWeightCache::load( } if (!pairs.empty()) { - // Full validation stays on with off-graph KV present: those constants are - // AOTI buffers, which the check permits to be absent, while every weight - // must still be supplied. + // Off-graph KV constants are AOTI buffers bound later by the KV pool, so + // they are allowed to be absent here, while every weight must be supplied. + const bool check_here = can_check_constants_bound(*handle); + if (check_here) { + std::unordered_set bound_names; + for (const auto& pair : pairs) { + bound_names.emplace(pair.name); + } + ET_CHECK_OK_OR_RETURN_ERROR( + check_constants_bound(*handle, num_constants, bound_names), + "CUDA FQN weights do not cover the AOTI library"); + } ET_CHECK_OK_OR_RETURN_ERROR( handle->update_user_managed_constant_buffer_pairs( handle->container_handle, pairs.data(), pairs.size(), /*use_inactive=*/false, - /*validate_full_update=*/true), + /*validate_full_update=*/!check_here), "Failed to bind CUDA FQN weights"); } ET_LOG(