colabfold-kernels 0.1.0__tar.gz

colabfold_kernels-0.1.0/LICENSE

@@ -0,0 +1,26 @@
+The MIT License (MIT)
+=====================
+
+Copyright © 2026 The ColabFold Development Team
+
+Permission is hereby granted, free of charge, to any person
+obtaining a copy of this software and associated documentation
+files (the “Software”), to deal in the Software without
+restriction, including without limitation the rights to use,
+copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the
+Software is furnished to do so, subject to the following
+conditions:
+
+The above copyright notice and this permission notice shall be
+included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND,
+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
+OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
+HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
+WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
+OTHER DEALINGS IN THE SOFTWARE.
+

colabfold_kernels-0.1.0/PKG-INFO

@@ -0,0 +1,34 @@
+Metadata-Version: 2.4
+Name: colabfold-kernels
+Version: 0.1.0
+Summary: Fast fused kernels for ColabFold
+Author-email: Milot Mirdita <milot@mirdita.de>
+License: MIT
+Project-URL: Homepage, https://github.com/mirditalab/colabfold-kernels
+Project-URL: Issues, https://github.com/mirditalab/colabfold-kernels/issues
+Keywords: alphafold,colabfold,pallas,triton,jax,attention,bioinformatics
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Dynamic: license-file
+
+# colabfold-kernels
+
+Fast fused kernels for ColabFold, written in
+[Pallas](https://docs.jax.dev/en/latest/pallas/index.html) with native CUDA-fallback
+for Volta- and Turing-generation GPUs.
+
+```python
+import jax.numpy as jnp
+from colabfold_kernels import attention, dispatch
+
+kernel = attention(dispatch(), jnp.bfloat16, key_dim=96, value_dim=96, layout="seq")
+if kernel is not None:
+    out = kernel(q, k, v, mask_bias, nonbatched_bias, scale)
+```
+
+`attention`, `layer_norm` and `gated_dual_proj` each return `None` if not supported.

colabfold_kernels-0.1.0/README.md

@@ -0,0 +1,16 @@
+# colabfold-kernels
+
+Fast fused kernels for ColabFold, written in
+[Pallas](https://docs.jax.dev/en/latest/pallas/index.html) with native CUDA-fallback
+for Volta- and Turing-generation GPUs.
+
+```python
+import jax.numpy as jnp
+from colabfold_kernels import attention, dispatch
+
+kernel = attention(dispatch(), jnp.bfloat16, key_dim=96, value_dim=96, layout="seq")
+if kernel is not None:
+    out = kernel(q, k, v, mask_bias, nonbatched_bias, scale)
+```
+
+`attention`, `layer_norm` and `gated_dual_proj` each return `None` if not supported.

colabfold_kernels-0.1.0/pyproject.toml

@@ -0,0 +1,28 @@
+[build-system]
+requires = ["setuptools>=64", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "colabfold-kernels"
+version = "0.1.0"
+description = "Fast fused kernels for ColabFold"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "MIT" }
+authors = [{ name = "Milot Mirdita", email = "milot@mirdita.de" }]
+keywords = ["alphafold", "colabfold", "pallas", "triton", "jax", "attention", "bioinformatics"]
+classifiers = [
+  "Intended Audience :: Science/Research",
+  "License :: OSI Approved :: MIT License",
+  "Operating System :: POSIX :: Linux",
+  "Topic :: Scientific/Engineering :: Bio-Informatics",
+]
+[project.urls]
+Homepage = "https://github.com/mirditalab/colabfold-kernels"
+Issues = "https://github.com/mirditalab/colabfold-kernels/issues"
+
+[tool.setuptools]
+package-dir = { "" = "src" }
+
+[tool.setuptools.packages.find]
+where = ["src"]

colabfold_kernels-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

colabfold_kernels-0.1.0/src/colabfold_kernels/__init__.py

@@ -0,0 +1,9 @@
+"""Fused Pallas/Triton kernels for ColabFold's AlphaFold2 and AlphaFold3 backends."""
+
+from colabfold_kernels.device import compute_capability, dispatch, shared_memory_limit
+from colabfold_kernels.fused_ops import attention, gated_dual_proj, layer_norm
+
+__all__ = ["attention", "compute_capability", "dispatch", "gated_dual_proj",
+           "layer_norm", "shared_memory_limit", "__version__"]
+
+__version__ = "0.1.0"

colabfold_kernels-0.1.0/src/colabfold_kernels/device.py

@@ -0,0 +1,29 @@
+"""What the GPUs here can run, so a caller need not work it out itself."""
+
+
+def compute_capability():
+  """The NVIDIA compute capability as an integer, e.g. 75, or None."""
+  try:
+    import jax
+    cc = str(jax.devices()[0].compute_capability)
+    return int(round(float(cc) * 10)) if "." in cc else int(cc)
+  except Exception:
+    return None
+
+
+def shared_memory_limit():
+  """The smallest shared memory budget across the GPUs here, or None."""
+  import jax
+  limits = [getattr(d, "shared_memory_per_block_optin", 0) for d in jax.local_devices()
+            if d.platform == "gpu"]
+  limits = [limit for limit in limits if limit]
+  return min(limits) if limits else None
+
+
+def dispatch(backend="auto"):
+  """The mapping the fused_ops selectors read. 'auto' picks what this GPU supports."""
+  cc = compute_capability()
+  if backend == "auto":
+    # XLA gates Pallas/Triton to sm_80+, so older NVIDIA parts take the CUDA kernels
+    backend = "cuda_legacy" if cc is not None and cc < 80 else "pallas"
+  return {"use_pallas": True, "kernel_backend": backend, "compute_capability": cc}

colabfold_kernels-0.1.0/src/colabfold_kernels/fused_ops.py

@@ -0,0 +1,72 @@
+"""Pick the fused kernel (volta, Pallas, or None for XLA) for each op."""
+import functools
+
+
+def _legacy_cc(global_config):
+  """Give the compute capability if the CUDA kernels were asked for, else None."""
+  if global_config.get('kernel_backend', 'pallas') != 'cuda_legacy':
+    return None
+  return global_config.get('compute_capability')
+
+
+def _volta():
+  from colabfold_kernels import volta
+  return volta
+
+
+def attention(global_config, dtype, key_dim, value_dim, layout='heads'):
+  """-> f(q, k, v, mask_bias, nonbatched_bias, scale) or None.
+
+  layout 'heads' takes q/k/v as [b, h, S, c], 'seq' as [b, S, h, c].
+  """
+  import jax.numpy as jnp
+  if not global_config.get('use_pallas', False):
+    return None
+  cc = _legacy_cc(global_config)
+  if dtype == jnp.float16 and cc is not None and layout == 'heads':
+    va = _volta()
+    # The kernels only handle a fixed set of head dims, and need key == value.
+    if (va.available(cc) and key_dim == value_dim
+        and va.supports(int(key_dim), cc)):
+      return functools.partial(va.volta_attention, cc=cc)
+  if (global_config.get('kernel_backend', 'pallas') == 'pallas'
+      and dtype in (jnp.bfloat16, jnp.float16)):
+    from colabfold_kernels.tri_flash import pallas_attention
+    if layout == 'seq':
+      return functools.partial(pallas_attention, seq_major=True)
+    return pallas_attention
+  return None
+
+
+def layer_norm(global_config, dtype):
+  """-> f(x, scale, offset, *, eps) or None. Last-axis, scale+offset only."""
+  import jax.numpy as jnp
+  if not global_config.get('use_pallas', False):
+    return None
+  cc = _legacy_cc(global_config)
+  if dtype == jnp.float16 and cc is not None:
+    va = _volta()
+    if va.ops_available(cc):
+      return va.volta_layer_norm
+  if (global_config.get('kernel_backend', 'pallas') == 'pallas'
+      and dtype in (jnp.bfloat16, jnp.float32)):
+    from colabfold_kernels.tri_mul import pallas_layer_norm
+    return pallas_layer_norm
+  return None
+
+
+def gated_dual_proj(global_config, dtype):
+  """-> f(x, wp, bp, wg, bg, mask, *, split) or None."""
+  import jax.numpy as jnp
+  if not global_config.get('use_pallas', False):
+    return None
+  cc = _legacy_cc(global_config)
+  if dtype == jnp.float16 and cc is not None:
+    va = _volta()
+    if va.ops_available(cc):
+      return functools.partial(va.volta_gated_dual_proj, cc=cc)
+  if (global_config.get('kernel_backend', 'pallas') == 'pallas'
+      and dtype in (jnp.bfloat16, jnp.float16)):
+    from colabfold_kernels.tri_mul import gated_dual_proj as gdp
+    return gdp
+  return None

colabfold_kernels-0.1.0/src/colabfold_kernels/tri_flash.py

@@ -0,0 +1,167 @@
+"""Pallas/Triton flash attention for the AF2 Evoformer (MSA + triangle).
+
+Inference (forward) only, no backward pass
+
+Computes, per (batch n, head h):
+  softmax_k( scale * Q_q . K_k + bias_{h,q,k}  if mask_{n,k} else -inf ) . V_k
+Layout: q/k/v [N, H, S, D]; bias [H, Sq, Sk] (shared across N); kmask [N, Sk] bool.
+"""
+import functools, math
+import numpy as np
+import jax, jax.numpy as jnp
+from jax.experimental import pallas as pl
+from jax.experimental.pallas import triton as plgpu
+
+# jnp.dot is stable across jax versions, contrary to pl.dot/plgpu.dot
+_dot = functools.partial(jnp.dot, preferred_element_type=jnp.float32)
+
+NEG = -0.7 * float(np.finfo(np.float32).max)
+LOG2E = math.log2(math.e)
+
+
+def _pow2_tile(d):
+  """Triton tiles are powers of two, and its dot needs at least 16 columns."""
+  return max(16, 1 << (int(d) - 1).bit_length())
+
+
+def _chunks(d):
+  """d as power-of-two dot widths: narrow dots beat one wider tile of padding."""
+  out, off, rem = [], 0, int(d)
+  while rem > 0:
+    w = 1 << (rem.bit_length() - 1) if rem >= 16 else 16
+    out.append((off, w))
+    off, rem = off + w, rem - w
+  tile = _pow2_tile(d)
+  return tuple(out) if off < tile else ((0, tile),)
+
+
+def _kernel(q_ref, k_ref, v_ref, bias_ref, kmask_ref, o_ref, *,
+            sm_scale, block_q, block_k, sq, d, chunks):
+  # q_ref [block_q, W]; k_ref/v_ref [Sk, W]; bias_ref [block_q, Sk]; kmask_ref [Sk],
+  # W being d split into power-of-two dot widths. Sq/Sk are NOT padded to block
+  # multiples; boundary blocks are handled with masked loads (offs < Sk) and a
+  # masked output store (offs_q < Sq) so no jnp.pad copy of q/k/v is needed.
+  sk = k_ref.shape[0]
+  # only a last chunk reaching past d needs a lane mask, and those lanes read zero
+  cols = [None if off + w <= d else (off + jnp.arange(w)) < d for off, w in chunks]
+  qs = [q_ref[:, pl.dslice(off, w)] if c is None else
+        plgpu.load(q_ref.at[:, pl.dslice(off, w)], mask=c[None, :], other=0.0)
+        for (off, w), c in zip(chunks, cols)]
+  m_i = jnp.full(block_q, -float('inf'), jnp.float32)
+  l_i = jnp.zeros(block_q, jnp.float32)
+  os = [jnp.zeros((block_q, w), jnp.float32) for _, w in chunks]
+
+  def body(j, carry):
+    o_prev, m_prev, l_prev = carry
+    start = j * block_k
+    kb = (start + jnp.arange(block_k)) < sk        # [block_k] in-bounds keys
+    sl = pl.dslice(start, block_k)
+    qk = None
+    for qc, (off, w), c in zip(qs, chunks, cols):
+      m = kb[:, None] if c is None else kb[:, None] & c[None, :]
+      kc = plgpu.load(k_ref.at[sl, pl.dslice(off, w)], mask=m, other=0.0)
+      part = _dot(qc, kc.T)                    # [block_q, block_k]
+      qk = part if qk is None else qk + part
+    bias = plgpu.load(bias_ref.at[:, sl], mask=kb[None, :], other=0.0)
+    qk = (qk * sm_scale + bias) * LOG2E
+    km = plgpu.load(kmask_ref.at[sl], mask=kb, other=False)  # OOB keys -> masked
+    qk = jnp.where(km[None, :], qk, NEG)
+    m_curr = jnp.max(qk, axis=-1)
+    m_next = jnp.maximum(m_prev, m_curr)
+    corr = jnp.exp2(m_prev - m_next)
+    s = jnp.exp2(qk - m_next[:, None])
+    l_next = corr * l_prev + s.sum(axis=-1)
+    o_next = []
+    for oc, (off, w), c in zip(o_prev, chunks, cols):
+      m = kb[:, None] if c is None else kb[:, None] & c[None, :]
+      vc = plgpu.load(v_ref.at[sl, pl.dslice(off, w)], mask=m, other=0.0)
+      o_next.append(corr[:, None] * oc + _dot(s.astype(vc.dtype), vc))
+    return o_next, m_next, l_next
+
+  os, m_i, l_i = jax.lax.fori_loop(0, pl.cdiv(sk, block_k), body, (os, m_i, l_i))
+  q_valid = (pl.program_id(0) * block_q + jnp.arange(block_q)) < sq
+  for oc, (off, w), c in zip(os, chunks, cols):
+    m = q_valid[:, None] if c is None else q_valid[:, None] & c[None, :]
+    plgpu.store(o_ref.at[:, pl.dslice(off, w)],
+                (oc / l_i[:, None]).astype(o_ref.dtype), mask=m)
+
+
+@functools.partial(jax.jit, static_argnames=(
+    "sm_scale", "block_q", "block_k", "num_warps", "num_stages", "seq_major"))
+def tri_flash(q, k, v, bias, kmask, *, sm_scale, block_q=64, block_k=64,
+              num_warps=4, num_stages=2, seq_major=False):
+  # q/k/v [N,H,Sq,D], or [N,Sq,H,D] when seq_major; bias [H,Sq,Sk]; kmask [N,Sk]
+  # bool. No padding: Sq/Sk need not be block multiples; the kernel masks the
+  # partial q-block (output) and the partial k-block (loads), so the big q/k/v
+  # tensors are never copied.
+  if seq_major:
+    N, Sq, H, D = q.shape
+    Sk = k.shape[1]
+  else:
+    N, H, Sq, D = q.shape
+    Sk = k.shape[2]
+  chunks = _chunks(D)
+  dk = chunks[-1][0] + chunks[-1][1]           # the tile the chunks span
+  # grid (q_block, head, row): row N is innermost so the shared bias[h,q_block,:]
+  # stays hot in L2 while sweeping rows (the bias is re-read for every row).
+  grid = (pl.cdiv(Sq, block_q), H, N)
+  if seq_major:
+    bs_qo = pl.BlockSpec((None, block_q, None, dk), lambda i, h, j: (j, i, h, 0))
+    bs_kv = pl.BlockSpec((None, Sk, None, dk),      lambda i, h, j: (j, 0, h, 0))
+    out_shape = jax.ShapeDtypeStruct((N, Sq, H, D), q.dtype)
+  else:
+    bs_qo = pl.BlockSpec((None, None, block_q, dk), lambda i, h, j: (j, h, i, 0))
+    bs_kv = pl.BlockSpec((None, None, Sk, dk),      lambda i, h, j: (j, h, 0, 0))
+    out_shape = jax.ShapeDtypeStruct((N, H, Sq, D), q.dtype)
+  bs_bias = pl.BlockSpec((None, block_q, Sk),    lambda i, h, j: (h, i, 0))
+  bs_mask = pl.BlockSpec((None, Sk),             lambda i, h, j: (j, 0))
+  return pl.pallas_call(
+      functools.partial(_kernel, sm_scale=sm_scale, block_q=block_q,
+                        block_k=block_k, sq=Sq, d=D, chunks=chunks),
+      grid=grid,
+      in_specs=[bs_qo, bs_kv, bs_kv, bs_bias, bs_mask],
+      out_specs=bs_qo,
+      out_shape=out_shape,
+      compiler_params=plgpu.CompilerParams(num_warps=num_warps, num_stages=num_stages),
+      name="tri_flash_fwd",
+  )(q, k, v, bias, kmask)
+
+
+def pallas_attention(q, k, v, mask_bias, nonbatched_bias, scale, *, seq_major=False):
+  """AF2 attention via the Pallas flash kernel.
+
+  q/k/v heads-major [b, h, S, c], or seq-major [b, S, h, c] when seq_major.
+  mask_bias additive [b,1,1,S_kv] (~-1e9 = invalid); nonbatched_bias
+  [h,S_qo,S_kv] shared, or None.
+  Returns q's layout.
+  """
+  if seq_major:
+    b, sq, h, c = q.shape
+    sk = k.shape[1]
+  else:
+    b, h, sq, c = q.shape
+    sk = k.shape[2]
+  # The kernel reuses q's channel dim for v (one BlockSpec), so key_dim must equal value_dim
+  # True for AF2's attention
+  assert q.shape[-1] == v.shape[-1], (
+      f'tri_flash needs key_dim == value_dim, got {q.shape[-1]} vs {v.shape[-1]}')
+  # mask_bias is 0 or big_neg(dtype): -1e9, or -1e4 in float16. Template pointwise
+  # attention shares one row across the batch, which the row grid index cannot index.
+  kmask = jnp.broadcast_to(mask_bias[:, 0, 0, :] > -1e3, (b, sk))   # [b, S_kv] bool
+  bias = (jnp.zeros((h, sq, sk), q.dtype)
+          if nonbatched_bias is None else nonbatched_bias.astype(q.dtype))
+  # Key tile matches the head: a wider tile stages smem the kernel never reads.
+  block_k = _pow2_tile(min(64, c))
+  # a head dim that is not a power of two, or below 16, is tiled wider and masked
+  out = tri_flash(q, k, v, bias, kmask, sm_scale=float(scale),
+                  block_q=64, block_k=block_k, seq_major=seq_major)
+  return out.astype(q.dtype)
+
+
+def ref_attn(q, k, v, bias, kmask, sm_scale):
+  logits = jnp.einsum('nhqd,nhkd->nhqk', q, k) * sm_scale + bias[None]
+  logits = jnp.where(kmask[:, None, None, :], logits, NEG)
+  w = jax.nn.softmax(logits.astype(jnp.float32), axis=-1).astype(v.dtype)
+  return jnp.einsum('nhqk,nhkd->nhqd', w, v)
+
+

colabfold_kernels-0.1.0/src/colabfold_kernels/tri_mul.py

@@ -0,0 +1,213 @@
+"""Pallas/Triton fused gated dual projection for AF2 triangle-multiplication.
+
+Inference (forward) only, no backward pass
+
+Computes the masked, sigmoid-gated dual projection in one kernel so the [M, 2c]
+projection/gate intermediates are never materialized. The triangle matmul itself
+stays a native XLA batched GEMM.
+
+  out[m, p] = mask[m] * (x[m]·wp[:,p] + bp[p]) * sigmoid(x[m]·wg[:,p] + bg[p])
+where x = layer-normed pair [M=N*N, K=c_z], wp/wg [K, P=2*c_i].
+"""
+import functools
+import jax, jax.numpy as jnp
+from jax.experimental import pallas as pl
+from jax.experimental.pallas import triton as plgpu
+
+# jnp.dot is stable across jax versions, contrary to pl.dot/plgpu.dot
+_dot = functools.partial(jnp.dot, preferred_element_type=jnp.float32)
+
+
+def _gdp_kernel(x_ref, wp_ref, bp_ref, wg_ref, bg_ref, mask_ref, o_ref, *,
+                activation=jax.nn.sigmoid):
+  x = x_ref[...]                                      # [BM, K] bf16
+  proj = _dot(x, wp_ref[...]) + bp_ref[...][None, :].astype(jnp.float32)
+  gate = _dot(x, wg_ref[...]) + bg_ref[...][None, :].astype(jnp.float32)
+  o = mask_ref[...][:, None].astype(jnp.float32) * proj * activation(gate)
+  o_ref[...] = o.astype(o_ref.dtype)
+
+
+def _half(x, wp, bp, wg, bg, mask):
+  proj = _dot(x, wp) + bp[None, :].astype(jnp.float32)
+  gate = _dot(x, wg) + bg[None, :].astype(jnp.float32)
+  return mask[:, None].astype(jnp.float32) * proj * jax.nn.sigmoid(gate)
+
+
+def _gdp_kernel_split_cm(x_ref, wpl_ref, bpl_ref, wgl_ref, bgl_ref,
+                         wpr_ref, bpr_ref, wgr_ref, bgr_ref, mask_ref,
+                         ol_ref, or_ref):
+  # _gdp_kernel_split, stored channel-major [ci, BM] for the triangle einsum.
+  x = x_ref[...]
+  m = mask_ref[...]
+  ol_ref[...] = _half(x, wpl_ref[...], bpl_ref[...], wgl_ref[...], bgl_ref[...],
+                      m).T.astype(ol_ref.dtype)
+  or_ref[...] = _half(x, wpr_ref[...], bpr_ref[...], wgr_ref[...], bgr_ref[...],
+                      m).T.astype(or_ref.dtype)
+
+
+def _gdp_kernel_split(x_ref, wpl_ref, bpl_ref, wgl_ref, bgl_ref,
+                      wpr_ref, bpr_ref, wgr_ref, bgr_ref, mask_ref,
+                      ol_ref, or_ref):
+  # Computes the left [BM, ci] and right [BM, ci] gated projections as two
+  # separate dots (weights pre-split on the host) and writes them to two
+  # contiguous outputs -- so the triangle einsum consumes them directly with no
+  # post-kernel strided slice/copy of the combined [N,N,2ci] tensor.
+  x = x_ref[...]
+  m = mask_ref[...]
+  ol_ref[...] = _half(x, wpl_ref[...], bpl_ref[...], wgl_ref[...], bgl_ref[...],
+                      m).astype(ol_ref.dtype)
+  or_ref[...] = _half(x, wpr_ref[...], bpr_ref[...], wgr_ref[...], bgr_ref[...],
+                      m).astype(or_ref.dtype)
+
+
+@functools.partial(jax.jit,
+                   static_argnames=("block_m", "split", "channel_major", "activation"))
+def gated_dual_proj(x, wp, bp, wg, bg, mask, *, block_m=64, split=False,
+                    channel_major=False, activation=jax.nn.sigmoid):
+  # x [M,K]; wp/wg [K,P]; bp/bg [P]; mask [M].
+  # split=False -> gated [M,P]; split=True -> (left [M,P/2], right [M,P/2]),
+  # the two halves written contiguously so the caller needs no slice.
+  M, K = x.shape
+  P = wp.shape[1]
+  Mp = -(-M // block_m) * block_m
+  x = jnp.pad(x, [(0, Mp - M), (0, 0)])
+  mask = jnp.pad(mask, [(0, Mp - M)])
+  in_specs = [
+      pl.BlockSpec((block_m, K), lambda i: (i, 0)),
+      pl.BlockSpec((K, P), lambda i: (0, 0)),
+      pl.BlockSpec((P,), lambda i: (0,)),
+      pl.BlockSpec((K, P), lambda i: (0, 0)),
+      pl.BlockSpec((P,), lambda i: (0,)),
+      pl.BlockSpec((block_m,), lambda i: (i,)),
+  ]
+  cp = plgpu.CompilerParams(num_warps=4, num_stages=1)
+  if split:
+    ci = P // 2
+    wpl, wpr = wp[:, :ci], wp[:, ci:]
+    wgl, wgr = wg[:, :ci], wg[:, ci:]
+    bpl, bpr = bp[:ci], bp[ci:]
+    bgl, bgr = bg[:ci], bg[ci:]
+    wspec = pl.BlockSpec((K, ci), lambda i: (0, 0))
+    bspec = pl.BlockSpec((ci,), lambda i: (0,))
+    split_in = [in_specs[0], wspec, bspec, wspec, bspec, wspec, bspec,
+                wspec, bspec, in_specs[5]]
+    if channel_major:
+      bs = pl.BlockSpec((ci, block_m), lambda i: (0, i))
+      left, right = pl.pallas_call(
+          _gdp_kernel_split_cm,
+          grid=(Mp // block_m,), in_specs=split_in, out_specs=[bs, bs],
+          out_shape=[jax.ShapeDtypeStruct((ci, Mp), x.dtype),
+                     jax.ShapeDtypeStruct((ci, Mp), x.dtype)],
+          compiler_params=cp, name="gated_dual_proj_split_cm",
+      )(x, wpl, bpl, wgl, bgl, wpr, bpr, wgr, bgr, mask)
+      return left[:, :M], right[:, :M]
+    bs = pl.BlockSpec((block_m, ci), lambda i: (i, 0))
+    left, right = pl.pallas_call(
+        _gdp_kernel_split,
+        grid=(Mp // block_m,), in_specs=split_in, out_specs=[bs, bs],
+        out_shape=[jax.ShapeDtypeStruct((Mp, ci), x.dtype),
+                   jax.ShapeDtypeStruct((Mp, ci), x.dtype)],
+        compiler_params=cp, name="gated_dual_proj_split",
+    )(x, wpl, bpl, wgl, bgl, wpr, bpr, wgr, bgr, mask)
+    return left[:M], right[:M]
+  out = pl.pallas_call(
+      functools.partial(_gdp_kernel, activation=activation),
+      grid=(Mp // block_m,), in_specs=in_specs,
+      out_specs=pl.BlockSpec((block_m, P), lambda i: (i, 0)),
+      out_shape=jax.ShapeDtypeStruct((Mp, P), x.dtype),
+      compiler_params=cp, name="gated_dual_proj",
+  )(x, wp, bp, wg, bg, mask)
+  return out[:M]
+
+
+def ref_gdp(x, wp, bp, wg, bg, mask):
+  proj = x @ wp + bp
+  gate = x @ wg + bg
+  return (mask[:, None] * proj * jax.nn.sigmoid(gate))
+
+
+# Fused LayerNorm: bf16 in/out, fp32 internal. One kernel replaces the
+# bf16->fp32 upcast + layernorm + fp32->bf16 downcast (kills the convert ops
+# and never materializes the fp32 tensor), while keeping fp32 accuracy.
+def _ln_kernel(x_ref, scale_ref, offset_ref, o_ref, *, eps, m, c, block_m):
+  bc = x_ref.shape[-1]
+  rows = (pl.program_id(0) * block_m + jnp.arange(block_m)) < m     # [block_m]
+  if bc == c:
+    x = plgpu.load(x_ref, mask=rows[:, None], other=0.0).astype(jnp.float32)
+    mean = jnp.mean(x, axis=-1, keepdims=True)
+    d = x - mean
+    var = jnp.mean(d * d, axis=-1, keepdims=True)
+    scale, offset = scale_ref[...], offset_ref[...]
+    valid = rows[:, None]
+  else:
+    # c is not a power of two: tile to the next one, mask the lanes past c
+    cols = jnp.arange(bc) < c
+    valid = rows[:, None] & cols[None, :]
+    x = plgpu.load(x_ref, mask=valid, other=0.0).astype(jnp.float32)
+    mean = jnp.sum(x, axis=-1, keepdims=True) / c
+    d = jnp.where(cols[None, :], x - mean, 0.0)
+    var = jnp.sum(d * d, axis=-1, keepdims=True) / c
+    scale = plgpu.load(scale_ref, mask=cols, other=0.0)
+    offset = plgpu.load(offset_ref, mask=cols, other=0.0)
+  y = d * jax.lax.rsqrt(var + eps)
+  y = y * scale[None, :].astype(jnp.float32) + offset[None, :].astype(jnp.float32)
+  plgpu.store(o_ref, y.astype(o_ref.dtype), mask=valid)
+
+
+# Triton keeps the tile in registers, so a wider one only spills and compiles slower
+_TILE_BYTES = 32 * 1024
+
+
+@functools.partial(jax.jit, static_argnames=("eps", "block_m"))
+def pallas_layer_norm(x, scale, offset, *, eps=1e-5, block_m=64):
+  # x [..., c]; normalize over last axis; scale/offset [c]. Returns x's shape/dtype.
+  c = x.shape[-1]
+  bc = max(16, 1 << (c - 1).bit_length())        # Triton tiles are powers of two
+  xf = x.reshape(-1, c)
+  m = xf.shape[0]
+  rows = max(1, _TILE_BYTES // (bc * x.dtype.itemsize))
+  block_m = min(block_m, 1 << (rows.bit_length() - 1))
+  out = pl.pallas_call(
+      functools.partial(_ln_kernel, eps=eps, m=m, c=c, block_m=block_m),
+      grid=(pl.cdiv(m, block_m),),
+      in_specs=[pl.BlockSpec((block_m, bc), lambda i: (i, 0)),
+                pl.BlockSpec((bc,), lambda i: (0,)),
+                pl.BlockSpec((bc,), lambda i: (0,))],
+      out_specs=pl.BlockSpec((block_m, bc), lambda i: (i, 0)),
+      out_shape=jax.ShapeDtypeStruct((m, c), x.dtype),
+      compiler_params=plgpu.CompilerParams(num_warps=4, num_stages=2),
+      name="pallas_layer_norm",
+  )(xf, scale, offset)
+  return out.reshape(x.shape)
+
+
+def ref_layer_norm(x, scale, offset, eps=1e-5):
+  xf = x.astype(jnp.float32)
+  mean = jnp.mean(xf, -1, keepdims=True)
+  var = jnp.mean((xf - mean) ** 2, -1, keepdims=True)
+  y = (xf - mean) * jax.lax.rsqrt(var + eps) * scale + offset
+  return y.astype(x.dtype)
+
+
+if __name__ == "__main__":
+  import time
+  key = jax.random.PRNGKey(0)
+  M, K, P = 1000 * 1000, 128, 256
+  x = (jax.random.normal(key, (M, K)) * 0.3).astype(jnp.bfloat16)
+  wp = (jax.random.normal(key, (K, P)) * 0.1).astype(jnp.bfloat16)
+  wg = (jax.random.normal(key, (K, P)) * 0.1).astype(jnp.bfloat16)
+  bp = (jax.random.normal(key, (P,)) * 0.1).astype(jnp.bfloat16)
+  bg = (jax.random.normal(key, (P,)) * 0.1).astype(jnp.bfloat16)
+  mask = (jax.random.uniform(key, (M,)) > 0.1).astype(jnp.bfloat16)
+  o = gated_dual_proj(x, wp, bp, wg, bg, mask)
+  r = ref_gdp(x.astype(jnp.float32), wp.astype(jnp.float32), bp.astype(jnp.float32),
+              wg.astype(jnp.float32), bg.astype(jnp.float32), mask.astype(jnp.float32))
+  err = float(jnp.max(jnp.abs(o.astype(jnp.float32) - r)))
+  print(f"gated_dual_proj bf16 err={err:.3e}")
+  f = jax.jit(gated_dual_proj)
+  o = f(x, wp, bp, wg, bg, mask); o.block_until_ready()
+  t = time.time()
+  for _ in range(50): o = f(x, wp, bp, wg, bg, mask)
+  o.block_until_ready()
+  print(f"per-call {(time.time()-t)/50*1000:.2f} ms")
+

colabfold_kernels-0.1.0/src/colabfold_kernels/volta.py

@@ -0,0 +1,169 @@
+"""sm_70/75 CUDA kernels (colabfold-legacy-kernels) used instead of Pallas."""
+import ctypes
+import functools
+import os
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+# Head dims each kernel instantiates: mma.sync (sm_75+) and wmma (sm_70).
+_HEAD_DIMS_MMA = (8, 16, 32, 64)
+_HEAD_DIMS_WMMA = (8, 16, 32, 64)
+
+_LOADED = {}
+
+
+def _package_path(kernel, cc):
+    """Give the library path from colabfold_legacy_kernels, or None."""
+    try:
+        import colabfold_legacy_kernels as clk
+    except ImportError:
+        return None
+    try:
+        return clk.library_path(kernel, cc)
+    except (FileNotFoundError, KeyError):
+        return None
+
+
+def _load(kernel, cc, symbols):
+    """dlopen the library for one kernel and register its FFI targets."""
+    key = (kernel, cc, tuple(symbols))
+    if key in _LOADED:
+        return _LOADED[key]
+    path = _package_path(kernel, cc)
+    if not path or not os.path.exists(path):
+        _LOADED[key] = False
+        return False
+    lib = ctypes.cdll.LoadLibrary(path)
+    for sym in symbols:
+        jax.ffi.register_ffi_target(
+            sym, jax.ffi.pycapsule(getattr(lib, sym)), platform="CUDA")
+    _LOADED[key] = True
+    return True
+
+
+def _attn_symbol(cc):
+    """sm_75+ gets the CUTLASS kernel; sm_70 the wmma one."""
+    return "VoltaMma" if cc >= 75 else "VoltaWmma"
+
+
+def available(cc):
+    """True if the attention library for this device loads."""
+    return _load("attention", cc, (_attn_symbol(cc),))
+
+
+def _native_dims(cc):
+    return _HEAD_DIMS_MMA if cc >= 75 else _HEAD_DIMS_WMMA
+
+
+def _pad_target(head_dim, cc):
+    """Smallest supported head dim >= head_dim, or None."""
+    for d in sorted(_native_dims(cc)):
+        if d >= head_dim:
+            return d
+    return None
+
+
+def supports(head_dim, cc):
+    if head_dim in _native_dims(cc):
+        return True
+    # wmma: zero-pad to the next supported head dim (exact, beats XLA on V100).
+    if cc < 75:
+        return _pad_target(head_dim, cc) is not None
+    return False
+
+
+def ops_available(cc):
+    """LayerNorm is Volta-capable; the CUTLASS gdp needs sm_75+."""
+    syms = ("VoltaLayerNorm", "VoltaGdp") if cc >= 75 else ("VoltaLayerNorm",)
+    ok = _load("layer_norm", cc, syms)
+    if ok and cc < 75:
+        ok = _load("gated_dual_proj", cc, ("VoltaGdpWmma",))
+    return ok
+
+
+@functools.partial(jax.jit, static_argnames=("sym", "sm_scale", "block_q", "block_k"))
+def _attn_call(q, k, v, bias, kmask, *, sym, sm_scale, block_q, block_k):
+    n, h, sq, d = q.shape
+    # sequential: the template stack vmaps over templates, and the handlers take a fixed rank.
+    return jax.ffi.ffi_call(sym, jax.ShapeDtypeStruct((n, h, sq, d), jnp.float16),
+                            vmap_method="sequential")(
+        q, k, v, bias, kmask,
+        scale=np.float32(sm_scale),
+        block_q=np.int64(block_q), block_k=np.int64(block_k))
+
+
+def volta_attention(q, k, v, mask_bias, nonbatched_bias, scale, cc,
+                    block_q=64, block_k=32):
+    """Drop-in for tri_flash.pallas_attention. q/k/v [b,h,S,c] float16."""
+    in_dtype = q.dtype
+    b, h, sq, c = q.shape
+    sk = k.shape[2]
+    assert c == v.shape[-1], "kernel needs key_dim == value_dim"
+    pad_to = None
+    if c not in _native_dims(cc):
+        pad_to = _pad_target(c, cc)
+        assert pad_to is not None, f"unsupported head_dim {c}"
+        pad = [(0, 0)] * 3 + [(0, pad_to - c)]
+        q, k, v = (jnp.pad(t.astype(jnp.float16), pad) for t in (q, k, v))
+    if cc >= 75 and (pad_to or c) == 64 and (block_q, block_k) == (64, 32):
+        block_k = 64    # the mma kernel has no (64, 64, 32) instantiation
+    # Template pointwise attention shares one mask row across the batch; the kernel
+    # indexes the mask per row, so broadcast instead of reading past its end.
+    kmask = jnp.broadcast_to(mask_bias[:, 0, 0, :] > -1e3, (b, sk)).astype(jnp.uint8)
+    bias = (jnp.zeros((h, sq, sk), jnp.float16) if nonbatched_bias is None
+            else nonbatched_bias.astype(jnp.float16))
+    out = _attn_call(q.astype(jnp.float16), k.astype(jnp.float16),
+                     v.astype(jnp.float16), bias, kmask,
+                     sym=_attn_symbol(cc), sm_scale=float(scale),
+                     block_q=block_q, block_k=block_k)
+    if pad_to is not None:
+        out = out[..., :c]
+    return out.astype(in_dtype)
+
+
+@functools.partial(jax.jit, static_argnames=("eps",))
+def _ln_call(x, scale, offset, *, eps):
+    m, c = x.shape
+    return jax.ffi.ffi_call(
+        "VoltaLayerNorm", jax.ShapeDtypeStruct((m, c), jnp.float16),
+        vmap_method="sequential")(
+            x, scale, offset, eps=np.float32(eps))
+
+
+def volta_layer_norm(x, scale, offset, *, eps=1e-5):
+    """Drop-in for tri_mul.pallas_layer_norm. Normalises the last axis."""
+    c = x.shape[-1]
+    out = _ln_call(x.reshape(-1, c).astype(jnp.float16),
+                   scale.astype(jnp.float32), offset.astype(jnp.float32), eps=eps)
+    return out.reshape(x.shape).astype(x.dtype)
+
+
+@functools.partial(jax.jit, static_argnames=("sym",))
+def _gdp_call(x, wp, bp, wg, bg, mask, *, sym):
+    m = x.shape[0]
+    n = wp.shape[1]
+    return jax.ffi.ffi_call(sym, jax.ShapeDtypeStruct((m, n), jnp.float16),
+                            vmap_method="sequential")(
+        x, wp, bp, wg, bg, mask)
+
+
+def volta_gated_dual_proj(x, wp, bp, wg, bg, mask, cc, *, split=True,
+                          channel_major=False):
+    """Drop-in for tri_mul.gated_dual_proj. x [M,K]; wp/wg [K,P]; mask [M]."""
+    sym = "VoltaGdp" if cc >= 75 else "VoltaGdpWmma"
+    f16 = lambda t: t.astype(jnp.float16)
+    if not split:
+        return _gdp_call(f16(x), f16(wp), f16(bp), f16(wg), f16(bg),
+                         f16(mask), sym=sym).astype(x.dtype)
+    ci = wp.shape[1] // 2
+    halves = []
+    for lo, hi in ((0, ci), (ci, 2 * ci)):
+        halves.append(_gdp_call(f16(x), f16(wp[:, lo:hi]), f16(bp[lo:hi]),
+                                f16(wg[:, lo:hi]), f16(bg[lo:hi]), f16(mask),
+                                sym=sym).astype(x.dtype))
+    if channel_major:
+        # XLA would transpose row-major halves for the einsum anyway.
+        halves = [h.T for h in halves]
+    return halves[0], halves[1]

colabfold_kernels-0.1.0/src/colabfold_kernels.egg-info/PKG-INFO

@@ -0,0 +1,34 @@
+Metadata-Version: 2.4
+Name: colabfold-kernels
+Version: 0.1.0
+Summary: Fast fused kernels for ColabFold
+Author-email: Milot Mirdita <milot@mirdita.de>
+License: MIT
+Project-URL: Homepage, https://github.com/mirditalab/colabfold-kernels
+Project-URL: Issues, https://github.com/mirditalab/colabfold-kernels/issues
+Keywords: alphafold,colabfold,pallas,triton,jax,attention,bioinformatics
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Dynamic: license-file
+
+# colabfold-kernels
+
+Fast fused kernels for ColabFold, written in
+[Pallas](https://docs.jax.dev/en/latest/pallas/index.html) with native CUDA-fallback
+for Volta- and Turing-generation GPUs.
+
+```python
+import jax.numpy as jnp
+from colabfold_kernels import attention, dispatch
+
+kernel = attention(dispatch(), jnp.bfloat16, key_dim=96, value_dim=96, layout="seq")
+if kernel is not None:
+    out = kernel(q, k, v, mask_bias, nonbatched_bias, scale)
+```
+
+`attention`, `layer_norm` and `gated_dual_proj` each return `None` if not supported.

colabfold_kernels-0.1.0/src/colabfold_kernels.egg-info/SOURCES.txt

@@ -0,0 +1,13 @@
+LICENSE
+README.md
+pyproject.toml
+src/colabfold_kernels/__init__.py
+src/colabfold_kernels/device.py
+src/colabfold_kernels/fused_ops.py
+src/colabfold_kernels/tri_flash.py
+src/colabfold_kernels/tri_mul.py
+src/colabfold_kernels/volta.py
+src/colabfold_kernels.egg-info/PKG-INFO
+src/colabfold_kernels.egg-info/SOURCES.txt
+src/colabfold_kernels.egg-info/dependency_links.txt
+src/colabfold_kernels.egg-info/top_level.txt

colabfold_kernels-0.1.0/src/colabfold_kernels.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

colabfold_kernels-0.1.0/src/colabfold_kernels.egg-info/top_level.txt

@@ -0,0 +1 @@
+colabfold_kernels