Initial: extracted Escha packed AQLM codebook + audit + reproducers

LAYOUT_NOTES.md

@@ -0,0 +1,110 @@
+# Escha packed layout — reverse-engineering notes (updated Aug 2026, post-Option-A)
+
+## Summary — where we are now
+
+**Modal audit + smart probe: SUCCESS.** The full (K, k_slot, code_value) →
+(row, col, weight_value) codebook has been extracted in a compact form (120 MB,
+~1024 A10G op calls total, ~2 min GPU wall time) and verified for isolated
+single-code lookups.
+
+**End-to-end packed inference: PARTIALLY WORKING.** MLX `escham_reconstruct`
+correctly reproduces the CUDA op's output for any single-slot code (max abs
+diff = fp16/bf16 rounding, ~0.01). For a real expert whose codes activate all
+262 K slots simultaneously, my reconstruction picks up an unresolved
+density-dependent bias equal in norm to the raw delta (~4 kOh). This bias is
+NOT expert-independent, contradicting the simple `w_bare = w0 + delta` model
+that the linearity audit (`docs/escha_op_signature.md`) predicts.
+
+Root cause candidates (all unverified — Modal workspace hit spend limit before
+the follow-up probe could run):
+  1. Escha's `escha_t128` implementation differs from the plain-Hadamard
+     `t128` in `transform.py` by a scaling / permutation / bias term. If so,
+     our compose+invert self-round-trip stays consistent but never reaches
+     the reference op's true `w_bare`.
+  2. The op has a per-tile bias (function of `(bi, bj)` alone, but non-uniform
+     across blocks) that our extraction folded into the "baseline"
+     via the `w0 = op(all-zeros)` subtraction.
+  3. The op is not truly linear at high code density (superposition test only
+     verified for 100 random slots out of 262 K).
+
+Fix path — one more Modal probe would resolve it: dump `op(all-zeros, in_p,
+out_p, K)` as an fp16 tensor (in_p × out_p, ~6 MB total for both K values),
+plus one `op(real_expert_code)` reference for cross-check. Blocked on Modal
+spend limit at 2026-08-01 17:20 CST.
+
+## What we extracted (Modal smart probe — `codebooks/modal_smart_probe.py`)
+
+  Runtime: 46.5 s for K=2 (256 op calls) + ~90 s for K=3 (768 op calls).
+  Total ≈ 2 min A10G, ~$0.04.
+
+- `codebooks/layout_v2/compact.pkl` (120 MB): compact sparse codebook.
+  Keys `K{2,3}_positions[k]` = (n_nz, 2) int8 row/col positions;
+       `K{2,3}_values[k]`    = (65536, n_nz) fp16 codebook values.
+- `codebooks/layout_v2/cb_K{2,3}.npy` on the `escha-codebooks` Modal Volume:
+  fully dense (k_max, 65536, 16, 16) fp16, ~1 GB + 1.5 GB.
+
+Verified properties (see `docs/escha_op_signature.md`):
+  - Op is exactly LINEAR in codes across up to 100 random slot activations
+    (superposition |diff|=0).
+  - Codebook is (bi, bj)-INVARIANT across all tile positions (up to (bi*16,
+    bj*16) offset), including corners like (bi=127, bj=63) vs (bi=0, bj=0).
+  - Op supports leading batch dimensions on `packed` — enabling further
+    parallelism if needed.
+
+## Structural regularities in the k_slot layout
+
+For K=2 (32 slots per 16×16 block), the (row, col) support of each k_slot is:
+  - Row pattern cycles with period 4: {[4,5,11,12,13], [2,3,9,10,11],
+    [0,1,8,9,15]+extra, [6,7,13,14,15]}
+  - Col pattern: `col_c = k_slot // 4`, cols = `[col_c, col_c + 8]`
+    (with an extra col at `k_slot % 4 == 2`)
+
+For K=3 (48 slots), similar cycle with period 6.
+
+These patterns are extractable in ~1 s of GPU time (1 op call per k_slot);
+they encode which 8-10 output positions each `cb[k, v]` writes to.
+
+## Bug in the earlier `cb_K2.npy` / `cb_K3.npy` extraction
+
+The prior `modal_extract_v2.py` only captured `d[first_nz_row, :32]` — a single
+32-value slice of the first non-zero row of the delta. This missed the other
+4-5 rows of each pattern. That's why plugging cb_K2/cb_K3 into `eschamoe.py`
+produced L2 diff 1.16 on the first row: the codebook was under-specified by 5x.
+
+The new `compact.pkl` extraction fixes this — it captures every non-zero
+position across all 65 K code values.
+
+## Files on disk
+
+- `codebooks/modal_op_audit.py`    — signature + linearity audit (Modal)
+- `codebooks/modal_smart_probe.py` — batched 1024-op codebook extraction
+- `codebooks/layout_v2/compact.pkl` — 120 MB codebook artifact
+- `codebooks/layout_v2/cb_K{2,3}.npy` — dense form (only on Modal Volume)
+- `codebooks/extract_baseline.py` — attempt at recovering `w0` from Option B
+  M matrices (produces per-expert-varying estimate; see § summary above)
+- `eschamoe.py`                    — MLX decoder using the new codebook
+- `docs/escha_op_signature.md`     — full op audit report
+
+## What would unblock full packed inference
+
+One more Modal call, ~30 s A10G, cost ~$0.02:
+
+```python
+@app.function(gpu="A10G")
+def dump_baselines():
+    # For each (in_f, out_f, K) shape used by Escha-W2, dump op(all-zeros).
+    import torch, escha
+    op = torch.ops.escha.escham_reconstruct
+    for (in_f, out_f, K, cshape) in [(2048, 1024, 2, (128, 64, 32)),
+                                      (512, 2048, 3, (32, 128, 48))]:
+        p0 = torch.zeros(cshape, dtype=torch.int16, device="cuda")
+        w0 = op(p0, in_f, out_f, K, True, False).cpu().numpy()
+        # Also dump op(one_random_expert_code) for cross-check
+        code = ...  # load one real expert code
+        w_ref = op(code, in_f, out_f, K, True, False).cpu().numpy()
+```
+
+Plus: verify the linearity assumption at HIGH code density (dense 262 K slots)
+by comparing `op(code)` against `w0 + my_reconstruct(code)`. If they differ
+by more than fp16 noise, the op has hidden non-linear structure that would
+need one more targeted probe.

OP_SIGNATURE_AUDIT.md

@@ -0,0 +1,201 @@
+# Escha `escham_reconstruct` — signature + linearity audit
+
+## 1. Introspection
+
+escha module: `/usr/local/lib/python3.12/site-packages/escha/__init__.py`
+escha._C: `<module 'escha._C' from '/usr/local/lib/python3.12/site-packages/escha/_C.cpython-312-x86_64-linux-gnu.so'>`
+
+### `dir(escha._C)`
+  - `escha_aqlm_auto` ? — escha_aqlm_auto(ta: torch.Tensor, codes: torch.Tensor, codebooks: torch.Tensor, scales: torch.Tensor, codes_T: torch.Tensor | None = None, k_tile: typing.SupportsInt = 64) -> torch.Tensor
+  - `escha_aqlm_fused_hmma` ? — escha_aqlm_fused_hmma(ta: torch.Tensor, codes_T: torch.Tensor, codebooks: torch.Tensor, scales: torch.Tensor, k_tile: typing.SupportsInt = 64) -> torch.Tensor
+  - `escha_aqlm_gemv` ? — escha_aqlm_gemv(ta: torch.Tensor, codes: torch.Tensor, codebooks: torch.Tensor, scales: torch.Tensor) -> torch.Tensor
+  - `escha_aqlm_prepare_codes_transposed` ? — escha_aqlm_prepare_codes_transposed(codes: torch.Tensor) -> torch.Tensor
+  - `escha_binary_gemv` ? — escha_binary_gemv(arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: torch.Tensor) -> torch.Tensor
+  - `escha_binary_gemv_reg` ? — escha_binary_gemv_reg(arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: torch.Tensor) -> torch.Tensor
+  - `escha_decgemv` ? — escha_decgemv(x: torch.Tensor, packed_codes: torch.Tensor, scale: torch.Tensor, transform_left: torch.Tensor, transform_right: torch.Tensor, a1: torch.Tensor, a2: torch.Tensor, td1: typing.SupportsInt, td2: typing.SupportsInt, ic: typing.SupportsInt, expic: typing.SupportsInt, gain: torch.Tensor | None = None, block_size: typing.SupportsInt = 0) -> torch.Tensor
+  - `escha_decgemv_inplace` ? — escha_decgemv_inplace(arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: torch.Tensor, arg4: torch.Tensor, arg5: torch.Tensor, arg6: torch.Tensor, arg7: typing.SupportsInt, arg8: typing.SupportsInt, arg9: typing.SupportsInt, arg10: typing.SupportsInt, arg11: torch.Tensor, arg12: torch.Tensor, arg13: torch.Tensor, arg14: torch.Tensor) -> None
+  - `escha_decgemv_reg` ? — escha_decgemv_reg(arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: torch.Tensor, arg4: torch.Tensor, arg5: torch.Tensor, arg6: torch.Tensor, arg7: typing.SupportsInt, arg8: typing.SupportsInt, arg9: typing.SupportsInt, arg10: typing.SupportsInt) -> torch.Tensor
+  - `escha_dequant` ? — escha_dequant(packed_codes: torch.Tensor, scale: torch.Tensor, transform_left: torch.Tensor, transform_right: torch.Tensor, a1: torch.Tensor, a2: torch.Tensor, td1: typing.SupportsInt, td2: typing.SupportsInt, ic: typing.SupportsInt, expic: typing.SupportsInt, gain: torch.Tensor | None = None, block_size: typing.SupportsInt = 0) -> torch.Tensor
+  - `escha_fused_dequant_gemm` ? — escha_fused_dequant_gemm(ta_fp16: torch.Tensor, packed_codes: torch.Tensor, scale: torch.Tensor) -> torch.Tensor
+  - `escha_fused_dequant_gemm_auto` ? — escha_fused_dequant_gemm_auto(ta_fp16: torch.Tensor, packed_codes: torch.Tensor, scale: torch.Tensor, packed_T: torch.Tensor | None = None) -> torch.Tensor
+  - `escha_fused_dequant_gemm_v2` ? — escha_fused_dequant_gemm_v2(ta_fp16: torch.Tensor, packed_codes: torch.Tensor, scale: torch.Tensor) -> torch.Tensor
+  - `escha_fused_dequant_gemm_v3` ? — escha_fused_dequant_gemm_v3(ta_fp16: torch.Tensor, packed_T: torch.Tensor, scale: torch.Tensor) -> torch.Tensor
+  - `escha_fused_dequant_gemm_v4` ? — escha_fused_dequant_gemm_v4(ta_fp16: torch.Tensor, packed_T: torch.Tensor, scale: torch.Tensor) -> torch.Tensor
+  - `escha_fused_dequant_gemm_v5` ? — escha_fused_dequant_gemm_v5(ta_fp16: torch.Tensor, packed_codes: torch.Tensor, scale: torch.Tensor) -> torch.Tensor
+  - `escha_init` ? — escha_init() -> None
+  - `escha_lut_binary_gemv` ? — escha_lut_binary_gemv(arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: torch.Tensor) -> torch.Tensor
+  - `escha_prepare_packed_transposed` ? — escha_prepare_packed_transposed(packed_codes: torch.Tensor) -> torch.Tensor
+  - `escha_transform` ? — escha_transform(arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: torch.Tensor, arg4: torch.Tensor, arg5: typing.SupportsInt, arg6: typing.SupportsInt, arg7: typing.SupportsInt, arg8: typing.SupportsInt) -> list[torch.Tensor]
+  - `escha_transform_fast` ? — escha_transform_fast(arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: torch.Tensor, arg4: torch.Tensor, arg5: typing.SupportsInt, arg6: typing.SupportsInt, arg7: typing.SupportsInt, arg8: typing.SupportsInt) -> list[torch.Tensor]
+  - `escha_transform_fp16` ? — escha_transform_fp16(arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: torch.Tensor, arg4: torch.Tensor, arg5: typing.SupportsInt, arg6: typing.SupportsInt, arg7: typing.SupportsInt, arg8: typing.SupportsInt) -> list[torch.Tensor]
+  - `escha_transform_fp16_no_bias` ? — escha_transform_fp16_no_bias(arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: torch.Tensor, arg4: torch.Tensor, arg5: typing.SupportsInt, arg6: typing.SupportsInt, arg7: typing.SupportsInt, arg8: typing.SupportsInt) -> torch.Tensor
+  - `eschax_binary_search` ? — eschax_binary_search(arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: typing.SupportsInt, arg4: typing.SupportsInt, arg5: typing.SupportsInt) -> torch.Tensor
+  - `eschax_dequant` ? — eschax_dequant(arg0: torch.Tensor, arg1: typing.SupportsFloat) -> torch.Tensor
+  - `eschax_eschax_decode` ? — eschax_eschax_decode(arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: torch.Tensor, arg4: typing.SupportsInt, arg5: typing.SupportsInt) -> torch.Tensor
+  - `eschax_eschax_gemv` ? — eschax_eschax_gemv(arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: torch.Tensor, arg4: torch.Tensor, arg5: typing.SupportsFloat, arg6: torch.Tensor, arg7: typing.SupportsInt, arg8: typing.SupportsInt) -> torch.Tensor
+  - `eschax_gemv` ? — eschax_gemv(arg0: torch.Tensor, arg1: torch.Tensor, arg2: typing.SupportsFloat, arg3: torch.Tensor) -> torch.Tensor
+  - `eschax_huffman_decode` ? — eschax_huffman_decode(arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: typing.SupportsInt, arg4: typing.SupportsInt) -> torch.Tensor
+  - `eschax_huffman_gemv` ? — eschax_huffman_gemv(arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: torch.Tensor, arg4: typing.SupportsFloat, arg5: torch.Tensor, arg6: typing.SupportsInt, arg7: typing.SupportsInt) -> torch.Tensor
+
+### torch.ops.escha.escham_reconstruct
+  escha.escham_reconstruct
+  overloads: ['default']
+  default: escha::escham_reconstruct(Tensor packed, int in_features, int out_features, int K, bool cbA, bool mul1) -> Tensor
+
+## 2. Shape acceptance test
+
+  min-K2 128x128: cshape=(8, 8, 32) in_f=128 out_f=128 K=2 -> OK, w.shape=(128, 128) dtype=torch.float16
+  min-K3 128x128: cshape=(8, 8, 48) in_f=128 out_f=128 K=3 -> OK, w.shape=(128, 128) dtype=torch.float16
+  escha gate_up (K=2, in=2048/out=1024): cshape=(128, 64, 32) in_f=2048 out_f=1024 K=2 -> OK, w.shape=(2048, 1024) dtype=torch.float16
+  escha down (K=3, in=512/out=2048): cshape=(32, 128, 48) in_f=512 out_f=2048 K=3 -> OK, w.shape=(512, 2048) dtype=torch.float16
+  leading batch (2, 8, 8, 32) K=2: cshape=(2, 8, 8, 32) in_f=128 out_f=128 K=2 -> OK, w.shape=(128, 128) dtype=torch.float16
+  leading batch (16, 8, 8, 32) K=2: cshape=(16, 8, 8, 32) in_f=128 out_f=128 K=2 -> OK, w.shape=(128, 128) dtype=torch.float16
+  leading batch (4, 128, 64, 32) K=2: cshape=(4, 128, 64, 32) in_f=2048 out_f=1024 K=2 -> OK, w.shape=(2048, 1024) dtype=torch.float16
+
+## 3. Full delta pattern at slot (0,0,0)
+
+For (in=2048, out=1024, K=2) tile: what is the full (row, col) support
+of the delta when we set exactly code[0,0,0] = v, for various v?
+
+  v=     1:   8 nonzero positions, rows=[4, 5, 11, 12, 13], cols=[0, 8]
+  v=     2:   8 nonzero positions, rows=[4, 5, 11, 12, 13], cols=[0, 8]
+  v=     3:   8 nonzero positions, rows=[4, 5, 11, 12, 13], cols=[0, 8]
+  v=     4:   8 nonzero positions, rows=[4, 5, 10, 11, 12, 13], cols=[0, 8]
+  v=     5:   9 nonzero positions, rows=[4, 5, 10, 11, 12, 13], cols=[0, 8]
+  v=     7:   9 nonzero positions, rows=[4, 5, 10, 11, 12, 13], cols=[0, 8]
+  v=    10:   9 nonzero positions, rows=[4, 5, 10, 11, 12, 13], cols=[0, 8]
+  v=    16:   8 nonzero positions, rows=[3, 4, 5, 10, 11, 12, 13], cols=[0, 8]
+  v=    64:   8 nonzero positions, rows=[2, 3, 4, 5, 10, 11, 12, 13], cols=[0, 8]
+  v=   256:   8 nonzero positions, rows=[2, 3, 4, 5, 10, 11, 12], cols=[0, 8]
+  v=  1024:   8 nonzero positions, rows=[2, 3, 4, 5, 10, 11], cols=[0, 8]
+  v=  4096:   8 nonzero positions, rows=[2, 3, 4, 10, 11], cols=[0, 8]
+  v= 16384:   8 nonzero positions, rows=[2, 3, 10, 11], cols=[0, 8]
+  v= 32767:  15 nonzero positions, rows=[2, 3, 4, 5, 10, 11, 12, 13], cols=[0, 8]
+  v=    -1:  15 nonzero positions, rows=[2, 3, 4, 5, 10, 11, 12, 13], cols=[0, 8]
+  v=  -100:  14 nonzero positions, rows=[2, 3, 4, 5, 10, 11, 12, 13], cols=[0, 8]
+  v=-32768:   8 nonzero positions, rows=[2, 3, 10, 11], cols=[0, 8]
+
+## 4. Superposition test (LINEARITY in codes)
+
+Test: op(all-zeros with code[bi_a, bj_a, k_a]=v_a AND code[bi_b, bj_b, k_b]=v_b)
+      == op(only code[bi_a, bj_a, k_a]=v_a) + op(only code[bi_b, bj_b, k_b]=v_b) - op(zeros)
+If yes, we can probe many (bi, bj) slots simultaneously in ONE op call.
+
+  2-pos, distinct (bi,bj), same k: |combined|=1.034e+01 |diff|=0.000e+00 rel=0.000e+00
+  2-pos, same (bi,bj), diff k: |combined|=1.034e+01 |diff|=0.000e+00 rel=0.000e+00
+  2-pos, same (bi,bj), diff K-slice: |combined|=1.034e+01 |diff|=0.000e+00 rel=0.000e+00
+  8-pos random: |combined|=2.231e+01 |diff|=0.000e+00 rel=0.000e+00
+  100-pos random: |combined|=8.165e+01 |diff|=0.000e+00 rel=0.000e+00
+
+## 5. Slot invariance test — is the codebook shared across (bi, bj)?
+
+Compare delta patterns for the SAME value v at DIFFERENT (bi, bj) with the same k_slot.
+If they are identical up to a (bi*16, bj*16) offset, the codebook is (bi, bj)-invariant.
+
+  v=     1 (bi=0,bj=0) vs (bi=1,bj=0): |diff|=0.000e+00 rel=0.000e+00
+  v=     1 (bi=0,bj=0) vs (bi=0,bj=1): |diff|=0.000e+00 rel=0.000e+00
+  v=     1 (bi=0,bj=0) vs (bi=1,bj=1): |diff|=0.000e+00 rel=0.000e+00
+  v=     1 (bi=0,bj=0) vs (bi=5,bj=3): |diff|=0.000e+00 rel=0.000e+00
+  v=     1 (bi=0,bj=0) vs (bi=127,bj=63): |diff|=0.000e+00 rel=0.000e+00
+  v=   100 (bi=0,bj=0) vs (bi=1,bj=0): |diff|=0.000e+00 rel=0.000e+00
+  v=   100 (bi=0,bj=0) vs (bi=0,bj=1): |diff|=0.000e+00 rel=0.000e+00
+  v=   100 (bi=0,bj=0) vs (bi=1,bj=1): |diff|=0.000e+00 rel=0.000e+00
+  v=   100 (bi=0,bj=0) vs (bi=5,bj=3): |diff|=0.000e+00 rel=0.000e+00
+  v=   100 (bi=0,bj=0) vs (bi=127,bj=63): |diff|=0.000e+00 rel=0.000e+00
+  v= 32767 (bi=0,bj=0) vs (bi=1,bj=0): |diff|=0.000e+00 rel=0.000e+00
+  v= 32767 (bi=0,bj=0) vs (bi=0,bj=1): |diff|=0.000e+00 rel=0.000e+00
+  v= 32767 (bi=0,bj=0) vs (bi=1,bj=1): |diff|=0.000e+00 rel=0.000e+00
+  v= 32767 (bi=0,bj=0) vs (bi=5,bj=3): |diff|=0.000e+00 rel=0.000e+00
+  v= 32767 (bi=0,bj=0) vs (bi=127,bj=63): |diff|=0.000e+00 rel=0.000e+00
+  v=-32768 (bi=0,bj=0) vs (bi=1,bj=0): |diff|=0.000e+00 rel=0.000e+00
+  v=-32768 (bi=0,bj=0) vs (bi=0,bj=1): |diff|=0.000e+00 rel=0.000e+00
+  v=-32768 (bi=0,bj=0) vs (bi=1,bj=1): |diff|=0.000e+00 rel=0.000e+00
+  v=-32768 (bi=0,bj=0) vs (bi=5,bj=3): |diff|=0.000e+00 rel=0.000e+00
+  v=-32768 (bi=0,bj=0) vs (bi=127,bj=63): |diff|=0.000e+00 rel=0.000e+00
+
+## 6. k_slot pattern
+
+For each k_slot, what is the row/col support of the (0, 0, k)+v=1 delta?
+
+### K=2, cshape=(128, 64, 32)
+  k= 0: rows=[4, 5, 11, 12, 13] cols=[0, 8] n_pos=8
+  k= 1: rows=[2, 3, 9, 10, 11] cols=[0, 8] n_pos=8
+  k= 2: rows=[0, 1, 8, 9, 15] cols=[1, 8, 9] n_pos=8
+  k= 3: rows=[6, 7, 13, 14, 15] cols=[0, 8] n_pos=8
+  k= 4: rows=[4, 5, 11, 12, 13] cols=[1, 9] n_pos=8
+  k= 5: rows=[2, 3, 9, 10, 11] cols=[1, 9] n_pos=8
+  k= 6: rows=[0, 1, 8, 9, 15] cols=[2, 9, 10] n_pos=8
+  k= 7: rows=[6, 7, 13, 14, 15] cols=[1, 9] n_pos=8
+  k= 8: rows=[4, 5, 11, 12, 13] cols=[2, 10] n_pos=8
+  k= 9: rows=[2, 3, 9, 10, 11] cols=[2, 10] n_pos=8
+  k=10: rows=[0, 1, 8, 9, 15] cols=[3, 10, 11] n_pos=8
+  k=11: rows=[6, 7, 13, 14, 15] cols=[2, 10] n_pos=8
+  k=12: rows=[4, 5, 11, 12, 13] cols=[3, 11] n_pos=8
+  k=13: rows=[2, 3, 9, 10, 11] cols=[3, 11] n_pos=8
+  k=14: rows=[0, 1, 8, 9, 15] cols=[4, 11, 12] n_pos=8
+  k=15: rows=[6, 7, 13, 14, 15] cols=[3, 11] n_pos=8
+  k=16: rows=[4, 5, 11, 12, 13] cols=[4, 12] n_pos=8
+  k=17: rows=[2, 3, 9, 10, 11] cols=[4, 12] n_pos=8
+  k=18: rows=[0, 1, 8, 9, 15] cols=[5, 12, 13] n_pos=8
+  k=19: rows=[6, 7, 13, 14, 15] cols=[4, 12] n_pos=8
+  k=20: rows=[4, 5, 11, 12, 13] cols=[5, 13] n_pos=8
+  k=21: rows=[2, 3, 9, 10, 11] cols=[5, 13] n_pos=8
+  k=22: rows=[0, 1, 8, 9, 15] cols=[6, 13, 14] n_pos=8
+  k=23: rows=[6, 7, 13, 14, 15] cols=[5, 13] n_pos=8
+  k=24: rows=[4, 5, 11, 12, 13] cols=[6, 14] n_pos=8
+  k=25: rows=[2, 3, 9, 10, 11] cols=[6, 14] n_pos=8
+  k=26: rows=[0, 1, 8, 9, 15] cols=[7, 14, 15] n_pos=8
+  k=27: rows=[6, 7, 13, 14, 15] cols=[6, 14] n_pos=8
+  k=28: rows=[4, 5, 11, 12, 13] cols=[7, 15] n_pos=8
+  k=29: rows=[2, 3, 9, 10, 11] cols=[7, 15] n_pos=8
+  k=30: rows=[0, 1, 8, 9, 15] cols=[0, 8, 15] n_pos=8
+  k=31: rows=[6, 7, 13, 14, 15] cols=[7, 15] n_pos=8
+### K=3, cshape=(32, 128, 48)
+  k= 0: rows=[2, 3, 10, 11] cols=[0, 8] n_pos=5
+  k= 1: rows=[1, 2, 3, 8, 9] cols=[0, 8] n_pos=5
+  k= 2: rows=[5, 6, 7, 12, 13] cols=[0, 8] n_pos=5
+  k= 3: rows=[4, 5, 11, 12, 13] cols=[0, 8] n_pos=6
+  k= 4: rows=[0, 1, 8, 9, 15] cols=[1, 8, 9] n_pos=6
+  k= 5: rows=[6, 7, 14, 15] cols=[0, 8] n_pos=5
+  k= 6: rows=[2, 3, 10, 11] cols=[1, 9] n_pos=5
+  k= 7: rows=[1, 2, 3, 8, 9] cols=[1, 9] n_pos=5
+  k= 8: rows=[5, 6, 7, 12, 13] cols=[1, 9] n_pos=5
+  k= 9: rows=[4, 5, 11, 12, 13] cols=[1, 9] n_pos=6
+  k=10: rows=[0, 1, 8, 9, 15] cols=[2, 9, 10] n_pos=6
+  k=11: rows=[6, 7, 14, 15] cols=[1, 9] n_pos=5
+  k=12: rows=[2, 3, 10, 11] cols=[2, 10] n_pos=5
+  k=13: rows=[1, 2, 3, 8, 9] cols=[2, 10] n_pos=5
+  k=14: rows=[5, 6, 7, 12, 13] cols=[2, 10] n_pos=5
+  k=15: rows=[4, 5, 11, 12, 13] cols=[2, 10] n_pos=6
+  k=16: rows=[0, 1, 8, 9, 15] cols=[3, 10, 11] n_pos=6
+  k=17: rows=[6, 7, 14, 15] cols=[2, 10] n_pos=5
+  k=18: rows=[2, 3, 10, 11] cols=[3, 11] n_pos=5
+  k=19: rows=[1, 2, 3, 8, 9] cols=[3, 11] n_pos=5
+  k=20: rows=[5, 6, 7, 12, 13] cols=[3, 11] n_pos=5
+  k=21: rows=[4, 5, 11, 12, 13] cols=[3, 11] n_pos=6
+  k=22: rows=[0, 1, 8, 9, 15] cols=[4, 11, 12] n_pos=6
+  k=23: rows=[6, 7, 14, 15] cols=[3, 11] n_pos=5
+  k=24: rows=[2, 3, 10, 11] cols=[4, 12] n_pos=5
+  k=25: rows=[1, 2, 3, 8, 9] cols=[4, 12] n_pos=5
+  k=26: rows=[5, 6, 7, 12, 13] cols=[4, 12] n_pos=5
+  k=27: rows=[4, 5, 11, 12, 13] cols=[4, 12] n_pos=6
+  k=28: rows=[0, 1, 8, 9, 15] cols=[5, 12, 13] n_pos=6
+  k=29: rows=[6, 7, 14, 15] cols=[4, 12] n_pos=5
+  k=30: rows=[2, 3, 10, 11] cols=[5, 13] n_pos=5
+  k=31: rows=[1, 2, 3, 8, 9] cols=[5, 13] n_pos=5
+  k=32: rows=[5, 6, 7, 12, 13] cols=[5, 13] n_pos=5
+  k=33: rows=[4, 5, 11, 12, 13] cols=[5, 13] n_pos=6
+  k=34: rows=[0, 1, 8, 9, 15] cols=[6, 13, 14] n_pos=6
+  k=35: rows=[6, 7, 14, 15] cols=[5, 13] n_pos=5
+  k=36: rows=[2, 3, 10, 11] cols=[6, 14] n_pos=5
+  k=37: rows=[1, 2, 3, 8, 9] cols=[6, 14] n_pos=5
+  k=38: rows=[5, 6, 7, 12, 13] cols=[6, 14] n_pos=5
+  k=39: rows=[4, 5, 11, 12, 13] cols=[6, 14] n_pos=6
+  k=40: rows=[0, 1, 8, 9, 15] cols=[7, 14, 15] n_pos=6
+  k=41: rows=[6, 7, 14, 15] cols=[6, 14] n_pos=5
+  k=42: rows=[2, 3, 10, 11] cols=[7, 15] n_pos=5
+  k=43: rows=[1, 2, 3, 8, 9] cols=[7, 15] n_pos=5
+  k=44: rows=[5, 6, 7, 12, 13] cols=[7, 15] n_pos=5
+  k=45: rows=[4, 5, 11, 12, 13] cols=[7, 15] n_pos=6
+  k=46: rows=[0, 1, 8, 9, 15] cols=[0, 8, 15] n_pos=6
+  k=47: rows=[6, 7, 14, 15] cols=[7, 15] n_pos=5

README.md

@@ -0,0 +1,131 @@
+---
+license: apache-2.0
+tags:
+  - reverse-engineering
+  - aqlm
+  - quantization
+  - mlx
+  - escha
+base_model: EschaLabs/Qwen3.6-35B-A3B-Escha-W2
+---
+
+# Escha-W2 packed AQLM codebook — reverse-engineered reference dump
+
+This repository contains the **first public extraction** of the `escham_reconstruct`
+codebook lattice used by EschaLabs' 2-bit AQLM+Hadamard quantized checkpoint
+[`EschaLabs/Qwen3.6-35B-A3B-Escha-W2`](https://huggingface.co/EschaLabs/Qwen3.6-35B-A3B-Escha-W2).
+
+The Escha packed format stores each MoE expert's `gate_up_proj` (K=2) and
+`down_proj` (K=3) projections as (in/16, out/16, 16·K) int16 codes plus per-row
+and per-column scales. Decoding those codes into a dense fp16 weight matrix
+requires two fixed codebook tables that ship **inside the CUDA `.so`** rather
+than in the safetensors — they were previously inaccessible outside a Linux
+GPU environment running the reference `escha` wheel.
+
+This repo makes the codebook portable.
+
+## Contents
+
+| File | Size | Description |
+|---|---|---|
+| `compact.pkl` | 120 MB | The extracted codebook in sparse-compact form (fp16). See "Format" below. |
+| `OP_SIGNATURE_AUDIT.md` | 15 KB | Full Modal-side introspection of `escha._C` — every operator, its schema, and the linearity / (bi, bj)-invariance proofs. |
+| `LAYOUT_NOTES.md` | 5 KB | Notes on the structural regularities of the k_slot layout, the residual "baseline" question, and known limitations. |
+| `modal_op_audit.py` | 12 KB | Reproducible Modal script (~1 min A10G) that produces the audit report. |
+| `modal_smart_probe.py` | 11 KB | Reproducible Modal script (~2 min A10G) that produces `compact.pkl`. |
+
+## Format
+
+```python
+import pickle
+d = pickle.load(open("compact.pkl", "rb"))
+# For each K in {2, 3}:
+for K in (2, 3):
+    positions = d[f"K{K}_positions"]  # list of (n_nz, 2) int8 (row, col) positions
+    values    = d[f"K{K}_values"]     # list of (65536, n_nz) fp16 codebook values
+    # Reconstruct dense (k_max, 65536, 16, 16) fp16:
+    import numpy as np
+    k_max = len(positions)
+    dense = np.zeros((k_max, 65536, 16, 16), dtype=np.float16)
+    for k, (pos, val) in enumerate(zip(positions, values)):
+        for i, (r, c) in enumerate(pos):
+            dense[k, :, r, c] = val[:, i]
+```
+
+To decode a packed expert weight tile back into fp16, sum the per-slot codebook
+lookups placed at each (bi, bj) block:
+
+```python
+# code: int16 (in_f/16, out_f/16, 16*K)
+in_f = 2048  # or 512 for down_proj
+out_f = 1024  # or 2048 for down_proj
+K = 2  # or 3 for down_proj
+w = np.zeros((in_f, out_f), dtype=np.float32)
+bi_max, bj_max = in_f // 16, out_f // 16
+for k in range(16 * K):
+    idx = code[:, :, k].astype(np.int32) & 0xFFFF  # int16 -> uint16
+    blocks = dense[k, idx]  # (bi_max, bj_max, 16, 16)
+    w += blocks.transpose(0, 2, 1, 3).reshape(in_f, out_f)
+```
+
+**IMPORTANT — known limitation.** For a real expert whose codes activate all
+262 K slots simultaneously, the above summation matches the CUDA op only up to
+an unresolved additive term (per-projection norm ~4e3). This term is _not_
+captured in the codebook (which stores deltas from `op(all-zeros code)`) and
+could not be extracted in this session — the Modal workspace hit its spend
+limit after the codebook extraction completed. See `LAYOUT_NOTES.md` for
+the 30-second follow-up probe that would resolve it.
+
+For a working end-to-end port that skips `escham_reconstruct` entirely (pre-
+dequantized to fp16 on Modal, no runtime decode needed), see
+[`KaedeTai/Qwen3.6-35B-A3B-Escha-W2-MLX`](https://huggingface.co/KaedeTai/Qwen3.6-35B-A3B-Escha-W2-MLX).
+
+## Verified properties
+
+- **Linearity** (superposition): `op(A+B) - op(0) = (op(A) - op(0)) + (op(B) - op(0))`
+  holds exactly for up to 100 random slot activations.
+- **(bi, bj)-invariance**: same code at any tile position produces the same
+  16x16 delta (offset by (bi*16, bj*16)). Tested for corners including
+  (bi=127, bj=63) vs (bi=0, bj=0).
+- **Structural regularity**: the per-k_slot (row, col) support cycles with
+  period 4 in k_slot (K=2).
+- **Op signature**: `escham_reconstruct(Tensor packed, int in_f, int out_f,
+  int K, bool cbA, bool mul1) -> Tensor` — one default overload, accepts
+  leading batch dims on `packed`.
+
+## Reproducing the extraction
+
+Requires a Modal account and the `EschaLabs/escha-runtime-qwen3moe` wheel on
+Hugging Face (public):
+
+```bash
+modal run modal_op_audit.py          # ~1 min A10G, produces OP_SIGNATURE_AUDIT.md
+modal run modal_smart_probe.py       # ~2 min A10G, produces compact.pkl
+```
+
+The smart probe uses **~1024 op calls total** across both K values —
+compared to the naive one-code-at-a-time sweep which would require
+**~328 million op calls (91 h A10G, ~$100)**. The speedup comes from three
+observations, each verified by the audit:
+
+1. Op is exactly linear in codes -> many perturbations can be superposed
+   in a single op call.
+2. The codebook is (bi, bj)-invariant -> each of the 8192 tile-blocks in
+   the (128, 64, 32) code tensor is a free "test bed" for a different
+   codebook entry.
+3. Different k_slots at the same (bi, bj) overlap in output positions ->
+   we must use different (bi, bj) for different (k, v) probes, but that's
+   fine since we have 8192 of them.
+
+Net: 65,536 codes x 32 k_slots / 8192 blocks-per-op = 256 op calls for K=2.
+
+## Credits
+
+- **EschaLabs** — for open-weight Qwen3.6-35B-A3B-Escha-W2 and the reference
+  runtime.
+- **AQLM** (Egiazarian et al., 2024) — the residual-codebook quantization
+  scheme that Escha builds on.
+
+## License
+
+Apache 2.0. Same as the base model.

compact.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:af17e0808291cf7d7d9e5f5e50debbd838833771dcd2d6b4e41be7b529001cd7
+size 127933716

modal_op_audit.py

@@ -0,0 +1,303 @@
+"""Signature + linearity audit of torch.ops.escha.escham_reconstruct.
+
+Answers:
+1. Full op signature (what argument shapes/dtypes are accepted?)
+2. Is the op LINEAR in codes? (i.e., op(A + B) == op(A) + op(B) - op(0))
+   - If yes -> superposition holds, we can probe many (bi, bj) simultaneously.
+3. If we set exactly one code slot to value v, what is the FULL delta pattern?
+   - Which (row, col) positions in the output are nonzero?
+4. Is the (codebook entry) slot-invariant across (bi, bj)?
+   - I.e., does slot (0, 0, 0) v=v produce the same block pattern (offset by 16)
+     as slot (1, 1, 0) v=v ?
+5. Does the pattern depend on the *tile shape* passed to the op?
+   - Compare 128x128 tile vs 2048x1024 real tile.
+
+Writes /vol/op_audit.pkl and /vol/op_audit_report.md.
+"""
+
+from __future__ import annotations
+
+import pickle
+from pathlib import Path
+
+import modal
+
+
+WHEEL_REVISION = "1.0.2+qwen3moe"
+
+image = (
+    modal.Image.from_registry(
+        "nvidia/cuda:12.8.1-cudnn-runtime-ubuntu24.04",
+        add_python="3.12",
+    )
+    .apt_install("curl", "binutils", "git", "ca-certificates")
+    .pip_install("wheel", "pip", "setuptools")
+    .pip_install(
+        "torch==2.9.*",
+        index_url="https://download.pytorch.org/whl/cu128",
+    )
+    .pip_install("numpy", "safetensors", "huggingface_hub[cli]")
+    .run_commands(
+        f"echo escha wheel revision: {WHEEL_REVISION}",
+        "mkdir -p /escha",
+        "hf download EschaLabs/escha-runtime-qwen3moe --include 'sglang/*' --local-dir /escha",
+        "pip install --no-deps /escha/sglang/escha-*.whl",
+    )
+)
+
+vol = modal.Volume.from_name("escha-codebooks", create_if_missing=True)
+app = modal.App("escha-op-audit", image=image)
+
+
+@app.function(gpu="A10G", timeout=1800, memory=16 * 1024, volumes={"/vol": vol})
+def audit() -> dict:
+    import inspect
+    import time
+    import numpy as np
+    import torch
+    import escha  # noqa: F401
+
+    report_lines: list[str] = []
+
+    def log(msg: str) -> None:
+        print(msg, flush=True)
+        report_lines.append(msg)
+
+    log("# Escha `escham_reconstruct` — signature + linearity audit")
+    log("")
+
+    # ---- (1) Introspection ----
+    log("## 1. Introspection")
+    log("")
+    C = escha._C
+    log(f"escha module: `{escha.__file__}`")
+    log(f"escha._C: `{C}`")
+    log("")
+    log("### `dir(escha._C)`")
+    for name in sorted(dir(C)):
+        if name.startswith("_"):
+            continue
+        obj = getattr(C, name)
+        try:
+            sig = str(inspect.signature(obj))
+        except Exception:
+            sig = "?"
+        try:
+            doc = (getattr(obj, "__doc__", None) or "").strip().splitlines()[:1]
+            doc = doc[0] if doc else ""
+        except Exception:
+            doc = ""
+        log(f"  - `{name}` {sig} — {doc}")
+    log("")
+    op = torch.ops.escha.escham_reconstruct
+    log(f"### torch.ops.escha.escham_reconstruct")
+    log(f"  {op}")
+    try:
+        log(f"  overloads: {op.overloads()}")
+    except Exception:
+        pass
+    try:
+        log(f"  default schema: {op._schema}")
+    except Exception:
+        pass
+    for ov_name in ("default",):
+        try:
+            ov = getattr(op, ov_name)
+            log(f"  {ov_name}: {ov._schema}")
+        except Exception as e:
+            log(f"  {ov_name}: {e}")
+    log("")
+
+    # ---- (2) Try different code tensor shapes ----
+    log("## 2. Shape acceptance test")
+    log("")
+    device = "cuda"
+
+    def try_shape(cshape: tuple[int, ...], in_f: int, out_f: int, K: int, tag: str) -> None:
+        try:
+            p = torch.zeros(cshape, dtype=torch.int16, device=device)
+            w = op(p, in_f, out_f, K, True, False)
+            log(f"  {tag}: cshape={cshape} in_f={in_f} out_f={out_f} K={K} -> OK, w.shape={tuple(w.shape)} dtype={w.dtype}")
+        except Exception as e:
+            log(f"  {tag}: cshape={cshape} in_f={in_f} out_f={out_f} K={K} -> {type(e).__name__}: {e}")
+
+    # Minimum: 128x128
+    try_shape((8, 8, 32), 128, 128, 2, "min-K2 128x128")
+    try_shape((8, 8, 48), 128, 128, 3, "min-K3 128x128")
+    # Escha actual
+    try_shape((128, 64, 32), 2048, 1024, 2, "escha gate_up (K=2, in=2048/out=1024)")
+    try_shape((32, 128, 48), 512, 2048, 3, "escha down (K=3, in=512/out=2048)")
+    # Batched leading dim?
+    try_shape((2, 8, 8, 32), 128, 128, 2, "leading batch (2, 8, 8, 32) K=2")
+    try_shape((16, 8, 8, 32), 128, 128, 2, "leading batch (16, 8, 8, 32) K=2")
+    try_shape((4, 128, 64, 32), 2048, 1024, 2, "leading batch (4, 128, 64, 32) K=2")
+    log("")
+
+    # ---- (3) Full delta pattern at slot (0,0,0) for various values ----
+    log("## 3. Full delta pattern at slot (0,0,0)")
+    log("")
+    log("For (in=2048, out=1024, K=2) tile: what is the full (row, col) support")
+    log("of the delta when we set exactly code[0,0,0] = v, for various v?")
+    log("")
+    in_f, out_f, K = 2048, 1024, 2
+    cshape = (128, 64, 32)
+    p0 = torch.zeros(cshape, dtype=torch.int16, device=device)
+    w0 = op(p0, in_f, out_f, K, True, False).detach().cpu().numpy().astype(np.float32)
+    full_deltas = {}
+    for v in [1, 2, 3, 4, 5, 7, 10, 16, 64, 256, 1024, 4096, 16384, 32767, -1, -100, -32768]:
+        p = torch.zeros(cshape, dtype=torch.int16, device=device)
+        p[0, 0, 0] = np.int16(np.uint16(v & 0xFFFF)) if v >= 0 else np.int16(v)
+        w = op(p, in_f, out_f, K, True, False).detach().cpu().numpy().astype(np.float32)
+        d = w - w0
+        # Find all non-zero (row, col) positions:
+        nz_pos = np.argwhere(np.abs(d) > 1e-6)
+        # Should be within rows [0..16) × cols [0..16) (or thereabouts) if layout is
+        # bi*16 offset. Show ALL positions:
+        vals = [(int(r), int(c), float(d[r, c])) for r, c in nz_pos]
+        nz_rows_unique = sorted(set(r for r, _, _ in vals))
+        nz_cols_unique = sorted(set(c for _, c, _ in vals))
+        full_deltas[v] = {
+            "positions": vals,
+            "nz_rows": nz_rows_unique,
+            "nz_cols": nz_cols_unique,
+            "n_positions": len(vals),
+        }
+        log(f"  v={v:6d}: {len(vals):3d} nonzero positions, rows={nz_rows_unique}, cols={nz_cols_unique}")
+    log("")
+
+    # ---- (4) Superposition test at real tile shape ----
+    log("## 4. Superposition test (LINEARITY in codes)")
+    log("")
+    log("Test: op(all-zeros with code[bi_a, bj_a, k_a]=v_a AND code[bi_b, bj_b, k_b]=v_b)")
+    log("      == op(only code[bi_a, bj_a, k_a]=v_a) + op(only code[bi_b, bj_b, k_b]=v_b) - op(zeros)")
+    log("If yes, we can probe many (bi, bj) slots simultaneously in ONE op call.")
+    log("")
+
+    def linearity_probe(pairs: list[tuple[tuple[int, int, int], int]], tag: str) -> None:
+        # Sum of individual perturbations
+        cumulative_delta = np.zeros_like(w0)
+        for (pos, v) in pairs:
+            pi = torch.zeros(cshape, dtype=torch.int16, device=device)
+            pi[pos] = np.int16(np.uint16(v & 0xFFFF)) if v >= 0 else np.int16(v)
+            wi = op(pi, in_f, out_f, K, True, False).detach().cpu().numpy().astype(np.float32)
+            cumulative_delta += (wi - w0)
+        # Combined
+        pc = torch.zeros(cshape, dtype=torch.int16, device=device)
+        for (pos, v) in pairs:
+            pc[pos] = np.int16(np.uint16(v & 0xFFFF)) if v >= 0 else np.int16(v)
+        wc = op(pc, in_f, out_f, K, True, False).detach().cpu().numpy().astype(np.float32)
+        combined_delta = wc - w0
+        # Compare
+        diff = np.abs(cumulative_delta - combined_delta)
+        norm_combined = np.linalg.norm(combined_delta)
+        norm_diff = np.linalg.norm(diff)
+        rel = norm_diff / max(norm_combined, 1e-9)
+        log(f"  {tag}: |combined|={norm_combined:.3e} |diff|={norm_diff:.3e} rel={rel:.3e}")
+        return rel
+
+    # Two positions, different (bi, bj), same k_slot
+    linearity_probe([((0, 0, 0), 100), ((1, 1, 0), 200)], "2-pos, distinct (bi,bj), same k")
+    # Same block, different k_slot
+    linearity_probe([((0, 0, 0), 100), ((0, 0, 5), 200)], "2-pos, same (bi,bj), diff k")
+    # Same block, different k_slot, same k-group
+    linearity_probe([((0, 0, 0), 100), ((0, 0, 16), 200)], "2-pos, same (bi,bj), diff K-slice")
+    # 8 random positions (dense superposition test)
+    torch.manual_seed(42)
+    pairs = []
+    for i in range(8):
+        pos = (int(torch.randint(0, 128, (1,))), int(torch.randint(0, 64, (1,))), int(torch.randint(0, 32, (1,))))
+        v = int(torch.randint(1, 65536, (1,)))
+        pairs.append((pos, v))
+    linearity_probe(pairs, "8-pos random")
+
+    # 100 positions (many-slot dense superposition)
+    torch.manual_seed(1)
+    pairs = []
+    used_positions = set()
+    for _ in range(100):
+        while True:
+            pos = (int(torch.randint(0, 128, (1,))), int(torch.randint(0, 64, (1,))), int(torch.randint(0, 32, (1,))))
+            if pos not in used_positions:
+                used_positions.add(pos)
+                break
+        v = int(torch.randint(1, 65536, (1,)))
+        pairs.append((pos, v))
+    linearity_probe(pairs, "100-pos random")
+    log("")
+
+    # ---- (5) Slot invariance: does (0,0,0)+v produce same shape as (1,1,0)+v (offset by 16 rows/cols)? ----
+    log("## 5. Slot invariance test — is the codebook shared across (bi, bj)?")
+    log("")
+    log("Compare delta patterns for the SAME value v at DIFFERENT (bi, bj) with the same k_slot.")
+    log("If they are identical up to a (bi*16, bj*16) offset, the codebook is (bi, bj)-invariant.")
+    log("")
+    for v in [1, 100, 32767, -32768]:
+        p_a = torch.zeros(cshape, dtype=torch.int16, device=device)
+        p_a[0, 0, 0] = np.int16(np.uint16(v & 0xFFFF)) if v >= 0 else np.int16(v)
+        w_a = op(p_a, in_f, out_f, K, True, False).detach().cpu().numpy().astype(np.float32)
+        d_a = (w_a - w0)[:16, :16]  # extract block at (0..16, 0..16)
+
+        for (bi_test, bj_test) in [(1, 0), (0, 1), (1, 1), (5, 3), (127, 63)]:
+            p_b = torch.zeros(cshape, dtype=torch.int16, device=device)
+            p_b[bi_test, bj_test, 0] = np.int16(np.uint16(v & 0xFFFF)) if v >= 0 else np.int16(v)
+            w_b = op(p_b, in_f, out_f, K, True, False).detach().cpu().numpy().astype(np.float32)
+            d_b = (w_b - w0)[bi_test*16:(bi_test+1)*16, bj_test*16:(bj_test+1)*16]
+            diff = np.abs(d_a - d_b)
+            rel = float(np.linalg.norm(diff) / max(np.linalg.norm(d_a), 1e-9))
+            log(f"  v={v:6d} (bi=0,bj=0) vs (bi={bi_test},bj={bj_test}): |diff|={float(np.linalg.norm(diff)):.3e} rel={rel:.3e}")
+    log("")
+
+    # ---- (6) k_slot invariance: does (0,0,k)+v have consistent pattern for different k? ----
+    log("## 6. k_slot pattern")
+    log("")
+    log("For each k_slot, what is the row/col support of the (0, 0, k)+v=1 delta?")
+    log("")
+    ks_patterns: dict = {}
+    for K_test in (2, 3):
+        cs = (128, 64, 32) if K_test == 2 else (32, 128, 48)
+        inf, outf = (2048, 1024) if K_test == 2 else (512, 2048)
+        p0k = torch.zeros(cs, dtype=torch.int16, device=device)
+        w0k = op(p0k, inf, outf, K_test, True, False).detach().cpu().numpy().astype(np.float32)
+        log(f"### K={K_test}, cshape={cs}")
+        for k in range(16 * K_test):
+            p = torch.zeros(cs, dtype=torch.int16, device=device)
+            p[0, 0, k] = 1
+            w = op(p, inf, outf, K_test, True, False).detach().cpu().numpy().astype(np.float32)
+            d = (w - w0k)[:16, :16]
+            nz_pos = np.argwhere(np.abs(d) > 1e-6)
+            rows = sorted(set(int(r) for r, _ in nz_pos))
+            cols = sorted(set(int(c) for _, c in nz_pos))
+            ks_patterns[(K_test, k)] = {"rows": rows, "cols": cols, "n_pos": len(nz_pos),
+                                        "dense_block": d.tolist()}
+            log(f"  k={k:2d}: rows={rows} cols={cols} n_pos={len(nz_pos)}")
+    log("")
+
+    # Save report + full audit data
+    Path("/vol").mkdir(parents=True, exist_ok=True)
+    result = {
+        "full_deltas": full_deltas,
+        "ks_patterns": ks_patterns,
+        "report_md": "\n".join(report_lines),
+    }
+    with open("/vol/op_audit.pkl", "wb") as f:
+        pickle.dump(result, f)
+    with open("/vol/op_audit_report.md", "w") as f:
+        f.write("\n".join(report_lines))
+    vol.commit()
+
+    return result
+
+
+@app.local_entrypoint()
+def main() -> None:
+    result = audit.remote()
+    out_pkl = Path("/Users/kaede/mlx-video/mlx_video/models/qwen3_5_moe_escha/codebooks/op_audit.pkl")
+    out_md = Path("/Users/kaede/mlx-video/docs/escha_op_signature.md")
+    out_pkl.parent.mkdir(parents=True, exist_ok=True)
+    out_md.parent.mkdir(parents=True, exist_ok=True)
+    with open(out_pkl, "wb") as f:
+        pickle.dump(result, f)
+    with open(out_md, "w") as f:
+        f.write(result["report_md"])
+    print(f"[local] wrote {out_pkl}")
+    print(f"[local] wrote {out_md}")

modal_smart_probe.py

@@ -0,0 +1,257 @@
+"""Extract the FULL layout of escham_reconstruct in a handful of op calls.
+
+Findings from `modal_op_audit.py`:
+  1. Op is EXACTLY linear in codes: op(A+B) = op(A) + op(B) - op(0)  (|diff|=0)
+  2. Codebook is SLOT-INVARIANT across (bi, bj): same delta at any (bi, bj)
+     up to a (bi*16, bj*16) offset.
+  3. Each (k, v) pair produces a fixed (16, 16) block pattern with ~5-9 nonzero
+     positions.
+
+Extraction algorithm:
+  For each K ∈ {2, 3}:
+    Use cshape (128, 64, 32) for K=2  → 8192 blocks per op
+    Use cshape (32, 128, 48) for K=3  → 4096 blocks per op
+    For each k_slot ∈ [0, 16*K):
+      Loop v ∈ [0, 65536) in batches of `blocks_per_op`:
+        Build ONE code tensor where slot[bi_i, bj_i, k_slot] = v_i for a batch
+        of `blocks_per_op` distinct (bi, bj) pairs.
+        Call op once. Read out each (16, 16) block. Store as codebook entry.
+
+Result:
+  cb_K2 : (32, 65536, 16, 16) fp16  — 4.3 GB
+  cb_K3 : (48, 65536, 16, 16) fp16  — 6.4 GB
+Total ~11 GB.
+
+But we know most codes produce SPARSE patterns (5-9 nonzeros). We can store
+compact: for each (K, k_slot) collect ONCE a mask of which (row, col) positions
+are nonzero, then store only the sparse values as (K, 65536, n_nz) fp16.
+This shrinks to <100 MB total. See `_compact_layout` below.
+"""
+
+from __future__ import annotations
+
+import pickle
+import time
+from pathlib import Path
+
+import modal
+
+
+WHEEL_REVISION = "1.0.2+qwen3moe"
+
+image = (
+    modal.Image.from_registry(
+        "nvidia/cuda:12.8.1-cudnn-runtime-ubuntu24.04",
+        add_python="3.12",
+    )
+    .apt_install("curl", "binutils", "git", "ca-certificates")
+    .pip_install("wheel", "pip", "setuptools")
+    .pip_install(
+        "torch==2.9.*",
+        index_url="https://download.pytorch.org/whl/cu128",
+    )
+    .pip_install("numpy", "safetensors", "huggingface_hub[cli]", "hf_transfer")
+    .run_commands(
+        f"echo escha wheel revision: {WHEEL_REVISION}",
+        "mkdir -p /escha",
+        "hf download EschaLabs/escha-runtime-qwen3moe --include 'sglang/*' --local-dir /escha",
+        "pip install --no-deps /escha/sglang/escha-*.whl",
+    )
+)
+
+vol = modal.Volume.from_name("escha-codebooks", create_if_missing=True)
+app = modal.App("escha-smart-probe", image=image)
+
+
+@app.function(gpu="A10G", timeout=3600 * 2, memory=32 * 1024, volumes={"/vol": vol})
+def extract_layout() -> dict:
+    import numpy as np
+    import torch
+    import escha  # noqa: F401
+
+    op = torch.ops.escha.escham_reconstruct
+    device = "cuda"
+
+    result: dict = {}
+
+    for K, cshape, in_f, out_f in [
+        (2, (128, 64, 32), 2048, 1024),
+        (3, (32, 128, 48), 512, 2048),
+    ]:
+        bi_max, bj_max, k_max = cshape
+        blocks_per_op = bi_max * bj_max
+        print(f"\n=== K={K} cshape={cshape} in_f={in_f} out_f={out_f} blocks_per_op={blocks_per_op} ===", flush=True)
+
+        # Baseline (all-zeros).
+        p0 = torch.zeros(cshape, dtype=torch.int16, device=device)
+        w0 = op(p0, in_f, out_f, K, True, False).detach().cpu().numpy().astype(np.float32)
+        print(f"  baseline w0 shape={w0.shape} norm={np.linalg.norm(w0):.3e}", flush=True)
+
+        # Assign a distinct (bi, bj) to each of blocks_per_op probes.
+        block_positions = [(bi, bj) for bi in range(bi_max) for bj in range(bj_max)]
+
+        # For each k_slot, we'll store a full (65536, 16, 16) fp16 tensor of block
+        # deltas. This is 32 MB per k_slot. Total: 32 × 32 MB = 1 GB (K=2), 48 × 32 MB = 1.5 GB (K=3).
+        # We'll compact to sparse form at the end.
+        cb_full = np.zeros((k_max, 65536, 16, 16), dtype=np.float16)
+        t_k_start = time.time()
+        for k in range(k_max):
+            t0 = time.time()
+            # Loop v in batches of `blocks_per_op`.
+            n_ops = 0
+            for v_start in range(0, 65536, blocks_per_op):
+                v_end = min(v_start + blocks_per_op, 65536)
+                n_probes = v_end - v_start
+                # Build code tensor.
+                p = torch.zeros(cshape, dtype=torch.int16, device=device)
+                for i in range(n_probes):
+                    v = v_start + i
+                    bi, bj = block_positions[i]
+                    # int16 wrap: values > 32767 become negative
+                    p[bi, bj, k] = np.int16(np.uint16(v & 0xFFFF)) if v < 32768 else np.int16(v - 65536)
+                # Call op ONCE for these n_probes probes.
+                w = op(p, in_f, out_f, K, True, False)
+                # Extract each (bi, bj) block as the codebook entry for (k, v).
+                w_np = w.detach().cpu().numpy().astype(np.float32)
+                for i in range(n_probes):
+                    v = v_start + i
+                    bi, bj = block_positions[i]
+                    block = w_np[bi*16:(bi+1)*16, bj*16:(bj+1)*16] - w0[bi*16:(bi+1)*16, bj*16:(bj+1)*16]
+                    cb_full[k, v] = block.astype(np.float16)
+                n_ops += 1
+            dt = time.time() - t0
+            total = time.time() - t_k_start
+            print(f"  K={K} k_slot={k:2d} done: {n_ops} op calls, {dt:.1f}s (cumulative {total:.1f}s)", flush=True)
+
+        result[f"cb_K{K}"] = cb_full   # (k_max, 65536, 16, 16) fp16
+        print(f"  K={K} total: {time.time() - t_k_start:.1f}s")
+
+    # === Sanity check: reproduce a real expert weight from the extracted layout. ===
+    print("\n=== sanity check: reproduce w_bare for gate_up L0/E0 ===", flush=True)
+    from safetensors import safe_open
+    from huggingface_hub import snapshot_download
+    import json
+
+    print("  downloading model snapshot (cached in volume)...", flush=True)
+    model_dir = snapshot_download(
+        "EschaLabs/Qwen3.6-35B-A3B-Escha-W2",
+        cache_dir="/vol/hf_cache",
+    )
+    idx = json.load(open(f"{model_dir}/model.safetensors.index.json"))
+    wm = idx["weight_map"]
+
+    def _load(name, expert=0):
+        shard = wm[name]
+        with safe_open(f"{model_dir}/{shard}", framework="pt") as f:
+            return f.get_tensor(name)[expert]
+
+    for tag, pfx, in_f, out_f, K, cshape in [
+        ("gate_up", "model.language_model.layers.0.mlp.experts.gate_up_proj", 2048, 1024, 2, (128, 64, 32)),
+        ("down",    "model.language_model.layers.0.mlp.experts.down_proj",    512,  2048, 3, (32, 128, 48)),
+    ]:
+        code = _load(f"{pfx}.escha_code", expert=0).cuda()
+        w_ref = op(code, in_f, out_f, K, True, False).detach().cpu().numpy().astype(np.float32)
+
+        # Reconstruct via extracted layout.
+        cb = result[f"cb_K{K}"]        # (k_max, 65536, 16, 16) fp16
+        # Baseline for this cshape:
+        p0 = torch.zeros(cshape, dtype=torch.int16, device=device)
+        w0_np = op(p0, in_f, out_f, K, True, False).detach().cpu().numpy().astype(np.float32)
+
+        w_recon = w0_np.copy()   # start from baseline
+        code_np = code.cpu().numpy()   # (bi_max, bj_max, k_max) int16
+        # Convert int16 → uint16 index
+        code_u = code_np.astype(np.int32) & 0xFFFF   # (bi_max, bj_max, k_max)
+        bi_max, bj_max, k_max = cshape
+        for k in range(k_max):
+            # Gather cb[k, code_u[:,:,k]] → (bi_max, bj_max, 16, 16)
+            blocks = cb[k, code_u[:, :, k]]   # (bi_max, bj_max, 16, 16) fp16
+            # Add to w_recon at the right positions.
+            # blocks[bi, bj] → w_recon[bi*16:(bi+1)*16, bj*16:(bj+1)*16]
+            # Reshape trick:
+            blocks_reshaped = blocks.astype(np.float32).transpose(0, 2, 1, 3).reshape(bi_max * 16, bj_max * 16)
+            w_recon += blocks_reshaped
+
+        diff = np.abs(w_recon - w_ref)
+        max_diff = float(diff.max())
+        mean_diff = float(diff.mean())
+        rel = float(np.linalg.norm(diff) / np.linalg.norm(w_ref))
+        print(f"  {tag}: |w_ref|={np.linalg.norm(w_ref):.3e} max_diff={max_diff:.3e} mean_diff={mean_diff:.3e} rel={rel:.3e}", flush=True)
+        result[f"sanity_{tag}"] = {"max_diff": max_diff, "mean_diff": mean_diff, "rel": rel}
+
+    # === Save extracted layout to volume ===
+    print("\n=== saving to /vol ===", flush=True)
+    Path("/vol/layout_v2").mkdir(parents=True, exist_ok=True)
+    for K in (2, 3):
+        cb = result[f"cb_K{K}"]
+        np.save(f"/vol/layout_v2/cb_K{K}.npy", cb)
+        print(f"  wrote /vol/layout_v2/cb_K{K}.npy ({cb.nbytes/1e9:.2f} GB)", flush=True)
+
+    # === Also compute a compact-sparse form ===
+    print("\n=== computing sparse-compact form ===", flush=True)
+    compact: dict = {}
+    for K in (2, 3):
+        cb = result[f"cb_K{K}"]   # (k_max, 65536, 16, 16) fp16
+        k_max = cb.shape[0]
+        # For each k_slot, find the union of (row, col) positions that are ever
+        # nonzero across all 65536 codes. Store cb[k, :, mask] compactly.
+        mask_per_k: list = []
+        vals_per_k: list = []
+        for k in range(k_max):
+            # Any code that has this position nonzero?
+            any_nz = np.any(cb[k].astype(np.float32) != 0, axis=0)   # (16, 16) bool
+            positions = np.argwhere(any_nz)   # (n_nz, 2)
+            n_nz = positions.shape[0]
+            # Extract the (65536, n_nz) values.
+            vals = cb[k, :, positions[:, 0], positions[:, 1]]   # (n_nz, 65536)
+            # Transpose to (65536, n_nz)
+            vals = vals.T.astype(np.float16)
+            mask_per_k.append(positions.astype(np.int8))
+            vals_per_k.append(vals)
+            if k < 3:
+                print(f"  K={K} k={k}: n_nz_positions={n_nz} vals_shape={vals.shape}", flush=True)
+        # Since n_nz may differ per k, store as a list.
+        compact[f"K{K}_positions"] = mask_per_k
+        compact[f"K{K}_values"] = vals_per_k
+    # Serialize.
+    with open("/vol/layout_v2/compact.pkl", "wb") as f:
+        pickle.dump(compact, f, protocol=pickle.HIGHEST_PROTOCOL)
+    print(f"  wrote /vol/layout_v2/compact.pkl", flush=True)
+
+    vol.commit()
+
+    # Return sanity metrics + sizes (avoid returning multi-GB tensors).
+    return {
+        "sanity_gate_up": result["sanity_gate_up"],
+        "sanity_down": result["sanity_down"],
+        "cb_K2_shape": result["cb_K2"].shape,
+        "cb_K3_shape": result["cb_K3"].shape,
+        "cb_K2_nbytes": int(result["cb_K2"].nbytes),
+        "cb_K3_nbytes": int(result["cb_K3"].nbytes),
+    }
+
+
+@app.function(image=image, volumes={"/vol": vol}, timeout=600)
+def fetch() -> dict:
+    """Fetch the compact codebook + optionally the full form."""
+    import os
+    result = {}
+    for name in ("compact.pkl",):
+        path = f"/vol/layout_v2/{name}"
+        if os.path.exists(path):
+            with open(path, "rb") as f:
+                result[name] = f.read()
+    return result
+
+
+@app.local_entrypoint()
+def main() -> None:
+    metrics = extract_layout.remote()
+    print(f"\n[local] extract metrics: {metrics}")
+
+    files = fetch.remote()
+    out_dir = Path("/Users/kaede/mlx-video/mlx_video/models/qwen3_5_moe_escha/codebooks")
+    for name, data in files.items():
+        target = out_dir / f"layout_v2_{name}"
+        target.write_bytes(data)
+        print(f"[local] wrote {target} ({len(data)/1e6:.2f} MB)")