🚧 Week 2 Day 2: Keep W4 Packed — Decode Waits on Weights

At the end of Day 1 your model reuses its K/V cache, so each decode step feeds only one new token through the network. That token is cheap to compute, but it still has to visit every projection weight in the model. Day 1’s roofline showed what that costs: for Qwen3-4B, one decode step streams about 8 GB of BF16 weights and does only one floating-point operation per byte it reads. Decode spends its time waiting on memory.

That is today’s hypothesis, and you can check it on your own model first. After Day 1 you saved a request benchmark and a decode replay of capacity-cache (see Find the Next Bottleneck). If decode takes most of the request and the projections group leads the decode replay, the hypothesis fits your machine: each decode step’s time is set by the weight bytes it reads, so reading fewer bytes should shrink the projections group and speed up decode. If another group leads, write down which one. You still need today’s packed weights for Days 3 to 5, and the comparison at the end of the day will show what packing actually did.

The checkpoint you download is already much smaller than that. The Qwen3-4B-MLX-4bit files store each weight as a 4-bit code. Up to now the course has unpacked those codes into BF16 when loading the model, then read the unpacked copy on every token. Today you stop doing that. The weights stay packed in memory, and a Metal kernel unpacks each one in registers right before multiplying it by the activation. Every token then reads about a quarter of the bytes.

By the end of the day, you will have:

1. an embedding lookup that unpacks only the rows it needs;
2. a quantized-matmul operator, backed by your own MLX primitive;
3. two Metal kernels for it: a simple one that is easy to check, and a fast one shaped for decode; and
4. a cached Qwen3 model whose projections and output head all run through that operator.

The starter already supplies the packed-weight container (QuantizedWeights and its from_mlx_layer loader), the C++ declaration and Python binding for the operator, CMake registration, and stubs that raise an error naming the missing function. Build both extensions, then run the first tests. They fail until you fill in the stubs:

pdm run build-ext
pdm run build-ext-ref
pdm run test --week 2 --day 2 -- -k task_1
pdm run test --week 2 --day 2 -- -k gpu

📚 Readings

• Model Compression and Quantization
• MLX Extensions Development Guide
• Quantized Matmul on GPU (Video)

Eight Weights, Four Bits Each

Start with a small example. Suppose a row of a weight matrix begins with these eight values:

w = [-0.50, -0.31, 0.02, 0.10, 0.33, 0.47, 0.61, 0.80]

Four bits give you 16 codes, 0 through 15. One simple way to map the values onto those codes is to spread the codes evenly between the smallest and largest value:

bias  b = min(w)                   = -0.50
scale s = (max(w) - min(w)) / 15   =  0.08667
code  q = round((w - b) / s)

q = [0, 2, 6, 7, 10, 11, 13, 15]

To get a weight back, reverse the mapping:

$$\hat{w}=qs+b$$

q * s + b = [-0.500, -0.327, 0.020, 0.107, 0.367, 0.453, 0.627, 0.800]

Each reconstructed value is within half a step (s / 2, about 0.043) of the original. That small error is the price of storing four bits per weight instead of sixteen.

You will never run this quantization step yourself. The checkpoint already contains the codes, scales, and biases, and the model may have chosen them differently from the min/max rule above. Your code only does the reconstruction, q * s + b, using the stored numbers exactly as they are.

Groups of 128

A single scale and bias for a whole matrix row would be a poor fit: one large weight would stretch the range and waste most of the 16 codes on values that never occur. Instead, each row is split into groups of consecutive weights, and every group gets its own scale and bias. The Qwen3 MLX 4-bit checkpoints use a group size of 128, and every required test in this chapter does too.

The format has a name: W4A16. Weights use 4-bit codes; activations stay 16-bit floating point. In this course, activations, scales, biases, and outputs are all BF16.

The Scale Can Be Negative

The codes are unsigned, but the stored scale is signed. With a positive scale, code 0 is the low end of the group’s range. With a negative scale, code 0 is the high end:

positive scale: s =  0.0867, b = -0.5  =>  q=0 is -0.5, q=15 is about  0.8
negative scale: s = -0.0867, b =  0.8  =>  q=0 is  0.8, q=15 is about -0.5

Because a stored scale can be negative, you cannot assume code 0 is the low end. If you rebuild the scale from an assumed min/max, every group with a negative scale comes out reversed. Use scale and bias as stored.

Eight Codes in One uint32

A 4-bit code is too small to address on its own, so eight codes share one 32-bit word. The first code sits in the lowest four bits:

codes  [a, b, c, d, e, f, g, h] = [0, 2, 6, 7, 10, 11, 13, 15]

word = (h << 28) | (g << 24) | (f << 20) | (e << 16)
     | (d << 12) | (c << 8)  | (b << 4)  | a
     = 0xFDBA7620

a = (word >> 0)  & 0xF = 0
b = (word >> 4)  & 0xF = 2
...
h = (word >> 28) & 0xF = 15

Read the hex digits from right to left and you get the codes in order: 0, 2, 6, 7, A, B, D, F.

For a matrix with K output rows and N input columns, the checkpoint stores three arrays:

packed codes: K × (N / 8)    uint32    eight codes per word
scales:       K × (N / 128)  bfloat16  one per group
biases:       K × (N / 128)  bfloat16  one per group

The QuantizedWeights container in src/tiny_llm/quantize.py holds exactly these fields:

QuantizedWeights.from_mlx_layer reads these fields from an MLX quantized layer and converts the scales and biases to BF16. Use it everywhere you load a quantized layer rather than writing a second loader.

How Much Faster Can Decode Get?

The packed codes are half a byte per weight, but each group of 128 also carries a 2-byte scale and a 2-byte bias. Add that metadata and recompute Day 1’s numbers for Qwen3-4B:

bytes per W4 weight  = 0.5 + (2 + 2) / 128   = 0.53125 bytes
streamed bytes/token = 4,022,272,000 × 0.53125 = 2.137 GB
arithmetic intensity = 8.045 GFLOPs / 2.137 GB = 3.765 FLOPs/byte

Decode reads 3.765 times fewer projection-weight bytes per token. If decode were limited only by reading those weights, that would also be the speedup. Real decode also reads K/V, runs normalization and attention, and launches many small operations, so expect a smaller gain.

A Bandwidth Ceiling for Each Chip

Apple publishes memory bandwidth for each chip, but not a comparable BF16 GPU throughput figure. Bandwidth alone still gives a useful upper bound:

best possible tokens/s = memory bandwidth / weight bytes per token

The table below uses the highest-bandwidth configuration of each chip in Apple’s published specifications, with decimal GB as Apple uses. These numbers are ceilings for the arithmetic above, not measurements.

Bandwidths come from Apple’s specifications for M1 Pro and Max, M1 Ultra, M2 Pro and Max, M2 Ultra, M3 Pro and Max, M3 Ultra in the 2025 Mac Studio, and M4 Pro and Max. The table stops at M4 Pro and Max; it makes no claim about later chips.

A real model falls short of these ceilings because it reads more than projection weights and cannot hold peak bandwidth for a whole step. The performance appendix keeps measured results separate from this exercise.

The roofline also applies only to decode, where one activation row meets each weight. Prefill sends many rows through the same weights, so each byte is reused many times and the bottleneck moves toward arithmetic. Day 3 gives prefill its own kernel.

The Math You Will Implement

The course writes a linear layer as C = A Bᵀ:

• A: activations, shape (M, N), BF16;
• B: weights, shape (K, N), stored as W4 codes;
• C: outputs, shape (M, K), BF16.

M is the number of activation rows after you flatten every leading dimension. During decode, M is 1. Each output is a dot product over the input dimension:

$$C[i,k]=\sum _{j=0}^{N-1}A[i,j]\cdot B[k,j]$$

With quantized weights, B[k, j] is reconstructed from its code and its group’s parameters, where g = floor(j / G):

$$B[k,j]=q[k,j]\cdot \text{scale}[k,g]+\text{bias}[k,g]$$

Split the sum by group. Inside a group the scale and bias are constants, so they come out of the inner sum:

$$C[i,k]=\sum _{g=0}^{N/G-1}\left(\text{scale}[k,g]\sum _{j^{\prime }=0}^{G-1}A[i,gG+j^{\prime }]q[k,gG+j^{\prime }]+\text{bias}[k,g]\sum _{j^{\prime }=0}^{G-1}A[i,gG+j^{\prime }]\right)$$

In words: for each group, multiply the scale by the dot product of activations and raw codes, add the bias times the plain sum of those activations, and add up the group results. The fast kernel later in this chapter uses exactly this form. For one group, the identity is:

$$\sum _j{a}_j(s{q}_j+b)=s\sum _j{a}_j{q}_j+b\sum _j{a}_j$$

Written as a loop, the direct version computes one output like this:

for each output C[i, k]:
  sum = 0.0                                  # float32 accumulator
  for each group g in 0 .. N/G - 1:
    scale = scales[k, g]
    bias  = biases[k, g]
    for each packed word p in 0 .. G/8 - 1:  # 16 words per group
      word = B_packed[k, g * (G/8) + p]
      for shift in [0, 4, 8, ..., 28]:
        q = (word >> shift) & 0xF
        a = A[i, g*G + p*8 + shift/4]
        sum += a * (q * scale + bias)
  C[i, k] = bfloat16(sum)

Task 1: Unpack in Python First

src/tiny_llm/quantize.py
src/tiny_llm/embedding.py

Before writing any GPU code, write the reconstruction in plain mlx.core. This version becomes the reference you trust when a kernel misbehaves.

The starter’s QuantizedWeights.from_mlx_layer already loads the packed fields; reuse it rather than writing another loader. Implement these functions:

• dequantize_weights in quantize.py. Unpack each uint32 with shifts and masks, repeat each group’s scale and bias across its 128 values, and compute q * scale + bias. Do not call mx.dequantize; the point is to write the unpacking yourself.
• QuantizedEmbedding.__init__, QuantizedEmbedding.__call__, and QuantizedEmbedding.as_linear in embedding.py.

The token embedding table is quantized too, and it is large: vocab_size rows of hidden_size values. A lookup should only unpack the rows it needs. embedding(input_ids) gathers the packed words, scales, and biases for the requested token ids, then calls dequantize_weights on just those rows. The ids may arrive as uint32 values from the sampler, so accept them as well as int32.

Qwen3 ties its output head to the embedding table: the same matrix that maps token ids to vectors also maps the final hidden state back to vocabulary logits. embedding.as_linear(h) is that output projection. Implement it as quantized_linear(h, weight) so it reuses the quantized-matmul path instead of expanding the whole vocab_size × hidden_size table. It starts working once Task 3 is done.

Also implement two thin wrappers in quantize.py:

• quantized_linear(x, w, bias=None) takes the same arguments as Week 1’s linear. On the course path, it chooses a kernel from the number of rows, as described in Task 3.
• mlx_quantized_linear(x, w, bias=None) computes the same result with mx.quantized_matmul(..., transpose=True). This is a control, not your implementation: the Task 4 model test builds a second copy of your model that routes through it, then compares the two. quantized_linear should call it only when w.use_mlx_quantized_linear is true.

The Task 1 tests check dequantize_weights and the embedding lookup against mx.dequantize:

pdm run test --week 2 --day 2 -- -k task_1

Task 2: Declare the Operator

src/extensions/src/tiny_llm_ext.h
src/extensions/bindings.cpp
src/extensions/src/quantized_matmul.cpp
src/extensions/CMakeLists.txt

An MLX extension operator has two parts. A C++ function checks its inputs and returns a lazy mx::array that records “this array will be produced by my primitive.” Later, when MLX evaluates the graph, it calls the primitive’s eval_gpu (or eval_cpu) to do the work. MLX owns the arrays and the command encoder; you supply the shape logic and the kernel launch. This is the same pattern as the axpby example in the MLX extension guide.

Most of the plumbing is already in place. Read the Day 2 quantized_matmul(...) declaration and the QuantizedMatmul class in tiny_llm_ext.h, and check that the existing m.def("quantized_matmul", ...) binding and the quantized_matmul.cpp entry in CMakeLists.txt match it. Do not add a second binding or a duplicate source entry. Keep everything in the tiny_llm_ext namespace.

Then fill in quantized_matmul.cpp:

• tiny_llm_ext::quantized_matmul(...) in quantized_matmul.cpp checks the inputs and returns the lazy output array of shape (M, K). Require group_size == 128 and bits == 4, BF16 scales, biases, and activations that all share one dtype, and shapes that agree with each other.
• QuantizedMatmul::eval_cpu(...) in quantized_matmul.cpp should throw a clear error. This operator runs only on the GPU; there is no CPU fallback to keep in sync.

On the Python side, quantized_matmul in quantize.py flattens the leading dimensions of a into M, calls mx.contiguous on each input, calls the extension, and reshapes the result back. Metal receives raw buffers, not strides, so the extension can only index correctly if the data is laid out row by row. Make that a rule at the boundary: call mx.contiguous once in Python, and check it again in C++ before encoding the kernel.

Build to catch signature and binding mismatches early:

pdm run build-ext
pdm run test --week 2 --day 2 -- -k task_1

Before You Write Metal: How to Debug a Kernel

You will not write a C++ CPU version of this operator, so GPU bugs need a plan. Use three levels, each checked against the one before it:

1. The equation in Python (dequantize_weights plus a normal matmul, or mx.quantized_matmul). This defines the right answer.
2. A simple Metal kernel where one thread computes one output. It is slow, but you can read it line by line against the equation.
3. The fast Metal kernel, using SIMD groups and wider loads. Compare it with the simple kernel, not with whole-model text output, which is too far from the bug to tell you anything.

Pick inputs whose answers you can predict: zeros, ones, a ramp, an identity matrix, and a fixed random seed. Try one shape that fits the kernel’s tiles exactly and one with a leftover tail, such as 8 and then 10 rows for an 8-row tile.

MLX is lazy, so an error may not appear where it happened. Force evaluation right after the operator you are testing:

expected = python_reference(*inputs)
actual = metal_operator(*inputs)
mx.eval(expected, actual)

assert actual.shape == expected.shape
assert actual.dtype == mx.bfloat16
assert mx.allclose(actual, expected, rtol=2e-2, atol=2e-2).item()

When something fails, check the boundary before the arithmetic: rank, shape, dtype, and contiguity in Python and C++, and whether each buffer index in the encoder matches the Metal function signature. The kind of failure narrows the search:

• Pipeline creation fails: the kernel name, its template instantiation, or the Metal source is wrong.
• The GPU faults or reads garbage: a grid size, bounds check, stride, or buffer binding is wrong.
• The result is finite but wrong: indexing, reduction, masking, unpacking, or the accumulator update is wrong.

For a wrong number, make the kernel simpler until it agrees: one output per thread, no cooperative loads. Then compare an intermediate, such as one unpacked group or one partial dot product, by writing it to a small debug output buffer. Printing from thousands of GPU threads is rarely readable. Put optimizations back one at a time, rerunning the aligned and tail shapes after each.

Task 3: Write the Metal Kernels

src/extensions/src/quantized_matmul.metal
src/extensions/src/quantized_matmul.cpp
src/tiny_llm/quantize.py

You will write two kernels that compute the same math, each shaped for a different M:

Implement these functions:

• quantized_matmul_vanilla_w4a16_g128 and quantized_matvec_x4_fast_w4a16_g128 in quantized_matmul.metal, instantiated for bfloat16_t;
• QuantizedMatmul::eval_gpu in quantized_matmul.cpp;
• quantized_matmul_vanilla and quantized_matvec_custom in quantize.py, so each kernel can be called directly in tests and comparisons.

A Quick Map of the GPU

Metal organizes work in four levels. You need all four to read the kernels below:

• Thread (lane). One instruction stream with its own registers.
• SIMD group. 32 lanes on Apple GPUs that run in lockstep. Lanes in one group can combine values directly with functions like simd_sum, without going through memory.
• Threadgroup. One or more SIMD groups scheduled together on one GPU core. They can share threadgroup memory and wait for each other with threadgroup_barrier.
• Grid. All the threadgroups in one launch. dispatchThreadgroups sets how many there are, and the GPU spreads them across its cores.

Stage 1: The Simple Kernel

Launch a 2D grid with one thread per output (i, k). Each thread walks all N inputs, unpacks eight codes from each uint32, applies the group’s scale and bias, and accumulates in float32. Check i < M and k < K before writing.

This kernel reloads the same activations in every thread and shares no work, so it is slow. Keep it anyway. It mirrors the loop in the previous section almost line for line, which makes it the kernel you compare everything else against. quantized_matmul_vanilla should always reach it.

Stage 2: A Matvec Built for Decode

With M = 1, there is only one activation row, and each output is a dot product of length N, which is 2,560 for Qwen3-4B’s hidden size. Giving that whole product to one thread leaves most of the GPU idle. Instead, give one output tile to a SIMD group:

• each lane takes a different slice of the N inputs and keeps a partial sum;
• at the end, simd_sum adds the 32 partial sums together;
• lane 0 writes the result.

That works for one output, but every lane still loads activations that the neighboring output will load again. So let each SIMD group compute four outputs at once. In the reference schedule, each lane loads two adjacent packed words of a row, which is 16 activations, and reuses those 16 activations for all four output rows. Two SIMD groups per threadgroup gives eight outputs per threadgroup. For M rows and K outputs, the grid has M × ceil(K / 8) threadgroups of 64 threads each.

Two more tricks make the inner loop cheaper.

Apply the bias once per group. Use the identity from the math section. Each lane keeps Σ a·q and Σ a separately, then adds scale * Σ a·q + bias * Σ a. The bias is applied once per group instead of once per weight.

Skip the shifts. Four codes fit in one 16-bit half-word. Masking that half-word with 0x00F0 gives 16 * q instead of q. If you divide the matching activation by 16 once, before the loop, the factor cancels. Call the pre-divided activations a0' to a3':

hw  = 0x7620                      # codes 0, 2, 6, 7
a0' = a0,  a1' = a1/16,  a2' = a2/256,  a3' = a3/4096

a0' * (hw & 0x000F) = a0        * 0     = a0 * 0
a1' * (hw & 0x00F0) = (a1/16)   * 32    = a1 * 2
a2' * (hw & 0x0F00) = (a2/256)  * 1536  = a2 * 6
a3' * (hw & 0xF000) = (a3/4096) * 28672 = a3 * 7

The kernel scales each lane’s 16 activations once, then uses only masks and multiply-adds for all four outputs. Fewer integer instructions do not guarantee a faster kernel, though. The extra live values compete for registers. Judge the trick by measuring the whole model, as described below.

A lane’s 16 values come from two adjacent words, and a group of 128 spans 16 words, so a lane’s slice never straddles two groups. That is why one scale and one bias per output per loop iteration are enough.

Read activations straight from device memory. One activation row is small and stays in cache. Copying it into threadgroup memory first would add a barrier to every projection. If you try that as an experiment, keep it only if the whole-model measurement improves.

Connect eval_gpu

Follow the axpby dispatch pattern:

1. Get the Metal device and command encoder from the stream.
2. Choose the kernel. The binding’s use_simdgroup flag defaults to true. When it is true and M <= 8, use the matvec kernel. When it is false, use the simple kernel. (Day 3 adds the case of use_simdgroup with M > 8.)
3. Bind the inputs, the output, and the constants M, N, and K, in the order of the kernel’s buffer indices.
4. Size the grid for the chosen kernel and call dispatchThreadgroups.

On the Python side:

• quantized_matvec_custom calls the extension with the default use_simdgroup and rejects inputs with more than 8 rows.
• quantized_matmul_vanilla calls quantized_matmul with use_simdgroup=False.
• quantized_linear counts the rows and calls quantized_matvec_custom when there are at most 8 and w.use_simdgroup_matvec is set; otherwise it calls quantized_matmul.

Then run the GPU tests:

pdm run build-ext
pdm run build-ext-ref
pdm run test --week 2 --day 2 -- -k gpu

These tests compare your kernels with mx.quantized_matmul: the simple kernel on a 128-row identity input and a 3-row random input; the matvec at M = 1, at M = 8, and at a real Qwen3-4B projection shape (one row of 2,560 inputs, 1,024 outputs); and the matvec against your simple kernel. The full Day 2 gate also runs output counts that do not fill the last tile (5 and 65 outputs).

Optional: Tune the Schedule

The four-output, two-SIMD-group layout is a measured starting point, not a rule. If you want to see why, change one choice at a time and measure the whole model each time:

• Outputs per SIMD group: try 2, 4, and 8. More outputs reuse each activation load more often but keep more accumulators alive, which uses registers.
• SIMD groups per threadgroup: try 2, 4, 8, and 16. More groups cover more outputs per threadgroup but can lower how many threadgroups fit on a core.

Adding SIMD groups to a threadgroup and adding threadgroups to the grid are different knobs. The first changes per-threadgroup register demand. The second changes how the outputs are divided. Keep the experiments separate so each number answers one question. If you add an experimental variant, keep the selected path’s kernel names unchanged.

Task 4: Use It in the Model

src/tiny_llm/qwen3_week2.py

Now make the cached model use the new operator. Update Qwen3ModelWeek2.__init__, Qwen3MultiHeadAttention.__call__, Qwen3MLP.__call__, and Qwen3ModelWeek2.__call__.

Add the packed path next to the dense one; do not replace it. The model takes a checkpoint name, and WEEK2_CHECKPOINT_FEATURES in qwen3_week2.py records what each checkpoint turns on. quantized_weights is false for kv-cache and capacity-cache and true from quantized-matvec onward. Keep the dense BF16 weights, embedding, and linear calls for the two Day 1 checkpoints, so capacity-cache stays a fair control. When quantized_weights is true:

• Load every attention projection (wq, wk, wv, wo) and MLP projection (w_gate, w_up, w_down) with QuantizedWeights.from_mlx_layer(...) instead of expanding it to BF16, and call quantized_linear(x, w) where you used to call linear(x, w).
• Load embed_tokens the same way and wrap it in your QuantizedEmbedding.
• If the model has a separate lm_head (that is, the embeddings are not tied), keep it as QuantizedWeights too and apply it with quantized_linear. It is a projection, not a lookup.
• Pass the constructor’s use_mlx_quantized_linear flag through to from_mlx_layer. The test builds a control model with this flag set and expects its projections to go through mlx_quantized_linear.
• Keep group_size and bits from each loaded layer, and let the extension reject anything other than 128 and 4.

Leave the Week 1 model alone. Its layers still take plain mx.array weights.

If the model outputs nan or nonsense, check dtypes first. Scales, biases, and activations must all be BF16, and the output must be BF16 too.

Run the whole Day 2 gate, then generate text with your model:

pdm run test --week 2 --day 2

pdm run main --solution tiny_llm --loader week2 \
  --week2-checkpoint quantized-matvec --model qwen3-0.6b --max-tokens 16

Check That the Model Really Uses Your Kernel

Passing tests is not quite enough. The tests compare outputs and call your operators directly, but they do not check which path the model takes internally. Make sure that:

• a decode step goes quantized_linear → quantized_matvec_custom → your primitive → the matvec kernel; and
• a prompt longer than eight tokens goes quantized_linear → quantized_matmul → your primitive → the simple kernel.

An easy way to check is to put a temporary exception in one branch of quantized_linear and confirm that the main command above hits it.

Then measure. Keep the model, prompt length, output length, and logit mode identical across every run you compare:

pdm run bench --solution tiny_llm --loader week2 \
  --week2-checkpoint quantized-matvec --model qwen3-0.6b \
  --num-seqs 1 --min-input-len 128 --max-input-len 128 \
  --min-output-len 65 --max-output-len 65 --warmup 2 \
  --prefill-logits last

Run it again with --solution tiny_llm_ref to compare with the reference solution, and again with --week2-checkpoint capacity-cache to compare with Day 1. Don’t change anything else between runs.

To see the whole progression, and how the replayed decode work changes:

pdm run bench-week2-progression --offline --solution tiny_llm --suite week2 --repeats 2 \
  --variant week2-capacity-cache --variant week2-quantized-matvec --variant mlx \
  --model qwen3-0.6b --input-len 128 --output-len 129 --warmup 2 \
  --prefill-logits last

pdm run profile-week2-kernels --solution tiny_llm --model qwen3-0.6b \
  --case capacity-cache:decode:128 --case quantized-matvec:decode:128 \
  --warmup 4 --iterations 12 \
  --json-output week2-day2-attribution.json

The profiler replays one decode step that ends at 128 cached tokens through the public operators each checkpoint selects, as described in Day 1. Both checkpoints use the same bounded cache, and in one command both cases share the same context, cache bound (C = S = 128 by default), and logit setting. The only thing that changes is the packed-weight path: projections and the embedding. If you change --cache-capacity or --prefill-logits, change it for both.

Then answer three questions from your own runs. Did packing change the replayed projections group, and by how much? Did the whole-request decode result in the progression run improve, and by how much? Does your model actually route its projections through quantized_linear (the wiring check earlier in this chapter)? Any of these can come out differently from what you expect: a kernel can be slower than the dense path, and a smaller replay group can still sit next to a worse whole-request result. Keep whatever you get.

Then look for the next bottleneck. Decode now runs your matvec, but a prompt with more than eight rows still goes to today’s readable matrix kernel. Replay one prefill as well:

pdm run profile-week2-kernels --solution tiny_llm --model qwen3-0.6b \
  --case quantized-matvec:prefill:128 --warmup 4 --iterations 12 \
  --json-output week2-day2-prefill.json

If the projections group leads that prefill replay, you have Day 3’s starting point. If something else does, note it before you go on.

For a sense of scale, here is one older measurement of this step. It used Qwen3-4B on an M4 Pro, with an earlier version of the course code and an earlier version of the profiler whose numbers are not comparable with today’s replay. Packing the weights cut the time spent in projections by 69.0%, and decode on a fixed workload went from 24.38 to 58.90 tokens per second. That older profile then showed normalization, positional encoding, and activation taking 33.5% of the time, which is what Day 4 goes after. Your 0.6B numbers will differ, and these are not targets to hit. The performance appendix has the full history.

To compare at 4B yourself, download the model with hf download Qwen/Qwen3-4B-MLX-4bit and rerun every row with --model qwen3-4b. The dense capacity-cache control still needs the memory listed in the model table.

Where Week 2 Ends Up

Today’s kernel is the first of several. Day 3 adds a SIMD matrix kernel for prefill, Day 4 adds RMSNorm, RoPE, and SwiGLU kernels, and Day 5 adds tiled prefill attention. The supplied reference adds further scheduling on top of those. The table compares that finished Week 2 reference with stock MLX running the same Qwen3-4B model. A value above 100% means the course model was faster.

Where these numbers come from: separate measurement runs on one Apple-silicon Mac, made on an earlier development revision of the reference before it was merged in PR #328. They are not a fresh timing of the merged code. Each case ran eight times on each model, with each course-model run paired with one MLX run. The throughput column is the ratio of the two medians (MLX’s median total time divided by the course model’s), not a median of the paired ratios. All runs used BF16 activations, W4 weights with group size 128, one request at a time, and the same fixed number of generated tokens. PR #328 records the remaining limits:

• in the 128 / 128 and 8,192 / 128 rows, several paired runs fell below 95% even though the medians did not;
• one long measurement campaign was interrupted twice, and the runs on both sides of each interruption are kept in the results;
• the course model’s outputs differ numerically from MLX’s;
• the course model failed a biology-answer check, and MLX’s answer to a factual check had its own problem;
• runs keep generating past a natural end-of-sequence token to reach the fixed count, so they say nothing about output quality over the full length.

These are whole-request times, including the CPU work of preparing and submitting GPU commands, not isolated kernel timings. The course model was slower than MLX to produce the first token in every case. Decode speed was mixed: faster at 2,048 / 16 and 8,192 / 8,192, slower in the rest. Its peak allocated memory was higher than MLX’s in every case except 128 / 128, by up to about 0.9 GB at 32,640 / 128. So the result is close to MLX, not uniformly faster. Your own runs will vary with your machine and model.

If You Get Stuck

If you need to move on before your kernels work, you can implement quantized_linear with mx.quantized_matmul, in the same way as mlx_quantized_linear, and leave the rest of your model unchanged. That swaps out one operator and keeps your cache and model wiring. It is not the same as --solution mlx, which runs a separate, complete MLX model and exercises none of your code.

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.