Mark scale dimensions to have the same batch size as input

Signed-off-by: Antoni Viros i Martin <[email protected]>

Author: ani300

fms_mo/aiu_addons/fp8/fp8_attn.py

@@ -317,6 +317,23 @@ def _spyre_scaled_paged_compute_op(
             attn_kwargs["left_padded_prompt_mask"],
             attn_kwargs["block_table"],
         )
+    
+    def __spyre_scaled_paged_validate_attn_kwargs_op(
+        input_ids: torch.Tensor,
+        position_ids: torch.Tensor,
+        past_key_value_states: Optional[list[tuple[torch.Tensor, torch.Tensor]]] = None,
+        **attn_kwargs,
+    ):
+        __spyre_paged_validate_attn_kwargs_op(input_ids, position_ids, past_key_value_states, **attn_kwargs)
+
+        if past_key_value_states is not None:
+            for k, v in past_key_value_states:
+                assert isinstance(k, ScaledTensor)
+                assert isinstance(v, ScaledTensor)
+
+                # assert that for each layer, the scales are per-sequence
+                assert k._scale.shape[0] == input_ids.shape[0]
+                assert v._scale.shape[0] == input_ids.shape[0]
 
     register_attention_op(
         "spyre_paged_attn_fp8",
@@ -325,5 +342,5 @@ def _spyre_scaled_paged_compute_op(
         is_prefill_op=lambda **attn_kwargs: attn_kwargs.get("block_table", None)
         is None,
         compute_decode_op=_spyre_scaled_paged_compute_op,
-        validate_attn_kwargs_op=__spyre_paged_validate_attn_kwargs_op,
+        validate_attn_kwargs_op=__spyre_scaled_paged_validate_attn_kwargs_op,
     )