De novo binder design via rfd3-mlx: the first method that generates a chain rather than folding one

[javierbq]

Summary

Add de novo binder design as an on-device capability, via the MLX-Swift port of
RFdiffusion3 in javierbq/rfd3-mlx (RFD3Kit).

Scope is deliberately one mode: given a target structure and a set of hotspot residues,
generate a new binder backbone against it. Sequence redesign, refolding, interface gating
and any run-until-N-pass loop are out of scope here — see Scope below.

The port is in better shape than protenix-mlx was at #309: the network is complete, the
featurizer is already in Swift, it designs end to end against a real target, and no MLX
type leaks through the public surface. The work is integration, not porting.

What already exists

javierbq/rfd3-mlx — a .library("RFD3Kit") SwiftPM product, macOS 14 / iOS 17, mlx-swift
pinned exact: 0.31.6 (the same pin MPNNKit and boltz-mlx already hold).

Layer	State
Swift engine (token encoder/init, diffusion transformer, decoder, sampler, sequence head)	working
Swift featurizer (Featurizer.binderDesign, motif/hotspot conditioning)	working, in Swift
End-to-end design (target PDB → binder backbone + sequence)	working
Measured peak-memory model + pre-GPU refusal (RFD3Budget, preflight)	working
Weight pack format + sha256 manifest + provenance (RFD3Pack)	working
Published tag / release pack	neither exists — see Package hygiene
Motif scaffolding, symmetry, small-molecule binders, partial diffusion	not exposed

The whole surface RayMol would consume:

public struct RFD3Model {
  public init(packDirectory: URL, verifyChecksums: Bool = true) throws
  public static func autoTarget(fromPDB: String) -> (residues: [Residue], hotspots: [Int])?
  public static func preflight(target: [Residue], options: Options = Options(),
                               budgetBytes: Int? = nil) throws -> FeatSet
  public func designBinder(target: [Residue], options: Options) throws -> Result
}
// Options: binderLength, hotspots, numTimesteps, nRecycle, seed, memoryBudgetBytes,
//          onProgress: ((Int, Int) -> Void)?
// Result:  pdb (chain A = target, held fixed; chain B = designed binder),
//          binderSequence, stats, seconds
// Stats:   binderCACAmeanA, backboneValidPct, radiusOfGyrationA, interfaceMinA,
//          contactsUnder8A, binderToHotspotMinA, targetDriftMaxA

Worth noting: RFD3 has a sequence head, so a design arrives with a sequence already.
That is what makes this mode useful on its own — the output can be refolded by the existing
cmd.predict without MPNN in the loop.

The architectural crux: this generates a chain, it does not fold one

Every method in docs/predictors.md maps a sequence to a structure. RFD3 does not fit that
contract anywhere:

	predictors today	RFD3
input	chain sequences (typed or from a selection)	a target structure (coordinates) + hotspot residue indices + a binder length
output	a fold of the chains you named	a chain that did not exist
spec	PredictionSpec(chains, name, alignments)	has no sequence to put in chains
wire	Request.chains: [Chain]	needs target coordinates and hotspots

So Predictor is the wrong base class, and predict rfd3, ... is the wrong surface. The
decision this issue has to make is where the seam goes. Options, in the order I'd rank them:

1. A sibling generators/ registry. weights.py (WeightBundle/WeightCache) and
   host.py (job transport) are already method-agnostic and get reused verbatim; only
   base.py's spec/contract is sequence-shaped. New surface, e.g. cmd.design_binder.
   Keeps Predictor honest and the predict UI untouched.
2. Widen the abstraction to a Method with predictor/generator variants. More churn in
   code that three shipped methods depend on, for one caller.
3. Special-case it. Cheapest now, worst later; motif scaffolding is one public entry
   point away in the featurizer, so there will be a second generator.

Input comes from a selection, not a typed sequence — which makes the natural UI
precedent Design mode (#217), not the predict bar (#288).

What RayMol already has that this reuses

Most of the hard infrastructure is done, which is the main argument for doing this now:

• Run MLX inference out-of-process so an OOM fails the job instead of the session #273 out-of-process MLX — load-bearing here. RFD3's OOM is an uncatchable
  std::terminate from a Metal completion handler; preflight covers the predictable case
  and Run MLX inference out-of-process so an OOM fails the job instead of the session #273 covers the rest.
• Weight download freezes the app with no UI: first predict blocks the main thread for 529 MB, so its own progress messages can't render #284 / Move MPNN.mpnnpack onto WeightCache instead of bundling it #275 WeightCache + non-blocking download UI — .rfd3pack becomes one more
  digest-pinned WeightBundle.
• Running predictions need an ambient progress indicator: today a job is invisible unless the object panel is open, and always invisible on iOS #291 progress tray — Options.onProgress(step, total) maps onto compose_progress
  directly. It needs it: a full-HSA-scale target is minutes per design.
• Objects can carry no data: RayMol needs generic per-object metrics, scoped to the object, a state, or a residue range #308 per-object metrics — the natural home for Stats, so a designed backbone carries
  its own geometry rather than printing it to the console.
• Toolbar is out of room: fold Move, Measure and Design into a single "Tools" menu #304 Tools menu — where the entry point lives.
• Interactive on-device design tool: logit-colored backbone + per-position AA editing with live repack #217 Design mode / MPNNKit — already ships; the eventual sequence-redesign step.

Package hygiene, before RayMol takes the dependency

• The root SPM manifest is not pushed. rfd3-mlx is 4 commits ahead of origin/main
  locally, and the commit that moves the manifest to the repo root is one of them. Until it
  lands, RayMol can only reference this as .package(path:), which cannot be checked in —
  it breaks CI, a clean clone, and the release/notarization and MAS archive flows.
• No tags. Nothing to pin upToNextMajor against. Cut 0.1.0; RayMol pins boltz-mlx at
  0.2.1, protenix-mlx at 0.1.3 and proteinmpnn-mlx at 0.1.2, and should not depend on a
  branch here.
• No published weight pack. There is no release asset to put a digest on, so no
  WeightBundle can be written yet.
• mlx-swift exact: 0.31.6 — a third exact pin on the same package. Not a blocker;
  worth relaxing all of them to a range eventually.
• xcodebuild needs -skipPackagePluginValidation -skipMacroValidation (mlx-swift's
  CudaBuild prebuild plugin), same as today.
• The engine sources are duplicated in-repo (RFD3Kit/Sources/RFD3Kit and
  swift/Sources/RFD3Core are byte-identical). Upstream cleanup, not RayMol's problem, but
  it is a divergence risk for anything we report bugs against.

Measured costs, and the lever that is still on the table

From rfd3-mlx's own benchmarking, peak GPU allocation is well modelled by atom count
A = Σ_target(resolved atoms) + 14 · binderLength:

peak_bytes ≈ 995.5·A² − 524,880·A + 1.226e9      (max residual 0.96%)

That decomposes as roughly 9% architecture floor (a [A, A, 16] fp32 pair tensor alive
for the whole run), 44% a one-time setup transient, and 47% MLX allocator slack.
RFD3Budget.defaultBudgetBytes() takes 60% of physical memory; preflight refuses above it
without touching the GPU. Runtime is minutes per design at 200 timesteps on an
A17/M-class GPU, which is what makes #291's progress tray and working cancellation
non-optional rather than polish.

One lever is already identified upstream and not yet taken: SinusoidalDistEmbed builds a
[A, A, 64] fp32 tensor whose output is identically zero outside the motif block, so at
realistic sizes ~99.9% of that work is multiplied by zero. Row-blocking it measured
bit-identical and takes the model from 995.5 to ~62 B/atom² — a ~16× cut in peak, and it
touches only setup, not the diffusion loop. Worth doing in rfd3-mlx before RayMol ships a
size guard fitted to the current number, or the guard will be refusing designs that would
now fit.

Where this sits in the pipeline

The product this is a stage of:

target + hotspots
  → [1] RFD3            backbone + sequence
  → [2] ProteinMPNN     resequence          (OPTIONAL — see below)
  → [3] Boltz refold    target MSA only, binder none
  → [4] filter          min_ipSAE + refold-vs-design RMSD
  → repeat until N pass

#	stage	input	output	in RayMol today
1	RFD3	target coords + hotspots + length	backbone and sequence	this issue
2	ProteinMPNN	backbone, target fixed	sequence(s)	MPNNKit ships (#217)
3	Boltz refold	2-chain complex + target MSA	structure + PAE	ships (cmd.predict)
4	filter	PAE + both structures	pass/fail	nothing
—	loop	—	N passing designs	nothing

This issue is box 1 only. But three things about the rest have to be settled here,
because they are cheap now and expensive to retrofit:

1. A design and its refold must stay associated. The honest test of a design is not
   min_ipSAE alone, it is whether the refold reproduces the designed pose —
   target-anchored RMSD between chain B as designed and chain B as refolded. That requires
   a refold to know which design it came from, which makes it a data-model decision in
   this issue, whatever representation designs get (objects, states, groups). Concrete
   reason to care: a measured run scoring min_ipSAE 0.70 had its binder docked 15.6 Å from
   the reference pose. The scalar alone would have passed it; the pairing is what catches
   it.

2. The handoff to Boltz is already free — don't invent a format for it. Result.pdb is
   chain A = target (held fixed) + chain B = designed binder, which is exactly the two-chain
   complex the refold takes, and the wire Request already supports a real alignment for
   the target and none for the binder. So the 1→3 edge costs nothing provided this issue
   emits that pair as one object rather than flattening it or dropping the target.

3. MPNN is a branch, not a stage. RFD3's sequence head means a design already carries a
   sequence, so box 2 is optional and off the critical path. Worth stating so the pipeline
   isn't built assuming it.

Deliberately not decided here: the loop policy (how many trajectories, when to stop),
the gate cutoff, and whether MPNN runs by default. Those belong in a (tracking) issue for
the pipeline — following #274's precedent — which this issue would be the first child of.

Scope

In: the RFD3Kit dependency; an rfd3 runtime and job manager reusing the #273
out-of-process host; a .rfd3pack weight bundle; a command surface taking a target
selection + hotspots + binder length + seed/steps; the designed backbone loaded as an object
whose Stats land as per-object metrics; a minimal entry point under Tools; macOS only.

Out: MPNN redesign of the generated sequence; refolding and any interface gate
(min_ipSAE); the run-until-N-pass loop and campaign management; motif scaffolding, symmetry,
small-molecule and nucleic-acid binders, partial diffusion; iOS (follow #249's pattern).

Naming guardrail

A generated chain is a designed backbone, never a "binder", anywhere in the UI, object
names, or metrics — until it has been refolded and passed an interface gate, neither of which
is in this issue. This is a deliberate product rule, not a wording preference: the whole
point of the later refold step is that generation alone does not license the word.

Suggested staging

1. Upstream, in rfd3-mlx: push the 4 pending commits, tag 0.1.0, publish a pack as a
   digest-pinned release asset with manifest.json + weights, and state the pack's member
   list so a WeightBundle can be written against it. Take the SinusoidalDistEmbed
   row-blocking win here too, so the size guard is fitted once to the number we ship.
2. Decide the seam (generators/ vs widening Predictor) — everything else follows from
   it, so this is the one thing worth settling before code.
3. RayMol, macOS only: the generator module + tests, RFD3JobManager with a per-runtime
   size guard fed by RFD3Budget, the runtime added to what
   PyMOLBridge.mm advertises in RAYMOL_PREDICT_RUNTIMES (so check_available can refuse
   before a download), measured memory table.
4. UI: target/hotspot picking (Design mode's residue picking, not the predict bar),
   length + count + seed, progress via Running predictions need an ambient progress indicator: today a job is invisible unless the object panel is open, and always invisible on iOS #291, results as objects carrying Objects can carry no data: RayMol needs generic per-object metrics, scoped to the object, a state, or a residue range #308 metrics.
5. Follow-ups, separately: refold + interface gate; MPNN redesign; the loop; motif
   scaffolding; iOS.

Acceptance

• Nothing in modules/pymol/predictors/ changes behaviour for the three shipped
  predictors; if the seam is a sibling registry, predicting.py is untouched.
• check_available refuses on a build without the runtime, before any weight
  download, and names the missing runtime.
• preflight refuses an over-budget design without touching the GPU, and the refusal
  names the measured budget — asserted by a test with no model pack present
  (RFD3Model.preflight is static and weight-free for exactly this reason).
• A design job that OOMs fails the job, not the session (Run MLX inference out-of-process so an OOM fails the job instead of the session #273).
• The target is provably unmoved: Stats.targetDriftMaxA is asserted ~0, since the
  contract is that chain A is held fixed.
• Cancel works mid-rollout, and progress ticks per diffusion step.
• Stats are attached to the object as metrics (Objects can carry no data: RayMol needs generic per-object metrics, scoped to the object, a state, or a residue range #308), not printed.
• A design is emitted as the target + designed chain together, not the binder alone —
  that pair is the refold's input, so splitting it here makes stage 3 re-derive it.
• A design carries a stable identity (seed + options + target) that a later refold can be
  keyed to, so refold-vs-design RMSD is computable without guessing which design a
  prediction came from. No refold is implemented here; the association is.
• No UI string, object name or metric label calls the output a "binder".
• Both macOS and iOS slices hand-compiled before merge (feat(ios/ipad): iOS/iPad UI polish — landscape-only movie dock, hugging inspector, show/hide tongues + top pane rail #174, Load ~/.raymolrc startup script, migrated from ~/.pymolrc #226/Fix iOS build: gate ~/.raymolrc migration prompt to macOS #238), even though
  the feature is macOS-gated.

Related: #273, #284, #291, #308, #217, #304, #249. Precedent for adding a method: #306, #309.