CoreWeave inference workloads often need customized attention behavior that cannot rely on a framework primitive. Implement single-head scaled dot-product attention using only Python lists and the standard library.
Given query matrix Q, key matrix K, value matrix V, an optional boolean mask, and a causal flag, compute:
Attention(Q, K, V) = softmax((Q × Kᵀ) / sqrt(d)) × V
where d is the key dimension. A mask value of true means the key may be attended to, while false excludes it. If causal is true, position i may attend only to key positions j <= i. Apply both restrictions when both are provided. Every query row has at least one permitted key.
Implement attention(q, k, v, mask, causal). q is an n x d list of numbers, k is an m x d list, and v is an m x dv list. mask is either None or an n x m boolean matrix. causal is a boolean. Return an n x dv list of floating-point numbers. Do not use NumPy or machine-learning libraries. Use a numerically stable softmax by subtracting the largest permitted score before exponentiation.
q has shape n x d, k has shape m x d, and v has shape m x dvmask is None or has shape n x mdef attention(q, k, v, mask, causal):