DYNAMIC_UNBOUND support for portable runtime: lazy KV cache allocation

CMakeLists.txt

@@ -215,6 +215,9 @@ endif()
 
 # Disable the deprecated constant_buffer path.
 add_definitions(-DET_ENABLE_DEPRECATED_CONSTANT_BUFFER=0)
+if(EXECUTORCH_ENABLE_DYNAMIC_ALLOCATOR)
+  add_definitions(-DET_DYNAMIC_ALLOCATOR_ENABLED)
+endif()
 
 if(EXECUTORCH_ENABLE_EVENT_TRACER)
   add_definitions(-DET_EVENT_TRACER_ENABLED)

examples/models/llama/export_llama_lib.py

@@ -422,6 +422,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,
@@ -808,6 +817,7 @@ def _prepare_for_llama_export(llm_config: LlmConfig) -> LLMEdgeManager:
             use_torchao_kernels_linear=llm_config.backend.torchao.use_torchao_kernels_linear,
             use_torchao_kernels_tied_embedding=llm_config.backend.torchao.use_torchao_kernels_tied_embedding,
             quantize_with_hqq=llm_config.quantization.use_hqq,
+            lazy_kv_cache=llm_config.export.lazy_kv_cache,
         )
     )
 
@@ -1430,6 +1440,12 @@ def _export_llama(llm_config: LlmConfig) -> LLMEdgeManager:  # noqa: C901
     # which corrupts the causal mask computation.
     if llm_config.model.use_attention_sink:
         additional_passes.append(InitializedMutableBufferPass(["cache_positions"]))
+    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)
@@ -1699,6 +1715,7 @@ def _get_source_transforms(  # noqa
     use_torchao_kernels_linear: bool = False,
     use_torchao_kernels_tied_embedding: bool = False,
     quantize_with_hqq: bool = True,
+    lazy_kv_cache: bool = False,
 ) -> List[Callable[[torch.nn.Module], torch.nn.Module]]:
     """
     Return a list of functions that transform a graph.
@@ -1825,6 +1842,24 @@ def _get_source_transforms(  # noqa
 
     if quantize_kv_cache:
         assert use_kv_cache, "quantize_kv_cache requires use_kv_cache=True"
+        # TODO: --lazy_kv_cache + --quantize_kv_cache is not yet
+        # supported. QuantizedKVCache.from_float reads the source cache shape
+        # to determine max_context_length; with lazy=True the shape is
+        # [B, 0, H, D], so it creates a permanently-zero-sized quantized cache
+        # that silently produces wrong results. Proper support requires:
+        #   1. Plumbing `lazy` into QuantizedKVCache (moderate: ~100 LOC across
+        #      custom_kv_cache.py and quantized_sdpa.py).
+        #   2. Making QuantizedKVCache register its k/v buffers with seq_len=0
+        #      and matching name patterns so MarkDynamicUnboundPass picks them up.
+        #   3. Ensuring the quantized update_cache C++ op also calls
+        #      maybe_resize_cache (or equivalent) to grow before dequant+write.
+        # Estimated effort: a few days of work, mostly in the quantized SDPA
+        # kernel which packs/unpacks int4 cache values in a layout-specific way
+        # that complicates in-place growth.
+        assert not lazy_kv_cache, (
+            "--lazy_kv_cache and --quantize_kv_cache cannot be used together yet. "
+            "QuantizedKVCache does not support DYNAMIC_UNBOUND lazy allocation."
+        )
         transforms.append(replace_kv_cache_with_quantized_kv_cache)
         # Right now
         transforms.append(replace_sdpa_with_quantized_sdpa)

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

@@ -13,6 +13,7 @@
 from executorch.exir.delegate import executorch_call_delegate
 from executorch.exir.dialects._ops import ops as exir_ops
 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
@@ -130,3 +131,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/custom_ops/custom_ops.py

@@ -207,16 +207,11 @@ def _validate_update_cache_params(
         ), f"Expected value and cache to have same size in dimension {i} but got {value.size(i)} and {cache.size(i)}"
 
     torch._check_is_size(start_pos)
-    if indices is None:
+    # Bounds checks are skipped for DYNAMIC_UNBOUND (lazy) caches where
+    # cache.size(1) starts at 0 and grows at runtime.
+    if indices is None and cache.size(1) > 0:
         torch._check(start_pos < cache.size(1))
-        assert start_pos < cache.size(
-            1
-        ), f"Start position {start_pos} must be less than sequence length {cache.size(1)}"
-
         torch._check((start_pos + seq_len) <= cache.size(1))
-        assert (start_pos + seq_len) <= cache.size(
-            1
-        ), f"Start position  + length = {start_pos + seq_len} must be less than sequence length {cache.size(1)}"
 
     if indices is not None:
         assert (

extension/llm/custom_ops/op_sdpa.cpp

@@ -593,6 +593,14 @@ Tensor& sdpa_with_kv_cache_out(
 
   ET_CHECK_MSG(q_projected.dim() == 4, "query must be a 4D tensor");
 
+  // For DYNAMIC_UNBOUND lazy KV cache: allocate if data is null.
+  if (key_cache.const_data_ptr() == nullptr) {
+    resize_tensor(key_cache, key_cache.sizes());
+  }
+  if (value_cache.const_data_ptr() == nullptr) {
+    resize_tensor(value_cache, value_cache.sizes());
+  }
+
   update_cache(k_projected, key_cache, start_pos, seq_len);
   update_cache(v_projected, value_cache, start_pos, seq_len);
 

extension/llm/custom_ops/op_update_cache.cpp

@@ -11,6 +11,7 @@
 #include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
 // @lint-ignore CLANGTIDY facebook-unused-include-check
 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
+#include <executorch/runtime/core/exec_aten/util/tensor_util.h>
 
 #include <executorch/extension/kernel_util/make_boxed_from_unboxed_functor.h>
 
@@ -207,6 +208,35 @@ Tensor& update_cache_impl(
 }
 } // anonymous namespace
 
+// Grow cache seq dimension if needed (for DYNAMIC_UNBOUND lazy KV cache).
+static bool maybe_resize_cache(
+    RuntimeContext& ctx,
+    const Tensor& value,
+    Tensor& cache,
+    int64_t start_pos) {
+  ET_CHECK_OR_RETURN_FALSE(cache.dim() == 4, "cache must be a 4D tensor");
+  ET_CHECK_OR_RETURN_FALSE(value.dim() == 4, "value must be a 4D tensor");
+  int64_t seq_len = value.size(1);
+  int64_t required_seq = start_pos + seq_len;
+  // Resize if cache is too small OR if data hasn't been allocated yet
+  // (lazy DYNAMIC_UNBOUND cache starts with capacity_bytes=0 and null data).
+  if (required_seq > cache.size(1) || cache.const_data_ptr() == nullptr) {
+    int64_t new_seq = std::max(required_seq, cache.size(1));
+    executorch::aten::SizesType new_sizes[] = {
+        static_cast<executorch::aten::SizesType>(cache.size(0)),
+        static_cast<executorch::aten::SizesType>(new_seq),
+        static_cast<executorch::aten::SizesType>(cache.size(2)),
+        static_cast<executorch::aten::SizesType>(cache.size(3)),
+    };
+    auto err =
+        resize_tensor(cache, {new_sizes, static_cast<size_t>(cache.dim())});
+    if (err != Error::Ok) {
+      return false;
+    }
+  }
+  return true;
+}
+
 // Original update_cache_out function without indices parameter
 Tensor& update_cache_out(
     RuntimeContext& ctx,
@@ -215,6 +245,11 @@ Tensor& update_cache_out(
     const int64_t start_pos,
     Tensor& output) {
   int64_t seq_len = value.size(1);
+  ET_KERNEL_CHECK(
+      ctx,
+      maybe_resize_cache(ctx, value, cache, start_pos),
+      InvalidArgument,
+      output);
   ET_KERNEL_CHECK(
       ctx,
       validate_cache_params(value, cache, start_pos, seq_len),
@@ -233,6 +268,11 @@ Tensor& update_cache_with_indices_out(
     const Tensor& indices,
     Tensor& output) {
   int64_t seq_len = value.size(1);
+  ET_KERNEL_CHECK(
+      ctx,
+      maybe_resize_cache(ctx, value, cache, start_pos),
+      InvalidArgument,
+      output);
   ET_KERNEL_CHECK(
       ctx,
       validate_cache_params(value, cache, start_pos, seq_len, indices),

extension/llm/custom_ops/test_lazy_kv_cache.py

@@ -0,0 +1,141 @@
+# 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.
+
+"""
+Tests for lazy KV cache (DYNAMIC_UNBOUND) support.
+
+Tests the update_cache custom op's ability to handle caches that start at
+seq_len=0 and grow on demand, which is the foundation for pay-as-you-go
+KV cache memory allocation.
+"""
+
+# pyre-unsafe
+
+import unittest
+
+import torch
+
+from executorch.extension.llm.custom_ops import custom_ops  # noqa
+
+
+class LazyKVCacheUpdateTest(unittest.TestCase):
+    """Test update_cache op with zero-sized initial caches (lazy KV cache)."""
+
+    def setUp(self):
+        torch.manual_seed(42)
+        self.batch_size = 1
+        self.num_heads = 4
+        self.head_dim = 8
+        self.max_seq_len = 64
+
+    def test_update_cache_grows_from_zero(self):
+        """Verify update_cache works when cache seq dim starts at full size
+        and tokens are appended sequentially."""
+        cache = torch.zeros(
+            (self.batch_size, self.max_seq_len, self.num_heads, self.head_dim),
+            dtype=torch.float32,
+        )
+
+        for pos in range(10):
+            value = torch.randn(
+                (self.batch_size, 1, self.num_heads, self.head_dim),
+                dtype=torch.float32,
+            )
+            torch.ops.llama.update_cache(value, cache, pos)
+            self.assertTrue(
+                torch.allclose(cache[:, pos : pos + 1, :, :], value),
+                f"Cache mismatch at position {pos}",
+            )
+
+    def test_custom_kv_cache_lazy_init(self):
+        """Verify CustomKVCache with lazy=True creates zero-sized buffers."""
+        from executorch.examples.models.llama.source_transformation.custom_kv_cache import (
+            CustomKVCache,
+        )
+
+        cache = CustomKVCache(
+            max_batch_size=1,
+            max_context_length=131072,  # 128K ceiling
+            n_heads=4,
+            head_dim=8,
+            dtype=torch.float32,
+            lazy=True,
+        )
+        self.assertEqual(cache.k_cache.shape[1], 0, "Lazy k_cache seq dim should be 0")
+        self.assertEqual(cache.v_cache.shape[1], 0, "Lazy v_cache seq dim should be 0")
+        self.assertEqual(cache.max_context_length, 131072)
+
+    def test_custom_kv_cache_non_lazy_init(self):
+        """Verify CustomKVCache without lazy=True creates full-sized buffers."""
+        from executorch.examples.models.llama.source_transformation.custom_kv_cache import (
+            CustomKVCache,
+        )
+
+        cache = CustomKVCache(
+            max_batch_size=1,
+            max_context_length=64,
+            n_heads=4,
+            head_dim=8,
+            dtype=torch.float32,
+            lazy=False,
+        )
+        self.assertEqual(cache.k_cache.shape[1], 64)
+        self.assertEqual(cache.v_cache.shape[1], 64)
+
+    def test_replace_kv_cache_with_lazy(self):
+        """Verify replace_kv_cache_with_custom_kv_cache passes lazy flag."""
+        from executorch.examples.models.llama.attention import KVCache
+        from executorch.examples.models.llama.source_transformation.custom_kv_cache import (
+            CustomKVCache,
+            replace_kv_cache_with_custom_kv_cache,
+        )
+
+        class FakeModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                # KVCache stores as [B, H, S, D]
+                self.kv = KVCache(
+                    max_batch_size=1,
+                    max_context_length=128,
+                    n_heads=4,
+                    head_dim=8,
+                    enable_dynamic_shape=False,
+                    dtype=torch.float32,
+                )
+
+            def forward(self, x):
+                return x
+
+        model = FakeModel()
+        replace_kv_cache_with_custom_kv_cache(model, lazy=True)
+        self.assertIsInstance(model.kv, CustomKVCache)
+        # CustomKVCache stores as [B, S, H, D], lazy means seq_dim=0
+        self.assertEqual(model.kv.k_cache.shape[1], 0)
+
+
+class LazyKVCacheMetaKernelTest(unittest.TestCase):
+    """Test that meta kernels work without upper-bound cache size checks."""
+
+    def test_meta_kernel_allows_start_pos_beyond_cache(self):
+        """Meta kernel should not reject start_pos >= cache.size(1)."""
+        value = torch.randn(1, 1, 4, 8)
+        # Cache with seq_len=0 (lazy)
+        cache = torch.zeros(1, 0, 4, 8)
+        # This should not raise — the runtime op handles resize
+        result = torch.ops.llama.update_cache(value, cache, 0)
+        self.assertIsNotNone(result)
+
+    def test_meta_kernel_allows_large_start_pos(self):
+        """Meta kernel should allow start_pos beyond current cache size for lazy caches."""
+        value = torch.randn(1, 1, 4, 8)
+        cache = torch.zeros(1, 0, 4, 8)
+        # Lazy cache (size(1)==0) skips bounds checks — runtime op handles resize
+        result = torch.ops.llama.update_cache(value, cache, 100)
+        self.assertIsNotNone(result)
+
+
+if __name__ == "__main__":
+    unittest.main()