Update Feather training runtime image

overlay/subsystems/htm.py

@@ -60,7 +60,15 @@ _HTM_HAS_CAI = _HTM_HAS_GPU and hasattr(_HTM_REGION_GPU, "step_many_cuda")
 # Opt-in via env var (default on when available).
 import os as _os_fused
 _HTM_HAS_FUSED = _HTM_HAS_GPU and hasattr(_HTM_REGION_GPU, "step_many_fused_cuda")
-_HTM_USE_FUSED = _HTM_HAS_FUSED and bool(int(_os_fused.environ.get("HYDRA_HTM_FUSED", "1")))
+_HTM_USE_FUSED = _HTM_HAS_FUSED and bool(int(_os_fused.environ.get("HYDRA_HTM_FUSED", "1")))
+
+
+def _is_fused_unavailable_error(exc: RuntimeError) -> bool:
+    message = str(exc)
+    return (
+        "Fused HTM kernel is unavailable" in message
+        or "fused HTM kernel disabled for this CUDA arch" in message
+    )
 
 
 class HTMLayer(nn.Module):
@@ -170,19 +178,30 @@ class HTMLayer(nn.Module):
         if _HTM_HAS_CAI and self._use_gpu and sdr.is_cuda:
             sdr_u8 = sdr.contiguous().to(torch.uint8) if sdr.dtype != torch.uint8 else sdr.contiguous()
             cols_out = torch.empty((B, T, self.n_columns), dtype=torch.uint8, device=sdr.device)
-            anom_out = torch.empty((B, T), dtype=torch.float32, device=sdr.device)
-            # Pick fused (1 launch) or legacy (12*T launches) path.
-            if _HTM_USE_FUSED:
-                for b in range(B):
-                    self._regions[b].step_many_fused_cuda(
-                        sdr_u8[b].__cuda_array_interface__,
-                        cols_out[b].__cuda_array_interface__,
-                        anom_out[b].__cuda_array_interface__,
-                        learn,
-                    )
-            else:
-                for b in range(B):
-                    self._regions[b].step_many_cuda(
+            anom_out = torch.empty((B, T), dtype=torch.float32, device=sdr.device)
+            # Pick fused (1 launch) or legacy (12*T launches) path.
+            if _HTM_USE_FUSED:
+                try:
+                    for b in range(B):
+                        self._regions[b].step_many_fused_cuda(
+                            sdr_u8[b].__cuda_array_interface__,
+                            cols_out[b].__cuda_array_interface__,
+                            anom_out[b].__cuda_array_interface__,
+                            learn,
+                        )
+                except RuntimeError as exc:
+                    if not _is_fused_unavailable_error(exc):
+                        raise
+                    for b in range(B):
+                        self._regions[b].step_many_cuda(
+                            sdr_u8[b].__cuda_array_interface__,
+                            cols_out[b].__cuda_array_interface__,
+                            anom_out[b].__cuda_array_interface__,
+                            learn,
+                        )
+            else:
+                for b in range(B):
+                    self._regions[b].step_many_cuda(
                         sdr_u8[b].__cuda_array_interface__,
                         cols_out[b].__cuda_array_interface__,
                         anom_out[b].__cuda_array_interface__,
@@ -280,30 +299,49 @@ class HTMLayer(nn.Module):
                         [anom_out[b].__cuda_array_interface__ for b in range(B)],
                         learn,
                     )
-                except RuntimeError as _e:
-                    if "COOPERATIVE_LAUNCH_TOO_LARGE" in str(_e):
-                        # Batch too large for cooperative grid. Fall back to
-                        # sequential per-region fused launches (each B=1).
-                        for b in range(B):
+                except RuntimeError as _e:
+                    if "COOPERATIVE_LAUNCH_TOO_LARGE" in str(_e):
+                        # Batch too large for cooperative grid. Fall back to
+                        # sequential per-region fused launches (each B=1).
+                        for b in range(B):
                             self._regions[b].step_many_fused_cuda(
                                 sdr_u8[b].__cuda_array_interface__,
                                 cols_out[b].__cuda_array_interface__,
-                                anom_out[b].__cuda_array_interface__,
-                                learn,
-                            )
-                    else:
-                        raise
-            elif _HTM_USE_FUSED:
-                for b in range(B):
-                    self._regions[b].step_many_fused_cuda(
-                        sdr_u8[b].__cuda_array_interface__,
-                        cols_out[b].__cuda_array_interface__,
-                        anom_out[b].__cuda_array_interface__,
-                        learn,
-                    )
-            else:
-                for b in range(B):
-                    self._regions[b].step_many_cuda(
+                                anom_out[b].__cuda_array_interface__,
+                                learn,
+                            )
+                    elif _is_fused_unavailable_error(_e):
+                        for b in range(B):
+                            self._regions[b].step_many_cuda(
+                                sdr_u8[b].__cuda_array_interface__,
+                                cols_out[b].__cuda_array_interface__,
+                                anom_out[b].__cuda_array_interface__,
+                                learn,
+                            )
+                    else:
+                        raise
+            elif _HTM_USE_FUSED:
+                try:
+                    for b in range(B):
+                        self._regions[b].step_many_fused_cuda(
+                            sdr_u8[b].__cuda_array_interface__,
+                            cols_out[b].__cuda_array_interface__,
+                            anom_out[b].__cuda_array_interface__,
+                            learn,
+                        )
+                except RuntimeError as exc:
+                    if not _is_fused_unavailable_error(exc):
+                        raise
+                    for b in range(B):
+                        self._regions[b].step_many_cuda(
+                            sdr_u8[b].__cuda_array_interface__,
+                            cols_out[b].__cuda_array_interface__,
+                            anom_out[b].__cuda_array_interface__,
+                            learn,
+                        )
+            else:
+                for b in range(B):
+                    self._regions[b].step_many_cuda(
                         sdr_u8[b].__cuda_array_interface__,
                         cols_out[b].__cuda_array_interface__,
                         anom_out[b].__cuda_array_interface__,