Fix GenAI regression: disable past_present_share_buffer for dual head_dim

Models with different head_dim across layers (e.g., Gemma4 with
head_dim=256 for sliding and global_head_dim=512 for full-attention)
cannot use shared KV buffers. Dynamically detect dual head_dim from
config and set past_present_share_buffer=false.

Implementation:
- GenaiConfigGenerator: add _search_overrides dict applied in generate()
- auto_export: check config.global_head_dim != config.head_dim

134 tests pass (15 gemma4 + 119 ort_genai).

Signed-off-by: Justin Chu <justinchu@microsoft.com>

Author: justinchuby

src/mobius/integrations/ort_genai/auto_export.py

@@ -655,6 +655,14 @@ def _write_genai_config(
             **audio_kwargs,
         )
 
+    # Disable past_present_share_buffer for models with dual head_dim
+    # (e.g., different head_dim for sliding vs full-attention layers).
+    # Shared buffers require uniform head_dim across all KV cache layers.
+    global_head_dim = getattr(config, "global_head_dim", None)
+    head_dim = getattr(config, "head_dim", None)
+    if global_head_dim and head_dim and global_head_dim != head_dim:
+        generator._search_overrides["past_present_share_buffer"] = False
+
     return generator.write(output_dir)
 
 

src/mobius/integrations/ort_genai/genai_config.py

@@ -178,6 +178,9 @@ def __init__(
         # Optional audio fields (set via with_audio())
         self._audio: dict[str, Any] | None = None
 
+        # Search config overrides applied in generate()
+        self._search_overrides: dict[str, Any] = {}
+
     @classmethod
     def from_config(
         cls,
@@ -407,9 +410,12 @@ def generate(self) -> dict[str, Any]:
             model["speech"] = self._audio
         model.update(self._vlm_token_ids)
 
+        search = _default_search_params(ep=self.ep, context_length=self.context_length)
+        search.update(self._search_overrides)
+
         return {
             "model": model,
-            "search": _default_search_params(ep=self.ep, context_length=self.context_length),
+            "search": search,
         }
 
     def write(self, output_dir: str) -> str: