Remove overly granular checks on compatability

Signed-off-by: John St. John <jstjohn@nvidia.com>

Author: jstjohn

bionemo-recipes/recipes/evo2_megatron/src/bionemo/evo2/models/megatron/hyena/engine.py

@@ -83,6 +83,7 @@ def parallel_fir(
     fir_length,
     compute_state,
     use_subquadratic_ops=False,
+    check_subquadratic_ops=True,
 ):
     """Compute parallel finite impulse response filtering with optional state computation."""
     L = u.shape[1]  # noqa: N806
@@ -94,7 +95,8 @@ def parallel_fir(
     if fir_length >= 128:
         if use_subquadratic_ops:
             # subq-ops fft_causal_conv1d expects [B, D, L] input and [D, L] filter; dtypes must match
-            ensure_subquadratic_fft_causal_conv1d_supported()
+            if check_subquadratic_ops and u.is_cuda:
+                ensure_subquadratic_fft_causal_conv1d_supported()
             k = weight[:, :, :L].squeeze(1) if weight.dim() == 3 else weight[:, :L]
             u_fp32 = u.to(torch.float32)
             z = _subq_fft_causal_conv1d(u_fp32, k.to(torch.float32))
@@ -113,7 +115,8 @@ def parallel_fir(
             if _subq_causal_conv1d is None:
                 raise ImportError(_subq_error_msg)
             # subq-ops causal_conv1d expects pre-padded [B, D, L+pad] input and [D, K] weight.
-            ensure_subquadratic_causal_conv1d_supported()
+            if check_subquadratic_ops and u.is_cuda:
+                ensure_subquadratic_causal_conv1d_supported()
             pad_size = fir_length - 1
             x_padded = F.pad(u.to(torch.float32), (pad_size, 0))
             w = weight.squeeze(1) if weight.dim() == 3 else weight

bionemo-recipes/recipes/evo2_megatron/src/bionemo/evo2/models/megatron/hyena/hyena_utils.py

@@ -60,20 +60,9 @@ def causal_conv1d_fn(*args, **kwargs):
     from subquadratic_ops_torch.fft_causal_conv1d import fft_causal_conv1d as _subq_fft_causal_conv1d
     from subquadratic_ops_torch.implicit_filter import implicit_filter
 
-    def causal_conv1d(*args, **kwargs):
-        """Run guarded subquadratic causal_conv1d."""
-        ensure_subquadratic_causal_conv1d_supported()
-        return _subq_causal_conv1d(*args, **kwargs)
-
-    def b2b_causal_conv1d(*args, **kwargs):
-        """Run guarded subquadratic b2b_causal_conv1d."""
-        ensure_subquadratic_b2b_causal_conv1d_supported()
-        return _subq_b2b_causal_conv1d(*args, **kwargs)
-
-    def fft_causal_conv1d(*args, **kwargs):
-        """Run guarded subquadratic fft_causal_conv1d."""
-        ensure_subquadratic_fft_causal_conv1d_supported()
-        return _subq_fft_causal_conv1d(*args, **kwargs)
+    causal_conv1d = _subq_causal_conv1d
+    b2b_causal_conv1d = _subq_b2b_causal_conv1d
+    fft_causal_conv1d = _subq_fft_causal_conv1d
 except ImportError as e:
     msg_causal_conv1d = f"Problem importing subquadratic_ops: {e}. causal_conv1d is not available."
     msg_b2b_causal_conv1d = f"Problem importing subquadratic_ops: {e}. b2b_causal_conv1d is not available."
@@ -471,7 +460,17 @@ def hyena_no_weight_decay_cond_with_embeddings(name, param):
     return ("embedding" in name) or hyena_no_weight_decay_cond(name, param)
 
 
-def fftconv_func(u, k, D, dropout_mask, gelu=True, k_rev=None, bidirectional=False, use_subquadratic_ops=False):  # noqa: N803
+def fftconv_func(
+    u,
+    k,
+    D,  # noqa: N803
+    dropout_mask,
+    gelu=True,
+    k_rev=None,
+    bidirectional=False,
+    use_subquadratic_ops=False,
+    check_subquadratic_ops=True,
+):
     """Apply a 1D convolution to the input sequence u using the filter k and the shortcut D."""
     seqlen = u.shape[-1]
     fft_size = 2 * seqlen
@@ -504,6 +503,8 @@ def fftconv_func(u, k, D, dropout_mask, gelu=True, k_rev=None, bidirectional=Fal
     # causal
     else:
         if use_subquadratic_ops:
+            if check_subquadratic_ops and u.is_cuda:
+                ensure_subquadratic_fft_causal_conv1d_supported()
             y = fft_causal_conv1d(u, k.squeeze(0))
         else:
             fft_size = max(fft_size, 2 * k.shape[-1])
@@ -902,6 +903,7 @@ def __init__(
         self.zigzag = zigzag
 
         self.use_subquadratic_ops = transformer_config.use_subquadratic_ops
+        self._subquadratic_ops_checked = False
 
         self.model_parallel_size = self.pg_collection.tp.size() if self.pg_collection.tp is not None else 1
         self.model_parallel_rank = self.pg_collection.tp.rank() if self.pg_collection.tp is not None else 0
@@ -984,6 +986,16 @@ def reset_parameters(self):
             bounds = math.sqrt(1 / self.kernel_size)
             torch.nn.init.uniform_(self.conv_bias, a=-bounds, b=bounds)
 
+    def _ensure_subquadratic_ops_supported(self):
+        """Run expensive subquadratic-op CUDA self-tests once per operator instance."""
+        if self._subquadratic_ops_checked or not self.use_subquadratic_ops:
+            return
+        if self.operator_type == "hyena_medium_conv" and self.kernel_size < 128:
+            ensure_subquadratic_causal_conv1d_supported()
+        else:
+            ensure_subquadratic_fft_causal_conv1d_supported()
+        self._subquadratic_ops_checked = True
+
     def forward_long(self, *, x1, x2, v, h, bias, inference_context):
         """Forward pass long."""
         import bionemo.evo2.models.megatron.hyena.engine as engine
@@ -1074,6 +1086,7 @@ def get_filter_state(filter_name):
                 fir_length=self.kernel_size,  # self.short_filter_length,
                 compute_state=inference_context is not None,
                 use_subquadratic_ops=self.use_subquadratic_ops,
+                check_subquadratic_ops=False,
             )
             y = rearrange(y, "b d l -> b l d")
             y = y * x1
@@ -1099,6 +1112,8 @@ def forward(self, x1, x2, v, _hyena_use_cp=True, inference_context=None):
         Input shapes: bs, (num_groups, group_size), seq_length
         Output shapes: bs, (num_groups, group_size), seq_length
         """
+        if x1.is_cuda:
+            self._ensure_subquadratic_ops_supported()
         B, GDG, L = x1.shape  # noqa: N806
         x1, x2, v = x1[..., :L], x2[..., :L], v[..., :L]
 
@@ -1189,6 +1204,7 @@ def forward(self, x1, x2, v, _hyena_use_cp=True, inference_context=None):
                     gelu=False,
                     bidirectional=self.bidirectional,
                     use_subquadratic_ops=self.use_subquadratic_ops,
+                    check_subquadratic_ops=False,
                 )
                 z = z.to(v.dtype)
 
@@ -1388,6 +1404,7 @@ def __init__(
         self.num_groups = num_groups
         self.transformer_config = transformer_config
         self.use_subquadratic_ops = transformer_config.use_subquadratic_ops
+        self._subquadratic_ops_checked = False
         self.short_conv_L = hyena_config.short_conv_L
         self.local_init = local_init
         if pg_collection is None:
@@ -1543,6 +1560,7 @@ def __init__(
         """
         super().__init__()
         self.b2b_causal_conv1d_fn = b2b_causal_conv1d
+        self._check_subquadratic_ops = b2b_causal_conv1d is globals()["b2b_causal_conv1d"]
         if pg_collection is None:
             pg_collection = ProcessGroupCollection.use_mpu_process_groups()
         self.pg_collection = pg_collection
@@ -1567,6 +1585,14 @@ def __init__(
             raise ValueError(f"Operator type {operator_type} not supported")
 
         self.effective_pad_size = (self._mixer_kernel_size - 1) + (self._proj_conv_kernel_size - 1)
+        self._subquadratic_ops_checked = False
+
+    def _ensure_subquadratic_ops_supported(self):
+        """Run the B2B CUDA self-test once per wrapper instance."""
+        if self._subquadratic_ops_checked or not self._check_subquadratic_ops:
+            return
+        ensure_subquadratic_b2b_causal_conv1d_supported()
+        self._subquadratic_ops_checked = True
 
     def forward(self, x, _use_cp=True):
         """Forward pass for the B2BCausalConv1dModule.
@@ -1580,6 +1606,8 @@ def forward(self, x, _use_cp=True):
         # Validate input dimensions
         if x.dim() != 3:
             raise ValueError("Input tensor must be 3D [batch_size, hidden_dim, seq_len]")
+        if x.is_cuda:
+            self._ensure_subquadratic_ops_supported()
 
         # Extract weights at runtime to avoid parameter registration
         proj_weight = self._proj_conv_module.short_conv_weight
@@ -1713,6 +1741,9 @@ def get_filter_state(filter_name):
         L = u.shape[1]  # noqa: N806
         fir_state = get_filter_state("fir")
         if fir_state is None:
+            if self.use_subquadratic_ops and u.is_cuda and not self._subquadratic_ops_checked:
+                ensure_subquadratic_causal_conv1d_supported()
+                self._subquadratic_ops_checked = True
             z_pre, fir_state = engine.parallel_fir(
                 u=u,
                 weight=torch.tensor(weight),  # self.short_filter_weight,
@@ -1722,6 +1753,7 @@ def get_filter_state(filter_name):
                 fir_length=self.kernel_size,  # self.short_filter_length,
                 compute_state=inference_context is not None,
                 use_subquadratic_ops=self.use_subquadratic_ops,
+                check_subquadratic_ops=False,
             )
         else:
             if len(u.shape) > 2:

bionemo-recipes/recipes/evo2_megatron/tests/bionemo/evo2/models/megatron/hyena/test_hyena_utils.py

@@ -278,6 +278,29 @@ def test_b2b_causal_conv1d_module_device_handling():  # noqa: D103
         assert result_cuda.device == x_cuda.device, "Device mismatch on CUDA"
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for subquadratic guard test")
+@patch("bionemo.evo2.models.megatron.hyena.hyena_utils.ensure_subquadratic_b2b_causal_conv1d_supported")
+@patch("bionemo.evo2.models.megatron.hyena.hyena_utils.b2b_causal_conv1d")
+def test_b2b_causal_conv1d_module_checks_subquadratic_kernel_once(mock_b2b, mock_ensure):  # noqa: D103
+    mock_b2b.side_effect = mock_b2b_causal_conv1d
+    proj_conv = MockProjConv(kernel_size=3)
+    mixer = MockMixer(kernel_size=5)
+    b2b_module = B2BCausalConv1dModule(
+        proj_conv,
+        mixer,
+        operator_type="hyena_short_conv",
+        b2b_causal_conv1d=mock_b2b,
+        pg_collection=MockProcessGroupCollection(),
+    )
+
+    x = torch.randn(2, 96, 32, device="cuda")
+    b2b_module(x)
+    b2b_module(x)
+
+    assert mock_ensure.call_count == 1
+    assert mock_b2b.call_count == 2
+
+
 def test_b2b_causal_conv1d_effective_padding_size():
     """Test the zigzag pattern for data distribution in context parallel mode."""
     proj_conv = MockProjConv(kernel_size=3)