Cached Transformer Decoder Block

Problem

Implement this precisely specified single-head pre-norm decoder block. First RMS-normalize each row of x with unit gain and eps, then project Q, K, V using weights["q"], ["k"], ["v"]. Append new K/V rows to past_keys/past_values. Use scaled causal attention with cached position offsets, project through weights["o"], and add x.

RMS-normalize this residual, compute ReLU(norm @ weights["up"]) @ weights["down"], and add the residual again. Return {"output": final rows, "keys": appended keys, "values": appended values}. This exercise has no bias, positional embeddings, dropout, or learned normalization gain; its feed-forward activation is ReLU.

def cached_decoder_block(x, past_keys, past_values, weights, eps):

Constraints

• x is T × D; cached K/V are P × D or [] when P=0; 1 ≤ T,D ≤ 8, 0 ≤ P ≤ 8.
• q/k/v/o weights are D × D, up is D × F, down is F × D; 1 ≤ F ≤ 8.
• eps ≥ 0 and both normalization denominators are positive.
• Inputs satisfy these constraints; invalid-input handling is not graded. Use only Python standard library and NumPy.

Examples

Arguments: [
  [
    [
      1
    ]
  ],
  [],
  [],
  {
    "q": [
      [
        1
      ]
    ],
    "k": [
      [
        1
      ]
    ],
    "v": [
      [
        1
      ]
    ],
    "o": [
      [
        1
      ]
    ],
    "up": [
      [
        1
      ]
    ],
    "down": [
      [
        1
      ]
    ]
  },
  0
]
Expected result: {
  "output": [
    [
      3
    ]
  ],
  "keys": [
    [
      1
    ]
  ],
  "values": [
    [
      1
    ]
  ]
}

Arguments: [
  [
    [
      1
    ]
  ],
  [
    [
      0
    ]
  ],
  [
    [
      4
    ]
  ],
  {
    "q": [
      [
        1
      ]
    ],
    "k": [
      [
        1
      ]
    ],
    "v": [
      [
        1
      ]
    ],
    "o": [
      [
        1
      ]
    ],
    "up": [
      [
        1
      ]
    ],
    "down": [
      [
        1
      ]
    ]
  },
  0
]
Expected result: {
  "output": [
    [
      3.806824264109985
    ]
  ],
  "keys": [
    [
      0
    ],
    [
      1
    ]
  ],
  "values": [
    [
      4
    ],
    [
      1
    ]
  ]
}