feat(escha): serve Qwen3.6-35B-A3B-Escha-W2 natively (2-bit eschamoe + int8), lossless AR ~58 tok/s

[davidtai]

Summary

Native serving for EschaLabs/Qwen3.6-35B-A3B-Escha-W2 — a 2-bit Qwen3.6-A3B checkpoint
(qwen3_5_moe hybrid: 30 GDN + 10 full-attn, 256 experts / top-8). It loads and serves through the
ordinary runtime (mtplx.runtime.load(path, mtp=True)), decoding the 2-bit eschamoe experts and
int8 non-experts on the fly — no transcode, no dense-weight materialization, weights never cached.
Recommended path is autoregressive; MTP spec-decode is experimental (see end).

Improvements (before → after)

change	before	after
int8 non-experts, fused matvec (no dequant-at-load)	25.1 tok/s	59.6 tok/s (2.37×, −2 GiB)
eschamoe compute fp32 → bf16 (native dtype)	fp32 cast storm	identical output, faster
mx.compile decode chain (ESCHA_COMPILE)	56.5	58.5 tok/s (bit-identical)
async-eval AR decode (MTPLX_AR_ASYNC_PIPELINE)	49.4	56.5 tok/s (bit-identical)
chunked prefill (bounded memory)	125 GiB @32k	22.4 GiB @32k
base qwen3_5_mtp MTP injector object level	broken (AttributeError)	fixed (also fixes A3B)

Decode @1024, lossless throughout: 49.4 (sync) → 56.5 (async) → 58.5 tok/s (async + compile).

Benchmark — prefill / decode / memory (1024 → 128k)

context	prefill tok/s	decode tok/s	peak memory
1,024	428	58.1	16.2 GiB
16,384	446	53.6	20.8 GiB
32,768	425	49.5	22.4 GiB
49,152	414	46.5	25.8 GiB
65,536	399	43.2	29.1 GiB
81,920	384	40.8	32.4 GiB
98,304	371	38.7	35.8 GiB
114,688	358	36.7	39.0 GiB
131,072	345	34.8	42.4 GiB

Example generation

Prompt: Write is_palindrome(s) (ignore case/spaces/punctuation), with a docstring and two asserts.

def is_palindrome(s):
    """Returns True if the string is a palindrome, ignoring case, spaces, and punctuation."""
    cleaned = ''.join(char.lower() for char in s if char.isalnum())
    return cleaned == cleaned[::-1]

assert is_palindrome("A man, a plan, a canal: Panama") == True
assert is_palindrome("racecar") == True

The base-path fix

inject_qwen3_5_mtp_support attached the MTP surface at the wrong object level, so
validate_mtp_support failed and forward raised AttributeError. Fixed to the dual-level pattern
(TextModel carries .mtp; the outer model delegates), matching inject_mtp_support. Covered by a
new hermetic CPU test. This repairs the base (non-escha) qwen3_5_mtp path too.

Files

mtplx/eschamoe.py + eschamoe_gather.npz (2-bit decode kernel + compiled decode path) ·
mtplx/int8_linear.py (fused int8 matvec / GEMM-prefill) · mtplx/escha_load.py (loader) ·
mtplx/qwen3_5_mtp_patch.py (Escha head + object-level fix) · mtplx/runtime.py (dispatch/load hook) ·
mtplx/generation.py (opt-in async AR pipeline, default off) · tests/ (object-level + bit-exact
decode) · NOTICE (attribution).

MTP status

The native MTP head binds and drafts; short-context spec-decode is bit-exact. Long-context
spec-decode is experimental — its multi-token verify uses a different scaled_dot_product_attention
kernel than single-token decode, which can flip a near-tie greedy argmax. Serve AR for now.

Attribution

escha_qmv and the int8 tiling are ported/adapted from
dusterbloom/higgs; see NOTICE.

🤖 Generated with Claude Code

[youssofal]

Good work. The reverse-engineering writeup (PTX + goldens, cross-validated, 0/3.1M mismatches) is thorough, the vendor-golden fixture tests are the right approach, and the attribution in NOTICE/int8_linear is how borrowing should look. Nothing here is a refusal; this is sequencing plus a set of asks.

Verified on our side: your hermetic injector test fails on current main exactly as you claimed. .mtp lands on the outer wrapper while validate_mtp_support inspects the TextModel, so it is a real pre-existing break for qwen3_5_mtp-type checkpoints (our shipping 35B is unaffected via a different load path, which is why nobody noticed). We'd like to take the injector fix and its two tests as a standalone cherry-pick so it rides the imminent release. Can you split that commit out, or bless us doing the cherry-pick with attribution?

The eschamoe serving lane itself we'd land in the release after. The diff is 1.7k lines into two of the hottest files (generation.py, runtime.py), which have moved ~100 commits since your base, so it needs a rebase plus our own A/B regression gate at native sampling settings before it ships.

Asks for that round:

1. Rebase onto the new main once it's tagged.
2. is_escha_checkpoint runs at the head of every _load_base_model. Check model_type before touching the filesystem so non-A3B loads don't pay the index read, and lose the bare except on the index parse.
3. Promote the bench/validate_int8.py cases into tests/ (the strided-slice case earned permanent battery coverage). bench/validate_qmv.py is referenced but not in the diff.
4. The evidence ladder (49.4 to 56.5 to 58.5) is greedy-measured. Headline numbers need to be at the model's shipped sampling settings, with greedy as a labeled extra. Related: MTPLX_AR_ASYNC_PIPELINE currently helps only greedy requests. Plain temperature/top-p/top-k sampling can stay on-device the same way (penalties genuinely can't), and that follow-up is what makes the +14% real for product traffic.
5. ESCHA_MTP_SHARE_LAYER=20 cites an acceptance sweep; put the sweep numbers in the evidence doc.
6. Your own doc recommends AR until the verify-kernel near-tie flip is fixed. Make the code enforce that (default the escha path to AR) rather than trusting readers.

[youssofal]

The qwen3_5_mtp validation fix from this PR (MTP surface at the TextModel level) shipped in 2.6.0 with your regression test, credited in the notes. The Escha W2 serving lane itself is still under review here.