Arm backend: Add DeepSeek-R1-Distill-Qwen FP8 layer tests

backends/arm/test/models/DeepSeek_R1_Distill_Qwen/test_deepseek_r1_distill_qwen_layers.py

@@ -10,10 +10,12 @@
 
 import pytest
 import torch
+from executorch.backends.arm.ao_ext import MXFPOpConfig
 from executorch.backends.arm.test import common
 from executorch.backends.arm.test.models.DeepSeek_R1_Distill_Qwen.deepseek_r1_distill_qwen_test_config import (
     get_deepseek_r1_distill_qwen_1_5b_checkpoint_config,
 )
+from executorch.backends.arm.test.ops.mxfp.common import MXFPTosaPipelineFP
 from executorch.backends.arm.test.tester.test_pipeline import (
     TosaPipelineFP,
     VgfPipeline,
@@ -32,6 +34,7 @@
 )
 
 input_t = Tuple[torch.Tensor, ...]
+aten_op_mxfp_linear = "torch.ops.tosa_mxfp.linear.default"
 
 
 def _make_deepseek_r1_distill_qwen_1_5b_layer_config():
@@ -84,6 +87,10 @@ def _to_bfloat16(
     )
 
 
+def _is_linear(module: torch.nn.Module, _fqn: str) -> bool:
+    return isinstance(module, torch.nn.Linear)
+
+
 class RotaryEmbeddingModel(DeepSeekR1DistillQwenTestModule):
     def __init__(self, config) -> None:
         super().__init__()
@@ -326,6 +333,12 @@ class DeepSeekR1DistillQwenTestCase:
     TOSA_BF16_TEST_CASES
 )
 
+TOSA_MXFP8_TEST_CASES: dict[str, DeepSeekR1DistillQwenTestCase] = {
+    "attention": DeepSeekR1DistillQwenTestCase(model_cls=AttentionModel),
+    "mlp": DeepSeekR1DistillQwenTestCase(model_cls=MLPModel),
+    "decoder_layer": DeepSeekR1DistillQwenTestCase(model_cls=DecoderLayerModel),
+}
+
 
 @common.parametrize(
     "test_case",
@@ -369,6 +382,59 @@ def test_deepseek_r1_distill_qwen_tosa_FP_bf16(
         pipeline.run()
 
 
+@common.parametrize(
+    "test_case",
+    TOSA_MXFP8_TEST_CASES,
+)
+def test_deepseek_r1_distill_qwen_tosa_mxfp8_fp32(
+    test_case: DeepSeekR1DistillQwenTestCase,
+):
+    model, inputs = test_case.model_cls.prepare_model_and_inputs()
+    mxfp_config = MXFPOpConfig(weight_dtype=torch.float8_e4m3fn)
+
+    with torch.no_grad():
+        pipeline = MXFPTosaPipelineFP[input_t](
+            model,
+            inputs,
+            aten_op=aten_op_mxfp_linear,
+            exir_op=[],
+            filter_fn=_is_linear,
+            frobenius_threshold=0.05,
+            cosine_threshold=0.995,
+            mxfp_config=mxfp_config,
+            tosa_version="1.1",
+            tosa_extensions=["mxfp"],
+        )
+        pipeline.run()
+
+
+@common.parametrize(
+    "test_case",
+    TOSA_MXFP8_TEST_CASES,
+)
+def test_deepseek_r1_distill_qwen_tosa_mxfp8_bf16(
+    test_case: DeepSeekR1DistillQwenTestCase,
+):
+    model, inputs = test_case.model_cls.prepare_model_and_inputs()
+    model, inputs = _to_bfloat16(model, inputs)
+    mxfp_config = MXFPOpConfig(weight_dtype=torch.float8_e4m3fn)
+
+    with torch.no_grad():
+        pipeline = MXFPTosaPipelineFP[input_t](
+            model,
+            inputs,
+            aten_op=aten_op_mxfp_linear,
+            exir_op=[],
+            filter_fn=_is_linear,
+            frobenius_threshold=0.05,
+            cosine_threshold=0.995,
+            mxfp_config=mxfp_config,
+            tosa_version="1.1",
+            tosa_extensions=["bf16", "mxfp"],
+        )
+        pipeline.run()
+
+
 @common.SkipIfNoModelConverter
 @common.parametrize(
     "test_case",