Remove kv_shared from ModelConfig, and pass it as layer input.

haozha111 · copybara-github · commit f9c49af39245 · 2025-05-29T15:00:12.000-07:00
PiperOrigin-RevId: 764898438
diff --git a/ai_edge_torch/generative/layers/model_config.py b/ai_edge_torch/generative/layers/model_config.py
@@ -134,8 +134,6 @@ class AttentionConfig:
   value_norm_config: NormalizationConfig = dataclasses.field(
       default_factory=NormalizationConfig
   )
-  # Whether the KV cache is shared with the previous attention block.
-  kv_shared: bool = False
   relative_attention_num_buckets: int = 0
   relative_attention_max_distance: int = 0
   # Softcap on the output logits.

Original file line number	Diff line number	Diff line change
`@@ -134,8 +134,6 @@ class AttentionConfig:`
`134`	`134`	`value_norm_config: NormalizationConfig = dataclasses.field(`
`135`	`135`	`default_factory=NormalizationConfig`
`136`	`136`	`)`
`137`		`- # Whether the KV cache is shared with the previous attention block.`
`138`		`- kv_shared: bool = False`
`139`	`137`	`relative_attention_num_buckets: int = 0`
`140`	`138`	`relative_attention_max_distance: int = 0`
`141`	`139`	`# Softcap on the output logits.`