DYNAMIC_UNBOUND support for portable runtime: lazy KV cache allocation

Enable DYNAMIC_UNBOUND tensors in the portable runtime, allowing KV cache
buffers to be dynamically managed rather than statically memory-planned.
This is the architectural foundation for pay-as-you-go memory allocation
in ExecuTorch LLM inference.

Core changes:
- DynamicAllocator interface with allocate/reallocate/free
- PalDynamicAllocator default impl (PAL-backed, 2x growth policy)
- TrackingDynamicAllocator for memory stats observability
- MemoryManager gains 4th slot for DynamicAllocator (backward compatible)
- TensorImpl gains dynamic_allocator_ and capacity_bytes_ fields
- TensorImpl::internal_resize_contiguous handles DYNAMIC_UNBOUND resize
- tensor_parser_portable.cpp: remove DYNAMIC_UNBOUND rejection, wire up
  allocator at load time for tensors with no memory-planned data
- method.cpp: FreeCall frees dynamic memory; destructor cleans up all
- Module API auto-creates PalDynamicAllocator (DYNAMIC_UNBOUND just works)

Export changes:
- MarkDynamicUnboundPass marks KV cache buffers as DYNAMIC_UNBOUND
- --lazy_kv_cache flag for Llama export

Co-authored-by: Claude <noreply@anthropic.com>

Author: psiddh

examples/models/llama/export_llama_lib.py

@@ -418,6 +418,15 @@ def build_args_parser() -> argparse.ArgumentParser:
         help="maximum length of context for model to remember",
     )
 
+    parser.add_argument(
+        "--lazy_kv_cache",
+        action="store_true",
+        default=False,
+        help="Mark KV cache buffers as DYNAMIC_UNBOUND so they are allocated "
+        "lazily at runtime instead of at load time. Reduces initial memory "
+        "usage when max_context_length is large.",
+    )
+
     parser.add_argument(
         "--local_global_attention",
         type=parse_list_of_ints,
@@ -1362,6 +1371,13 @@ def _export_llama(llm_config: LlmConfig) -> LLMEdgeManager:  # noqa: C901
     if llm_config.base.model_class.value in TORCHTUNE_DEFINED_MODELS:
         additional_passes = [InitializedMutableBufferPass(["kv_cache_pos"])]
 
+    if llm_config.export.lazy_kv_cache:
+        from executorch.exir.passes.mark_dynamic_unbound_pass import (
+            MarkDynamicUnboundPass,
+        )
+
+        additional_passes.append(MarkDynamicUnboundPass())
+
     # export_to_edge
     builder_manager = _prepare_for_llama_export(llm_config)
     if (

exir/passes/mark_dynamic_unbound_pass.py

@@ -0,0 +1,36 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import List, Optional
+
+from executorch.exir.pass_base import ExportPass
+
+
+class MarkDynamicUnboundPass(ExportPass):
+    """
+    Marks matching placeholder nodes with ``et_dynamic_unbound`` metadata.
+
+    After ``SpecPropPass`` creates ``TensorSpec`` for each placeholder,
+    ``update_placeholder_tensor_specs`` reads this flag and sets the spec's
+    ``shape_dynamism`` to ``DYNAMIC_UNBOUND``.  The memory planner then skips
+    those tensors, and the runtime allocates their memory lazily via
+    ``DynamicAllocator``.
+
+    Typical usage: mark KV cache buffers so they start unallocated and grow
+    on demand, avoiding the full upfront memory cost of max_context_length.
+    """
+
+    def __init__(
+        self,
+        name_patterns: Optional[List[str]] = None,
+    ) -> None:
+        super().__init__()
+        self.name_patterns = name_patterns or ["k_cache", "v_cache"]
+
+    def placeholder(self, name: str, arg, meta):
+        if any(pattern in name for pattern in self.name_patterns):
+            meta["et_dynamic_unbound"] = True
+        return super().placeholder(name, arg, meta)

exir/passes/spec_prop_pass.py

@@ -12,6 +12,7 @@
 import torch
 from executorch.exir.delegate import executorch_call_delegate
 from executorch.exir.pass_base import ExportPass, ProxyValue
+from executorch.exir.schema import TensorShapeDynamism
 from executorch.exir.tensor import TensorSpec
 from torch.export.exported_program import ExportGraphSignature
 from torch.fx.node import Node
@@ -121,3 +122,7 @@ def update_placeholder_tensor_specs(
                 in exported_program.graph_signature.inputs_to_lifted_tensor_constants
             ):
                 spec.const = True
+            if isinstance(spec, TensorSpec) and node.meta.get(
+                "et_dynamic_unbound", False
+            ):
+                spec.shape_dynamism = TensorShapeDynamism.DYNAMIC_UNBOUND

extension/llm/export/config/llm_config.py

@@ -256,6 +256,7 @@ class ExportConfig:
     export_only: bool = False
     foundation_weights_file: Optional[str] = None
     lora_weights_file: Optional[str] = None
+    lazy_kv_cache: bool = False
 
     def __post_init__(self):
         if self.max_context_length < self.max_seq_length:
@@ -695,6 +696,8 @@ def from_args(cls, args: argparse.Namespace) -> "LlmConfig":  # noqa: C901
             llm_config.export.foundation_weights_file = args.foundation_weights_file
         if hasattr(args, "lora_weights_file"):
             llm_config.export.lora_weights_file = args.lora_weights_file
+        if hasattr(args, "lazy_kv_cache"):
+            llm_config.export.lazy_kv_cache = args.lazy_kv_cache
 
         # QuantizationConfig
         if hasattr(args, "quantization_mode"):

extension/module/module.cpp

@@ -13,6 +13,7 @@
 #include <executorch/extension/flat_tensor/flat_tensor_data_map.h>
 #include <executorch/extension/memory_allocator/malloc_memory_allocator.h>
 #include <executorch/extension/named_data_map/merged_data_map.h>
+#include <executorch/runtime/executor/pal_dynamic_allocator.h>
 #include <executorch/runtime/platform/runtime.h>
 
 namespace executorch {
@@ -389,8 +390,13 @@ runtime::Error Module::load_method(
       planned_memory = method_holder.planned_memory->planned_memory.get();
     }
 
+    method_holder.dynamic_allocator =
+        std::make_unique<runtime::PalDynamicAllocator>();
     method_holder.memory_manager = std::make_unique<runtime::MemoryManager>(
-        memory_allocator_.get(), planned_memory, temp_allocator_.get());
+        memory_allocator_.get(),
+        planned_memory,
+        temp_allocator_.get(),
+        method_holder.dynamic_allocator.get());
     auto res_method = program_->load_method(
         method_name.c_str(),
         method_holder.memory_manager.get(),

extension/module/module.h

@@ -14,6 +14,7 @@
 #include <unordered_set>
 #include <vector>
 
+#include <executorch/runtime/executor/dynamic_allocator.h>
 #include <executorch/runtime/executor/program.h>
 
 #ifdef USE_ATEN_LIB
@@ -694,6 +695,7 @@ class Module {
 
   struct MethodHolder {
     std::unique_ptr<PlannedMemory> planned_memory;
+    std::unique_ptr<runtime::DynamicAllocator> dynamic_allocator;
     std::unique_ptr<runtime::MemoryManager> memory_manager;
     std::unique_ptr<Method> method;
   };

extension/module/targets.bzl

@@ -25,6 +25,7 @@ def define_common_targets():
                 "//executorch/extension/data_loader:mmap_data_loader",
                 "//executorch/extension/flat_tensor:flat_tensor_data_map" + aten_suffix,
                 "//executorch/extension/named_data_map:merged_data_map" + aten_suffix,
+                "//executorch/runtime/executor:pal_dynamic_allocator",
             ],
             exported_deps = [
                 "//executorch/runtime/executor:program_no_prim_ops" + aten_suffix,

runtime/core/portable_type/targets.bzl

@@ -41,6 +41,7 @@ def define_common_targets():
             "//executorch/runtime/core/exec_aten/util:scalar_type_util",
             "//executorch/runtime/core/exec_aten/util:tensor_shape_to_c_string",
             "//executorch/runtime/core:tag",
+            "//executorch/runtime/executor:dynamic_allocator",
         ],
     )
 

runtime/core/portable_type/tensor_impl.cpp

@@ -113,12 +113,59 @@ Error TensorImpl::internal_resize_contiguous(ArrayRef<SizesType> new_sizes) {
       }
 
       break;
-    case TensorShapeDynamism::DYNAMIC_BOUND:
-      // TODO(T175194371): Unbounded dynamic tensor resizing is not yet
-      // supported: treat them as upper-bounded.
     case TensorShapeDynamism::DYNAMIC_UNBOUND: {
       const auto new_numel = compute_numel(new_sizes.data(), dim_);
 
+      ET_CHECK_OR_RETURN_ERROR(
+          static_cast<size_t>(new_numel) <= numel_bound_,
+          NotSupported,
+          "Attempted to resize a dynamic unbound tensor beyond its ceiling of %zu elements to %zu elements.",
+          numel_bound_,
+          new_numel);
+
+      const size_t needed_bytes =
+          static_cast<size_t>(new_numel) * elementSize(type_);
+      // If capacity_bytes_ is 0 but data_ is non-null, the buffer is
+      // externally managed (e.g., stack-allocated in tests). Use the
+      // original numel bound as the effective capacity.
+      const size_t effective_capacity = capacity_bytes_ > 0
+          ? capacity_bytes_
+          : static_cast<size_t>(numel_bound_) * elementSize(type_);
+      if (needed_bytes > effective_capacity) {
+        ET_CHECK_OR_RETURN_ERROR(
+            dynamic_allocator_ != nullptr,
+            NotSupported,
+            "DYNAMIC_UNBOUND tensor needs reallocation but has no DynamicAllocator");
+        size_t actual_size = 0;
+        void* new_data = dynamic_allocator_->reallocate(
+            data_,
+            capacity_bytes_,
+            needed_bytes,
+            alignof(std::max_align_t),
+            &actual_size);
+        ET_CHECK_OR_RETURN_ERROR(
+            new_data != nullptr,
+            MemoryAllocationFailed,
+            "Failed to reallocate DYNAMIC_UNBOUND tensor to %zu bytes",
+            needed_bytes);
+        data_ = new_data;
+        capacity_bytes_ = actual_size;
+      }
+
+      if (strides_ && dim_order_) {
+        auto error =
+            dim_order_to_stride(new_sizes.data(), dim_order_, dim_, strides_);
+        if (error != Error::Ok) {
+          return error;
+        }
+      }
+      numel_ = new_numel;
+      std::copy(new_sizes.begin(), new_sizes.end(), sizes_);
+    } break;
+
+    case TensorShapeDynamism::DYNAMIC_BOUND: {
+      const auto new_numel = compute_numel(new_sizes.data(), dim_);
+
       ET_CHECK_OR_RETURN_ERROR(
           static_cast<size_t>(new_numel) <= numel_bound_,
           NotSupported,

runtime/core/portable_type/tensor_impl.h

@@ -12,6 +12,7 @@
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/portable_type/scalar_type.h>
 #include <executorch/runtime/core/tensor_shape_dynamism.h>
+#include <executorch/runtime/executor/dynamic_allocator.h>
 
 // Forward declaration of a helper that provides access to internal resizing
 // methods of TensorImpl. Real definition is in
@@ -203,6 +204,26 @@ class TensorImpl {
     data_ = ptr;
   }
 
+  /// Returns the dynamic allocator for DYNAMIC_UNBOUND tensors, or nullptr.
+  DynamicAllocator* dynamic_allocator() const {
+    return dynamic_allocator_;
+  }
+
+  /// Sets the dynamic allocator for lazy allocation.
+  void set_dynamic_allocator(DynamicAllocator* allocator) {
+    dynamic_allocator_ = allocator;
+  }
+
+  /// Returns the capacity in bytes of the current dynamic allocation.
+  size_t capacity_bytes() const {
+    return capacity_bytes_;
+  }
+
+  /// Sets the capacity in bytes of the current dynamic allocation.
+  void set_capacity_bytes(size_t capacity) {
+    capacity_bytes_ = capacity;
+  }
+
   /*
    * DEPRECATED: Use torch::executor::resize_tensor() or
    * torch::executor::resize_tensor_impl().
@@ -261,6 +282,13 @@ class TensorImpl {
 
   /// Specifies the mutability of the shape of the tensor.
   const TensorShapeDynamism shape_dynamism_;
+
+  /// Allocator for DYNAMIC_UNBOUND tensors. nullptr for other dynamism types.
+  DynamicAllocator* dynamic_allocator_ = nullptr;
+
+  /// Capacity in bytes of the buffer pointed to by data_, when managed by
+  /// dynamic_allocator_. 0 means no allocation yet.
+  size_t capacity_bytes_ = 0;
 };
 
 /**