Implement Attention in Python

NVIDIA TensorRT-LLM uses attention during transformer inference. Implement a reference version of scaled dot product attention that supports arbitrary query, key, and value sequences, an optional attention mask, and optional causal masking.

For each query, compute score = query · key / sqrt(d), where d is the key dimension. Apply masks before softmax, then return the weighted sum of value vectors and the attention-weight matrix.

Formal Specification

Implement attention(query, key, value, mask=None, causal=False).

• query is a q_len x d list of floats.
• key is a k_len x d list of floats.
• value is a k_len x value_dim list of floats.
• mask, when provided, is a q_len x k_len Boolean matrix. True means the key may be attended to.
• If causal is True, query position i may attend only to key positions j <= i.
• Return [output, weights], where output is q_len x value_dim and weights is q_len x k_len.
• If a query has no permitted keys, return an all-zero weight row and output row.

Use a numerically stable softmax by subtracting the maximum permitted score. Do not use NumPy or deep-learning libraries.