fn attention_f16(q: Tensor<f32, [D]>,
                 key_cache: Tensor<f16, [L, SeqLen, KvDim]>,
                 value_cache: Tensor<f16, [L, SeqLen, KvDim]>,
                 layer: i32, pos: i32, H: i32, scale: f32) -> Tensor<f32, [D]>
    with tile=[8, 64], parallel=[h, t] {

  var att: Tensor<f32, [H, SeqLen]> = zeros([H, SeqLen])

  # Compute attention scores
  att[h, t] = if t > pos { -inf } else {
    sum(i) q[h * Dh + i] * (key_cache[layer, t, h * Dh + i] as f32) * scale
  }

  var weights: Tensor<f32, [H, SeqLen]> = softmax(att)
  # ... weighted sum over values
}