temp hack

Differential Revision: [D93870396](https://our.internmc.facebook.com/intern/diff/D93870396/)

[ghstack-poisoned]

Author: kimishpatel

examples/models/llama/source_transformation/custom_kv_cache.py

@@ -325,6 +325,7 @@ def _replace_kv_cache_with_quantized_kv_cache(module):
                     child,
                     QuantizedCacheType.AffineAsymmetric,
                     use_custom_update_cache_op=True,
+                    is_seq_at_dim_2=child.is_seq_at_dim_2,
                 ),
             )
         else:
@@ -421,6 +422,7 @@ def _replace_kv_cache_with_custom_kv_cache(module):
                     n_heads,
                     head_dim,
                     dtype=cache_dtype,
+                    is_seq_at_dim_2=True, # hacking temporarily
                 ),
             )
         else:

examples/models/llama/source_transformation/sdpa.py

@@ -92,6 +92,7 @@ def _replace_sdpa_with_custom_op(
                 SDPACustom(
                     child.dim,
                     use_attention_mask=use_attention_mask,
+                    is_seq_at_dim_2=True, # hacking temporarily
                 ),
             )
         else: