StarCoder¶

StarCoder is a family of code-generation language models developed by the BigCode open-science collaboration. Released in 2023, StarCoder demonstrated that specialized code models trained on permissively licensed data could match or exceed proprietary alternatives. Its architecture introduces two key ideas relevant to code generation: Multi-Query Attention (MQA) for efficient inference during long code completions, and Fill-in-the-Middle (FIM) training for non-left-to-right code generation¹.

1. Architecture Overview¶

StarCoder is a decoder-only autoregressive transformer optimized for code. It uses Multi-Query Attention (a single KV head shared across all query heads), learned absolute position embeddings, and a context length of 8,192 tokens -- sufficient for most single-file code completions.

2. Key Innovations¶

2.1 Multi-Query Attention (MQA)¶

MQA, introduced by Shazeer (2019)², uses a single key and value head shared across all query heads:

2.2 Fill-in-the-Middle (FIM)¶

Standard left-to-right training cannot generate code that fills a gap between existing code. FIM transforms training examples by splitting them into prefix, middle, and suffix, then reordering:

2.3 Code-Specific Tokenizer¶

StarCoder uses a byte-level BPE tokenizer with special handling for code:

• Whitespace grouping: Sequences of spaces are tokenized as single tokens (e.g., 4 spaces = one token), critical for indentation-sensitive languages
• Repository context tokens: Special tokens encode file paths and repository structure
• FIM special tokens: <fim_prefix>, <fim_middle>, <fim_suffix>

3. Architecture Diagram¶

flowchart TD
    INPUT["Code Token IDs"] --> EMB

    subgraph EMB["Embedding"]
        TOK_EMB["Token Embedding"]
        POS_EMB["Learned Position Embedding"]
        SUM["Sum"]
        TOK_EMB --> SUM
        POS_EMB --> SUM
    end

    SUM --> BLOCKS

    subgraph BLOCKS["N x StarCoderBlock"]
        LN1["LayerNorm"] --> ATTN
        subgraph ATTN["Multi-Query Attention"]
            Q_PROJ["Q Projection\n(n_heads x d_head)"]
            KV_PROJ["K,V Projection\n(1 x d_head each)"]
            Q_PROJ --> DOT["Scaled Dot-Product\n(all Q heads share K,V)"]
            KV_PROJ --> DOT
            DOT --> MASK["Causal Mask"]
            MASK --> SM["Softmax"]
            SM --> OUT["Output Projection"]
        end
        ATTN --> RES1["Residual Add"]
        RES1 --> LN2["LayerNorm"]
        LN2 --> FFN["FFN (GELU)"]
        FFN --> RES2["Residual Add"]
    end

    BLOCKS --> FINAL_LN["Final LayerNorm"]
    FINAL_LN --> LM_HEAD["LM Head"]

4. Configuration Parameters¶

5. Mathematical Formulation¶

5.1 MQA Attention Scores¶

For query head \( i \) with shared key and value:

The final output concatenates all heads:

5.2 KV Cache Memory Comparison¶

5.3 FIM Probability Factorization¶

Standard autoregressive:

FIM reorders the factorization:

This allows the model to condition on both prefix and suffix when generating the middle segment.

6. Zig Implementation¶

6.1 StarCoderConfig¶

pub const StarCoderConfig = struct {
    n_layers: u32,
    d_model: u32,
    n_heads: u32,
    n_kv_heads: u32 = 1,        // MQA: always 1
    d_ff: u32,
    vocab_size: u32 = 49152,
    max_seq_len: u32 = 8192,
    norm_eps: f32 = 1e-5,
    activation: ActivationType = .gelu,
    use_bias: bool = true,

    pub fn headDim(self: StarCoderConfig) u32 {
        return self.d_model / self.n_heads;
    }

    /// KV cache is dramatically smaller with MQA
    pub fn kvCacheSize(self: StarCoderConfig, seq_len: u32) u64 {
        // Only 1 KV head per layer
        return 2 * @as(u64, self.n_layers) * seq_len
            * self.headDim() * @sizeOf(f32);
    }
};

6.2 Multi-Query Attention¶

pub const MultiQueryAttention = struct {
    wq: Linear,    // [d_model, n_heads * d_head]
    wk: Linear,    // [d_model, d_head]  -- single head
    wv: Linear,    // [d_model, d_head]  -- single head
    wo: Linear,    // [n_heads * d_head, d_model]
    n_heads: u32,
    d_head: u32,

    pub fn forward(
        self: *MultiQueryAttention,
        x: Tensor(f32),
        pos: u32,
        kv_cache: *KVCache,
    ) !Tensor(f32) {
        // Q: [seq, n_heads * d_head]
        const q = self.wq.forward(x);

        // K, V: [seq, d_head] -- single head, shared by all query heads
        const k = self.wk.forward(x);
        const v = self.wv.forward(x);

        // Update KV cache (only 1 head to cache)
        kv_cache.update(k, v, pos);

        // Each query head attends to the SAME K, V
        var outputs: [MAX_HEADS]Tensor(f32) = undefined;
        for (0..self.n_heads) |h| {
            const q_h = q.headSlice(h, self.d_head);
            const cached_k = kv_cache.keys(pos);
            const cached_v = kv_cache.values(pos);

            const scores = scaledDotProduct(q_h, cached_k, self.d_head);
            applyCausalMask(scores, pos);
            const attn = softmax(scores);
            outputs[h] = matmul(attn, cached_v);
        }

        const concat = try concatenateHeads(outputs[0..self.n_heads]);
        return self.wo.forward(concat);
    }
};

6.3 FIM Prompt Construction¶

pub const FIMFormatter = struct {
    prefix_token: u32,    // <fim_prefix>
    suffix_token: u32,    // <fim_suffix>
    middle_token: u32,    // <fim_middle>

    /// Transform code into FIM format for infilling
    pub fn formatInfill(
        self: FIMFormatter,
        prefix: []const u32,
        suffix: []const u32,
        allocator: Allocator,
    ) ![]u32 {
        var tokens = std.ArrayList(u32).init(allocator);

        try tokens.append(self.prefix_token);
        try tokens.appendSlice(prefix);
        try tokens.append(self.suffix_token);
        try tokens.appendSlice(suffix);
        try tokens.append(self.middle_token);
        // Model generates the middle portion autoregressively

        return tokens.toOwnedSlice();
    }
};

7. Variants¶

8. Educational Value¶

9. References¶