From 9103f540315604dad00bec57cd58c7c143ffa0da Mon Sep 17 00:00:00 2001 From: Antoni Viros i Martin Date: Fri, 1 Aug 2025 18:51:32 +0000 Subject: [PATCH 1/3] Change paged FP8 prefill back to regular sdpa Signed-off-by: Antoni Viros i Martin --- fms_mo/aiu_addons/fp8/fp8_attn.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/fms_mo/aiu_addons/fp8/fp8_attn.py b/fms_mo/aiu_addons/fp8/fp8_attn.py index e4e4224..d0a98d5 100644 --- a/fms_mo/aiu_addons/fp8/fp8_attn.py +++ b/fms_mo/aiu_addons/fp8/fp8_attn.py @@ -21,6 +21,7 @@ import torch # Local +from fms.modules.attention import _sdpa_compute_op from fms_mo.aiu_addons.fp8.fp8_utils import ScaledTensor from fms_mo.prep import available_packages import fms_mo.aiu_addons.fp8.fp8_spyre_op # pylint: disable=unused-import @@ -340,7 +341,7 @@ def __spyre_scaled_paged_validate_attn_kwargs_op( register_attention_op( "spyre_paged_attn_fp8", _spyre_scaled_paged_store_op, - compute_op=_math_fp8_compute_op, + compute_op=_sdpa_compute_op, is_prefill_op=lambda **attn_kwargs: attn_kwargs.get("block_table", None) is None, compute_decode_op=_spyre_scaled_paged_compute_op, From a01a9169ca9f03da661790402c29da247890068b Mon Sep 17 00:00:00 2001 From: Antoni Viros i Martin Date: Fri, 1 Aug 2025 20:25:39 +0000 Subject: [PATCH 2/3] linting Signed-off-by: Antoni Viros i Martin --- fms_mo/aiu_addons/fp8/fp8_attn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fms_mo/aiu_addons/fp8/fp8_attn.py b/fms_mo/aiu_addons/fp8/fp8_attn.py index d0a98d5..c07a4cb 100644 --- a/fms_mo/aiu_addons/fp8/fp8_attn.py +++ b/fms_mo/aiu_addons/fp8/fp8_attn.py @@ -18,10 +18,10 @@ import math # Third Party +from fms.modules.attention import _sdpa_compute_op import torch # Local -from fms.modules.attention import _sdpa_compute_op from fms_mo.aiu_addons.fp8.fp8_utils import ScaledTensor from fms_mo.prep import available_packages import fms_mo.aiu_addons.fp8.fp8_spyre_op # pylint: disable=unused-import From f86feedbb155148cd3d2034578721f47d914dc88 Mon Sep 17 00:00:00 2001 From: Antoni Viros i Martin Date: Fri, 1 Aug 2025 20:53:01 +0000 Subject: [PATCH 3/3] more linting Signed-off-by: Antoni Viros i Martin --- fms_mo/aiu_addons/fp8/fp8_attn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fms_mo/aiu_addons/fp8/fp8_attn.py b/fms_mo/aiu_addons/fp8/fp8_attn.py index c07a4cb..5a091ae 100644 --- a/fms_mo/aiu_addons/fp8/fp8_attn.py +++ b/fms_mo/aiu_addons/fp8/fp8_attn.py @@ -18,7 +18,6 @@ import math # Third Party -from fms.modules.attention import _sdpa_compute_op import torch # Local @@ -30,6 +29,7 @@ # Third Party from fms.modules.attention import ( AttentionKwargs, + _sdpa_compute_op, _sdpa_update_attn_kwargs, register_attention_op, )