The Eight-Headed Bug That Hid Behind a System Prompt

24 April 2026 · Discuss on GitHub

When I added system prompt support to Mila’s chat CLI, something unexpected happened. Prompts that had previously produced perfect responses started generating fluent, coherent, and completely context-oblivious text. Ask for a story about rabbits living on the moon eating green carrots — get a charming tale about Hungarian meadow rabbits instead. The model wasn’t broken. It was just ignoring everything I said.

This is the story of how a single innocent constant — kPrefillChunkSize — quietly corrupted eight different stages of the GQA prefill pipeline, and why it took a system prompt to wake the beast.

The Setup

Mila’s prefill pipeline processes long prompts in chunks of 64 tokens (kPrefillChunkSize). For each chunk, the GQA attention op writes K/V entries into the cache, expands them for grouped query attention, computes Q @ K^T via cuBLASLt, applies a causal softmax, computes att @ V, and unpermutes the output. All pre-allocated buffers are sized to kPrefillChunkSize rows.

The key insight — which should have been obvious but wasn’t — is that kPrefillChunkSize serves two distinct roles:

1. The buffer allocation size — how many rows were reserved at build time
2. The actual row count — how many rows the current chunk is using

For full chunks these are identical. For partial chunks — the final chunk when T_prompt is not a multiple of 64 — they diverge. And every place in the pipeline that confused one for the other became a bug.

The Trigger

Without a system prompt, typical user messages fit comfortably within 64 tokens — a single full chunk. All bugs were dormant. The pipeline had never actually executed a partial chunk followed by meaningful downstream processing.

Adding a system prompt pushed T_prompt past 64, activating the partial chunk path for the first time. Eight bugs woke up simultaneously.

The Bugs

Every bug was a variation of the same mistake — kPrefillChunkSize used as a stride where chunk_len was required:

cuBLASLt wrote preatt_ with 64-row head strides regardless of actual chunk size. The softmax then read the wrong row of preatt_ for every head beyond head 0. The model was computing attention over garbage for every layer, every head, every chunk beyond the first.

The Symptom

The corrupted attention heads didn’t produce noise. They produced fluent, grammatically correct, contextually irrelevant text. The model was responding from its language prior — “what statistically follows a system prompt” — rather than from the actual prompt content.

• Ask about rabbits on the moon → get meadow rabbits
• Ask about the weather in Paris → get a recitation of the system prompt persona
• Ask a question with tool calls active → get chain-of-thought reasoning about how to use a tool

This made the bug particularly deceptive. A model producing garbage output is obviously broken. A model producing polished, confident, wrong answers requires careful prompting to diagnose. The tell was specificity: the model correctly responded to the shape of the prompt (a story request, a question) but ignored every specific detail (the moon, green carrots, Paris).

The Diagnosis

The debugging session worked through the pipeline systematically, clearing each stage in turn:

• ✅ KV cache write — correct
• ✅ KV cache expand — correct (incremental, persistent, right index math)
• ✅ RoPE absolute positioning — correct (abs_pos = position_offset + t)
• ✅ Block residual stream — correct
• ✅ Chunk loop hidden state — correct (prefill is stateless across chunks by design)
• ❌ prefill_softmax row offset — first fix (chunk_stride → chunk_len)
• ❌ cuBLASLt partial plan strides — remaining seven fixes

The softmax fix was necessary but not sufficient. After it was applied, the model produced better output on small partial chunks (71 tokens: one full + seven-token partial worked correctly) but still failed on larger inputs with multiple preceding full chunks. The debug print:

[GQA::prefill_softmax] B=1 NH=24 T_stride=4096 chunk_stride=64 chunk_len=11 position_offset=64

confirmed the softmax parameters were correct — which redirected suspicion to the cuBLASLt plans themselves.

The smoking gun was strideC in build_qk_prefill_plan:

// Before — bug
const long long strideC = static_cast<long long>(prefill_window_size) * max_seq_length;

// After — fix
const long long strideC = static_cast<long long>(chunk_rows) * max_seq_length;

cuBLASLt was writing preatt_ with 64-row head strides. The softmax (now correctly striding by chunk_len) was reading with 11-row head strides. Every head beyond head 0 was reading from a completely different location in the buffer.

The Fix

Eight coordinated one-line fixes, all of the form: replace kPrefillChunkSize with chunk_len in stride and offset calculations.

// fix(gqa): correct partial-chunk prefill pipeline for multi-chunk prompts
//
// Eight coordinated bugs caused all heads beyond head 0 to read garbage
// data during partial-chunk prefill (T_prompt > kPrefillChunkSize with
// a non-zero remainder). All bugs shared the same root cause: every
// stage used kPrefillChunkSize as the row stride instead of chunk_len.

A comment has been added to kPrefillChunkSize at its declaration:

// Buffer allocation size only — never use as a stride in matmul plans
// or kernel index arithmetic. Always use the actual chunk_len for the
// current iteration.

Lessons

Single-chunk testing is insufficient for chunked pipelines. A test suite that never exercises the partial chunk path cannot catch partial chunk bugs — no matter how thorough it is otherwise. Mila’s validation suite now includes explicit multi-chunk prefill cases with non-multiples of kPrefillChunkSize.

The most dangerous bugs produce fluent wrong answers. Numerical corruption that manifests as coherent text is harder to catch than corruption that manifests as garbage. When a model ignores specific details of a prompt while responding correctly to its general shape, suspect attention score corruption before anything else.

Buffer allocation size and active row count are different things. In a chunked pipeline with pre-allocated buffers, the temptation to use the allocation constant as a stride is always present and always wrong on partial chunks. Name them differently, document the distinction at the declaration site, and never let them meet in index arithmetic.

Mila is a solo CUDA-accelerated LLM inference framework targeting local inference. Built in C++23 with cuBLASLt and custom CUDA kernels.