Skip to content

atatess/manifold

Repository files navigation

Manifold

A local, type-aware diffusion code editor.
Whole-region code edits with an on-device diffusion LLM, steered toward type-correctness by a type checker inside the denoising loop.

tests MIT MLX on-device


The idea

Autoregressive assistants (Copilot and friends) are great at left-to-right completion and weak at "rewrite this whole region under a constraint while preserving the signature." That is a constrained denoising problem — what diffusion language models are good at.

Manifold leans into that, and adds the piece nobody ships: it closes a feedback loop between a type checker and the diffusion sampler. Between denoising steps, it runs the type checker on the partial code and uses the diagnostics to decide what to re-mask and re-write next. The model denoises toward type-correctness instead of being linted after the fact.

Three properties, combined — and as of mid-2026 no shipping product combines all three:

  1. Local / on-device. Runs on Apple Silicon via MLX. No cloud inference — works in airplane mode.
  2. Type-aware denoising. Type-checker feedback is injected between denoising steps, not as a post-hoc pass.
  3. Multi-region, cross-file. Joint denoising over several highlighted spans at once.

Model: ByteDance-Seed/Stable-DiffCoder-8B-Instruct (MIT, block diffusion), 4-bit MLX weights (~4.6 GB) — measured on an M2 Max at ~4.7 GB resident, ~3 s to load, and ~1 s per forward at interactive window sizes. (Per-forward latency is ~linear in sequence length: ~200 ms at a 16-token toy prompt, ~1 s at a ~220-token windowed edit; a prefix KV cache reclaims ~2.3x of it.)

It works today

Type feedback in the loop fixes the model's own type errors. Blank a function body, and the naive denoise fills a value that type-checks wrong; the type-aware loop runs pyright between steps, re-masks the offending span, and re-denoises toward a correct fill:

def this_year() -> str:
    """Return the current year as a string."""
    return ____

# naive denoise:      2023  # Example year        ← int, but the signature says -> str
# type-aware denoise: str(datetime.now().year)    ← type-correct AND right

Across the acceptance eval, type-aware denoising cut pyright errors 12 → 6 over the erroring cases (4/7 fully fixed), and — the point the raw count misses — it produces the semantically right fix where a checker-blind re-roll does not (year_str: 2023str(datetime.now().year)). A blind re-roll with the same budget but no diagnostics can lower the raw error count more by regenerating into degenerate fills; the type feedback is what steers toward code that is both type-correct and right. (Phase 4's ablation makes the same point more sharply: joint visibility without the checker scored worse than editing regions independently — see below.) The whole thing runs on-device and streams into a VS Code diff preview you accept or reject; scripts/phase3_e2e.py drives it end-to-end over the real stdio protocol (transcript).

The steering is a general hook, not a type-checker special case. Ban the % token mid-denoise and the same mechanism reroutes to a different, still-correct implementation:

prompt:  "Write a Python function is_even(n) that returns True if n is even."

denoise (no intervention):     return n % 2 == 0
denoise (ban "%" mid-loop):    return n // 2 == n / 2     ← rerouted, still correct

How it works

Stable-DiffCoder generates block-by-block: the edit region is a run of <[MASK_TOKEN]> tokens, and each block is refined over several denoising steps, committing the highest-confidence positions first. Attention is block-causal (bidirectional within a block, causal across blocks). Manifold reimplements this reverse-diffusion loop so it can pause between steps, read the partial code, and intervene — which is exactly where the type checker goes.

flowchart LR
    A[Masked region<br/>id 5 = MASK] --> B[Forward pass<br/>block-causal mask]
    B --> C{Pause · inspect}
    C -->|run pyright on<br/>partial decode| D[Diagnostics →<br/>token spans]
    D -->|re-mask / bias<br/>type-error spans| E[Commit top-k<br/>by confidence]
    E -->|holes remain| B
    E -->|done| F[Type-correct edit]
Loading

The forward pass is an injected dependency, so the loop is runtime-agnostic and unit-tested against a fake model — and the type-feedback policy is just one implementation of the step hook, never a separate lint pass.

Architecture

Two processes:

  • Model server (Python) — loads the model, runs the custom denoising loop with step-level hooks, and runs the type checker between steps. All model logic lives here. manifold serve speaks stdio JSON-RPC (NDJSON): one warm process keeps the 8B model loaded; stdout is protocol-only at the file- descriptor level so native library chatter can never corrupt a frame. The checker itself is kept warm too — a persistent pyright-langserver turns each in-loop check from a ~365 ms cold subprocess into a ~5 ms LSP round-trip, with an automatic fall-back to the subprocess checker if it is unavailable.
  • VS Code extension (TypeScript) — UI only, zero runtime dependencies: select a region, stream the live denoise into a diff preview, accept/reject. The model sees a token window around the region (per-forward latency is ~linear in sequence length; a prefix KV cache skips recomputing the stable prefix each step — ~2.3x on a prefix-heavy edit, byte-identical output); pyright always checks the full reconstructed file, with the original file's diagnostics subtracted as a baseline so only errors the fill introduced steer repair. See extension/.

Status & roadmap

Phase What State
0 Runnable MLX forward over a masked snippet ✅ done
1 Controllable block-diffusion loop (pause / inspect / intervene) ✅ done
2 pyright in the loop (the core demo) ✅ done — gate table
2.5 Prefix KV cache + warm pyright-langserver (perf) ✅ done — 2.3x, byte-identical; ~365 ms → ~5 ms/check
3 VS Code extension (stdio JSON-RPC), incl. natural-language instruction ✅ built — E2E transcript
4 Multi-region joint denoise ✅ done — gate table
4+ Cross-file joint denoise ✅ engine + RPC + demo; extension UI deferred

Every layer is validated at the cheapest honest level: the denoising loop, schedule, offsets, engine, and RPC server against an injected fake forward and fake checkers (219 pytest, plus a live pyright-langserver round-trip skipped when the binary is absent); the extension's codec/client/UI helpers and a cross-process transport test against the real manifold serve --fake (55 vitest); and the real model via scripted end-to-end runs whose transcripts are committed. Phase 2's acceptance eval: type-aware denoising cut pyright errors 12 → 6 over the erroring cases, 4/7 fully fixed. Phase 4's acceptance eval (two-region cross-dependency cases): joint + type-aware denoise cut cross-region errors 11 → 5, 3/6 fully fixed — while joint visibility without the checker scored 16 (worse than independent), isolating the type feedback as the source of the gain. Phase 3's scripted E2E drives the real server over the real protocol (stream, cancel, clean exit, frame purity). The one remaining Phase 3 gate is the manual VS Code checklist in extension/README.md — the human half, run in the Extension Development Host.

Quickstart

# Apple Silicon (M-series), ~5 GB free for weights
uv sync --extra dev
uv run pytest                       # full suite, no model needed

hf download hunterbown/Stable-DiffCoder-8B-Instruct-mlx-4Bit \
  --local-dir models/stable-diffcoder-4bit       # ~4.6 GB, one-time

uv run python scripts/phase0_forward.py          # forward-pass sanity (latency, finite logits)
uv run python scripts/phase1_denoise.py          # streaming denoise
uv run python scripts/phase1_denoise.py --ban " %"   # the intervention demo above
uv run python scripts/eval_typeaware.py          # Phase 2 acceptance eval (naive vs type-aware)
uv run python scripts/phase3_e2e.py              # Phase 3 E2E: real server over stdio JSON-RPC
uv run python scripts/eval_multiregion.py        # Phase 4 acceptance eval (multi-region joint denoise)
uv run python scripts/eval_crossfile.py          # Phase 4 cross-file demo (per-file vs workspace feedback)
uv run python scripts/bench_kv_cache.py          # Phase 2.5 prefix KV cache: parity + speedup

For the editor experience, see extension/README.md (npm ci, then F5 in VS Code).

Layout

src/manifold/
  masks.py       block-causal additive attention mask          (pure, tested)
  schedule.py    per-step transfer counts + confidence select  (pure, tested)
  sampler.py     the controllable denoising loop + repair      (runtime-agnostic, tested)
  model_mlx.py   MLX forward backend (threads a custom mask)
  offsets.py     token positions <-> char ranges                (pure, tested)
  typecheck.py   pyright in the loop: parse, re-mask, score    (tested)
  langserver.py  warm pyright-langserver (LSP): ~5ms/check vs ~365ms cold (tested)
  engine.py      text-in/text-out edits: windowing, baseline-delta, streaming, instruction (tested)
  rpc.py         stdio JSON-RPC: codec, dispatch, server loop  (tested)
  cli.py         `manifold serve [--fake]`
scripts/         phase runners + latency/E2E checkpoints
tests/           pytest suite (fake forward, fake checkers, pipe-driven server)
eval/            typed cases + committed acceptance artifacts
extension/       the VS Code extension (TS, zero runtime deps, vitest + tsc)
docs/plans/      design

See docs/plans/2026-07-01-manifold-design.md for the full design and the engineering decisions behind it.

License

MIT — see LICENSE. Stable-DiffCoder and its weights are MIT-licensed by ByteDance Seed.

About

Local, type-aware diffusion code editor — on-device whole-region code edits steered by type-checker feedback inside the denoising loop (Apple Silicon / MLX)

Topics

Resources

License

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors