From bac2de5bb2bb4eb5c0a639747ef8b5f19bce25db Mon Sep 17 00:00:00 2001
From: tangchengxiang <2064027004@qq.com>
Date: Thu, 17 Sep 2026 08:36:50 +0000
Subject: [PATCH] fix(metax): allow strict FP32 GEMM precision

Keep TF32 as the default and support INFINIOP_METAX_ALLOW_TF32=0 at descriptor creation. Validate the precision choice in isolated processes.
---
 README.md                                     |  6 ++
 src/infiniop/ops/gemm/metax/gemm_metax.cc     | 10 +++-
 .../ops/test_metax_gemm_precision.py          | 59 +++++++++++++++++++
 3 files changed, 73 insertions(+), 2 deletions(-)
 create mode 100644 test/infinicore/ops/test_metax_gemm_precision.py

diff --git a/README.md b/README.md
index 4569c0157..b84c41fde 100644
--- a/README.md
+++ b/README.md
@@ -109,6 +109,12 @@ python scripts/install.py [XMAKE_CONFIG_FLAGS]
 | `--ccl=[y\|n]`           | 是否编译 InfiniCCL 通信库接口实现 | n
 | `--graph=[y\|n]`         | 是否编译 cuda graph 接口实现      | n
 
+MetaX FP32 GEMM permits TF32 by default. Set `INFINIOP_METAX_ALLOW_TF32=0`
+before starting the process to request `MCBLAS_COMPUTE_32F` instead of
+`MCBLAS_COMPUTE_32F_FAST_TF32`. The policy is captured when each descriptor
+is created and remains fixed during graph replay; restart the process when
+changing it. FP16 and BF16 GEMM retain their existing FP32 accumulation.
+
 ##### 手动安装底层库
 
 0. 生成九齿算子（可选）
diff --git a/src/infiniop/ops/gemm/metax/gemm_metax.cc b/src/infiniop/ops/gemm/metax/gemm_metax.cc
index 9d45099dc..0e3f758d8 100644
--- a/src/infiniop/ops/gemm/metax/gemm_metax.cc
+++ b/src/infiniop/ops/gemm/metax/gemm_metax.cc
@@ -1,11 +1,14 @@
 #include "gemm_metax.h"
 #include "../../../devices/metax/metax_common.h"
 #include "../../../devices/metax/metax_handle.h"
+#include <cstdlib>
+#include <cstring>
 
 namespace op::gemm::metax {
 
 struct Descriptor::Opaque {
     std::shared_ptr<device::metax::Handle::Internal> internal;
+    bool allow_tf32;
 };
 
 Descriptor::~Descriptor() {
@@ -26,9 +29,12 @@ infiniStatus_t Descriptor::create(
     auto result = MatmulInfo::create(c_desc, a_desc, b_desc, MatrixLayout::COL_MAJOR);
     CHECK_RESULT(result);
 
+    // Capture the precision policy with the descriptor, including graph replay.
+    const char *allow_tf32 = std::getenv("INFINIOP_METAX_ALLOW_TF32");
+    const bool use_tf32 = allow_tf32 == nullptr || std::strcmp(allow_tf32, "0") != 0;
     *desc_ptr = new Descriptor(
         dtype, result.take(), 0,
-        new Opaque{handle->internal()},
+        new Opaque{handle->internal(), use_tf32},
         handle->device, handle->device_id);
     return INFINI_STATUS_SUCCESS;
 }
@@ -57,7 +63,7 @@ infiniStatus_t Descriptor::calculate(
         break;
     case INFINI_DTYPE_F32:
         a_type = b_type = c_type = HPCC_R_32F;
-        compute_type = HCBLAS_COMPUTE_32F_FAST_TF32;
+        compute_type = _opaque->allow_tf32 ? HCBLAS_COMPUTE_32F_FAST_TF32 : HCBLAS_COMPUTE_32F;
         break;
 
     default:
diff --git a/test/infinicore/ops/test_metax_gemm_precision.py b/test/infinicore/ops/test_metax_gemm_precision.py
new file mode 100644
index 000000000..6a179debe
--- /dev/null
+++ b/test/infinicore/ops/test_metax_gemm_precision.py
@@ -0,0 +1,59 @@
+"""Check MetaX GEMM precision in processes with independent descriptor caches."""
+
+import os
+import subprocess
+import sys
+
+import pytest
+import torch
+from infinicore.lib import _infinicore
+
+import infinicore
+
+
+@pytest.mark.skipif(
+    _infinicore.get_device_count(_infinicore.Device.Type.METAX) == 0,
+    reason="A MetaX device is required.",
+)
+@pytest.mark.parametrize("allow_tf32", [None, "0"])
+def test_metax_gemm_precision(allow_tf32):
+    env = os.environ.copy()
+    if allow_tf32 is None:
+        env.pop("INFINIOP_METAX_ALLOW_TF32", None)
+    else:
+        env["INFINIOP_METAX_ALLOW_TF32"] = allow_tf32
+    result = subprocess.run(
+        [sys.executable, __file__], env=env, capture_output=True, text=True
+    )
+    assert result.returncode == 0, result.stdout + result.stderr
+
+
+if __name__ == "__main__":
+    generator = torch.Generator().manual_seed(20260917)
+    for dtype in (torch.float32, torch.float16, torch.bfloat16):
+        for rows in (1, 3, 257):
+            x = torch.randn(rows, 768, generator=generator).to(dtype)
+            weight = torch.randn(96, 768, generator=generator).to(dtype)
+            expected = (x.double() @ weight.double().T).to(dtype)
+            x_device, weight_device = x.cuda(), weight.cuda().T
+            output = torch.empty((rows, 96), device="cuda", dtype=dtype)
+            torch.cuda.synchronize()
+            infinicore.matmul(
+                infinicore.from_torch(x_device),
+                infinicore.strided_from_blob(
+                    weight_device.data_ptr(),
+                    list(weight_device.shape),
+                    list(weight_device.stride()),
+                    dtype=infinicore.utils.to_infinicore_dtype(dtype),
+                    device=infinicore.device("cuda", 0),
+                ),
+                out=infinicore.from_torch(output),
+            )
+            infinicore.sync_device()
+            strict = (
+                dtype == torch.float32 and os.getenv("INFINIOP_METAX_ALLOW_TF32") == "0"
+            )
+            tolerance = (2e-4, 1e-5) if strict else (0.1, 1e-2)
+            torch.testing.assert_close(
+                output.cpu(), expected, atol=tolerance[0], rtol=tolerance[1]
+            )