Inkling-512k-NVFP4-AQLM-hybrid

A memory-optimized hybrid-quantized build of thinkingmachines/Inkling (MoE, 256 routed experts/layer, 6 active). Each MoE layer's experts are split into a small hot set kept at NVFP4 (4 bpw) and a large cold set stored as AQLM 2-bit. This is the 512k tier: 14.8% of expert slots are NVFP4 (hot) and 85.2% are AQLM-2bit (cold), averaging **2.66 bits per weight** over the routed experts. It serves at 512k context (with MTP speculative decoding) on 4x96GB GPUs — verified, numbers below.

This checkpoint needs a CUSTOM FUSED MoE KERNEL that dispatches NVFP4 hot experts and AQLM cold experts within the same fused grouped-GEMM. It does NOT load on stock vLLM / Transformers — use the SM120 fork below. The on-disk format and dequant math are fully specified so other kernels can consume it directly.

Serving (VERIFIED: 512k + MTP on 4x96GB)

The 14.8%-hot tier is ~358 GB weights, and it serves at 512k context WITH MTP speculative decoding on 4x96GB = 384 GB (and up to 640k without MTP). Serving needs the SM120 vLLM fork jarrelscy/vllm-glm52-sm120 (branch feature/inkling), which carries the custom fused hybrid MoE kernels (NVFP4 hot + AQLM cold in one grouped GEMM / gemv, occupancy-tuned V2 decode gemv), Triton paged rel-bias attention (split-KV decode with fused combine, varlen prefill, sliding-window), the restored MTP multi-prefill speculator, per-request adaptive draft suspension, and FULL_DECODE_ONLY CUDA-graph support.

git clone -b feature/inkling https://github.com/jarrelscy/vllm-glm52-sm120
cd vllm-glm52-sm120
docker build -f Dockerfile.inkling-sm120 -t inkling-sm120 .
docker run --gpus all --ipc=host -p 8001:8001 \
  -v /path/to/this/repo/snapshot:/models/inkling:ro \
  inkling-sm120                    # default MODE=512k-mtp; -e MODE=640k for long-ctx

Reproducible build: this image pins flashinfer to a known-good commit and forces a NotImplementedError stub for nvidia-cutlass-dsl's cute.experimental (the 4.5.2 wheel was republished in place with a CuteExperimentalDSL ref that fatally AttributeErrors at import flashinfer.comm on CUDA<13.1). A plain docker build is therefore index-independent.

Both modes: TP4, bf16 KV cache (lossless — no KV quantization), --gpu-memory-utilization 0.97, --max-num-batched-tokens 2048, FULL_DECODE_ONLY CUDA graphs (capture sizes [1,2,4]). MTP uses the model's own speculative heads ({"method": "mtp", "num_speculative_tokens": 2}); drafts are verified against the base model, so greedy output is unchanged (lossless). 512k-mtp mode also enables per-request adaptive draft suspension (INKLING_ADAPTIVE_SPEC=1): requests whose content drafts poorly (e.g. free prose) automatically fall back to plain decode, so MTP can no longer make a request slower than the non-MTP stack.

mode	context	spec decode	needle-in-haystack
512k-mtp (default)	524,288	MTP ns=2	PASS @ 520,759 tok
640k	655,360	off	PASS @ 653,695 tok

Measured throughput (4x RTX PRO 6000 Blackwell 96GB, TP4)

(2026-07-20: all decode numbers below re-measured after the decode kernel trio — occupancy-tuned gemv V2, fused split-KV combine, verify-via-gemv — and adaptive speculation landed default-on; the trio alone cut MTP round cost ~19–20% at both short context and 512k depth.)

Decode — OpenAI-API server mode, greedy, 600 new tokens, short prompts, 512k-mtp mode. Ranges over repeated passes: MTP acceptance (and therefore speed) depends on how predictable the content is.

workload	decode tok/s	MTP tokens/round (max 3)
counting	74–80	3.00
code	65–77	~2.7–3.0
prose	49–67	adaptive (suspends when drafting doesn't pay)

Reference: the same stack without MTP decodes 50.0–51.7 tok/s (engine-measured, roughly workload-independent; was 40.8 pre-trio). At 516k-token depth, 512k-mtp runs 48.7 ms/round (−20% vs the previous build), measuring 33–35 tok/s on the depth benchmark content and scaling up with acceptance as at short context; 640k mode decodes 34.5 tok/s at 640k depth (pre-trio measurement).

Prefill — 1,780 tok/s on 4k–16k prompts (engine-measured, flat); **1,174 tok/s sustained filling to 512k**; ~1,075 tok/s sustained to 640k (pre-trio measurement).

Multimodal — audio + image understanding smoke-gated on the shipped container (512k-mtp mode).

Quality

Held-out perplexity (teacher-forced, 5-domain mix disjoint from calibration; gate-proxy = cold experts requantized to NVFP4 so it loads on stock vLLM, whose per-expert error tracks the true hybrid). Lower is better.

domain	Inkling-NVFP4	prior hybrid	this build (v3, 14.8% hot)
code	2.222	2.250	2.244
agentic	1.930	1.938	1.934
instruction	4.674	5.169	4.892
medical	5.464	5.964	5.817
prose	5.399	6.691	6.563
aggregate	3.406	3.678	3.608

This v3 build re-tiers and re-calibrates on a broad multimodal corpus (code / agentic / instruction / medical / prose, plus image + audio routing). It beats the prior text-only-calibrated hybrid by −1.9% aggregate, with the gains concentrated where the prior build was starved (instruction −5.4%, medical −2.5%, prose −1.9%); code/agentic are unchanged within noise. This 14.8% tier edges the 10%-hot tier (−1.1% vs the prior build).

Provenance

• Hot experts (NVFP4): sliced verbatim from the ModelOpt NVFP4 base (thinkingmachines/Inkling-NVFP4), no requantization.
• Cold experts (AQLM): quantized from the Inkling BF16 teacher. Both w13 and w2 cold books are activation-aware converged then output-preserving PV-tuned against the BF16 teacher on multimodal calibration activations (text + image + audio routing), with a disjoint held-out early-stop.
• Assignment of hot vs cold per layer is in hf_quant_config.json['aqlm_hybrid']['aqlm_layer_books'] and reproduced from the multimodal REAP importance study (assign_512k_v3.json, 14.8% hot). Reproduction scripts are in code/.

Layout summary

• 63 NVFP4-packed MoE layers carry the hybrid expert format below (NVFP4 hot slabs + AQLM cold).
• 1 bf16-special MoE layer(s) (layer 2) have no NVFP4 base, so they use the SAME hybrid format EXCEPT the hot slabs are stored as bf16 (experts.{proj}_hot_bf16, no .scale/.scale2); their cold experts are AQLM like every other layer. The fused uncompressed bf16 experts.w13_weight / experts.w2_weight are dropped. Layers 0-1 are dense MLP (not experts).
• All non-expert tensors (attention, embeddings, norms, shared_experts, gate, vision, audio, MTP) are copied verbatim from the NVFP4 base.
• Total hot slots 2417 / 16384 expert slots.

On-disk tensor format (per NVFP4-packed MoE layer L)

prefix = model.llm.layers.{L}.mlp.experts. Routed experts (256 total) are partitioned into n_hot hot + n_cold cold. For each projection proj in {w13, w2} (w13 = fused gate|up, out=6144, in=6144; w2 = down, out=6144, in=3072):

Hot — NVFP4-packed layers (ModelOpt, hot_format=nvfp4):

tensor	dtype	shape
{prefix}.{proj}_hot_weight	uint8	[n_hot, out, in/2] (2 fp4 codes/byte)
{prefix}.{proj}_hot_weight.scale	float8_e4m3	[n_hot, out, in/16] (per-block-16 scale)
{prefix}.{proj}_hot_weight.scale2	float32	[n_hot] (per-expert global scale)

Hot — bf16-special layer(s) (hot_format=bf16, e.g. layer 2):

tensor	dtype	shape
{prefix}.{proj}_hot_bf16	bfloat16	[n_hot, out, in] (dense bf16, no scales)

Cold (AQLM) — same for every layer. One code tensor PER BOOK, smallest lossless dtype:

tensor	dtype	shape
{prefix}.{proj}_cold_codes.{b}	uint8 if book has <=256 entries else int16	[n_cold, out, in/8] (group_size=8)
{prefix}.{proj}_cold_codebook.{b}	float16	[entries_b, 8]
{prefix}.{proj}_cold_scales	float16	[n_cold, out] (per-output-row scale)

Per-book code dtypes: w13 book0 = int16 (65536 entries); w2 book0 = int16 (65536), w2 book1 = uint8 (256 entries, codes 0..255 — lossless). {proj}_cold_codes.{b} replaces the old single {proj}_cold_codes tensor (the book axis is now the trailing .{b} suffix), saving ~38 GB losslessly.

Index maps:

tensor	dtype	shape	meaning
{prefix}.hot_ids	int32	[n_hot]	routed-expert index (0..255) of each hot slab
{prefix}.cold_ids	int32	[n_cold]	routed-expert index (0..255) of each cold slab

hot_ids[i] gives the original routed-expert id for hot slab i (and likewise cold_ids); together they partition 0..255. group_size = 8.

• w13: 1 book, entries = [65536].
• w2: 2 books, entries = [65536, 256].

Dequant math

Cold (AQLM), expert e, output row r, group g (8 weights):

acc[8] = 0
for b in range(n_books):
    acc += codebook_b[ codes_b[e, r, g] ]           # codes_b = cold_codes.{b}; gather an 8-vector
W[e, r, g*8:(g+1)*8] = acc * cold_scales[e, r]        # per-row scale

Codes are read from the per-book tensor {proj}_cold_codes.{b}. uint8 books (<=256 entries) index the codebook directly. For the int16 65536-entry book, negative int16 values wrap to the intended unsigned row (idx & 0xFFFF).

Hot (NVFP4, standard ModelOpt e2m1):

codes = unpack_nibbles(hot_weight)                    # 0..15 -> e2m1 LUT value
W = LUT[codes] * scale.repeat_interleave(16, dim=-1) * scale2

e2m1 LUT = [0, .5, 1, 1.5, 2, 3, 4, 6, -0, -.5, -1, -1.5, -2, -3, -4, -6]; low nibble is the first weight of each pair.

Files

• model-*.safetensors + model.safetensors.index.json — the sharded weights.
• hf_quant_config.json — original NVFP4 config plus the aqlm_hybrid block (quant_method, per-layer {n_nvfp4, n_cold, packed}, book entry counts).
• config.json — base config plus a quantization_config block advertising quant_method = inkling_nvfp4_aqlm_hybrid.
• code/ — the exact repackaging scripts (build_hybrid.py, ink_common.py, ...).