Update core/mamba_block.py

core/mamba_block.py

@@ -7,7 +7,7 @@ import math
 class MambaBlock(nn.Module):
     """
     Production-ready Mamba block for graph processing
-    Based on official Mamba implementation with graph optimizations
+    Device-safe implementation
     """
     def __init__(self, d_model, d_state=16, d_conv=4, expand=2, dt_rank="auto", bias=False):
         super().__init__()
@@ -56,6 +56,11 @@ class MambaBlock(nn.Module):
         Returns: (batch, length, d_model)
         """
         batch, length, _ = x.shape
+        device = x.device
+        
+        # Ensure all parameters are on correct device
+        self.A_log = self.A_log.to(device)
+        self.D = self.D.to(device)
         
         # Input projection and split
         xz = self.in_proj(x)  # (batch, length, 2 * d_inner)
@@ -81,6 +86,7 @@ class MambaBlock(nn.Module):
     def selective_scan(self, u):
         """Selective scan operation - core of Mamba"""
         batch, length, d_inner = u.shape
+        device = u.device
         
         # Compute ∆, B, C
         x_dbl = self.x_proj(u)  # (batch, length, dt_rank + 2*d_state)
@@ -92,15 +98,20 @@ class MambaBlock(nn.Module):
         return self._selective_scan_pytorch(u, delta, B, C)
     
     def _selective_scan_pytorch(self, u, delta, B, C):
-        """PyTorch implementation of selective scan"""
+        """PyTorch implementation of selective scan - device safe"""
         batch, length, d_inner = u.shape
+        device = u.device
+        
+        # Ensure A_log and D are on correct device
+        A_log = self.A_log.to(device)
+        D = self.D.to(device)
         
         # Discretize
-        deltaA = torch.exp(delta.unsqueeze(-1) * (-torch.exp(self.A_log)))  # (batch, length, d_inner, d_state)
+        deltaA = torch.exp(delta.unsqueeze(-1) * (-torch.exp(A_log)))  # (batch, length, d_inner, d_state)
         deltaB_u = delta.unsqueeze(-1) * B.unsqueeze(2) * u.unsqueeze(-1)  # (batch, length, d_inner, d_state)
         
         # Initialize state
-        x = torch.zeros((batch, d_inner, self.d_state), device=u.device, dtype=u.dtype)
+        x = torch.zeros((batch, d_inner, self.d_state), device=device, dtype=u.dtype)
         ys = []
         
         for i in range(length):
@@ -111,6 +122,6 @@ class MambaBlock(nn.Module):
         y = torch.stack(ys, dim=1)  # (batch, length, d_inner)
         
         # Add skip connection
-        y = y + u * self.D
+        y = y + u * D
         
         return y