add Dao-style standard FHT accuracy tests （新增 Dao 风格标准 FHT 精度测试） (flagos-ai#4863)

Author: dawnyang050-debug

benchmark/test_hadamard_transform.py

@@ -12,15 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import math
+
 import pytest
 import torch
+import torch.nn.functional as F
 import triton
 from packaging.version import Version
 
 import flag_gems
 
 from . import base, consts
 
+try:
+    from scipy.linalg import hadamard as scipy_hadamard
+except ImportError:  # pragma: no cover
+    scipy_hadamard = None
+
 _TRITON_VERSION = Version(triton.__version__.split("+")[0])
 _SKIP_JOIN_BUG = _TRITON_VERSION < Version("3.5.0")
 _skip_if_join_bug = pytest.mark.skipif(
@@ -37,8 +45,8 @@
     (1024, 512),
     (1024, 1024),
     (1024, 4096),
-    (1024, 16384),
-    (1024, 32768),
+    # (1024, 16384),  # scipy full matrix too slow; temporarily commented
+    # (1024, 32768),
     (8192, 256),
     (8192, 512),
     (8192, 1024),
@@ -56,20 +64,32 @@ def ht_input_fn(shape, dtype, device):
     yield (torch.randn(batch, dim, dtype=dtype, device=device),)
 
 
-def _hadamard_matrix(n: int, device) -> torch.Tensor:
-    H = torch.tensor([[1.0]], device=device)
-    while H.shape[0] < n:
-        H = torch.cat([torch.cat([H, H], dim=1), torch.cat([H, -H], dim=1)], dim=0)
-    return H
+def _hadamard_transform_ref(x, scale=1.0):
+    """Reference matching tests/test_hadamard_transform.py (Dao scipy matrix multiply)."""
+    if scipy_hadamard is None:
+        raise ImportError("Please install scipy")
+    x_shape = x.shape
+    dim = x.shape[-1]
+    x = x.reshape(-1, dim)
+    log_dim = math.ceil(math.log2(dim)) if dim > 0 else 0
+    dim_padded = 1 << log_dim if dim > 0 else 1
+    if dim != dim_padded:
+        x = F.pad(x, (0, dim_padded - dim))
+    out = F.linear(
+        x,
+        torch.tensor(
+            scipy_hadamard(dim_padded, dtype=float),
+            dtype=x.dtype,
+            device=x.device,
+        ),
+    )
+    out = out * scale
+    return out[..., :dim].reshape(*x_shape)
 
 
 def torch_ht(x):
-    """Reference: Hadamard matrix multiply in fp32."""
-    dim = x.shape[-1]
-    padded = 1 << (dim - 1).bit_length() if dim > 1 else 1
-    H = _hadamard_matrix(padded, x.device).to(x.dtype)
-    x_padded = torch.nn.functional.pad(x, (0, padded - dim))
-    return (x_padded @ H.T)[..., :dim]
+    """Benchmark baseline: same scipy Hadamard ref as correctness tests."""
+    return _hadamard_transform_ref(x)
 
 
 class HadamardBenchmark(base.GenericBenchmark2DOnly):
@@ -79,6 +99,10 @@ class HadamardBenchmark(base.GenericBenchmark2DOnly):
     def set_more_shapes(self):
         return []
 
+    def set_shapes(self, *args, **kwargs):
+        # Force _FHT_SHAPES; do not fall back to GenericBenchmark2DOnly in core_shapes.yaml
+        self.shapes = self.DEFAULT_SHAPES
+
 
 @pytest.mark.hadamard_transform
 @_skip_if_join_bug
@@ -145,11 +169,11 @@ def ht_mn_input_fn(shape, dtype, device):
 
 
 def torch_ht_mn(x):
-    """Reference: pad to next power of 2 and run standard FHT."""
+    """Benchmark baseline: pad to next power of 2, then same scipy ref (not gems)."""
     dim = x.shape[-1]
     padded = 1 << (dim - 1).bit_length()
-    x_padded = torch.nn.functional.pad(x, (0, padded - dim))
-    return flag_gems.hadamard_transform(x_padded)[..., :dim]
+    x_padded = F.pad(x, (0, padded - dim))
+    return _hadamard_transform_ref(x_padded)[..., :dim]
 
 
 def gems_ht_mn(x):

tests/test_hadamard_transform.py

@@ -12,8 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import math
+
 import pytest
 import torch
+import torch.nn.functional as F
 import triton
 from packaging.version import Version
 
@@ -22,13 +25,40 @@
 from . import accuracy_utils as utils
 from . import conftest as cfg
 
+try:
+    from scipy.linalg import hadamard as scipy_hadamard
+except ImportError:  # pragma: no cover
+    scipy_hadamard = None
+
 _TRITON_VERSION = Version(triton.__version__.split("+")[0])
 _SKIP_JOIN_BUG = _TRITON_VERSION < Version("3.5.0")
 _skip_if_join_bug = pytest.mark.skipif(
     _SKIP_JOIN_BUG,
     reason=f"triton {triton.__version__} has tt.join layout bug (fixed in 3.5.0)",
 )
 
+# Dao-style dims for standard FHT (exclude 16384/32768: scipy full matrix too large)
+if cfg.QUICK_MODE:
+    HADAMARD_DIMS = [64, 137, 256, 512, 1024]
+else:
+    HADAMARD_DIMS = [
+        1,
+        2,
+        4,
+        8,
+        16,
+        32,
+        64,
+        128,
+        256,
+        512,
+        137,
+        1024,
+        2048,
+        4096,
+        8192,
+    ]
+
 if cfg.QUICK_MODE:
     HADAMARD_MN_CASES = [
         (1536, 3, "12N"),
@@ -57,6 +87,72 @@
 }
 
 
+def _hadamard_transform_ref(x, scale=1.0):
+    """Reference matching Dao-AILab fast_hadamard_transform_interface.hadamard_transform_ref."""
+    if scipy_hadamard is None:
+        raise ImportError("Please install scipy")
+    x_shape = x.shape
+    dim = x.shape[-1]
+    x = x.reshape(-1, dim)
+    log_dim = math.ceil(math.log2(dim)) if dim > 0 else 0
+    dim_padded = 1 << log_dim if dim > 0 else 1
+    if dim != dim_padded:
+        x = F.pad(x, (0, dim_padded - dim))
+    out = F.linear(
+        x,
+        torch.tensor(
+            scipy_hadamard(dim_padded, dtype=float),
+            dtype=x.dtype,
+            device=x.device,
+        ),
+    )
+    out = out * scale
+    return out[..., :dim].reshape(*x_shape)
+
+
+@pytest.mark.hadamard_transform
+@_skip_if_join_bug
+@pytest.mark.parametrize("dtype", utils.FLOAT_DTYPES)
+@pytest.mark.parametrize("dim", HADAMARD_DIMS)
+def test_hadamard_transform(dim, dtype):
+    """Dao-style accuracy test vs scipy Hadamard matrix multiply (fp32 ground truth)."""
+    if scipy_hadamard is None:
+        pytest.skip("scipy is required for hadamard_transform_ref")
+
+    atol = 3e-3 if dtype == torch.float32 else 5e-3
+    if dtype == torch.bfloat16:
+        atol = 5e-2
+
+    torch.random.manual_seed(0)
+    batch_size = 15
+    device = flag_gems.device
+    x = torch.randn(batch_size, dim, device=device, dtype=dtype).requires_grad_()
+    x_ref = x.detach().clone().requires_grad_()
+    x_pt = x.detach().clone().requires_grad_()
+    scale = 1 / math.sqrt(dim)
+
+    out = flag_gems.hadamard_transform(x, scale=scale)
+    out_ref = _hadamard_transform_ref(x_ref.float(), scale=scale)
+    out_pt = _hadamard_transform_ref(x_pt, scale=scale)
+
+    out_err = (out.float() - out_ref).abs().max().item()
+    pt_err = (out_pt.float() - out_ref).abs().max().item()
+    assert (
+        out_err < 2 * pt_err + atol
+    ), f"forward dim={dim} dtype={dtype}: out_err={out_err}, pt_err={pt_err}, atol={atol}"
+
+    g = torch.randn_like(out)
+    out.backward(g)
+    out_ref.backward(g)
+    out_pt.backward(g)
+
+    dx_err = (x.grad.float() - x_ref.grad.float()).abs().max().item()
+    dx_pt_err = (x_pt.grad.float() - x_ref.grad.float()).abs().max().item()
+    assert (
+        dx_err < 2 * dx_pt_err + atol
+    ), f"backward dim={dim} dtype={dtype}: dx_err={dx_err}, dx_pt_err={dx_pt_err}, atol={atol}"
+
+
 def _ref_mn(x: torch.Tensor, M: int) -> torch.Tensor:
     """Reference: 2-kernel version (H_M column transform in fp32 + standard FHT)."""
     *leading, dim = x.shape
@@ -93,7 +189,10 @@ def _ref_mn(x: torch.Tensor, M: int) -> torch.Tensor:
 
     ym = torch.stack(rows, dim=1).reshape(batch * M, n_cols)  # keep fp32
     ym = flag_gems.hadamard_transform(ym)  # FHT in fp32
-    return ym.to(orig_dtype).reshape(*leading, dim)
+    out = ym.to(orig_dtype).reshape(*leading, dim)
+    if cfg.TO_CPU:
+        out = out.cpu()
+    return out
 
 
 @pytest.mark.hadamard_transform_mn