hyper-connections 0.4.4__tar.gz → 0.4.6__tar.gz

{hyper_connections-0.4.4 → hyper_connections-0.4.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyper-connections
-Version: 0.4.4
+Version: 0.4.6
 Summary: Hyper-Connections
 Project-URL: Homepage, https://pypi.org/project/hyper-connections/
 Project-URL: Repository, https://github.com/lucidrains/hyper-connections

{hyper_connections-0.4.4 → hyper_connections-0.4.6}/hyper_connections/mHCv2.py

@@ -47,6 +47,10 @@ def l1norm(t, dim):
     return F.normalize(t, p = 1, dim = dim)
 
 def sinkhorn_knopps(log_alpha, iters = 20):
+
+    if iters <= 0:
+        return log_alpha
+
     assert log_alpha.shape[-2] == log_alpha.shape[-1]
 
     dtype = log_alpha.dtype
@@ -63,6 +67,10 @@ def sinkhorn_knopps(log_alpha, iters = 20):
     return alpha.to(dtype)
 
 def log_domain_sinkhorn_knopps(log_alpha, iters = 20):
+
+    if iters <= 0:
+        return log_alpha
+
     assert log_alpha.shape[-2] == log_alpha.shape[-1]
 
     dtype = log_alpha.dtype
@@ -109,6 +117,7 @@ def get_init_and_expand_reduce_stream_functions(
     add_attn_pool_reduce_stream = False,
     disable = None,
     sinkhorn_iters = 20,
+    use_triton_sinkhorn = False,
     **kwargs
 ):
     disable = default(disable, num_streams == 1 and num_fracs == 1)
@@ -116,7 +125,7 @@ def get_init_and_expand_reduce_stream_functions(
     hyper_conn_klass = ManifoldConstrainedHyperConnections if not disable else Residual
 
     kwargs.pop('add_attn_pool_reduce_stream', None)
-    init_hyper_conn_fn = partial(hyper_conn_klass, num_streams, num_fracs = num_fracs, sinkhorn_iters = sinkhorn_iters, **kwargs)
+    init_hyper_conn_fn = partial(hyper_conn_klass, num_streams, num_fracs = num_fracs, sinkhorn_iters = sinkhorn_iters, use_triton_sinkhorn = use_triton_sinkhorn, **kwargs)
     expand_reduce_fns = get_expand_reduce_stream_functions(
         num_streams,
         add_stream_embed = add_stream_embed,
@@ -231,7 +240,7 @@ class ManifoldConstrainedHyperConnections(Module):
         residual_mix_constraint_fn: Callable | None = None,
         forward_method_names: tuple[str, ...] = (),
         num_dynamic_alpha_proposals = 1,
-
+        use_triton_sinkhorn = False,
     ):
         """
         Appendix J, Algorithm2 in - https://arxiv.org/abs/2409.19606
@@ -306,10 +315,22 @@ class ManifoldConstrainedHyperConnections(Module):
         # Hres constraint related
         # by default is sinkhorn
 
-        self.residual_mix_constraint_fn = default(
-            residual_mix_constraint_fn,
-            partial(sinkhorn_knopps if not log_domain_sinkhorn else log_domain_sinkhorn_knopps, iters = sinkhorn_iters)
-        )
+        use_triton_sinkhorn_and_available = False
+
+        if use_triton_sinkhorn:
+            try:
+                from hyper_connections.triton_sinkhorn import triton_sinkhorn, is_triton_available
+                use_triton_sinkhorn_and_available = is_triton_available()
+            except ImportError:
+                use_triton_sinkhorn_and_available = False
+
+        if use_triton_sinkhorn_and_available:
+            self.residual_mix_constraint_fn = partial(triton_sinkhorn, iters = sinkhorn_iters)
+        else:
+            self.residual_mix_constraint_fn = default(
+                residual_mix_constraint_fn,
+                partial(sinkhorn_knopps if not log_domain_sinkhorn else log_domain_sinkhorn_knopps, iters = sinkhorn_iters)
+            )
 
         # dropouts
 

hyper_connections-0.4.6/hyper_connections/triton_sinkhorn.py

@@ -0,0 +1,160 @@
+import torch
+import triton
+import triton.language as tl
+from torch.autograd import Function
+
+@triton.jit
+def sinkhorn_kernel_forward_log(
+    input_ptr,
+    output_ptr,
+    M, N,
+    stride_b, stride_m, stride_n,
+    iters: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid_b = tl.program_id(0)
+    offs_m = tl.arange(0, BLOCK_SIZE)
+    offs_n = tl.arange(0, BLOCK_SIZE)
+    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
+    
+    curr_input_ptr = input_ptr + pid_b * stride_b
+    # Use a large negative value for log-space padding to avoid interference
+    log_alpha = tl.load(curr_input_ptr + offs_m[:, None] * stride_m + offs_n[None, :] * stride_n, mask=mask, other=-1e10)
+    
+    # Use static_range to force unrolling and avoid compiler bugs with dynamic loops in this environment
+    for _ in tl.static_range(iters):
+        # Column-wise Log-Softmax (dim=-2)
+        col_max = tl.max(tl.where(mask, log_alpha, -1e10), axis=0)
+        exp_weights_col = tl.exp(log_alpha - col_max[None, :])
+        exp_weights_col = tl.where(mask, exp_weights_col, 0.0)
+        col_lse = col_max + tl.log(tl.sum(exp_weights_col, axis=0))
+        log_alpha = log_alpha - col_lse[None, :]
+        log_alpha = tl.where(mask, log_alpha, -1e10)
+        
+        # Row-wise Log-Softmax (dim=-1)
+        row_max = tl.max(tl.where(mask, log_alpha, -1e10), axis=1)
+        exp_weights_row = tl.exp(log_alpha - row_max[:, None])
+        exp_weights_row = tl.where(mask, exp_weights_row, 0.0)
+        row_lse = row_max + tl.log(tl.sum(exp_weights_row, axis=1))
+        log_alpha = log_alpha - row_lse[:, None]
+        log_alpha = tl.where(mask, log_alpha, -1e10)
+    
+    result_alpha = tl.exp(log_alpha)
+    result_alpha = tl.where(mask, result_alpha, 0.0)
+    
+    curr_output_ptr = output_ptr + pid_b * stride_b
+    tl.store(curr_output_ptr + offs_m[:, None] * stride_m + offs_n[None, :] * stride_n, result_alpha, mask=mask)
+
+@triton.jit
+def sinkhorn_kernel_backward_log(
+    grad_output_ptr,
+    output_ptr,
+    grad_input_ptr,
+    M, N,
+    stride_b, stride_m, stride_n,
+    iters: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid_b = tl.program_id(0)
+    offs_m = tl.arange(0, BLOCK_SIZE)
+    offs_n = tl.arange(0, BLOCK_SIZE)
+    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
+    
+    curr_output_ptr = output_ptr + pid_b * stride_b
+    curr_grad_output_ptr = grad_output_ptr + pid_b * stride_b
+    
+    alpha = tl.load(curr_output_ptr + offs_m[:, None] * stride_m + offs_n[None, :] * stride_n, mask=mask, other=0.0)
+    grad_alpha = tl.load(curr_grad_output_ptr + offs_m[:, None] * stride_m + offs_n[None, :] * stride_n, mask=mask, other=0.0)
+    
+    # Ensure they are truly zeroed in padded areas for sum robustness
+    alpha = tl.where(mask, alpha, 0.0)
+    grad_alpha = tl.where(mask, grad_alpha, 0.0)
+    
+    for _ in tl.static_range(iters):
+        # Backward of Row-wise Normalization
+        # Sum only over valid elements
+        row_sum_grad_alpha = tl.sum(tl.where(mask, grad_alpha * alpha, 0.0), axis=1)
+        grad_alpha = grad_alpha - row_sum_grad_alpha[:, None]
+        grad_alpha = tl.where(mask, grad_alpha, 0.0)
+        
+        # Backward of Column-wise Normalization
+        col_sum_grad_alpha = tl.sum(tl.where(mask, grad_alpha * alpha, 0.0), axis=0)
+        grad_alpha = grad_alpha - col_sum_grad_alpha[None, :]
+        grad_alpha = tl.where(mask, grad_alpha, 0.0)
+    
+    grad_input = alpha * grad_alpha
+    
+    curr_grad_input_ptr = grad_input_ptr + pid_b * stride_b
+    tl.store(curr_grad_input_ptr + offs_m[:, None] * stride_m + offs_n[None, :] * stride_n, grad_input, mask=mask)
+
+class TritonSinkhornFunction(Function):
+    @staticmethod
+    def forward(ctx, log_alpha, iters=20):
+        # Handle matrix size limits to avoid register spilling/SRAM overflow
+        M, N = log_alpha.shape[-2:]
+        if max(M, N) > 256:
+             from hyper_connections.mHCv2 import log_domain_sinkhorn_knopps
+             return log_domain_sinkhorn_knopps(log_alpha, iters)
+
+        batch_shape = log_alpha.shape[:-2]
+        log_alpha_flat = log_alpha.view(-1, M, N).contiguous()
+        B = log_alpha_flat.shape[0]
+        
+        output = torch.empty_like(log_alpha_flat)
+        BLOCK_SIZE = max(32, triton.next_power_of_2(max(M, N)))
+        
+        sinkhorn_kernel_forward_log[(B,)](
+            log_alpha_flat,
+            output,
+            M, N,
+            log_alpha_flat.stride(0), log_alpha_flat.stride(1), log_alpha_flat.stride(2),
+            iters=iters,
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=4
+        )
+        
+        ctx.save_for_backward(output)
+        ctx.iters = iters
+        return output.view(*batch_shape, M, N)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        output, = ctx.saved_tensors
+        iters = ctx.iters
+        B, M, N = output.shape
+        BLOCK_SIZE = max(32, triton.next_power_of_2(max(M, N)))
+        
+        # Explicit contiguity for grad_output
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty_like(output)
+        
+        sinkhorn_kernel_backward_log[(B,)](
+            grad_output.view(B, M, N),
+            output,
+            grad_input,
+            M, N,
+            grad_input.stride(0), grad_input.stride(1), grad_input.stride(2),
+            iters=iters,
+            BLOCK_SIZE=BLOCK_SIZE,
+            num_warps=4
+        )
+        
+        return grad_input.view_as(grad_output), None
+
+def triton_sinkhorn(log_alpha, iters=20):
+    if log_alpha.is_cuda:
+        try:
+            return TritonSinkhornFunction.apply(log_alpha, iters)
+        except Exception:
+            pass
+    
+    # fallback
+    from hyper_connections.mHCv2 import sinkhorn_knopps
+    return sinkhorn_knopps(log_alpha, iters = iters)
+
+def is_triton_available():
+    try:
+        import triton
+        return torch.cuda.is_available()
+    except ImportError:
+        return False

{hyper_connections-0.4.4 → hyper_connections-0.4.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "hyper-connections"
-version = "0.4.4"
+version = "0.4.6"
 description = "Hyper-Connections"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{hyper_connections-0.4.4 → hyper_connections-0.4.6}/tests/test_hyper_connections.py

@@ -280,3 +280,38 @@ def test_mhcv2(
     residual = reduce_stream(residual)
 
     assert residual.shape == (2, 1024, 512)
+
+def test_triton_sinkhorn():
+    import torch
+    if not torch.cuda.is_available():
+        pytest.skip('CUDA not available')
+
+    from hyper_connections.triton_sinkhorn import triton_sinkhorn
+    from hyper_connections.mHCv2 import sinkhorn_knopps, log_domain_sinkhorn_knopps
+
+    B, M, N = 2, 16, 16
+    log_alpha = torch.randn(B, M, N, device = 'cuda', requires_grad = True, dtype = torch.float32)
+    iters = 20
+
+    # 1. Forward equivalence with sinkhorn_knopps
+    out_triton = triton_sinkhorn(log_alpha, iters = iters)
+    out_torch = sinkhorn_knopps(log_alpha, iters = iters)
+    torch.testing.assert_close(out_triton, out_torch, atol = 1e-4, rtol = 1e-4)
+
+    # 2. Forward equivalence with log_domain_sinkhorn_knopps
+    out_log_torch = log_domain_sinkhorn_knopps(log_alpha, iters = iters)
+    torch.testing.assert_close(out_triton, out_log_torch, atol = 1e-4, rtol = 1e-4)
+
+    # 3. Backward parity check
+    out_triton.backward(torch.ones_like(out_triton))
+    grad_triton = log_alpha.grad.clone()
+
+    log_alpha.grad.zero_()
+    out_torch.backward(torch.ones_like(out_torch))
+    grad_torch = log_alpha.grad.clone()
+
+    torch.testing.assert_close(grad_triton, grad_torch, atol = 1e-3, rtol = 1e-3)
+
+    log_alpha_double = torch.randn(1, 4, 4, device = 'cuda', requires_grad = True, dtype = torch.float64)
+
+    torch.autograd.gradcheck(triton_sinkhorn, (log_alpha_double, 10), eps = 1e-6, atol = 1e-5)