Handle fallback gate and dense Wigner inputs

The fast GPU backend can be selected for channel widths and autocast dtypes that the packed FP32 Triton gate does not support. Keep the packed GEMM outputs, but route their activation through the existing PyTorch block implementation unless the actual outputs are FP32 with a power-of-two channel width. This avoids rejecting valid fast-backend configurations without repeating the GEMMs.

Compact Wigner storage also made the exported fused operations reject their prior dense [E, 9, 9] inputs. Pack dense block-diagonal inputs before entering custom autograd so the public interface, dense gradient shape, and zero off-block gradients are preserved.

Test Plan:
```bash
PYTHONPATH="$PWD/src" pytest -q tests/core/models/uma/uma_fast/test_execution_backends.py -k "gate_activation_fallback or compact_edge_degree_matches_dense"
PYTHONPATH="$PWD/src" pytest -q tests/core/models/uma/uma_fast/test_fused_edgewise.py -k "exported_fused_ops_accept_dense_wigner or wigner_conv1_fused_dynamic_compile"
pre-commit run --files src/fairchem/core/models/uma/escn_md_block.py src/fairchem/core/models/uma/nn/execution_backends.py src/fairchem/core/models/uma/triton/fused_wigner.py tests/core/models/uma/uma_fast/test_execution_backends.py tests/core/models/uma/uma_fast/test_fused_edgewise.py
```

Authored with assistance from an AI coding tool.

Author: mlazos

src/fairchem/core/models/uma/escn_md_block.py

@@ -269,18 +269,12 @@ def _forward_chunk_fused(
             x_full, edge_index, wigner, x_edge, sphere_channels
         )
         if self.act_type == "gate":
-            if x_full.dtype == torch.float32:
-                x0, x1, x2 = self.so2_conv_1.gemm_outputs_from_packed(
-                    m0_buf, m1_buf, m2_buf
-                )
-                x_blocks = self.backend.fused_gate_activation(
-                    x0, x1, x2, self.hidden_channels
-                )
-            else:
-                x_blocks, x_0_gating = self.so2_conv_1.gemm_blocks_from_packed(
-                    m0_buf, m1_buf, m2_buf
-                )
-                x_blocks = self.act.forward_m_blocks(x_0_gating, x_blocks)
+            x0, x1, x2 = self.so2_conv_1.gemm_outputs_from_packed(
+                m0_buf, m1_buf, m2_buf
+            )
+            x_blocks = self.backend.gate_activation(
+                x0, x1, x2, self.hidden_channels, self.act
+            )
             g0, g1, g2 = self.so2_conv_2.gemms_from_blocks(x_blocks)
         else:
             x_message, x_0_gating = self.so2_conv_1.gemms_from_packed(

src/fairchem/core/models/uma/nn/execution_backends.py

@@ -16,6 +16,7 @@
 from fairchem.core.models.uma.nn.unified_radial import UnifiedRadialMLP
 
 if TYPE_CHECKING:
+    from fairchem.core.models.uma.nn.activation import GateActivation
     from fairchem.core.units.mlip_unit.api.inference import (
         InferenceSettings,
     )
@@ -508,15 +509,26 @@ def fused_node_to_edge_conv1_pack(
         )
 
     @staticmethod
-    def fused_gate_activation(
+    def gate_activation(
         x0_full: torch.Tensor,
         x1: torch.Tensor,
         x2: torch.Tensor,
         channels: int,
+        activation: GateActivation,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        from fairchem.core.models.uma.triton import packed_gate_op
-
-        return packed_gate_op(x0_full, x1, x2, channels)
+        if (
+            channels > 0
+            and channels & (channels - 1) == 0
+            and x0_full.dtype == torch.float32
+            and x1.dtype == torch.float32
+            and x2.dtype == torch.float32
+        ):
+            from fairchem.core.models.uma.triton import packed_gate_op
+
+            return packed_gate_op(x0_full, x1, x2, channels)
+
+        gating, x0 = x0_full.split((2 * channels, 3 * channels), dim=-1)
+        return activation.forward_m_blocks(gating, (x0, x1, x2))
 
     @staticmethod
     def fused_conv2_inv_edge_to_node(

src/fairchem/core/models/uma/triton/fused_wigner.py

@@ -43,11 +43,26 @@
 
 
 def _compact_l2_wigner(wigner: Tensor, num_edges: int) -> Tensor:
-    if wigner.ndim != 2 or wigner.shape[1] != 35:
+    if wigner.ndim != 2 or wigner.shape != (num_edges, 35):
         raise ValueError("wigner must have shape [E, 35]")
     return wigner.reshape(num_edges, 35)
 
 
+def _prepare_l2_wigner(wigner: Tensor, num_edges: int) -> Tensor:
+    if wigner.ndim == 2 and wigner.shape == (num_edges, 35):
+        return wigner
+    if wigner.ndim == 3 and wigner.shape == (num_edges, 9, 9):
+        return torch.cat(
+            (
+                wigner[:, :1, :1].flatten(1),
+                wigner[:, 1:4, 1:4].flatten(1),
+                wigner[:, 4:9, 4:9].flatten(1),
+            ),
+            dim=1,
+        )
+    raise ValueError("wigner must have shape [E, 35] or [E, 9, 9]")
+
+
 # =============================================================================
 # Producer-side fused wigner -> conv1 (emits conv1's GEMM-ready packed buffers)
 # =============================================================================
@@ -268,13 +283,14 @@ def wigner_conv1_fused_op(
     Args:
         x_full: Node features [N, 9, C] (L-major).
         edge_index: Edge indices [2, E].
-        wigner: Compact Wigner blocks [E, 35].
+        wigner: Compact Wigner blocks [E, 35] or a dense matrix [E, 9, 9].
         radial: Per-layer conv1 radial embedding [E, 6*2C] (rad_func applied).
         C: sphere_channels.
 
     Returns:
         (m0, m1, m2) GEMM-ready packed buffers.
     """
+    wigner = _prepare_l2_wigner(wigner, edge_index.shape[1])
     return WignerConv1FusedFunction.apply(x_full, edge_index, wigner, radial, C)
 
 
@@ -534,12 +550,13 @@ def wigner_inv_conv2_fused_op(
         g0: conv2 fc_m0 output [E, 3C] (rows M0,M1,M2).
         g1: conv2 m=1 block-GEMM output [E, 4C] (rows M3,M4,M5,M6).
         g2: conv2 m=2 block-GEMM output [E, 2C] (rows M7,M8).
-        wigner: Compact inverse Wigner blocks [E, 35].
+        wigner: Compact inverse Wigner blocks [E, 35] or a dense matrix [E, 9, 9].
         C: sphere_channels.
 
     Returns:
         x_rotated [E, 9, C] (L-major).
     """
+    wigner = _prepare_l2_wigner(wigner, g0.shape[0])
     return WignerInvConv2FusedFunction.apply(g0, g1, g2, wigner, C)
 
 
@@ -621,6 +638,7 @@ def wigner_inv_conv2_scatter_op(
     Uses direct atomic accumulation by default and preserves PyTorch's
     deterministic-algorithm behavior through a materialized fallback.
     """
+    wigner = _prepare_l2_wigner(wigner, g0.shape[0])
     return WignerInvConv2ScatterFunction.apply(
         g0, g1, g2, wigner, scatter_target, num_nodes, C
     )

tests/core/models/uma/uma_fast/test_execution_backends.py

@@ -26,6 +26,7 @@
 from fairchem.core.datasets.ase_datasets import AseDBDataset
 from fairchem.core.datasets.atomic_data import AtomicData
 from fairchem.core.datasets.collaters.simple_collater import data_list_collater
+from fairchem.core.models.uma.nn.activation import GateActivation
 from fairchem.core.models.uma.nn.execution_backends import UMASFastGPUBackend
 from fairchem.core.models.uma.triton.constants import M_TO_L_GATHER_IDX
 from fairchem.core.models.uma.triton.node_to_edge_wigner_permute import (
@@ -142,6 +143,65 @@ def test_umas_fast_gpu_validation_accepts_hessian_loop():
     UMASFastGPUBackend.validate(lmax=2, mmax=2, settings=settings)
 
 
+@pytest.mark.gpu()
+@pytest.mark.parametrize(
+    ("channels", "dtype"),
+    [(96, torch.float32), (128, torch.bfloat16)],
+)
+def test_umas_fast_gpu_gate_activation_fallback(channels, dtype):
+    torch.manual_seed(42)
+    num_edges = 16
+    inputs = (
+        torch.randn(num_edges, 5 * channels, device="cuda", dtype=dtype),
+        torch.randn(num_edges, 4 * channels, device="cuda", dtype=dtype),
+        torch.randn(num_edges, 2 * channels, device="cuda", dtype=dtype),
+    )
+    activation = GateActivation(2, 2, channels, m_prime=True).cuda()
+    gating, x0 = inputs[0].split((2 * channels, 3 * channels), dim=-1)
+    expected = activation.forward_m_blocks(gating, (x0, inputs[1], inputs[2]))
+    actual = UMASFastGPUBackend.gate_activation(*inputs, channels, activation)
+
+    for actual_block, expected_block in zip(actual, expected, strict=True):
+        torch.testing.assert_close(actual_block, expected_block, rtol=0, atol=0)
+
+
+@pytest.mark.gpu()
+def test_umas_fast_gpu_gate_activation_fallback_dynamic_compile(
+    compile_reset_state,
+):
+    torch.manual_seed(42)
+    channels = 96
+    activation = GateActivation(2, 2, channels, m_prime=True).cuda()
+
+    def fn(x0_full, x1, x2):
+        return UMASFastGPUBackend.gate_activation(x0_full, x1, x2, channels, activation)
+
+    compiled = torch.compile(fn, fullgraph=True, dynamic=True)
+    for num_edges in (17, 31):
+        inputs = (
+            torch.randn(num_edges, 5 * channels, device="cuda", requires_grad=True),
+            torch.randn(num_edges, 4 * channels, device="cuda", requires_grad=True),
+            torch.randn(num_edges, 2 * channels, device="cuda", requires_grad=True),
+        )
+        reference_inputs = tuple(
+            value.detach().clone().requires_grad_() for value in inputs
+        )
+        actual = compiled(*inputs)
+        expected = fn(*reference_inputs)
+        grad_outputs = tuple(torch.randn_like(value) for value in actual)
+        actual_grads = torch.autograd.grad(actual, inputs, grad_outputs)
+        expected_grads = torch.autograd.grad(expected, reference_inputs, grad_outputs)
+
+        for actual_block, expected_block in zip(actual, expected, strict=True):
+            torch.testing.assert_close(
+                actual_block, expected_block, rtol=1e-6, atol=1e-6
+            )
+        for actual_grad, expected_grad in zip(
+            actual_grads, expected_grads, strict=True
+        ):
+            torch.testing.assert_close(actual_grad, expected_grad, rtol=1e-6, atol=1e-6)
+
+
 @pytest.mark.gpu()
 def test_compact_edge_degree_matches_dense():
     torch.manual_seed(42)

tests/core/models/uma/uma_fast/test_fused_edgewise.py

@@ -217,6 +217,69 @@ def test_wigner_inv_conv2_fused_matches_pytorch(sphere_channels):
     ), f"Max diff: {(ref_out - triton_out).abs().max()}"
 
 
+@pytest.mark.gpu()
+def test_exported_fused_ops_accept_dense_wigner(torch_deterministic):
+    torch.manual_seed(42)
+    num_nodes, num_edges, channels = 8, 16, 128
+    edge_index = torch.randint(0, num_nodes, (2, num_edges), device="cuda")
+    dense = _create_block_diagonal_wigner(num_edges, "cuda").requires_grad_()
+    compact = _compact_l2_wigner(dense.detach()).requires_grad_()
+    x_dense = torch.randn(num_nodes, 9, channels, device="cuda", requires_grad=True)
+    x_compact = x_dense.detach().clone().requires_grad_()
+    radial_dense = torch.randn(
+        num_edges, 12 * channels, device="cuda", requires_grad=True
+    )
+    radial_compact = radial_dense.detach().clone().requires_grad_()
+
+    dense_outputs = wigner_conv1_fused_op(
+        x_dense, edge_index, dense, radial_dense, channels
+    )
+    compact_outputs = wigner_conv1_fused_op(
+        x_compact, edge_index, compact, radial_compact, channels
+    )
+    grad_outputs = tuple(torch.randn_like(value) for value in dense_outputs)
+    dense_grads = torch.autograd.grad(
+        dense_outputs, (x_dense, dense, radial_dense), grad_outputs
+    )
+    compact_grads = torch.autograd.grad(
+        compact_outputs, (x_compact, compact, radial_compact), grad_outputs
+    )
+
+    for actual, expected in zip(dense_outputs, compact_outputs, strict=True):
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    torch.testing.assert_close(dense_grads[0], compact_grads[0], rtol=0, atol=0)
+    torch.testing.assert_close(
+        _compact_l2_wigner(dense_grads[1]), compact_grads[1], rtol=0, atol=0
+    )
+    torch.testing.assert_close(dense_grads[2], compact_grads[2], rtol=0, atol=0)
+
+    dense_inv = _create_block_diagonal_wigner(num_edges, "cuda").requires_grad_()
+    compact_inv = _compact_l2_wigner(dense_inv.detach()).requires_grad_()
+    dense_inputs = tuple(
+        torch.randn(num_edges, multiple * channels, device="cuda", requires_grad=True)
+        for multiple in (3, 4, 2)
+    )
+    compact_inputs = tuple(
+        value.detach().clone().requires_grad_() for value in dense_inputs
+    )
+    dense_output = wigner_inv_conv2_fused_op(*dense_inputs, dense_inv, channels)
+    compact_output = wigner_inv_conv2_fused_op(*compact_inputs, compact_inv, channels)
+    grad_output = torch.randn_like(dense_output)
+    dense_grads = torch.autograd.grad(
+        dense_output, (*dense_inputs, dense_inv), grad_output
+    )
+    compact_grads = torch.autograd.grad(
+        compact_output, (*compact_inputs, compact_inv), grad_output
+    )
+
+    torch.testing.assert_close(dense_output, compact_output, rtol=0, atol=0)
+    for actual, expected in zip(dense_grads[:-1], compact_grads[:-1], strict=True):
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    torch.testing.assert_close(
+        _compact_l2_wigner(dense_grads[-1]), compact_grads[-1], rtol=0, atol=0
+    )
+
+
 @pytest.mark.gpu()
 @pytest.mark.parametrize("sphere_channels", [128, 256])
 def test_wigner_inv_conv2_scatter_matches_materialized(sphere_channels):
@@ -362,17 +425,19 @@ def test_wigner_conv1_fused_deterministic_backward(
 
 
 @pytest.mark.gpu()
-def test_wigner_conv1_fused_dynamic_compile(compile_reset_state):
+@pytest.mark.parametrize("dense_wigner", [False, True])
+def test_wigner_conv1_fused_dynamic_compile(compile_reset_state, dense_wigner):
     torch.manual_seed(42)
     num_nodes, channels = 8, 128
     compiled = torch.compile(wigner_conv1_fused_op, fullgraph=True, dynamic=True)
 
     for num_edges in (17, 31):
         x = torch.randn(num_nodes, 9, channels, device="cuda", requires_grad=True)
         edge_index = torch.randint(0, num_nodes, (2, num_edges), device="cuda")
-        wigner = _compact_l2_wigner(
-            _create_block_diagonal_wigner(num_edges, "cuda")
-        ).requires_grad_()
+        wigner = _create_block_diagonal_wigner(num_edges, "cuda")
+        if not dense_wigner:
+            wigner = _compact_l2_wigner(wigner)
+        wigner.requires_grad_()
         radial = torch.randn(
             num_edges, 12 * channels, device="cuda", requires_grad=True
         )