Update src/llmcompressor/modifiers/quantization/cache.py

per tensor are same

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: Eviannn

src/llmcompressor/modifiers/quantization/cache.py

@@ -115,13 +115,14 @@ def update(
             q_value_states, KVCacheScaleType.VALUE, layer_idx
         )
 
-        # reshape for per channel scenario
-        # from [batch_size, seq_len - residual_length, num_heads * head_dim]
-        # to [batch_size, num_heads, seq_len - residual_length, head_dim]
-        qdq_key_states = qdq_key_states.view(
-            qdq_key_states.shape[0], qdq_key_states.shape[1], num_heads, head_dim).transpose(1, 2)
-        qdq_value_states = qdq_value_states.view(
-            qdq_value_states.shape[0], qdq_value_states.shape[1], num_heads, head_dim).transpose(1, 2)
+        if self.quantization_args.strategy == QuantizationStrategy.CHANNEL:
+            # reshape for per channel scenario
+            # from [batch_size, seq_len - residual_length, num_heads * head_dim]
+            # to [batch_size, num_heads, seq_len - residual_length, head_dim]
+            qdq_key_states = qdq_key_states.view(
+                qdq_key_states.shape[0], qdq_key_states.shape[1], num_heads, head_dim).transpose(1, 2)
+            qdq_value_states = qdq_value_states.view(
+                qdq_value_states.shape[0], qdq_value_states.shape[1], num_heads, head_dim).transpose(1, 2)
 
         keys_to_return, values_to_return = qdq_key_states, qdq_value_states