Your question is Custom Attention Mechanism. Start with the requirements on the right.
Run and submit as often as you like. When you're ready, talk me through your approach or go straight to the code.
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):