Implement Multi-Head Attention | Dataford Interview Questions

Implement Multi-Head Self-Attention (MHA) from scratch in PyTorch, ensuring your implementation supports batching and causal masking.

Accept x with shape [batch, sequence, model_dim] and projection matrices with shape [model_dim, model_dim]. Return the projected attention output as a nested list with the same shape as x; inputs are provided as JSON-compatible lists and must be converted to PyTorch tensors internally.

Use scaled dot-product attention, split the model dimension across num_heads, apply the upper-triangular future-token mask when causal is true, concatenate heads, and apply w_o. Assume model_dim is divisible by num_heads.

Constraints

x has shape [batch_size, sequence_length, model_dim]
Each projection matrix has shape [model_dim, model_dim]
1 <= batch_size <= 128
1 <= sequence_length <= 2048
1 <= num_heads <= model_dim
model_dim is divisible by num_heads
causal is a boolean

Function Signature

def mha(x, w_q, w_k, w_v, w_o, num_heads, causal):

Your question is Implement Multi-Head Attention. Start with the requirements in the Question tab.

Run and submit as often as you like. When you're ready, talk me through your approach or go straight to the code.

Problem

Implement Multi-Head Self-Attention (MHA) from scratch in PyTorch, ensuring your implementation supports batching and causal masking.

Constraints

x has shape [batch_size, sequence_length, model_dim]
Each projection matrix has shape [model_dim, model_dim]
1 <= batch_size <= 128
1 <= sequence_length <= 2048
1 <= num_heads <= model_dim
model_dim is divisible by num_heads
causal is a boolean

Function Signature

def mha(x, w_q, w_k, w_v, w_o, num_heads, causal):

Your question is Implement Multi-Head Attention. Start with the requirements in the Question tab.

Run and submit as often as you like. When you're ready, talk me through your approach or go straight to the code.