Attention with Online Softmax

Problem

Compute one query’s scaled dot-product attention while processing keys in blocks of at most block_size. Return a Dv-wide weighted value vector; mask=false excludes a key, and all-masked input returns zeros.

Maintain a running maximum m, exponential sum l, and weighted-value sum a. When the maximum rises, rescale the previous sums before adding the new block. Do not materialize an S-wide score or probability vector: additional working storage must be O(block_size*(D+Dv) + Dv), apart from input arrays.

def online_attention(query, keys, values, block_size, mask):

Constraints

• query has D entries; keys S × D, values S × Dv, mask S booleans.
• 1 ≤ S,D,Dv ≤ 32; 1 ≤ block_size ≤ 32; all entries are finite.
• Inputs satisfy these constraints; invalid-input handling is not graded. Use only Python standard library and NumPy.

Examples

Arguments: [
  [
    1
  ],
  [
    [
      1
    ]
  ],
  [
    [
      2
    ]
  ],
  1,
  [
    true
  ]
]
Expected result: [
  2
]

Arguments: [
  [
    0
  ],
  [
    [
      1
    ],
    [
      2
    ],
    [
      3
    ]
  ],
  [
    [
      1
    ],
    [
      2
    ],
    [
      6
    ]
  ],
  2,
  [
    true,
    true,
    true
  ]
]
Expected result: [
  3
]