Local AI Model Systems Advanced Lab

1.1 Config Archaeology and Shape Derivation

Objective

Learn to derive tensor shapes, cache sizes, and runtime constraints from model configuration before reading implementation code.

Why This Is Advanced

At the advanced level, model cards are not enough. You need to read config.json, infer the tensor program, and predict which assumptions the runtime will make.

The same nominal parameter count can hide different operational profiles:

• multi-head attention vs grouped-query attention
• full attention vs sliding-window attention
• dense feed-forward layers vs routed experts
• conventional transformer blocks vs hybrid recurrent or state-space layers
• static context limits vs RoPE scaling conventions

Shape Pass

For each model, derive:

head_dim = hidden_size / num_attention_heads
q_width = num_attention_heads * head_dim
kv_width = num_key_value_heads * head_dim
mlp_width = intermediate_size
kv_cache_width_per_token_per_layer = 2 * num_key_value_heads * head_dim

If any division is not integral, stop. The architecture assumptions are wrong.

Config Fields To Track

The minimum useful inspection list:

• hidden_size
• num_hidden_layers
• num_attention_heads
• num_key_value_heads
• intermediate_size
• max_position_embeddings
• rope_theta
• rope_scaling
• attention sliding-window fields
• MoE expert count, top-k, and active intermediate size

Exercise

Pick two target models and one DFlash draft model. For each, make a shape worksheet:

model id:
architecture family:
hidden size:
layers:
attention heads:
KV heads:
head dim:
KV width per layer:
MLP or expert width:
position encoding:
attention window:
shape assumptions that could break a loader:

Checkpoint

You should be able to say which dimensions determine KV cache memory and which dimensions determine MLP or expert compute.

1.2 Attention Variants and Position Scaling

Objective

Understand how attention variants and position encoding choices change inference cost, memory, and correctness.

Attention Variants

Multi-head attention gives each query head its own key and value head. Grouped-query attention uses fewer key/value heads than query heads. Multi-query attention is the extreme case where many query heads share one key/value head.

The practical effect:

• query projection cost can remain high
• KV cache memory can be much lower with GQA or MQA
• cache shape must match the model exactly
• kernels must know how to broadcast KV heads across query heads

Sliding Windows

Sliding-window attention limits each token's attention to a recent window for selected layers or all layers. This can reduce effective cache use, but it complicates verification and rollback when speculative decoding is involved.

Do not assume a serving engine handles every window pattern correctly. DFlash's README calls out backend changes around interleaved SWA, which is a reminder that architecture details leak into serving code.

Position Scaling

RoPE settings influence how positions are embedded. Extending context by changing RoPE scaling can be useful, but it is not a free correctness guarantee. Long-context behavior depends on training, scaling method, attention pattern, and runtime support.

Exercise

For one model with GQA and one model with sliding-window attention, answer:

What is the query head count?
What is the KV head count?
What is the cache memory reduction vs full MHA?
Does every layer use full attention?
Which runtime would I trust first and why?
What benchmark would reveal position-scaling failure?

Checkpoint

You should be able to explain why two models with the same hidden size and layer count can have different KV cache footprints.

1.3 MoE Routing and Capacity Lab

Objective

Reason about mixture-of-experts inference beyond the simple "active parameters" label.

The Routing Path

For each token, an MoE layer usually:

1. Computes router scores.
2. Selects top-k experts.
3. Dispatches token states to selected experts.
4. Runs expert MLPs.
5. Combines weighted expert outputs.

The active parameter count approximates the expert compute path, but implementation details decide whether the runtime is fast.

Advanced Failure Modes

Watch for:

• hot experts that dominate traffic
• expert imbalance across a batch
• small batches that underutilize expert kernels
• expert parallelism that needs multiple devices
• memory residency for experts that are rarely active
• router behavior changing across domains

Exercise

Build a tiny router simulator. Feed it synthetic router choices and compute:

• token count per expert
• maximum expert load
• mean expert load
• imbalance ratio
• number of experts touched

Then vary top-k and batch size. The goal is to see why "A3B" is not a complete performance model.

Checkpoint

You should be able to explain why an MoE model can have attractive active compute and still be awkward to serve efficiently.

2.1 KV Cache Layout, Paging, and Quantization

Objective

Understand KV cache engineering as a first-class inference system design problem.

Layout Questions

Every runtime has to decide how cache tensors are laid out and updated. Important questions:

• Is cache contiguous per request, paged, or block allocated?
• Are keys and values stored together or separately?
• Are cache elements fp16, bf16, int8, or lower precision?
• Can prompt cache be trimmed?
• Can cache pages be reused across requests?
• Can speculative verification roll back cache state cheaply?

Paging

Paged KV cache is useful because serving workloads have many requests with different sequence lengths. Instead of reserving one giant contiguous cache per request, the runtime can allocate fixed-size pages and map logical token positions to physical cache blocks.

Paging reduces waste but introduces bookkeeping. That bookkeeping matters when block verification accepts only part of a draft.

Quantized KV Cache

Quantizing weights and quantizing KV cache are separate decisions. KV quantization can reduce long-context memory pressure, but it can also affect quality and kernel support. Treat it as a measured tradeoff.

Exercise

Compare two cache strategies for a batch of requests:

requests = [512, 2048, 4096, 8192]
page size = 256 tokens
layers = 32
KV heads = 8
head dim = 128
cache dtype = fp16

Estimate reserved tokens for:

• contiguous allocation rounded to the max sequence length
• paged allocation rounded per request

Then calculate the memory difference.

Checkpoint

You should be able to explain why cache allocation strategy affects both maximum context and serving throughput.

2.2 Prefill, Decode, and Scheduler Design

Objective

Separate prompt prefill, token decode, batching, and scheduling decisions.

Prefill vs Decode

Prefill processes the prompt. It is parallel over prompt tokens and often looks like a large matrix-heavy workload. Decode generates new tokens. It is sequential per response and often limited by memory bandwidth, cache reads, and small-batch efficiency.

A tokens-per-second number is incomplete unless it tells you whether it measures prefill, decode, or both.

Scheduler Responsibilities

A serving scheduler decides:

• which requests enter a batch
• how much prefill work to run
• how many decode steps to interleave
• when to evict or compact cache
• how to preserve latency for short requests
• whether speculative decoding should run for this request

Speculation Interaction

Speculative decoding changes scheduler pressure. A request may need draft work, target verification, partial acceptance, and rollback. A good scheduler benefits from high acceptance. A bad scheduler can lose the speedup to coordination overhead.

Exercise

Design a scheduler policy for two workloads:

1. One interactive user with a long prompt and long response.
2. A server with many short prompts and mixed response lengths.

For each, decide:

batching policy:
prefill chunk size:
decode batch strategy:
speculation enabled:
cache eviction rule:
metric to optimize:
failure mode:

Checkpoint

You should be able to explain why a runtime that is excellent for offline throughput may feel worse for one local interactive session.

2.3 MLX Profiling and Memory Forensics

Objective

Design a profiling pass for MLX that distinguishes model cost, cache growth, lazy evaluation, and allocator behavior.

Profiling Layers

A useful MLX profiling pass records:

• prompt token count
• generated token count
• time to first token
• prompt processing tokens per second
• decode tokens per second
• peak MLX memory
• current memory after cleanup
• model and draft model IDs
• block size and acceptance lengths, if speculative decoding is enabled

Evaluation Boundaries

In MLX, lazy evaluation can make naive timings wrong. Force evaluation where the measurement needs a real value. If a benchmark spans too much work, split it into phases instead of guessing.

Memory Forensics

When memory rises unexpectedly, separate:

• loaded weights
• target KV cache
• draft KV cache
• hidden state capture
• temporary attention buffers
• tokenizer and Python overhead
• unreleased cached allocations

Exercise

Write a profiling plan for a DFlash MLX run:

target model:
draft model:
prompt tokens:
max generated tokens:
block size:
measure before load:
measure after load:
measure after prefill:
measure every N generated tokens:
record acceptance lengths:
cleanup step:

Checkpoint

You should be able to identify whether a slowdown is compute, cache growth, memory pressure, or benchmark contamination.

3.1 Speculative Decoding Correctness

Objective

Understand the correctness contract behind speculative decoding.

The Contract

A speculative decoder may use a cheaper model to propose tokens, but the target model controls what is committed. The implementation must preserve the intended target-model behavior while reducing wall-clock time.

The core invariants:

• proposed tokens are verified by the target
• accepted tokens match the target's sampled or deterministic path under the chosen algorithm
• rejected regions are replaced with target tokens
• target cache state reflects committed tokens only
• draft cache state does not corrupt future verification

Where Bugs Hide

Speculative decoding bugs often hide in:

• off-by-one verification windows
• cache rollback after partial acceptance
• EOS handling inside a proposed block
• sampling temperature mismatch between draft and target
• attention mask or position offset mismatch
• hidden-state capture from the wrong target layers

Exercise

Write a trace for one draft block:

current committed token:
draft proposal:
target verification tokens:
accepted prefix length:
committed new tokens:
target cache trim:
draft cache trim:
next hidden state slice:

Then deliberately create one off-by-one variant and explain how it would show up in output or metrics.

Checkpoint

You should be able to review a speculative decoding loop and identify the cache and token invariants it must preserve.

3.2 DFlash Block Diffusion Deep Dive

Objective

Read DFlash at the algorithm boundary: what the block-diffusion drafter consumes, proposes, and hands back to verification.

DFlash-Specific Questions

Focus on these questions while reading:

• Which target hidden states are captured?
• How are selected hidden states combined for the draft model?
• How does the draft represent masked proposal positions?
• What does block size control?
• How does sliding-window draft cache bound history?
• How are accepted tokens and hidden states sliced for the next step?

Implementation Anchors

In the MLX implementation, inspect:

• DFlashConfig
• DFlashAttention
• DFlashDraftModel.bind
• load_draft
• _patch_model
• stream_generate
• cache trimming and rollback branches

The purpose is to connect the paper's algorithm to concrete runtime mechanics.

Exercise

Create a table with three columns:

paper concept | MLX implementation location | runtime risk

Include at least:

• target hidden state extraction
• block proposal
• verification pass
• acceptance length
• cache rollback
• EOS handling
• sliding-window draft cache

Checkpoint

You should be able to explain which parts of DFlash are model architecture, which parts are decoding algorithm, and which parts are backend engineering.

3.3 Speedup Model and Benchmark Design

Objective

Build a sober speedup model for speculative decoding and design a benchmark that can falsify it.

Speedup Variables

The useful variables are:

• target baseline time per token
• draft time per proposed block
• target verification time per block
• average accepted tokens per block
• rollback overhead
• sampling and tokenizer overhead
• memory pressure from loading both models

The maximum speedup is constrained by accepted tokens per verification pass. If acceptance is low, speculation can be slower than baseline.

Benchmark Rules

A useful benchmark reports:

• exact target model and draft model
• hardware and memory
• runtime and version
• prompt length distribution
• generated token count
• prefill throughput
• decode throughput
• acceptance length histogram
• peak memory
• quality or task score if quality matters

Exercise

Design a benchmark matrix:

datasets:
prompt lengths:
generation lengths:
temperatures:
block sizes:
baseline runtime:
speculative runtime:
metrics:
stop condition:
decision rule:

Checkpoint

You should be able to reject a speedup claim that reports only one tokens-per-second number without acceptance or prompt-length context.

4.1 Serving Stack Comparison

Objective

Compare MLX, Transformers, vLLM, and SGLang using workload requirements instead of habit.

Comparison Axes

Use these axes:

• hardware target
• model architecture support
• quantization support
• KV cache strategy
• speculative decoding support
• batching and scheduling sophistication
• observability
• ease of hacking model internals
• production serving maturity

Runtime Profiles

MLX is attractive for Apple Silicon local experimentation and direct access to memory behavior. Transformers is useful for reference behavior and model hacking. vLLM and SGLang are serving-oriented runtimes where scheduler, cache management, and multi-request throughput matter more.

The correct choice changes with the question.

Exercise

Create a decision table for three workloads:

1. Local single-user research on an M-series machine.
2. Batch evaluation of 500 prompts.
3. Shared API serving with mixed prompt lengths.

For each, choose a runtime and defend the choice against one alternative.

Checkpoint

You should be able to explain why "fastest runtime" is not a complete requirement.

4.2 Reproducible Inference Experiments

Objective

Make local inference experiments reproducible enough that future you can trust the result.

Experiment Record

Every experiment should record:

• git commit or package versions
• target model ID and revision
• draft model ID and revision
• quantization format
• hardware
• memory pressure from other processes
• prompt dataset
• generation parameters
• random seed or deterministic mode
• warmup policy
• raw metrics
• failure logs

Reproducibility Traps

Common traps:

• model aliases moving under the same name
• silently changed tokenizer files
• warm runs compared against cold runs
• prompt formatting drift
• hidden thinking mode changes
• OS memory pressure changing between runs
• comparing different context lengths

Exercise

Write an experiment manifest for one DFlash run and one non-speculative baseline. Include enough detail that you could rerun it a month from now and know what changed.

Checkpoint

You should be able to tell the difference between a model improvement, a runtime improvement, and a benchmark artifact.

4.3 Capstone Local Inference Architecture Review

Objective

Produce an advanced architecture review for a local or near-local inference setup.

Capstone Scenario

Choose a demanding target such as:

• a large dense model with long context
• a large MoE model with small active parameter count
• a target-plus-draft DFlash pair
• a hybrid setup with local draft and remote target
• a local research machine plus a serving node

Required Memo

Write the review with these sections:

Goal:
Hardware and memory:
Target model:
Draft or auxiliary models:
Architecture-derived memory budget:
KV cache plan:
Attention and position-scaling risks:
MoE routing risks:
Runtime choice:
Scheduler and batching plan:
Speculation plan:
Benchmark matrix:
Quality guardrail:
Failure budget:
Decision:

Review Standard

A good memo identifies root causes before fixes. If the design fails, it should be clear whether the problem is weight memory, KV cache, architecture support, kernel support, scheduler overhead, or model quality.

Checkpoint

You should now be able to reason from model config to runtime choice to benchmark design without relying on vibes.