DLM Proposer + AR Verifier — runnable KV-cache-saving framework

[FluffyAIcode]

What this PR delivers

A runnable, end-to-end implementation of the speculative-decoding architecture from our prior design discussion, executing on real public weights — no mock, no fallback, no overfit — plus a forward-looking architecture document for the local inference engine that will wrap this core.

Role	Model	Params
Proposer	dllm-hub/Qwen3-0.6B-diffusion-mdlm-v0.1 (masked-diffusion DLM, same tokenizer as verifier)	0.75 B
Verifier	Qwen/Qwen3-1.7B (closest publicly-available stand-in for "Qwen 3.6")	1.72 B

No public Qwen 3.6 checkpoint exists; Qwen3-1.7B is the closest same-family AR model that shares the proposer's tokenizer (verified to produce identical prompt token-ids at startup) and is large enough for KV-cache savings to be meaningful. Swap-in is a one-flag change (--verifier-id).

Memory accounting

Metric is Net Bytes per Token (per-token persistent memory in steady-state long-context inference), defined as:

Net Bytes per Token (KV-only) = verifier_KV_per_token
                              + proposer_KV_per_token            (= 0 in this build; proposer recomputes per block)
                              + proposer_weight_bytes / (B * S)

Activation peak is not in Net Bytes per Token. A transient activation tensor is allocated when model(...) starts and freed when it returns; it does not accumulate across forwards and does not scale per session. It is a GPU capacity constraint (the forward must fit in HBM), not a per-token cost. The metric module reports it on a separate line.

Architecture

┌──────────────────┐     L tokens      ┌────────────────────────┐
│  DLM Proposer    │ ────────────────► │ AR Verifier            │
│  Qwen3-0.6B-MDLM │                   │ Qwen3-1.7B             │
│  K diffusion     │ ◄──────────────── │ DynamicCache trimmed   │
│  steps / block   │  accept / reject  │ to sink+window slots   │
└──────────────────┘                   └────────────────────────┘

• proposer.py — masked-diffusion block generator faithful to the model card's reference (low-confidence remasking, deterministic at temperature 0). The proposer recomputes per block — its persistent KV contribution to Net Bytes per Token is zero.
• verifier.py — SinkWindowVerifier slices each DynamicCache layer's K/V tensors after every step; new queries always use the global RoPE position (so RoPE on new K/Q is correct), and evicted tokens drop out of attention's view (StreamingLLM-style). Layer-shape invariants raise on mismatch.
• speculative.py — greedy speculative-decoding loop with rejection sampling. When sink + window >= full_seq_len, output is bit-equivalent to greedy AR — verified at runtime; the demo exits with code 2 on mismatch.
• baseline.py — reference greedy AR with full DynamicCache.
• metrics.py — KV byte counting; KV-only Net-Bytes-per-Token formula; capacity-constraint report; projection table.

Forward-looking: local inference engine architecture

docs/local-inference-engine.md describes the Mac/Ubuntu local engine that will wrap this algorithmic core. Highlights:

• Goals: ≤ 2 GB resident memory at S=128k for Qwen3-1.7B + 0.6B-DLM; ≥ 150 tok/s on M3 Max single-request; ≥ 400 / 1500 tok/s on RTX 4090 (single / aggregate).
• No PagedAttention. Under the sink+window invariant, every session's KV is a constant-size object. The three problems PagedAttention solves (fragmentation, prefix sharing, non-contiguity) all evaporate. A 30-LoC fixed-size slab pool replaces it and runs ~5–15% faster on contiguous KV.
• Backend: MLX on Mac (unified memory, fused 4-bit GEMM), CUDA + PyTorch on Linux (Flash-Attention 3, Marlin 4-bit GEMM). ~70% shared code.
• Phased plan P0 → P3 with concrete per-phase acceptance tests; phases are scoped by technical risk and dependency, not calendar time.
• The doc has a clear "what's already in this repo vs. what this doc describes" table so reviewers don't conflate the algorithmic core with the future engine.

Empirical results (from CPU runs)

1. Equivalence-regime self-test (sink+window covers full sequence)

prompt   : "Reply with exactly: OK."
config   : sink=4, window=64, block=8, K=8

baseline    output : "OK.<|im_end|>"   (peak KV =  3,584 KB)
speculative output : "OK.<|im_end|>"   (peak KV =  3,696 KB)
exact match        : True             ← "no intelligence loss" verified
acceptance rate    : 0.375

2. Compression-regime test (window ≪ sequence, real eviction)

prompt   : "Write a one-paragraph explanation of why prime numbers are infinite ..."
S        : 108 tokens (44 prompt + 64 generated)
config   : sink=4, window=24, block=16, K=16

Persistent (in Net Bytes per Token):
  verifier KV (full DynamicCache, baseline) =  12.10 MB total =  114,688 B/token
  verifier KV (sink+window,  speculative)   =   3.06 MB total =   29,734 B/token
                                                                 ── 3.86× verifier-side
  proposer KV                               =   0 B            (recomputed per block)
  proposer weights amortized at B=64,S=108  = 172,468 B/token  (small-S dominates here)
  Net Bytes per Token (KV-only)             = 202,202 B/token  (compression 0.57×)

Capacity (separate, NOT in Net Bytes per Token):
  proposer peak activation (single forward) =  31.30 MB
  verifier peak activation (single forward) =  12.75 MB

Projected Net Bytes per Token (KV-only) at canonical operating points

(per-slot KV measured = 114,688 B; cache_budget = 28 slots; proposer KV = 0)

B	S	Net Bytes per Token	compression vs full KV
1	8 192	145,912	0.79×
8	8 192	18,582	6.17×
8	32 768	4,646	24.69×
8	131 072	1,161	98.75×
8	1 048 576	145	790.02×
32	131 072	309	371.50×
64	131 072	167	688.36×
64	1 048 576	21	5,506.92×

These match the design's analytical prediction: at small B*S the proposer's weight bytes dominate (sub-unity ratios); at large B*S the only persistent cost is the bounded sink+window KV (28 slots × 114,688 B ≈ 3.06 MB total, amortized over S).

How to run

pip install -r requirements.txt
PYTHONPATH=. python3 scripts/smoke_test.py
PYTHONPATH=. python3 -m kv_cache_proposer.run_demo \
    --max-new-tokens 64 --block-size 16 --num-diffusion-steps 16 \
    --sink-size 4 --window-size 24 --batch-size-for-amortization 64 \
    --prompt "Write a one-paragraph explanation of why prime numbers are infinite, suitable for a high school student."

Sample logs and JSON results are committed under results/.

Honest caveats

1. Verifier model: Qwen3-1.7B carries KV on all 28 layers; Qwen 3.5/3.6 hybrid attention carries it on 16/64. Against a true Qwen 3.6 baseline of ~65 KB/token, the absolute compression numbers above would be smaller by a factor of ~1.75×; framework code is unchanged.
2. Acceptance rate is low (~0.12) because this proposer is trained with a different objective (Nemotron-SFT-Code masked diffusion) — not Repr-Align-aligned to Qwen3-1.7B. Low acceptance affects throughput, not correctness; the equivalence-regime self-test verifies output equivalence regardless.
3. Proposer activation memory uses dense logits [1, T, V]. The standard "compute logits at masked positions only" optimization is not applied; at long contexts it would be required to keep the forward fitting in HBM. Net Bytes per Token is independent of this optimization (activation is not in the metric); the capacity number reported is the value at the run's actual sequence length.
4. One environment fix: the dllm-hub modeling file has import dllm inside an if __name__ == "__main__": guard but transformers' static check_imports flags it; we install a no-op stub package to satisfy the check without altering any model logic. README has the portable command (works for any Python 3.x).

Out of scope (per design discussion, separate plumbing)

Multi-target verifier routing (Qwen / Gemma / DeepSeek), session-affinity scheduling, OTA, federated self-learning. Those are platform-level components and not part of this drop. The local-inference-engine doc covers the next layer up but its implementation is also separate work.