Sylla

Source-first course

Speculative and MoE Systems Lab

An advanced inference systems course for quantization internals, speculative decoding, MoE routing, paged cache allocation, and evidence-driven benchmarks.

Learning Outcomes

Source-First Materials

course.md

Speculative and MoE Systems Lab

Course Contract

Speculative and MoE Systems Lab is an advanced implementation course for techniques that change real inference economics: quantization, speculative decoding, MoE routing, paged cache allocation, and reproducible benchmark evidence.

The course assumes the learner can already build or read a small decoder path. It focuses on interventions that are easy to talk about and hard to implement correctly.

Audience

Engineers who completed NanoTransformer in MLX or have equivalent from-scratch transformer experience.

Prerequisites

  • NanoTransformer in MLX
  • Local AI Model Systems Advanced Lab
  • Comfort reading MLX, vLLM, or serving-runtime code

Course Outcome

By the end, the learner has implemented small versions of the systems ideas behind modern inference optimization and can write a defensible evidence report about when each idea helps or hurts.

Chapters

  1. Quantization Internals
  2. Speculative Decoding Beyond DFlash
  3. MoE Routing and Load Forensics
  4. Cache Allocation, Runtime Comparison, and Evidence

Source-First Rule

This course is assessed through runnable artifacts, traces, and systems memos. Platform lessons are optional wrappers around the artifact contract.

artifact-map.md

Artifact Map

Primary Artifacts

  • quant/group_quant.py: group-wise quantization implementation
  • quant/awq_probe.py: activation-aware quantization probe
  • spec/medusa_heads.py: simplified multi-head draft path
  • spec/verify_tree.py: verification and acceptance tracing
  • moe/router.py: top-k router and load statistics
  • moe/switch_layer.py: small Switch-style MoE block
  • cache/page_allocator.py: page allocator and fragmentation analyzer
  • bench/spec_bench.py: speculative decoding benchmark harness
  • bench/runtime_matrix.md: runtime comparison notes
  • reports/evidence_report.md: final capstone report

Evidence Standard

Every artifact must answer:

  • What does it optimize?
  • What correctness invariant can it violate?
  • What metric proves it helped?
  • What workload makes it worse?

Final Portfolio

The finished portfolio should include:

  • acceptance length histograms
  • routing imbalance traces
  • quantization error reports
  • cache fragmentation measurements
  • prefill vs decode timing
  • a recommendation memo for local MLX vs server runtime

references.md

References

Primary Reading Order

  1. Original speculative decoding paper and implementation notes.
  2. Medusa paper and reference implementation.
  3. EAGLE-style speculative decoding material for comparison.
  4. DFlash paper and MLX implementation as one case study.
  5. Switch Transformer and Mixtral-style MoE routing references.
  6. vLLM PagedAttention paper and block manager implementation notes.
  7. GPTQ and AWQ primary papers or reference implementations.
  8. MLX memory and compilation docs.

Reading Discipline

For every paper, extract:

  • the correctness contract
  • the claimed metric
  • the baseline
  • the hardware
  • the workload distribution
  • the implementation detail most likely to invalidate a naive reproduction

Deferral List

Do not spend this course on:

  • broad RLHF or DPO
  • full vLLM clone
  • multimodal inference unless it affects memory layout
  • agent frameworks
  • CUDA kernel writing beyond comparison notes

labs/medusa-verification-lab.md

Medusa Verification Lab

Goal

Implement a simplified speculative path with multiple draft heads and explicit target verification.

Task

  1. Produce multiple candidate tokens from auxiliary heads or a small simulated drafter.
  2. Verify candidates against the target path.
  3. Record accepted prefix length.
  4. Roll back or ignore unaccepted candidates.
  5. Compare against baseline decode.

Output

Create:

  • spec/medusa_heads.py
  • spec/verify_tree.py
  • reports/medusa_acceptance.md

Pass Condition

The report must include an acceptance histogram and a case where speculation loses.

labs/moe-routing-lab.md

MoE Routing Lab

Goal

Implement a small MoE layer and make expert imbalance visible.

Task

Build:

  1. Top-k router.
  2. Expert MLPs.
  3. Weighted combine step.
  4. Load statistics.
  5. Optional auxiliary load-balancing loss.

Run synthetic traffic and one real token batch.

Output

Create:

  • moe/router.py
  • moe/switch_layer.py
  • reports/moe_load_forensics.md

Pass Condition

You can show how top-k, batch size, and prompt distribution affect expert load.

labs/page-allocator-lab.md

Page Allocator Lab

Goal

Implement enough of a paged cache allocator to reason about fragmentation and rollback.

Task

Build an allocator that:

  • allocates fixed-size token pages
  • maps request IDs to logical pages
  • supports append
  • supports release
  • reports fragmentation
  • can undo a speculative allocation

Output

Create:

  • cache/page_allocator.py
  • tests/test_page_allocator.py
  • reports/page_allocator.md

Pass Condition

The report compares contiguous reservation with paged reservation on mixed sequence lengths.

labs/quantization-error-lab.md

Quantization Error Lab

Goal

Measure the difference between naive quantization and activation-aware choices.

Task

Implement group-wise quantization for a linear layer. Then collect activation samples and test whether preserving high-impact channels reduces output error.

Output

Create:

  • quant/group_quant.py
  • quant/awq_probe.py
  • reports/quant_error.md

Pass Condition

The report includes memory reduction, output error, and latency. It must state when quantization did not help.

rubrics/capstone.md

Capstone Rubric

Artifact

Build an evidence suite comparing at least two inference optimizations on one workload.

Examples:

  • Medusa-style speculation vs baseline decode
  • MoE layer vs dense layer with similar active compute
  • paged cache vs contiguous reservation
  • naive quantization vs activation-aware quantization

Rubric

  • Correctness, 3 points: invariants are stated and tested.
  • Measurement, 3 points: benchmark includes warmup, repeated runs, prompt length, generated length, and confidence or variance.
  • Root-cause analysis, 2 points: the report explains why an optimization helped or failed.
  • Reproducibility, 1 point: exact commands, model/config, and environment are recorded.
  • Judgment, 1 point: final recommendation states where this technique belongs and where it should be avoided.

Non-Passing Patterns

  • A speedup claim with no acceptance histogram.
  • A quantization claim with no output-error measurement.
  • A routing claim with no load distribution.
  • A cache claim with no fragmentation or allocation trace.

Audio Companion

Narration scripts, predefined question pairs, and walk-mode checks are available for this course.

2.1. Speculative Correctness Contract

Narration Script

Audio Companion Script: Speculative Correctness Contract

Speculative decoding is easy to explain badly.

The bad version is: use a smaller model to generate tokens, then hope they are good.

That is not the contract.

The useful version is: a cheaper draft path proposes tokens, and the target model verifies them before they become real output.

The draft is allowed to be fast. It is allowed to be imperfect. It is allowed to guess several tokens ahead. But it is not allowed to become the authority for the final sequence. The target model still controls what is committed.

That distinction matters because speculative decoding is an optimization. It should reduce wall-clock time without changing the behavior you meant to get from the target path.

The core loop has five ideas.

First, there is a committed prefix: the tokens that are already real.

Second, the draft proposes a block of future tokens.

Third, the target runs a verification pass over that proposed block.

Fourth, the algorithm accepts the longest prefix that is valid under the chosen verification rule.

Fifth, if the draft went too far, the rejected positions must not leak into future state.

That last point is where many bugs hide. The KV cache is not a passive performance detail. It is part of the model state used to produce future logits. If the cache contains keys and values for tokens that were never committed, the next token is being conditioned on a history that did not actually happen.

So after partial acceptance, cache state has to reflect committed tokens only. The target cache may need trimming. A draft cache may need trimming or may be rebuilt. Hidden state slices need to line up with the accepted region. Position offsets need to match the next real token.

There are also edge cases: an end-of-sequence token inside a proposed block, an off-by-one verification window, mismatched sampling temperature, or a RoPE offset mismatch. Any one of these can produce text that looks plausible while silently violating the target model's path.

The mental model is simple: draft tokens are proposals, not facts. Verification decides what becomes fact. Cache state must agree with the facts.

After this segment, you should be able to explain why speculative decoding can be fast without letting the draft model change the final behavior unchecked, and why cache rollback is a correctness issue, not just a speed issue.

Walking Questions

Why is speculative decoding an optimization rather than a change to the model's authority?

Expected concepts:

  • the draft proposes tokens
  • the target verifies proposed tokens
  • only accepted or target-produced tokens are committed

Hint: Separate proposal from commitment.

What must happen to cache state after a draft block is only partially accepted?

Expected concepts:

  • cache must represent committed tokens only
  • unaccepted draft positions must be trimmed or ignored
  • the next step must use the corrected target state

Hint: Ask what future logits will attend to.

What is one edge case that a speculative decoder has to handle carefully?

Expected concepts:

  • EOS inside a proposed block
  • off-by-one verification window
  • temperature or sampler mismatch
  • position offset mismatch

Hint: Name one place where token order, stopping, or position accounting can go wrong.

Modules

Quantization Internals

Move past stored-byte estimates into error, dequantization, and runtime behavior.

1.1 Group-wise Quantization

Objective

Implement group-wise quantization and measure its real tradeoffs.

Root Confusion

Smaller weights are not the same as faster inference or preserved quality.

Theory

Group-wise quantization stores compact values with per-group scales and possibly zero points. Smaller groups reduce error but increase metadata.

Hands-On Task

Quantize a linear layer by group. Compare dense output, quantized output, memory, and runtime.

Artifact

quant/group_quant.py and reports/group_quant.md

Checkpoint

What changed: stored weight size, temporary memory, latency, or output error?

Failure Mode

Reporting compression ratio without measuring error or runtime.

1.2 Activation-aware Quantization Probe

Objective

Understand why activation-aware quantization can outperform naive weight-only choices.

Root Confusion

Not every channel contributes equally to output error.

Theory

Activation-aware methods use representative activations to identify sensitive channels or groups. The goal is not just small weights, but low error on likely inputs.

Hands-On Task

Collect activation samples for a linear layer. Compare naive group quantization with a simple protected-channel variant.

Artifact

quant/awq_probe.py and reports/activation_aware_quant.md

Checkpoint

Which channels caused the largest output error?

Failure Mode

Calibrating on prompts that do not represent the target workload.

1.3 Fused Dequant and Linear

Objective

Reason about dequantization placement and kernel boundaries.

Root Confusion

Quantization can lose its benefit if dequantization creates large temporary tensors.

Theory

A fused dequant plus linear path can avoid materializing dense weights. Without fusion, memory and bandwidth may erase the compression advantage.

Hands-On Task

Implement a simple unfused path first. Then sketch or prototype a fused path and compare memory behavior.

Artifact

quant/dequant_linear.py and reports/dequant_placement.md

Checkpoint

Where was the dense tensor materialized?

Failure Mode

Counting quantized storage while ignoring dequantized temporaries.

Speculative Decoding Beyond DFlash

Study speculation as a correctness-preserving family of methods, then benchmark it.

2.1 Speculative Correctness Contract

Objective

State and test the correctness invariants of speculative decoding.

Root Confusion

Speculation is an optimization, not permission for a draft model to change committed output unchecked.

Theory

The target path verifies proposed tokens. Accepted tokens are committed. Rejected tokens are replaced by target output. Cache state must reflect committed tokens only.

Hands-On Task

Write a trace format for proposed tokens, target tokens, accepted prefix length, committed tokens, and cache trim.

Artifact

spec/trace.py and reports/spec_correctness_contract.md

Checkpoint

What cache state remains after a partial rejection?

Failure Mode

Accepting plausible draft text without preserving target-model behavior.

2.2 Medusa-style Draft Heads

Objective

Implement a simplified Medusa-style drafting path.

Root Confusion

Speculative decoding is a family of methods, not just DFlash.

Theory

Medusa-style methods attach extra heads that predict future tokens. Verification decides how many proposed tokens can be committed.

Hands-On Task

Build a toy multi-head drafter or simulate draft heads from a smaller model. Measure acceptance length across prompts.

Artifact

spec/medusa_heads.py and reports/medusa_heads.md

Checkpoint

How does acceptance change with temperature or prompt type?

Failure Mode

Judging a drafter by sample quality instead of accepted tokens per verification cost.

2.3 Speculative Benchmarking

Objective

Benchmark speculation in a way that can disprove the speedup claim.

Root Confusion

Speculation can be slower than baseline if acceptance does not amortize overhead.

Theory

The relevant metrics are baseline decode time, draft time, verification time, acceptance length, rollback cost, peak memory, and quality guardrails.

Hands-On Task

Run baseline decode and speculative decode on the same prompt set. Record acceptance histograms and per-block timing.

Artifact

bench/spec_bench.py and reports/spec_benchmark.md

Checkpoint

What workload made speculation lose?

Failure Mode

Reporting one aggregate tokens-per-second number without acceptance distribution.

MoE Routing and Load Forensics

Implement routing, compare MoE and dense layers, and diagnose imbalance.

3.1 Top-k Router

Objective

Implement a top-k router and understand expert assignment.

Root Confusion

Active parameter count hides routing and load distribution.

Theory

An MoE router scores experts for each token, selects top-k experts, and combines expert outputs with weights. The route distribution changes performance.

Hands-On Task

Implement top-k routing and report expert load for synthetic and real token batches.

Artifact

moe/router.py and reports/router_load.md

Checkpoint

Which expert was hottest, and why might that matter?

Failure Mode

Assuming evenly distributed expert usage without measuring it.

3.2 Switch-style MoE Layer

Objective

Build a small MoE layer and compare it to a dense layer.

Root Confusion

MoE efficiency depends on routing, dispatch, batch shape, and memory residency, not just active FLOPs.

Theory

A Switch-style layer routes tokens to expert MLPs and combines outputs. In small batches, dispatch overhead can dominate.

Hands-On Task

Implement a small MoE MLP and compare throughput and memory against a dense MLP with similar active compute.

Artifact

moe/switch_layer.py and reports/moe_vs_dense.md

Checkpoint

When did the dense layer win?

Failure Mode

Ignoring dispatch overhead and expert imbalance.

3.3 Load Balancing Forensics

Objective

Diagnose MoE load imbalance from traces.

Root Confusion

An MoE model can have good average compute but bad tail behavior.

Theory

Hot experts increase latency and reduce parallel efficiency. Load-balancing losses and capacity factors try to control this, but workload distribution still matters.

Hands-On Task

Generate routing traces for several prompt categories. Compute imbalance ratio, experts touched, and tail load.

Artifact

moe/load_forensics.py and reports/moe_load_forensics.md

Checkpoint

Which prompt class created the worst imbalance?

Failure Mode

Reporting mean throughput while hiding tail latency.

Cache Allocation, Runtime Comparison, and Evidence

Make memory allocation, runtime choice, and optimization claims falsifiable.

4.1 Page Allocator and Fragmentation

Objective

Implement a page allocator for cache-like token storage.

Root Confusion

Maximum context is not only a tensor formula. Allocation policy decides how much memory is wasted.

Theory

Paged allocation reserves fixed-size blocks and maps logical token positions to physical pages. It reduces waste for mixed sequence lengths but adds bookkeeping.

Hands-On Task

Build a page allocator and compare reservation waste against contiguous max-length allocation.

Artifact

cache/page_allocator.py and reports/page_allocator.md

Checkpoint

How much memory did paging save on mixed lengths?

Failure Mode

Optimizing allocation while breaking token order or rollback semantics.

4.2 Runtime Comparison Matrix

Objective

Compare MLX and serving runtimes by workload, not reputation.

Root Confusion

"Fast runtime" is not a requirement. Hardware, batching, cache, architecture support, and observability decide fit.

Theory

MLX is strong for Apple Silicon local work and implementation control. vLLM and SGLang are serving-oriented. Transformers is useful for reference behavior and model hacking.

Hands-On Task

Create a runtime decision matrix for local single-user, batch evaluation, and shared serving workloads.

Artifact

bench/runtime_matrix.md

Checkpoint

Which runtime would you choose for a single local research session, and why?

Failure Mode

Choosing the most fashionable runtime instead of the one that fits the workload.

4.3 Evidence Report

Objective

Produce a systems evidence report that rejects weak optimization claims.

Root Confusion

Advanced inference work is full of claims that sound plausible and fail under controlled measurement.

Theory

A good report states the claim, baseline, workload, metric, measurement method, result, root-cause explanation, and decision.

Hands-On Task

Choose one optimization from the course. Run the smallest benchmark that could prove it did not help. Write the report.

Artifact

reports/evidence_report.md

Checkpoint

What evidence would make you reverse your recommendation?

Failure Mode

Defending an optimization after the measurements say it does not fit the workload.