MEMTP — Multi-Ensemble Memory-Elastic Token Prediction

A clean numpy implementation of DeepSeek's DSpark speculative-decoding algorithm (Cheng et al. 2026) + the MEMTP extension: the draft-head ensemble grows to fill freed memory and breathes with the engine's coherence. Built and measured on an AMD Strix Halo APU (gfx1151) — no NVIDIA required.

memtp is the algorithm, decoupled from any trained drafter. The hardware-aware scheduler runs on any [0,1] confidence sequence (your engine's native-MTP / logprob confidences), so you can use it today without training anything. The trained drafter weights are a separate artifact (DeepSeek ship them; see Scope).

from memtp import memtp_depth_live, hardware_aware_prefix_scheduler, NgramMatryoshka

# MEMTP — scale the draft-head ENSEMBLE to the live free-memory headroom ("keep scaling width")
depth, telemetry = memtp_depth_live(max_depth=16)
# e.g. 21.5 GiB free -> ensemble depth 5 at R=1.0, breathing down to 1 under load

# DSpark Algorithm 1 — verify SMARTER not longer: admit each request's prefix by survival a=∏c,
# maximizing throughput Θ = τ·SPS(B). Lossless (non-anticipating early-stop).
sched = hardware_aware_prefix_scheduler([[0.95, 0.93, 0.90, 0.86], [0.55, 0.35, 0.25, 0.10]])
print(sched.lengths)   # the confident request verifies longer than the chatty one

What it is

DSpark accelerates LLM decoding by drafting a block of tokens and verifying it in one pass. Its two ideas:

1. Semi-autoregressive generation — a parallel backbone drafts the whole block in one forward pass, then a lightweight sequential head injects intra-block dependency so the suffix doesn't decay. (semi_ar_block_dist, MarkovHead, RnnHead.)
2. Confidence-scheduled verification — a confidence head estimates per-position survival; a hardware-aware scheduler spends the target's verify-budget only where it pays. (ConfidenceHead, analytical_acceptance, sequential_temperature_scaling, hardware_aware_prefix_scheduler.)

MEMTP is the extension this repo adds: the stacked draft heads are an ensemble, and the ensemble size is memory-elastic — as KV/cache compression frees tiles in the buffer, the ensemble grows to occupy them (each freed tile funds another head decoding deeper), breathing with the coherence order parameter R. The buffer is never under-utilized. (memtp_depth, memtp_depth_live.)

The n-gram matryoshka tower (NgramMatryoshka) realizes the spectrum from DSpark's Markov head (bigram) to its RNN head (full prefix) as a restriction tower where the n-gram order grows with block position (orders 1, 2, …, k+1), with the tower coherence R = ∏ Rᵢ.

Measured (live, on the box)

signal	value
free working memory	21.49 GiB
→ MEMTP ensemble depth @ R=1.0	5
→ @ R=0.3 (searching, leave headroom)	1
DSpark scheduler	confident request verifies longer; load trims the budget (the paper's Fig-8 curve)

Install

pip install -e .          # numpy-only core
pip install -e ".[dev]"   # + pytest
python -m memtp.core      # the self-checking demo
pytest                    # the test suite

Scope (honest)

• The algorithm is fully here (numpy). The hardware-aware scheduler works on any confidence sequence — no trained drafter needed.
• The trained drafter weights (the Markov / confidence head) need DeepSpec training (≈8-GPU + ~38TB target cache) or the shipped deepseek-ai/DeepSeek-V4-Pro-DSpark checkpoint. The heads here are structural (seeded) — drop in real weights and the same code runs.
• The serving speedup is not yet measured end-to-end here — there is no live vLLM wiring in this repo (the scheduler output isn't yet handed to num_speculative_tokens). The pieces are correct and tested; the production number is the next gate.

Credit

DSpark and DeepSpec are DeepSeek-AI's work (MIT) — github.com/deepseek-ai/DeepSpec. This repo is an independent clean-room numpy implementation of the published algorithm plus the MEMTP / matryoshka-tower extensions, by Intuition Labs.

License

MIT. _{(and yes — the name started as a "MEME-TP" joke before it earned the Multi-Ensemble backronym.)}