Decoding Cache

Task

This is a small, working decoder written with plain Python lists: two layers of causal self-attention with fixed weights. Generation is slow because every step runs the whole sequence through every layer again. Add a key and value cache, switched on with use_cache, so each step runs only the new token. The tokens and the logits must be the same as without the cache. Every result reports rows, the number of attention rows the model computed, so the saving can be checked.

Requirements

1. decode(request) takes one request object and returns a small object. tokens is a list of 1 to 8 integers from 0 to 5 (ValueError containing tokens otherwise). Any other op raises ValueError containing op. Do not change the arguments.
2. {op: "logits", tokens} returns {logits, rows}: the six scores for the token that follows, rounded to 6 decimals, and the attention rows computed.
3. {op: "generate", tokens, steps, use_cache} returns {tokens, logits, rows}: the input followed by steps greedy picks, the six logits behind each pick, and the attention rows computed. steps is a non-negative integer (ValueError containing steps otherwise). use_cache is true or false and defaults to false (ValueError containing use_cache otherwise). tokens plus steps must fit in 8 (ValueError containing length otherwise).
4. With use_cache true, tokens and logits are exactly the same as with use_cache false.
5. rows counts calls to attend(): one per layer for each position that is run. Without the cache, every step runs every position again. With the cache, every token is run through each layer once: 2 * (len(tokens) + steps - 1) rows when steps is at least 1, and 0 rows when steps is 0.
6. The weight tables, attend() and its row count, and the rounding must stay as they are.