🚧 Week 3 Day 1: Continuous Batching

You begin with the completed Week 2 single-request model: multi-offset RoPE and causal masking already have stable interfaces, and each request can own a dense KV cache. The Day 1 starter leaves four learner-owned slices behind those interfaces:

• dense batch assembly and masking in BatchingKvCache;
• mlx_quantized_linear plus the per-weight selector and the explicit dispatch_week3_batch_model factory;
• selector propagation through Qwen3ModelWeek2; and
• Request.try_prefill plus the request-admission/decode loop in batch_generate.

Complete and test those slices in order. The resulting continuous batch keeps several active requests on the device and replaces each request as soon as it finishes. Only quantized projections cross the Week 3 MLX seam; normalization, activation, RoPE, cache state, attention, and scheduling remain course-owned.

So far, each generation loop has processed only one request. That may not provide enough work to use the device efficiently, so Day 1 decodes several requests in each model call.

A static batch could select five prompts and run them together until every request finishes. However, generated sequences have different lengths. If four requests finish quickly while the fifth continues, most of the batch remains idle and queued requests cannot start.

Continuous batching instead sets a maximum number of active decode requests. When one finishes, the scheduler assigns its batch slot and KV-cache entry to a waiting request. This keeps the decode batch populated whenever work is queued.

The scheduler must also interleave prefill and decode work. We will use a simple policy: advance one pending prefill, then decode one token for every active request.

while requests_in_queue_or_in_progress:
    if prefill_request is not None:
        prefill_request.try_prefill()  # Day 1 processes the complete prompt
        if prefill_request.ready:
            if kv_cache.try_add(prefill_request):
                prefill_request = next(requests)
    if active_requests:
        tokens = decode(model, kv_cache)
        for request, token in zip(active_requests, tokens):
            request.append(token)

A complete prompt is admitted in one call on Day 1. This makes the scheduling policy easy to inspect and exposes an important limitation: one long prefill can delay every active request’s next decode step. Day 2 will add a bounded prefill budget to solve that fairness problem.

Task 1: Reuse RoPE and Causal Masking for Batched Requests

src/tiny_llm/week2_kernels.py::FastRoPE  (reuse unchanged)
src/tiny_llm/attention.py::causal_mask   (reuse unchanged)

Continuous batching requires one RoPE offset per batch element and a causal mask whose query and source lengths may differ. Verify those two Week 2 interfaces before adding the scheduler so the serving layer can use one model contract for every request position.

Verify multi-offset RoPE and the rectangular causal mask with:

pdm run test --week 3 --day 1 -- -k task_1

Task 2: Batch KV Cache

src/tiny_llm/kv_cache.py::BatchingKvCache

BatchingKvCache holds one request cache per decode slot. Because requests may have different sequence lengths, it must combine their keys and values into dense tensors and construct a matching B x 1 x L x S mask.

S = max(S_i across active requests)
L = mask_length (input parameter)
request_keys: H, S_i, D
request_values: H, S_i, D
batched_keys: B, H, S, D
batched_values: B, H, S, D
mask: B, 1, L, S

Right-align each active request in the common S dimension. The leading positions remain zero and masked out. Inactive slots remain fully masked.

keys_i, values_i = request_cache[i]
batched_keys[i, :, (S - S_i):S, :] = keys_i
batched_values[i, :, (S - S_i):S, :] = values_i
mask[i, :, 0:L, (S - S_i):S] = causal_mask(L, S_i)

You can verify your solution by running:

pdm run test --week 3 --day 1 -- -k task_2

Task 3: Add the Week 3 Projection Seam

src/tiny_llm/quantize.py::mlx_quantized_linear
src/tiny_llm/qwen3_week2.py::Qwen3ModelWeek2.__init__
src/tiny_llm/models.py::dispatch_week3_batch_model

Week 2 ends with a course-owned quantized matmul so you can inspect its loader, SIMD-matrix operations, and Split-K policy. Week 3 teaches serving mechanisms, so it should not make every cache and scheduler measurement depend on that teaching kernel’s remaining projection overhead.

Add a per-weight use_mlx_quantized_linear selector whose default remains False, preserving every Week 2 checkpoint. When selected, quantized_linear should call mx.quantized_matmul with the same packed weight, scales, biases, group size, bits, and transposed-weight convention. Then implement dispatch_week3_batch_model as the explicit construction seam: it builds the completed dense-cache Week 2 model with that selector enabled.

Call this batch-ready model with several requests, one offset per batch element, and the mask returned by BatchingKvCache. Exercise requests joining and leaving at different positions. Only quantized projections cross the MLX seam; normalization, RoPE, activation, attention, cache state, and scheduling remain in your solution. The model remains request-agnostic; slot ownership and lifecycle belong to the cache and scheduler.

Verify the projection seam and its batch-model factory with:

pdm run test --week 3 --day 1 -- -k task_3

Task 4: Batch Generate

src/tiny_llm/batch.py

First implement Request.try_prefill by prefilling the complete prompt in one call. The visible prefill_max_step input is reserved for Day 2; on Day 1, give this call a budget that covers the complete remaining prompt. Then complete the scheduler in batch_generate: move finished prefills into idle decode slots, collect the next token and offset for each slot, and remove requests that reach EOS or max_seq_len. A prompt longer than max_seq_len must fail before model or cache work, and a generated token that would cross the limit is not emitted.

Results are returned in completion order, because short requests can leave the batch before earlier long requests. Each result’s prompt_idx maps it back to the original input position; do not reorder completed results into input order.

Use the supplied scheduler checkpoint for full prefill, admission and slot reuse, immediate EOS, maximum-length termination, completion-order results, and their original prompt_idx values:

pdm run test --week 3 --day 1 -- -k task_4

Then run the complete scheduler against the real model:

pdm run batch-main --solution tiny_llm --loader week2

By default, batch-main uses Qwen3-0.6B with a batch size of five and a fixed prompt set. Treat it as a product smoke: watch requests enter, decode, finish, and release their slots. Its shuffled prompt order and cumulative wall-clock display are not the source of a decode-gap measurement.

Day 2’s deterministic bench-chunked-prefill runner owns that comparison. Its 512-token budget processes every prompt in the checked 64–512-token trace in one chunk, so that row is the reproducible Day 1 control. The runner records the exact token ids, output budget, seed, process order, and decode-completion gaps before Day 2 changes the prefill budget.

Your feedback is greatly appreciated. Join our Discord community.
Found an issue? Open an issue or pull request at github.com/skyzh/tiny-llm.
tiny-llm-book © 2025 by Alex Chi Z is licensed under CC BY-NC-SA 4.0.