Update Feather H200 training runtime image

overlay/htm_rust/src/gpu/mod.rs

@@ -38,6 +38,11 @@ use fused::FusedState;
 fn htm_gpu_fused_available() -> bool {
     matches!(option_env!("HTM_GPU_FUSED_AVAILABLE"), Some("1"))
 }
+
+#[pyfunction]
+fn gpu_fused_available() -> bool {
+    htm_gpu_fused_available()
+}
 
 /// Extract (device_ptr, shape, typestr) from a `__cuda_array_interface__` dict.
 /// Returns Err if the dict is malformed. Used by `step_many_cuda` to wrap
@@ -576,8 +581,9 @@ fn step_batch_fused_cuda(
     Ok(())
 }
 
-pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    m.add_class::<HTMRegionGpu>()?;
-    m.add_function(pyo3::wrap_pyfunction!(step_batch_fused_cuda, m)?)?;
-    Ok(())
-}
+pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<HTMRegionGpu>()?;
+    m.add_function(pyo3::wrap_pyfunction!(step_batch_fused_cuda, m)?)?;
+    m.add_function(pyo3::wrap_pyfunction!(gpu_fused_available, m)?)?;
+    Ok(())
+}

overlay/subsystems/htm.py

@@ -54,8 +54,15 @@ _HTM_HAS_CAI = _HTM_HAS_GPU and hasattr(htm_rust.HTMRegionGpu, "step_many_cuda")
 # inhibition (see htm_rust/docs/GPU_HTM.md §Fused Kernel).
 # Opt-in via env var (default on when available).
 import os as _os_fused
-_HTM_HAS_FUSED = _HTM_HAS_GPU and hasattr(htm_rust.HTMRegionGpu, "step_many_fused_cuda")
-_HTM_USE_FUSED = _HTM_HAS_FUSED and bool(int(_os_fused.environ.get("HYDRA_HTM_FUSED", "1")))
+_HTM_HAS_FUSED = _HTM_HAS_GPU and hasattr(htm_rust.HTMRegionGpu, "step_many_fused_cuda")
+_HTM_GPU_FUSED_RUNTIME = bool(
+    _HTM_HAS_FUSED and hasattr(htm_rust, "gpu_fused_available") and htm_rust.gpu_fused_available()
+)
+_HTM_USE_FUSED = (
+    _HTM_HAS_FUSED
+    and _HTM_GPU_FUSED_RUNTIME
+    and bool(int(_os_fused.environ.get("HYDRA_HTM_FUSED", "1")))
+)
 
 
 class HTMLayer(nn.Module):