Hand-Writing a Vision-Language Model, Part 1: Qwen3-4B's Attention Stack

What This Project Is

DeepSeek's progression over the past year has been hard to ignore. After reading their "Thinking with Visual Primitives" paper, it piqued my interest enough to want to understand it from the inside rather than just use it.

The goal is to recreate the core mechanism from that paper: interleaving coordinate tokens (<point>, <box>) directly into chain-of-thought reasoning, at small scale using Qwen3-4B as the language backbone and a pretrained ViT as the vision encoder.

The approach is deliberately the hard way: hand-write the Qwen3-4B architecture as PyTorch nn.Module classes, manually load the official safetensors weights, then bolt on a frozen ViT with a trainable projector. The LLaVA recipe, done by hand, for learning purposes.

This project is a graduation from the NLP journey (embeddings and attention from scratch), the CUDA kernel work (understanding GPU compute at the lowest level), and System Zero (a self-hosted agentic platform where the Reasoning Service runs Qwen). All of those threads converge here.

This is Part 1. It covers the first coding session: RMSNorm, RotaryEmbedding, and the GQA Attention block.

Why Qwen3-4B Specifically

Qwen3.5-4B on HuggingFace is already a multimodal model. It loads via AutoModelForImageTextToText and ships with vision preprocessor configs. Using it would defeat the entire point.

Qwen3-4B (the earlier generation) is confirmed text-only. Plain AutoModelForCausalLM, no vision configs, no preprocessor. This is the correct base: a clean text LLM that we will extend with vision ourselves.

Hardware constraint: GTX 1080, Pascal CC 6.1, 8GB VRAM. Qwen3-4B at bf16 full precision is approximately 8GB with no headroom for the ViT, projector, or activations. The plan is 4-bit quantisation via bitsandbytes to leave room for everything else.

Before Writing Any Code: Read the Weights First

The first step was not opening a Python file. It was opening config.json and model.safetensors.index.json and reading the actual architecture numbers and tensor name patterns before writing anything.

This matters because transformers have subtle implementation differences that are easy to get wrong if you assume rather than verify.

Key numbers from config.json:

398 tensors total = 36 layers × 11 + model.embed_tokens.weight + model.norm.weight. No lm_head.weight. The tied embeddings note confirms the LM head will need to reuse embed_tokens.weight transposed.

Two details that would have been easy to get wrong:

head_dim is not hidden_size / num_attention_heads. 2560 / 32 = 80. But the real head_dim is 128. Reading the actual tensor shapes off the safetensors file confirms this: q_proj.weight is [4096, 2560] (32 × 128 = 4096, not 2560). Q, K, V projections do not round-trip through hidden_size internally.

q_norm and k_norm per layer, shape [128] each. This is a Qwen3-specific detail not present in Qwen2: a shared RMSNorm applied independently to every attention head's 128-dim slice, right after Q/K projection and before RoPE is applied. Two extra norm operations per layer beyond the usual pre-attention and pre-MLP norms.

RMSNorm

class RMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, x):
        input_dtype = x.dtype
        x = x.to(torch.float32)
        variance = x.pow(2).mean(dim=-1, keepdim=True)
        x = x * torch.rsqrt(variance + self.eps)
        return self.weight * x.to(input_dtype)

RMSNorm is not a feature-learning layer. It is a stabilisation utility sitting between attention and MLP blocks in a 36-layer-deep stack. Without it, activation magnitudes drift across that depth, making training unstable. RMSNorm resets every token vector back to a consistent scale before each block.

The formula is output = weight * x / sqrt(mean(x², axis=-1) + eps). Note there is no mean-subtraction step, just mean(x²). That is the actual difference from LayerNorm, which also recenters. The original RMSNorm paper (Zhang & Sennrich, 2019) found recentering contributes little in transformers, so it is dropped for a simpler, cheaper operation.

RotaryEmbedding (RoPE)

class RotaryEmbedding(nn.Module):
    def __init__(self, head_dim, theta=1_000_000.0):
        super().__init__()
        inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def forward(self, x, position_ids):
        inv_freq = self.inv_freq[None, :, None].expand(position_ids.shape[0], -1, 1)
        positions = position_ids[:, None, :].float()
        freqs = (inv_freq.float() @ positions).transpose(1, 2)
        emb = torch.cat((freqs, freqs), dim=-1)
        return emb.cos().to(x.dtype), emb.sin().to(x.dtype)


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


def apply_rotary_pos_emb(q, k, cos, sin):
    cos, sin = cos.unsqueeze(1), sin.unsqueeze(1)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

Standard attention (Q·K) is position-blind. Dot products do not care about token order. RoPE injects position by rotating Q and K vectors by an angle proportional to token position, rather than adding a fixed positional vector to embeddings at the input layer.

The key property: rotating Q by position m and K by position n, then dotting them, produces a result that only depends on m - n. Relative distance falls out of the geometry automatically. Token 5 attending to token 3 gets the same relative-position signal as token 50 attending to token 48. The model never sees an explicit distance number.

RoPE applies to Q and K only, never V. It exists to shape attention scores. V is just the content being aggregated after scores are computed, so there is nothing for RoPE to do there.

Attention: Grouped-Query Attention

def repeat_kv(x, n_rep):
    if n_rep == 1:
        return x
    batch, num_kv_heads, seq_len, head_dim = x.shape
    x = x[:, :, None, :, :].expand(batch, num_kv_heads, n_rep, seq_len, head_dim)
    return x.reshape(batch, num_kv_heads * n_rep, seq_len, head_dim)


class Attention(nn.Module):
    def __init__(self, hidden_size, num_heads, num_kv_heads, head_dim, rms_norm_eps, bias=False):
        super().__init__()
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = head_dim
        self.n_rep = num_heads // num_kv_heads

        self.q_proj = nn.Linear(hidden_size, num_heads * head_dim, bias=bias)
        self.k_proj = nn.Linear(hidden_size, num_kv_heads * head_dim, bias=bias)
        self.v_proj = nn.Linear(hidden_size, num_kv_heads * head_dim, bias=bias)
        self.o_proj = nn.Linear(num_heads * head_dim, hidden_size, bias=bias)

        self.q_norm = RMSNorm(head_dim, eps=rms_norm_eps)
        self.k_norm = RMSNorm(head_dim, eps=rms_norm_eps)

    def forward(self, x, cos, sin, attn_mask=None):
        batch, seq_len, _ = x.shape

        q = self.q_norm(self.q_proj(x).view(batch, seq_len, self.num_heads, self.head_dim))
        k = self.k_norm(self.k_proj(x).view(batch, seq_len, self.num_kv_heads, self.head_dim))
        v = self.v_proj(x).view(batch, seq_len, self.num_kv_heads, self.head_dim)

        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
        q, k = apply_rotary_pos_emb(q, k, cos, sin)

        k, v = repeat_kv(k, self.n_rep), repeat_kv(v, self.n_rep)

        attn_weights = (q @ k.transpose(-2, -1)) / (self.head_dim ** 0.5)
        if attn_mask is not None:
            attn_weights = attn_weights + attn_mask
        attn_weights = torch.softmax(attn_weights.float(), dim=-1).to(q.dtype)
        attn_output = (attn_weights @ v).transpose(1, 2).contiguous().view(batch, seq_len, -1)
        return self.o_proj(attn_output)

Qwen3-4B uses Grouped-Query Attention: 32 query heads but only 8 key-value heads. Each physical KV head is shared by a group of 4 Q heads (n_rep = 32 / 8 = 4), cutting KV-cache size 4x compared to standard multi-head attention, at the cost of some representational flexibility in K and V.

The order of operations matters and is easy to get wrong:

1. Project input to Q, K, V
2. Reshape into per-head vectors
3. Apply q_norm and k_norm: normalise each head's 128-dim slice independently (Qwen3-specific, not present in Qwen2)
4. Transpose to (batch, heads, seq, head_dim)
5. Apply RoPE to Q and K
6. repeat_kv to expand K and V from 8 heads up to 32
7. Scaled dot-product attention with causal mask
8. Merge heads
9. o_proj back down to hidden_size

What Is Next

The attention stack is verified. Part 2 covers:

• SwiGLU MLP block
• DecoderLayer: wiring attention and MLP together with the two block-level RMSNorms
• Full Model class: embeddings, 36-layer stack, final norm, tied LM head
• Loading the real safetensors weights and verifying logits against AutoModelForCausalLM as ground truth

Once the full language model is verified, the ViT encoder and trainable projector get bolted on, and the vision-language model starts taking shape.