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, 1001 positions, 64 d_h] v [B batch, 12 heads, 1001 positions, 64 d_h]
This call computes the new token's K and V and concatenates them after the 1,000 cached positions, so attention reads 1,001. 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. In use after this call: 70.4 MiB per sequence across all 12 layers, at float32.
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)