From d6d412b9d50fc36a215b93e75ce2450fdec8cc4d Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Wed, 21 Oct 2026 01:32:39 -0700 Subject: [PATCH] [executorch][cuda] Check FQN weight coverage without AOTI's per-buffer warnings Every CUDA method load that binds FQN weights next to an off-graph KV cache printed one AOTI line per KV buffer, e.g. [WARNING] Found constant or module state buffer or empty module state parameter __et_offgraph_kv_layer_0_k in model, but not provided by user! The muse-glimmer batching CI log had 1672 of these lines out of 11650. CudaWeightCache::load binds weights with validate_full_update=true while it deliberately leaves the off-graph KV buffers unbound; the KV pool binds them later. AOTI's full update allows absent buffers but writes this warning to stderr for each one, on every load. When the AOTI library exports GetConstantType and GetConstantFromFolded, do the same completeness check ExecuTorch-side: folded constants, buffers, tensor constants and empty parameters may be absent, every other constant must be bound, else InvalidProgram. Then bind with validate_full_update=false. validate_full_update only gates that check in AOTI, so binding is otherwise unchanged. Libraries without the two APIs keep validate_full_update=true. Test plan: - g++ -fsyntax-only on cuda_weight_cache.cpp and cuda_backend.cpp against the CUDA 12.6 headers. - CI: CUDA unittest and muse-glimmer batching jobs pass, and their logs no longer contain the __et_offgraph_kv_ warnings. --- backends/cuda/runtime/cuda_backend.cpp | 5 ++ backends/cuda/runtime/cuda_delegate_handle.h | 14 ++++ backends/cuda/runtime/cuda_weight_cache.cpp | 70 ++++++++++++++++++-- 3 files changed, 85 insertions(+), 4 deletions(-) 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(