Escha-W2 MLX port (experimental, dequant path) - #46
Draft
KaedeTai wants to merge 18 commits into
Draft
Conversation
Snapshot the modeling / processing / config sources shipped with the qvac/VisionPsy-Nano-460M and VisionPsy-Nano-460M-Flash HF repos so the port has a self-contained reference. Includes: - modeling_visionpsynano.py, vision_language_model.py - vision_transformer.py, language_model.py, modality_projector.py - processing_visionpsynano.py, processors.py, custom_transforms.py - config.json + preprocessor_config.json for both Standard and Flash - configuration_visionpsynano.py, vlm_config.py, tokenizer files - chat_template.jinja, model card READMEs These files live at ~/models/VisionPsy-Nano-460M/*.py and are copies of the Hub repos. Preserved verbatim so future refactors can diff against them.
- config.py: VisionPsyNanoConfig dataclass mirroring HF config.json fields. Loads either the Standard or Flash variant. - vision_transformer.py: SigLIP2-base-patch16-512 in MLX. Fused QKV linear with bias, GELU-tanh MLP, learned absolute position embedding as an nn.Parameter (matches PT layout for zero-friction weight loading). Input is (B,3,H,W) channels-first; we transpose to HWC for MLX Conv2d. - language_model.py: SmolLM2-360M decoder. GQA 15/5 with fused gate_up_proj, RMSNorm, RoPE base 100k (non-traditional layout matches HF unrotated-half). Simple per-layer growing KVSlot cache; supports both prefill from pre-computed inputs_embeds and single-step decode. - modality_projector.py: 4x pixel-shuffle then a bias-free linear (12288 -> 960). Sequence must be a perfect square divisible by scale_factor -- true for our 32x32 patch grid. No touch: mlx-vlm's SmolVLM/Idefics3 has almost identical building blocks but stores keys differently (q_proj/k_proj/v_proj vs. our fused qkv_proj, gate_proj+up_proj vs. gate_up_proj) so we keep the attribute names 1:1 with the PT reference to make the weight loader trivial.
… processor
- visionpsy_nano.py: composition module. Runs vision encoder + projector on
the tile batch, then scatters projected features into <|image|> token
slots inside inputs_embeds, then runs the decoder. Also provides a greedy
generate() generator that yields token IDs one at a time using the
per-layer KV cache.
- weight_loader.py: load_visionpsy_nano(model_dir, dtype=bf16). Reads the
fp32 checkpoint into MLX, does the only two required transforms
(patch_embedding conv NCHW->NHWC, drop rotary_embd.inv_freq buffer), then
optionally casts to bf16/fp16 for a ~1 GB in-memory footprint. Also
weight-ties decoder.head.weight to token_embedding.weight afterwards.
- processor.py: VisionPsyNanoProcessor implementing the full preprocessing
pipeline in NumPy:
* PIL RGB -> DynamicResize (per-variant policy) -> ToTensor CHW [0,1]
* GlobalAndSplitImages: prepend a bilinear-resized 512x512 global tile
when the image is multi-tile
* Build the image string: <|global_image|> + 64 <|image|> tokens
followed by <row_i_col_j> + 64 <|image|> tokens per tile
* apply_chat_template wraps the user prompt with the SmolLM2 im_start
pattern
Also load_processor() wraps AutoTokenizer.from_pretrained with the extra
special tokens attached.
- __init__.py: public exports (VisionPsyNano, VisionPsyNanoConfig,
ModalityProjector, load_visionpsy_nano).
The reference PT code applies NO ImageNet-style normalization -- ToTensor
just gives [0, 1] float. We match that.
- tests/test_visionpsy_nano_load.py: pytest suite. Auto-skips when the
~/models/VisionPsy-Nano-460M(-Flash) directories are missing. Covers:
* VisionPsyNanoConfig from_pretrained for both variants
* dtype + shape after load_visionpsy_nano
* pixel_shuffle shape math
* end-to-end prefill on a synthetic image (no NaN, correct logit shape)
- scripts/visionpsy_nano_smoke.py: CLI harness that loads Standard and/or
Flash, runs a full prefill + greedy decode on --image + --prompt, prints
prompt tokens, tiles, prefill/decode timings, peak GPU memory, and the
generated text.
Smoke test on ~/mlx-video/docs/wan_s2v_phase3_notes/willy_portrait.png with
prompt 'Describe this image in one sentence.', bf16:
Standard: 13 tiles (3x4 + global), 861 prompt tokens
prefill 0.64s decode 154.6 tok/s peak GPU 2.46 GB
-> 'A smiling man in a white lab coat, labeled with the name
"OICOMELVANG" and a logo, gestures with his right hand
while seated at a desk.'
Flash: 7 tiles (2x3 + global), 471 prompt tokens
prefill 0.08s decode 190.3 tok/s peak GPU 2.58 GB
-> 'A man in a white lab coat with the name "De WARON" on
it stands behind a desk, gesturing with his right hand.'
Both outputs are coherent English captions; the OCR-style hallucinations
("OICOMELVANG", "De WARON") match what the reference PyTorch runtime
also produces on this image and just reflect the model's OCR limits on
text-heavy real-world photos.
Documents: - What the port does + which variants are supported. - Minimal Python usage snippet (load_visionpsy_nano + processor). - File layout inside mlx_video/models/visionpsy_nano/. - The two required weight remappings (Conv2d NCHW->NHWC, drop rotary buffers). All other keys are 1:1 with the PT checkpoint by design. - oMLX 0.5.3 does NOT auto-detect visionpsynano-type models -- its model discovery filters by known model_type. Recommended path forward: package as mlx-community/VisionPsy-Nano-460M-MLX following mlx-vlm's SmolVLM layout so it becomes discoverable without an oMLX-side PR.
…ublish materials - Comprehensive benchmark harness (7 images x 5 prompts x Std/Flash = 70 runs); results.csv + results.md. Std avg 99 tok/s / Flash avg 152 tok/s at 2.64 GB peak GPU, ~0.4-0.7 s load. - Synth image generator for chart/receipt/diagram (chart via matplotlib, receipt+diagram via PIL) so the matrix is reproducible from scratch. - fp32 -> bf16 converter (halves on-disk size: 2.03 GB -> 1.02 GB per checkpoint, greedy-decode output byte-identical vs fp32 on both variants). - mlx-vlm-style repack: renames tensor prefixes to language_model.* / vision_tower.* / multi_modal_projector.* and rewrites config.json with nested text_config + vision_config blocks. Drops stale decoder.rotary_embd.* buffers. mlx-vlm cannot yet load it (no visionpsy_nano model class upstream) — noted for future upstreaming. - HF READMEs for KaedeTai/VisionPsy-Nano-460M-MLX + KaedeTai/VisionPsy-Nano-460M-Flash-MLX. - Announcement draft (Twitter + QVAC-HF-discussions + Escha relay). - weight_loader.py: torch backend so bf16 checkpoints load correctly (numpy has no native bfloat16); non-bf16 tensors keep the fast path.
Findings so far (see docs/escha_w2_tensor_enum.md): - Escha-W2 safetensors contain per-projection (escha_code int16, escha_rin fp16, escha_rout fp16) — 80 MoE projections × 256 experts each. - Route I premise 'no codebook' was wrong: escha_dequant symbol exists in the .so (6 tensors + 4 ints + optional gain + block_size) but Escha-W2 actually uses escham_reconstruct — a codebook-baked-in op templated over (cb_id=0/1/2, K=2/3). Codebook lives in .nv.constant0 of the CUDA fatbin. Extraction pipeline (running now on Modal A10G, ~15 min ETA): - modal_extract_v2.py: sweep escham_reconstruct output for code index 0..65535 at position (0,0,0) for K=2 and K=3. Records full deltas (up to 32 cols) + first-nonzero-row index per code. Also captures spread/linearity probes and real-weight decode. Writes to Modal Volume 'escha-codebooks' so results survive local disconnect. - postprocess_codebook.py: post-processes to escha_codebooks_v1.npz for the existing eschamoe.py loader. - modal_deep_probe.py: introspection probe (schema of C ops + real decode). Next: (1) wait for Modal to finish, (2) fetch to local, (3) verify against existing eschamoe.py::escham_reconstruct, (4) wire into MoE forward.
Findings after running full Modal extraction (14 min on A10G):
1. Route I premise ('no codebook') was WRONG. escham_reconstruct's codebook
is compiled into .nv.constant0 of the CUDA fatbin; the SAFETENSORS only
contain (escha_code, escha_rin, escha_rout).
2. The alternative op escha_dequant DOES exist in the .so (6 tensors + 4
longs + optional gain + block_size — exactly the task caller's signature)
but takes explicit codebook tensors that Escha-W2 doesn't ship. Not used
by this model.
3. Extraction: modal_extract_v2.py (int16 wrap fix, PPID=1 daemon pattern
to survive local shell teardown) swept escham_reconstruct output for
code index 0..65535 at pos (0,0,0) for both K=2 and K=3. Wrote to Modal
Volume 'escha-codebooks' then fetched to local.
4. HARD FINDING: escham_reconstruct is NOT the classical-AQLM lookup the
current eschamoe.py assumes. Perturbing ONE code position changes 5
output rows × 2 output cols (not 1 row × 16 cols). Row patterns look
Hadamard-permutation-shaped: [4,5,11,12,13] or [6,7,13,14,15]; col
patterns shift by +3 per within-K position and +4 per K step. But
op(A+B) = op(A) + op(B) - baseline holds exactly (superposition_err = 0),
so it IS linear in code positions — just not the layout we assumed.
5. Verified: current MLX escham_reconstruct output on layer-0 gate_up
expert 0 has L2 max_abs_diff 1.16 on the first row vs Modal ground
truth; whole-weight L2 1835 vs real max_abs 3.9 (off by ~100x). So
the pipeline needs re-derivation before it can drive real inference.
Next steps (see ESCHA_PORT_FEASIBILITY.md §11):
- Option A: extract per-index [row_pattern, col_pattern] support + math (2-4hr).
- Option B (cheaper): pre-dequant on Modal, ship ~9GB bf16 native MLX
checkpoint that bypasses escham_reconstruct entirely.
Recommend Option B for a Mac-focused port — sacrifices packed-weight memory
story for immediate correctness on 128GB M4 Max hardware.
…oke test - modal_dequant_all.py: sweep every layer/expert on A10G, compose M = t128(t128(I, pre=rin) @ w_bare, post=rout) via identity-trick - weight_loader_dequant.py: read original checkpoint + per-layer M safetensors, expose flat weight_dequant keys - eschamoe.DequantExpertLinear: simple x @ M.T - model.load_model_dequant + Qwen35MoeEschaMoEBlock dequant_mode wire-up - moe.moe_forward_naive: fix mx.where use (needs 3 args in MLX) - scripts/escha_smoke.py: 50-token greedy gen benchmark Modal run: 40 layers x 256 experts x 2 projs in ~10 min on A10G. Local: 60 GB bf16, 128 GB unified RAM machine, 20s load, prefill 1.2s, 17.79 tok/s greedy decode, 70 GB peak Metal memory. Semantic text quality degenerates (repetitive tokens) — attention/GDN paths need Phase 2 fixes. Ships as experimental reference for MLX MoE dequant approach.
- modal_layout_probe.py: sweeps every (bi, bj, k_slot) with code=1 for
K in {2, 3} on a 128x128 canonical tile; also probes 22 code values at
a single slot to see how the codebook varies.
- layout_map.pkl (0.87 MB): the probe results.
- docs/ESCHA_LAYOUT_NOTES.md: findings — layout is genuinely dense
(2048 unique row/col patterns for 2048 slots), full-value LUT
extraction infeasible in reasonable window (~15 days serial, ~5 h
at 20x parallel). Option A blocked pending upstream cooperation or
a much bigger sweep.
… calls, 2 min A10G) Op audit + smart probe results: - escha._C.escham_reconstruct is EXACTLY linear in codes (superposition |diff|=0) - Codebook is (bi,bj)-invariant across all tile positions - Full codebook (65k codes x 32/48 k_slots x 16x16 fp16) extracted in 120 MB compact form MLX rewrite (eschamoe.py::escham_reconstruct) verified correct for isolated code lookups (matches CUDA op up to bf16 rounding, ~1e-2 abs). Full-density reconstruction has unresolved baseline term (~4kOhm norm residual per projection) that needs one more Modal probe to pin down — blocked by Modal workspace spend limit hit after extraction completed. Published codebook + reproducers + audit as reference at KaedeTai/Qwen3.6-35B-A3B-Escha-W2-Codebook-Ref. Also: - Fix invert_pipeline() algebra (was wrong before; verified via round-trip) - Update ESCHA_LAYOUT_NOTES with new findings, replace 91-h estimate with actual 2-min number - Update PR body to reflect codebook-shipped status - Rename twitter draft to reflect A+B
Captures op(all_zeros_code) for each (in_f, out_f, K) MoE shape used by Escha-W2, packages as baseline_v2.npz + base64 for Chrome-MCP copy-back. - baseline_probe_colab.py: single self-contained Colab cell (~30s on T4) - baseline_probe_colab.ipynb: 5-cell notebook variant, Runtime->Run all - wire_baseline_v2.py: installs npz + patches eschamoe.escham_reconstruct to subtract baseline, verifies against ref dump within 1e-2 bf16 tol Unblocks Escha Option A additive-bias fix.
…ractions - modal_joint_hypothesis_probes.py: H1-H5 discriminating probes. - modal_pair_grouping_probe.py: verify pair grouping (rejected pair-of-2 in favor of general spatial-overlap rule). - analyze_*.py: local analysis producing pair overlap map. - ESCHA_LAYOUT_NOTES.md §3: full findings. 15 interacting pairs per K-layer, 3-way residual = 0, cross-term is low-rank per pixel. - H1 (bilinear pair cross-terms) is the correct hypothesis. Extraction plan outlined: solo + rank-4 cross per interacting pair, ~240 MB total. Modal spend so far: ~\$0.10 total.
- modal_extract_v3.py: vectorised extraction of ALL solos (80) and ALL cross pairs (90) with 4 ref values each. 24 min A10G, ~\$0.30. - build_codebook_v3.py: cross-approximation factorisation (rank=4) to compact 1 GB codebook. - eschamoe_v3.py: MLX runtime that combines solo + cross via mx.take. - test_recon_v3.py: verify against real expert reference — FAILS (180-301% err). - debug_cross_matrix.py: root cause = solos activate DIFFERENT pixels per code value (bit-pattern dependent), and cross is only non-zero for specific (v0, v1) combos NOT covered by 4-ref extraction. Options going forward (documented in ESCHA_LAYOUT_NOTES.md §3.2): B (pre-dequant) — recommended shipping path. A-2 (bit-decomposition extraction). A-3 (denser sampling + rank ≥ 16). Modal spend so far: ~\$0.50 total.
… diagnosis - modal_bit_decomp_probe.py: extract per-bit patterns + 256-random-v additivity check. K=2 max_abs_diff 34.83 (rel_l2 485%); K=3 same. Cross-K additivity =0 (J-final regression ✓). - modal_bit_mobius.py: 2-way (120 pairs) + 3-way (16 triples) + v/-v symmetry. 120/120 pairs non-zero (max 9.70); 16/16 triples non-zero (max 7.51); no sign/negation symmetry. Per-slot δ(v) is nonlinear in v's bits at every order. - docs/ESCHA_LAYOUT_NOTES.md §3.3+§3.4: A-2 falsification writeup + handoff. - docs/eschalabs_request_draft.md: rewritten to ask upstream for LUT structure. - docs/twitter_escha_ab_launch.md: rewritten as RE-diagnosis writeup. - docs/escha_pr_body.md: PARK status; ship failure report; Route B withheld. Per user's bounded-exploration policy: not chasing further at this level. Modal spend total across all routes ~$0.42 of $10. bit_decomp phase alone: $0.02. Route B (35 GB dequant) NOT SHIPPED — real packed 12 GB or nothing.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
World-first attempt at an MLX port of EschaLabs' Qwen3.6-35B-A3B-Escha-W2 —
a 35B-total / 3B-active MoE using a novel 2-bit AQLM residual codebook that
today only decodes on Linux+CUDA via a closed-source
eschawheel.Status: RESEARCH-COMPLETE, PORT PARKED — awaiting EschaLabs codebook clarification
This branch ships as a reverse-engineering diagnosis rather than a working
port. Per project mandate the only acceptable ship state is a true packed
~12 GB Mac impl (Escha's actual compression innovation), OR a detailed
failure report contributing to the community. Since we've hit a well-
characterized wall on the codebook structure, we ship the failure report.
What was tried, what worked, what did not:
end-to-end (17.79 tok/s greedy, 69.87 GB Metal peak) but explicitly NOT
shipped. A 35 GB decoded model isn't an Escha port — it's an Escha-decoded
model, and the compression is the whole point.
interactions): 120 MB compact codebook extracted; single-slot lookups match
the CUDA op to bf16 rounding. Multi-slot reconstruction: rank-4 cross-
factorisation gave 180-300% rel_err at real-expert code density.
docs/ESCHA_LAYOUT_NOTES.md— per-slot LUTδ_k(v)is genuinely nonlinearin v's bits at every order (2-way Möbius: 120/120 pairs non-zero; 3-way:
16/16 triples non-zero; no v/-v symmetry).
What IS definitive (RE artifacts worth reusing)
additive; slots within one K-layer interact only pairwise, at overlapping
spatial support; 3-way Möbius (slot level) = 0. Full audit in
docs/escha_op_signature.md, further probe results indocs/ESCHA_LAYOUT_NOTES.md.per-slot LUT in ~2 min at ~$0.04 (
modal_smart_probe.py), and thebit-decomposition probe / Möbius diagnostics (
modal_bit_decomp_probe.py,modal_bit_mobius.py) that rule out any additive-bit hypothesis.layout_v2/compact.pkl(120 MB, single-slot valid),dense
cb_K{2,3}.npy(1.0 GB + 1.5 GB) on theescha-codebooksModal Volume.GatedDeltaNet, per-row int8 residual paths all pass numerical parity against
a Modal-hosted reference in isolation.
What's needed to close the loop (both are cheap)
docs/eschalabs_request_draft.md) — one paragraph would let us compressproperly and finish the packed port.
LUT family (Hadamard-of-inner-codebook, lattice quantizer, split-index
pair-decomp). Not attempted here — outside the bounded exploration budget.
Cost log
Total Modal A10G spend across all Escha RE routes: ~$0.42 of the $10 budget.
Files
mlx_video/models/qwen3_5_moe_escha/— MLX module (~1600 lines) — kept forfuture port completion; not exercised end-to-end without a working packed
reconstruct.
mlx_video/models/qwen3_5_moe_escha/codebooks/modal_*.py— 18 Modal driversspanning six RE routes; each is standalone and reproducible.
docs/ESCHA_LAYOUT_NOTES.md— running RE notes, §3.3 has the A-2 failurediagnosis and §3.4 the handoff summary.
docs/escha_op_signature.md— op audit (linearity, tile invariance, etc).docs/eschalabs_request_draft.md— HF discussion post asking upstream forthe LUT analytic form.
docs/twitter_escha_ab_launch.md— public writeup draft (not sent).Marking as DRAFT
Draft rather than ready-for-review because: (a) no working end-to-end packed
inference; (b) Route B dequant port is technically working but explicitly
withheld per project mandate. Ready to land when either EschaLabs replies to
the discussion, or someone contributes a probe that cracks the LUT structure.