Revert "Jm/bio 713 mxfp8 enum member (#285)"

cuequivariance_torch/cuequivariance_torch/primitives/triangle.py

@@ -22,21 +22,13 @@
     import enum
 
     class TriMulPrecision(enum.IntEnum):  # type: ignore
-        """Fallback precision enum when cuequivariance_ops_torch is not available.
-
-        Values must match ``cuequivariance_ops_torch.TriMulPrecision`` so that
-        ``precision=TriMulPrecision.MXFP8`` resolves identically whether or not
-        the backend ops package is installed.
-        """
+        """Fallback precision enum when cuequivariance_ops_torch is not available."""
 
         NONE = -1
         DEFAULT = 0
         TF32 = 1
         TF32x3 = 2
         IEEE = 3
-        # Tri-mul carrier-precision modes (opt-in; default is None = off).
-        MXFP8 = 4
-        BFLOAT16 = 5
 
 
 def triangle_attention(
@@ -205,9 +197,6 @@ def triangle_multiplicative_update(
             Available options:
             - None: Defaults to triton language dot's default for non-32b input and for 32b input, tf32/tf32x3 based on 1/0 value set in torch.backends.cuda.matmul.allow_tf32
             - IEEE: Use IEEE 754 precision
-            - MXFP8: Opt-in MXFP8 (E4M3 + UE8M0 block scales) acceleration for the
-              triangular projection. **Inference-only and off by default** — see note (7).
-              Equivalent to passing the string ``"mxfp8"``.
 
     Returns:
         Output tensor of shape (batch_size, seq_len, seq_len, hidden_dim)
@@ -219,18 +208,6 @@ def triangle_multiplicative_update(
         (4) We have moved away from the default round-towards-zero (RZ) implementation to round-nearest (RN) for better tf32 accuracy in cuex.triangle_multiplicative_update. In rare circumstances, this may cause minor differences in results observed.
         (5) When using torch compile, use `cueuivariance_ops_torch.init_triton_cache()` to initialize triton cache before calling torch compiled triangular multiplicative update.
         (6) Although the example demonstrates the most common case of one batch dimension, the API supports variable number of leading batch dimensions.
-        (7) **MXFP8 precision** (``precision=TriMulPrecision.MXFP8`` or ``"mxfp8"``) is an
-            opt-in, **inference-only** acceleration of the triangular projection. It is
-            **off by default** (``precision=None``). It engages only when *all* of the
-            following hold, and otherwise transparently falls back to bf16 (no error):
-            - **Inference only**: any input has ``requires_grad=True`` (training) falls back
-              to bf16 with a warning. The MXFP8 backward path is not yet implemented.
-            - **Sequence length** must be divisible by 32 (``seq_len % 32 == 0``); other
-              lengths warn and fall back to bf16.
-            - **Hardware**: a Blackwell GPU (CUDA compute capability >= 10.0, e.g.
-              sm_100 / sm_103) with the CUTLASS MXFP8 binding available; on other GPUs it
-              falls back to bf16.
-            Numerics are MXFP8-quantized, so expect small deviations from the bf16/None path.
 
     Example:
         >>> import torch

cuequivariance_torch/tests/primitives/triangle_test.py

@@ -181,41 +181,3 @@ def test_attention_pair_bias():
             b_ln_z=b_ln_z,
         )
         assert output.shape == torch.Size([batch_size, seq_len, hidden_dim])
-
-
-def test_trimul_precision_mxfp8_enum_member():
-    # BIO-714: MXFP8 must be discoverable on the public TriMulPrecision enum, whether it
-    # is imported from the cuequivariance_ops_torch backend or resolved from the in-module
-    # fallback enum (no-ops-installed case). Both must carry the same sentinel values.
-    assert hasattr(cuet.TriMulPrecision, "MXFP8")
-    assert hasattr(cuet.TriMulPrecision, "BFLOAT16")
-    assert int(cuet.TriMulPrecision.MXFP8.value) == 4
-    assert int(cuet.TriMulPrecision.BFLOAT16.value) == 5
-
-
-def test_triangle_multiplicative_update_mxfp8_inference():
-    # BIO-714: exercise the opt-in, inference-only MXFP8 path via the enum member.
-    # seq_len is > the eager bf16-fallback threshold (150) and divisible by 32 so the
-    # optimized path engages; on a Blackwell GPU (cc >= 10.0) this hits the MXFP8 CUTLASS
-    # kernel, and on other GPUs it transparently falls back to bf16 (same output shape).
-    if not torch.cuda.is_available():
-        return
-    try:
-        import cuequivariance_ops_torch  # noqa: F401
-    except ImportError:
-        return  # backend ops not installed; the MXFP8 path is unavailable
-    device = torch.device("cuda")
-    batch_size, seq_len, hidden_dim = 1, 160, 32
-    expected = torch.Size([batch_size, seq_len, seq_len, hidden_dim])
-    with torch.no_grad():
-        x = torch.randn(batch_size, seq_len, seq_len, hidden_dim, device=device)
-        mask = torch.ones(batch_size, seq_len, seq_len, device=device)
-        # Both the enum member and the legacy string spelling select the same path.
-        for precision in (cuet.TriMulPrecision.MXFP8, "mxfp8"):
-            output = cuet.triangle_multiplicative_update(
-                x=x,
-                direction="outgoing",
-                mask=mask,
-                precision=precision,
-            )
-            assert output.shape == expected