[KV cache] support-past-present-share-buffer (#268)

## Support past present share buffer

Add `supports_past_present_share_buffer=True`

Benefit: 

Save the `memcpy` overhead.

Author: rui-ren

src/mobius/_execution_providers.py

@@ -71,12 +71,21 @@ class EpCapabilities:
             quantization (0 = highest accuracy, 4 = fastest).
         provider_options: Default ORT GenAI provider options dict for this EP.
         enable_graph_capture: Whether this EP defaults to GPU graph capture.
-        supports_past_present_share_buffer: Whether this EP requires past and present
-            KV-cache tensors to share the same pre-allocated backing buffer.
-            ``True`` for WebGPU, which allocates the full KV-cache at model
-            load time and maps both past and present views into it.  ``False``
-            for all other EPs where the runtime manages KV-cache memory
-            dynamically.
+        supports_past_present_share_buffer: Whether past and present KV-cache
+            tensors alias the same pre-allocated buffer.  When ``True``, the
+            ORT GenAI runtime allocates a single KV-cache buffer at model load
+            and maps both past and present as views into it, avoiding a
+            per-step copy.  This is the recommended setting for every EP that
+            supports ``GroupQueryAttention`` (CPU, CUDA, DML, WebGPU,
+            TRT-RTX).  Set to ``False`` only for EPs that do not support GQA
+            or cannot handle aliased KV-cache buffers.
+        cap_kv_buffer_max_length: When ``True`` **and**
+            ``supports_past_present_share_buffer`` is also ``True``, the
+            generated ``max_length`` in genai_config is capped to avoid
+            pre-allocating huge KV-cache buffers on memory-constrained
+            devices.  ``True`` only for WebGPU (consumer GPU); ``False`` for
+            CUDA / CPU / DML / TRT-RTX where the runtime can handle large
+            pre-allocations.
     """
 
     name: str
@@ -90,6 +99,7 @@ class EpCapabilities:
     provider_options: dict[str, str] = dataclasses.field(default_factory=dict)
     enable_graph_capture: bool = False
     supports_past_present_share_buffer: bool = False
+    cap_kv_buffer_max_length: bool = False
 
     def __post_init__(self) -> None:
         if not self.supports_fused_rope and self.qkv_pack_dtypes:
@@ -98,6 +108,13 @@ def __post_init__(self) -> None:
                 f"supports_fused_rope=False — UnpackQKV lowering always fires for "
                 f"this EP, so packing would be immediately undone."
             )
+        if self.cap_kv_buffer_max_length and not self.supports_past_present_share_buffer:
+            raise ValueError(
+                f"EP '{self.name}': cap_kv_buffer_max_length=True requires "
+                f"supports_past_present_share_buffer=True — the cap only matters "
+                f"when the runtime pre-allocates the full KV-cache buffer at load "
+                f"time, which is what buffer sharing enables."
+            )
 
 
 class EpRegistry:
@@ -199,6 +216,7 @@ def _register_builtins() -> None:
             gqa_dtypes=frozenset({ir.DataType.FLOAT}),
             qkv_pack_dtypes=frozenset({ir.DataType.FLOAT}),
             default_int4_accuracy_level=4,
+            supports_past_present_share_buffer=True,
         ),
         EpCapabilities(
             name="cuda",
@@ -211,6 +229,7 @@ def _register_builtins() -> None:
                 "enable_cuda_graph": "0",
                 "enable_skip_layer_norm_strict_mode": "1",
             },
+            supports_past_present_share_buffer=True,
         ),
         EpCapabilities(
             name="dml",
@@ -221,6 +240,7 @@ def _register_builtins() -> None:
             qkv_pack_dtypes=frozenset(),
             supports_packed_multi_head_attention=True,
             supports_fused_rope=False,
+            supports_past_present_share_buffer=True,
         ),
         EpCapabilities(
             name="webgpu",
@@ -229,6 +249,7 @@ def _register_builtins() -> None:
             default_int4_accuracy_level=4,
             provider_options={"enableGraphCapture": "0", "validationMode": "basic"},
             supports_past_present_share_buffer=True,
+            cap_kv_buffer_max_length=True,
         ),
         EpCapabilities(
             name="trt-rtx",
@@ -239,12 +260,15 @@ def _register_builtins() -> None:
             supports_skip_layer_norm=False,
             enable_graph_capture=True,
             provider_options={"enable_cuda_graph": "1"},
+            supports_past_present_share_buffer=True,
         ),
         # onnx-standard: ONNX-only runtime — emits zero custom-domain ops.
         # All com.microsoft ops (SkipLayerNorm, PackedMHA) are expanded via
         # InlinePass to their standard-ONNX function bodies. No GQA or QKV
         # packing fusion is applied. Use this EP to produce models that run
         # on any conformant ONNX runtime without ORT extensions.
+        # KV buffer sharing is unsupported here: GQA isn't emitted, so
+        # standard Attention's concat-grow semantics handle the cache.
         EpCapabilities(
             name="onnx-standard",
             gqa_dtypes=frozenset(),  # no GroupQueryAttention

src/mobius/integrations/ort_genai/genai_config.py

@@ -61,11 +61,14 @@ def _default_search_params(*, ep: str, context_length: int) -> dict[str, Any]:
 
     caps = ep_registry.get(ep)
     share_buffer = caps.supports_past_present_share_buffer if caps is not None else False
-    if share_buffer:
-        # EPs that pre-allocate KV-cache for the full max_length at load time
-        # (e.g. WebGPU) need a capped default to avoid pre-allocating huge
-        # buffers (~8 GB for 128K-token models) on consumer hardware.  Users
-        # can raise the limit in genai_config.json for their target device.
+    cap_length = caps.cap_kv_buffer_max_length if caps is not None else False
+    if share_buffer and cap_length:
+        # Memory-constrained EPs (e.g. WebGPU on consumer GPUs) pre-allocate
+        # KV-cache for the full max_length at load time.  Cap the default to
+        # avoid pre-allocating huge buffers (~8 GB for 128K-token models).
+        # Users can raise the limit in genai_config.json for their device.
+        # The cap only applies when buffer sharing is also active — without
+        # sharing, the runtime grows the cache on demand and no cap is needed.
         max_length = min(context_length, _SHARE_BUFFER_MAX_LENGTH_CAP)
     else:
         max_length = context_length

src/mobius/integrations/ort_genai/genai_config_test.py

@@ -136,8 +136,8 @@ def test_search_params_defaults(self):
         assert search["top_p"] == pytest.approx(1.0)
         # max_length tracks the model's context window
         assert search["max_length"] == 8192
-        # CPU does not share past/present KV buffers
-        assert search["past_present_share_buffer"] is False
+        # CPU shares past/present KV buffers (all GQA-capable EPs do)
+        assert search["past_present_share_buffer"] is True
 
     def test_search_params_webgpu_sets_past_present_share_buffer(self):
         """WebGPU EP sets supports_past_present_share_buffer=True via EpCapabilities and caps max_length."""
@@ -174,8 +174,13 @@ def test_search_params_webgpu_small_context_not_capped(self):
         assert config["search"]["past_present_share_buffer"] is True
         assert config["search"]["max_length"] == 2048
 
-    def test_search_params_cuda_does_not_share_buffer(self):
-        """CUDA EP does not set past_present_share_buffer by default."""
+    def test_search_params_cuda_shares_buffer(self):
+        """CUDA EP sets past_present_share_buffer=True (all GQA-capable EPs do).
+
+        CUDA does NOT cap max_length — only memory-constrained EPs (WebGPU)
+        set ``cap_kv_buffer_max_length=True``.  Server-class GPUs handle
+        large pre-allocations.
+        """
         gen = GenaiConfigGenerator(
             "llama",
             vocab_size=32000,
@@ -185,15 +190,19 @@ def test_search_params_cuda_does_not_share_buffer(self):
             num_key_value_heads=8,
             head_dim=128,
             ep="cuda",
+            context_length=131072,
         )
         config = gen.generate()
-        assert config["search"]["past_present_share_buffer"] is False
+        assert config["search"]["past_present_share_buffer"] is True
+        # CUDA: full context_length, NOT capped at 4096
+        assert config["search"]["max_length"] == 131072
 
     def test_search_params_custom_ep_with_share_buffer(self):
         """A custom EP registered with supports_past_present_share_buffer=True gets the flag set.
 
         This proves the value comes from EpCapabilities, not from a hardcoded
-        'ep == webgpu' check.
+        'ep == webgpu' check.  The custom EP does NOT set
+        ``cap_kv_buffer_max_length``, so max_length is the full context.
         """
         from mobius._execution_providers import EpCapabilities, ep_registry
 
@@ -215,12 +224,43 @@ def test_search_params_custom_ep_with_share_buffer(self):
             )
             config = gen.generate()
             assert config["search"]["past_present_share_buffer"] is True
-            # Buffer-sharing EP: max_length capped at 4096
-            assert config["search"]["max_length"] == 4096
+            # No cap_kv_buffer_max_length: max_length = full context_length
+            assert config["search"]["max_length"] == 8192
         finally:
             # Clean up the test EP so it doesn't bleed into other tests
             ep_registry._entries.pop("test-custom-ep", None)
 
+    def test_search_params_custom_ep_with_max_length_cap(self):
+        """A custom EP with cap_kv_buffer_max_length=True caps max_length at 4096."""
+        from mobius._execution_providers import EpCapabilities, ep_registry
+
+        ep_registry.register(
+            EpCapabilities(
+                name="test-capped-ep",
+                supports_past_present_share_buffer=True,
+                cap_kv_buffer_max_length=True,
+            ),
+            overwrite=True,
+        )
+        try:
+            gen = GenaiConfigGenerator(
+                "llama",
+                vocab_size=32000,
+                hidden_size=4096,
+                num_hidden_layers=32,
+                num_attention_heads=32,
+                num_key_value_heads=8,
+                head_dim=128,
+                ep="test-capped-ep",
+                context_length=131072,
+            )
+            config = gen.generate()
+            assert config["search"]["past_present_share_buffer"] is True
+            # cap_kv_buffer_max_length=True: max_length capped at 4096
+            assert config["search"]["max_length"] == 4096
+        finally:
+            ep_registry._entries.pop("test-capped-ep", None)
+
     def test_session_options_present(self):
         """Decoder has session_options with log_id."""
         gen = GenaiConfigGenerator(