Attention in Qwen3.6-35B-A3B

QK, softmax, and the weighted blend of values.

📓 Companion notebook: EXDai/attention-mechanism — download and run every code snippet yourself. All plots in this article are generated by the notebook.

Episode 13 of EXD

---

A Note on the Real Model

Qwen3.6-35B-A3B has 40 layers. Only 10 of them — every 4th layer — use classic full attention. The other 30 use GatedDeltaNet, an efficient linear attention variant we'll cover in a later episode.

In this episode we focus only on the full-attention layers. Understanding the mechanism first makes the optimization make sense.

---

Where We Left Off

Ep12 ended with one 2048-dim hidden state per token, projected into Q, K, V, and a gate. Now we do it for the whole sentence at once — and crucially, on the real input to an attention layer, not a raw embedding.

To get there we run the model forward and keep every layer's output. The input to the first full-attention layer (layer 3) is hidden_states[3] — the embedding after layers 0, 1, and 2 have enriched it with context:

The →  cat  →  sat  →  on  →  the
 │      │      │      │      │
 └──────┴──────┴──────┴──────┘
             5 tokens × 2048 dims — the layer's input

What we're working with, per token:

Still no cross-token mixing. Every token's Q, K, V are independent projections. The mixing — the entire rest of this episode — happens when Q meets K.

---

Step 1 — QK Norm: Equalizing Volume

The attention score between two tokens is the dot product q · k. But Ep11 taught us the dot product measures both direction and magnitude — and Ep12 showed some heads are naturally louder than others. A loud key would dominate the softmax not because it's relevant, but because its vector is long.

Qwen3.6 fixes this with RMSNorm applied to Q and K, per head, before any scores exist. Similarity now depends on direction, not volume.

V is deliberately not normalized. The magnitude of V is the model's control knob for how much content a token carries — normalizing it would strip that control away.

---

Step 2 — RoPE: Position Enters the Score

Q and K are now direction-normalized but position-blind. Rotate a query to position 2 or position 4 and it would score identically against the same key. The sentence would read as a bag of tokens.

Ep10's mRoPE fixes this with rotation. Qwen3.6 rotates only the first 64 of the 256 dimensions per head (partial_rotary_factor = 0.25). The other 192 dimensions stay position-free, carrying pure semantic content.

For text-only input, the three mRoPE axes (temporal / height / width) all carry the same position, so the 3D rotation collapses to classic RoPE. The notebook applies it by hand — ten lines, no magic:

def rotate_half(x):
    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

def apply_rope(q, k, cos, sin):
    cos = cos.unsqueeze(1); sin = sin.unsqueeze(1)
    rotary_dim = cos.shape[-1]                       # 64
    q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
    k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
    q_rot = q_rot * cos + rotate_half(q_rot) * sin
    k_rot = k_rot * cos + rotate_half(k_rot) * sin
    return torch.cat([q_rot, q_pass], dim=-1), torch.cat([k_rot, k_pass], dim=-1)

The red dashed line at dimension 64 is a hard boundary: everything left of it rotates, everything right of it is untouched — exactly zero change. Position enters the score here, and nowhere else.

---

Step 3 — The Score Matrix (QKᵀ)

Now Q and K meet. Every query head asks its question of every key:

1. GQA expansion. 8 query heads share each KV head, so K is repeated 8× — heads 0–7 share KV head 0, heads 8–15 share KV head 1.
2. Dot products. Q @ Kᵀ produces, for every pair of tokens, how well the query at position i matches the key at position j.
3. Scaling. Divide by √256 = 1/16. Dot products over 256 dimensions have variance ~256 — un-scaled, softmax would saturate into a one-hot mess.

Each heatmap cell (i, j) is the raw score: "how much does the query at token i want to look at token j?" These are unnormalized — their absolute scale and offset are arbitrary, and negative values are normal. Softmax in the next step only cares about relative values within a row, so a low or negative cell can still win a large share of attention once the row is normalized.

---

Step 4 — The Causal Mask

The model generates one token at a time. Token i may only see tokens 0…i — the future must stay invisible. Everything above the diagonal is set to −inf, which softmax will turn into exactly 0.

This is why every attention heatmap has a black upper-right triangle: not a learned behavior, an enforced law of causality.

---

Step 5 — Softmax: Scores Become Probabilities

Each row is now a competition. Softmax converts the masked scores into probabilities that sum to 1 — a budget for how much of each key token's value the query token will absorb.

The −inf entries become exactly 0. The strongest score takes the largest share; the runner-ups still get slices. Softmax runs in float32 for stability — the scaling in Step 3 kept the inputs in softmax's working range.

---

Step 6 — The Weighted Sum: Context as a Blend

The probability row is the recipe. The output for token i is a weighted average of all value vectors — a mixture, not a choice. Take the last token, "the" (position 4): its output is the context that will predict the next word. It doesn't copy one token — it blends a mix of the five that came before it.

Each head does this independently with its own recipe, producing 16 different 256-dim context vectors. Sixteen interpretations of the same sentence — this is where multi-head attention gets its expressive power.

---

Step 7 — The Gate: Attention's Volume Knob

Qwen3.6's signature move (from Ep12): the Q projection also produces a gate — per head, per dimension. After attention, the context is multiplied by sigmoid(gate):

• gate → 0 ⇒ sigmoid ≈ 0 ⇒ the head is silenced, its attention didn't matter
• gate → 1 ⇒ sigmoid ≈ 1 ⇒ the head's context passes through unchanged

The gate is learned. The model decides per head, per token, how much it trusts that head's opinion right now — a volume knob on every head, turned in real time.

---

Step 8 — The Output Projection: Back to 2048

Sixteen heads, sixteen 256-dim contexts — 4096 dimensions of opinion. The residual stream is only 2048 wide, so W_O (2048 × 4096) squeezes the concatenated head outputs back down. This is the "project down" from Ep11: the model learns which information survives the bottleneck.

The result is added to the residual stream and the layer is done. 2048 in, 2048 out — always.

---

The Cost of Attention

The attention core has a weakness: every token pairs with every previous token. Doubling the sequence quadruples the QKᵀ and weighted-sum work — O(n²). The projections, by contrast, are O(n).

For a 5-token sentence, the quadratic part is tiny — less than a million FLOPs per layer. For a 32K-token context it's a wall:

This is exactly why Qwen3.6 runs full attention in only 10 of 40 layers — as "synchronization points" — and uses GatedDeltaNet, an O(n) linear attention variant, for the rest. GQA also helps: 2 KV heads shared across 16 Q heads means 8× less KV-cache memory than full multi-head attention.

---

What We Saw

One hidden state per token, one pipeline, eight steps:

• The whole mechanism is linear algebra plus one nonlinearity. Dot products, a mask, softmax, a weighted sum, two matrix multiplies.
• Position enters at RoPE, nowhere else. Without rotation, "cat" at position 2 and "cat" at position 4 would score identically.
• Context is a convex blend. Attention never copies — it mixes.
• The gate is a learned trust knob. Silenced heads still compute; the model just decides their opinion shouldn't count.
• The core is O(n²); the projections are O(n). This asymmetry explains why the model reserves full attention for 10 of 40 layers.

We deliberately stopped the sentence at "the". Ask the model what comes next — it says "mat", and when it produces that token, its attention lands on "cat" and "sat", the words it rhymes with. That's the setup for watching attention drive generation.

Next: Ep14 — GatedDeltaNet — why the other 30 layers can get away with a linear approximation, and what they trade away.