Transformer block › Attention
KV cache
×state · kept between callspast = tf.concat([past, next_outputs['presents']], axis=-2)
The keys and values of every position so far, kept between forward calls as each layer’s past.
- shape
- k [B batch, 12 heads, T positions, 64 d_h] v [B batch, 12 heads, T positions, 64 d_h]
Prefill computes K and V for the prompt and keeps them as this layer’s present. Nothing is allocated ahead: the cache grows one position per decoded token, up to the 1,024 learned positions, and the bars on the plate fill against those positions (sample.py 62–72). It holds every head (12); GPT-2 has one key/value head per query head.
Codemodel.py 58–66, 102–108 · sample.py 62–72
c = conv1d(x, 'c_attn', n_state*3)
q, k, v = map(split_heads, tf.split(c, 3, axis=2))
present = tf.stack([k, v], axis=1)
if past is not None:
pk, pv = tf.unstack(past, axis=1)
k = tf.concat([pk, k], axis=-2)
v = tf.concat([pv, v], axis=-2) def body(past, prev, output):
next_outputs = step(hparams, prev, past=past)
logits = next_outputs['logits'][:, -1, :] / tf.to_float(temperature)
logits = top_k_logits(logits, k=top_k)
logits = top_p_logits(logits, p=top_p)
samples = tf.multinomial(logits, num_samples=1, output_dtype=tf.int32)
return [
next_outputs['presents'] if past is None else tf.concat([past, next_outputs['presents']], axis=-2),
samples,
tf.concat([output, samples], axis=1)
]def attention_mask(nd, ns, *, dtype):
"""1's in the lower triangle, counting from the lower right corner.
Same as tf.matrix_band_part(tf.ones([nd, ns]), -1, ns-nd), but doesn't produce garbage on TPUs.
"""
i = tf.range(nd)[:,None]
j = tf.range(ns)
m = i >= j - ns + nd
return tf.cast(m, dtype)Sources3
Contextual evidence
Each call stacks only its own K and V into present (104); with past, the cached K and V are concatenated before the new ones for this call’s attention (105–108).
present · openai/gpt-2 / src/model.py
Pinned revision 9b63575ef427 ↗ · lines 102–108
Read code excerpt
c = conv1d(x, 'c_attn', n_state*3)
q, k, v = map(split_heads, tf.split(c, 3, axis=2))
present = tf.stack([k, v], axis=1)
if past is not None:
pk, pv = tf.unstack(past, axis=1)
k = tf.concat([pk, k], axis=-2)
v = tf.concat([pv, v], axis=-2)Contextual evidence
The sampling loop keeps past between calls: each step concatenates the call’s presents after it along the position axis (69); nothing is allocated ahead.
loop · openai/gpt-2 / src/sample.py
Pinned revision 9b63575ef427 ↗ · lines 62–72
Read code excerpt
def body(past, prev, output):
next_outputs = step(hparams, prev, past=past)
logits = next_outputs['logits'][:, -1, :] / tf.to_float(temperature)
logits = top_k_logits(logits, k=top_k)
logits = top_p_logits(logits, p=top_p)
samples = tf.multinomial(logits, num_samples=1, output_dtype=tf.int32)
return [
next_outputs['presents'] if past is None else tf.concat([past, next_outputs['presents']], axis=-2),
samples,
tf.concat([output, samples], axis=1)
]Contextual evidence
attention_mask(nd, ns) keeps the lower triangle counted from the lower right; with one query (nd = 1) every position is kept.
mask · openai/gpt-2 / src/model.py
Pinned revision 9b63575ef427 ↗ · lines 58–66
Read code excerpt
def attention_mask(nd, ns, *, dtype):
"""1's in the lower triangle, counting from the lower right corner.
Same as tf.matrix_band_part(tf.ones([nd, ns]), -1, ns-nd), but doesn't produce garbage on TPUs.
"""
i = tf.range(nd)[:,None]
j = tf.range(ns)
m = i >= j - ns + nd
return tf.cast(m, dtype)