Write HF static KV cache rows at the cache_position input

With transformers 5.11.0, StaticLayer.update ignores cache_position and writes
K/V at arange(n) + cumulative_length. Setting cumulative_length from
cache_position keeps eager results right, but the exported cache writes then
index rows with the output of an add. The ggml engine only accepts index_put_
row indices that come from a graph input or an integer arange, so
updated codes

Signed-off-by: jiseong.oh <jiseong.oh@samsung.com>

Author: Jiseong-oh

backends/arm/test/modules/test_moe.py

@@ -171,8 +171,8 @@ def _make_gpt_oss_moe(dtype: torch.dtype) -> tuple[GptOssMLP, GptOssConfig]:
         intermediate_size=16,
         num_local_experts=8,
         num_experts_per_tok=4,
-        experts_implementation="eager",
     )
+    config._experts_implementation = "eager"
     model = GptOssMLP(config)
     model.experts = TensorizedGptOssExperts(config)
     generator = torch.Generator().manual_seed(2648)

backends/native/examples/audio/export_hf_whisper.py

@@ -40,7 +40,10 @@
 
 import torch
 import torch.nn as nn
-from executorch.backends.native.examples.hf_llm.hf_utils import report_lowering
+from executorch.backends.native.examples.hf_llm.hf_utils import (
+    PositionedStaticLayer,
+    report_lowering,
+)
 from executorch.exir.native import to_native
 from torch.export import Dim, export, ExportedProgram
 from transformers import (
@@ -174,6 +177,8 @@ def __init__(self, model: WhisperForConditionalGeneration) -> None:
             dtype=torch.float32,
             device=torch.device("cpu"),
         )
+        for i, old in enumerate(self.self_cache.layers):
+            self.self_cache.layers[i] = PositionedStaticLayer.replacing(old, old.dtype)
         self.cross_cache = DynamicCache(config=config)
         self.cache = EncoderDecoderCache(self.self_cache, self.cross_cache)
         for layer_idx in range(config.decoder_layers):
@@ -195,8 +200,10 @@ def forward(
         cache_position: torch.Tensor,
         *cross_kv: torch.Tensor,
     ) -> torch.Tensor:
+        # transformers 5 takes positions from the cache's length, not cache_position.
         for layer in self.self_cache.layers:
             layer.cumulative_length = cache_position[:1].clone()
+            layer.positions = cache_position
         num_layers = len(self.cross_cache.layers)
         for layer_idx, layer in enumerate(self.cross_cache.layers):
             layer.keys = cross_kv[layer_idx]

backends/native/examples/hf_llm/BUCK

@@ -12,6 +12,7 @@ fbcode_target(
         "//caffe2:torch",
         "//executorch/backends/native:lib",
         "//executorch/exir:lib",
+        "fbsource//third-party/pypi/transformers:transformers",
     ],
 )
 

backends/native/examples/hf_llm/export_hf_llm.py

@@ -25,7 +25,10 @@
     gguf_architecture,
     gguf_to_hf,
 )
-from executorch.backends.native.examples.hf_llm.hf_utils import report_lowering
+from executorch.backends.native.examples.hf_llm.hf_utils import (
+    PositionedStaticLayer,
+    report_lowering,
+)
 from executorch.exir.native import Call, Sample, SampleSet, to_native, write_samples
 from executorch.extension.llm.export.load import assign_one, iter_checkpoint, RawIter
 from torch.export import Dim, export
@@ -118,15 +121,6 @@ def load_gguf(
 _CACHE_DTYPE_SDPA = "sdpa_in_cache_dtype"
 
 
-class _CastingStaticLayer(StaticLayer):
-    """A static cache layer that casts K/V writes to its own dtype."""
-
-    def update(self, key_states, value_states, cache_kwargs=None):
-        return super().update(
-            key_states.to(self.dtype), value_states.to(self.dtype), cache_kwargs
-        )
-
-
 def _sdpa_in_cache_dtype(module, query, key, value, attention_mask, **kwargs):
     output, weights = sdpa_attention_forward(
         module, query.to(key.dtype), key, value, attention_mask, **kwargs
@@ -155,22 +149,29 @@ def __init__(
             if hasattr(self.cached, "static_cache")
             else self.cached.cache
         )
-        if kv_cache_dtype != model.dtype:
-            self._store_kv_as(model, kv_cache_dtype)
+        change_dtype = kv_cache_dtype != model.dtype
+        self._write_cache_at_input_positions(kv_cache_dtype, change_dtype)
+        if change_dtype:
+            self._attend_in_cache_dtype(model)
 
-    def _store_kv_as(self, model: PreTrainedModel, dtype: torch.dtype) -> None:
+    def _write_cache_at_input_positions(
+        self, dtype: torch.dtype, change_dtype: bool
+    ) -> None:
         for i, old in enumerate(self.cache.layers):
             if type(old) is not StaticLayer:
-                raise ValueError(f"cannot change the KV dtype of {type(old).__name__}")
-            layer = _CastingStaticLayer(old.max_cache_len)
-            layer.lazy_initialization(
-                old.keys[:, :, :1].to(dtype), old.values[:, :, :1].to(dtype)
-            )
+                if change_dtype:
+                    raise ValueError(
+                        f"cannot change the KV dtype of {type(old).__name__}"
+                    )
+                continue
+            layer = PositionedStaticLayer.replacing(old, dtype)
             self.cache.layers[i] = layer
             self.cached.register_buffer(f"key_cache_{i}", layer.keys, persistent=False)
             self.cached.register_buffer(
                 f"value_cache_{i}", layer.values, persistent=False
             )
+
+    def _attend_in_cache_dtype(self, model: PreTrainedModel) -> None:
         # Keep the vmap-free mask TorchExportableModuleForDecoderOnlyLM selected.
         mask = ALL_MASK_ATTENTION_FUNCTIONS[model.config._attn_implementation]
         AttentionInterface.register(_CACHE_DTYPE_SDPA, _sdpa_in_cache_dtype)
@@ -180,9 +181,11 @@ def _store_kv_as(self, model: PreTrainedModel, dtype: torch.dtype) -> None:
     def forward(
         self, input_ids: torch.Tensor, cache_position: torch.Tensor
     ) -> torch.Tensor:
-        # StaticLayer writes at cumulative_length and ignores cache_position.
+        # transformers 5 takes positions from the cache's length, not cache_position.
         for layer in self.cache.layers:
             layer.cumulative_length = cache_position[:1].clone()
+            if isinstance(layer, PositionedStaticLayer):
+                layer.positions = cache_position
         return self.cached.model(
             input_ids=input_ids,
             cache_position=cache_position,

backends/native/examples/hf_llm/hf_utils.py

@@ -19,6 +19,31 @@
     get_lowered_submodules,
 )
 from torch.export import ExportedProgram
+from transformers.cache_utils import StaticLayer
+
+
+class PositionedStaticLayer(StaticLayer):
+    """A static cache layer that writes K/V, cast to its own dtype, at ``positions``.
+
+    transformers 5 writes at ``arange(n) + cumulative_length``. The native runtime
+    only accepts row indices that are a graph input or an arange, so the positions
+    are the ``cache_position`` input of the graph instead.
+    """
+
+    positions: torch.Tensor
+
+    @classmethod
+    def replacing(cls, old: StaticLayer, dtype: torch.dtype) -> "PositionedStaticLayer":
+        layer = cls(old.max_cache_len)
+        layer.lazy_initialization(
+            old.keys[:, :, :1].to(dtype), old.values[:, :, :1].to(dtype)
+        )
+        return layer
+
+    def update(self, key_states, value_states, *args, **kwargs):
+        self.keys.index_copy_(2, self.positions, key_states.to(self.dtype))
+        self.values.index_copy_(2, self.positions, value_states.to(self.dtype))
+        return self.keys, self.values
 
 
 def histogram(gm: torch.fx.GraphModule) -> collections.Counter:

backends/native/examples/hf_llm/test_export_hf_llm.py

@@ -316,6 +316,44 @@ def test_golden_samples_tolerance_follows_weight_quantization(self) -> None:
             samples = golden_samples(LastTokenLM(model, 16).eval(), [[3, 4]], 0)
             self.assertEqual(samples.tol, tol)
 
-    def test_fp16_kv_cache_rejects_sliding_window_layers(self) -> None:
-        with self.assertRaisesRegex(ValueError, "cannot change the KV dtype"):
-            LastTokenLM(_tiny("gemma3"), 32, torch.float16)
+    def test_cache_writes_index_the_cache_position_input(self) -> None:
+        seq = torch.export.Dim("seq", min=1, max=31)
+        for arch, kv_dtype in (("llama", torch.float32), ("gemma3", torch.float16)):
+            with self.subTest(arch=arch, kv_dtype=kv_dtype):
+                torch.manual_seed(7)
+                model = _tiny(arch)
+                fp32 = LastTokenLM(model, 32).eval()
+                wrapper = LastTokenLM(copy.deepcopy(model), 32, kv_dtype).eval()
+                ids = torch.tensor([[3, 4, 5, 6]])
+                with torch.no_grad():
+                    ep = torch.export.export(
+                        wrapper,
+                        (ids, torch.arange(4)),
+                        dynamic_shapes={
+                            "input_ids": {1: seq},
+                            "cache_position": {0: seq},
+                        },
+                        strict=True,
+                    )
+                    exported = ep.module()
+                    for step_ids, positions in (
+                        (ids, torch.arange(4)),
+                        (torch.tensor([[9]]), torch.tensor([4])),
+                        (ids, torch.arange(4)),
+                    ):
+                        self.assertTrue(
+                            torch.allclose(
+                                fp32(step_ids, positions),
+                                exported(step_ids, positions),
+                                atol=1e-2,
+                                rtol=2e-2,
+                            )
+                        )
+                writes = [
+                    n
+                    for n in ep.graph.nodes
+                    if n.target == torch.ops.aten.index_copy_.default
+                ]
+                self.assertEqual(len(writes), 2 * model.config.num_hidden_layers)
+                for n in writes:
+                    self.assertEqual(n.args[2].name, "cache_position")