Grouped-Query Attention

Problem

Compute scaled dot-product causal attention with shared KV heads. queries has shape Hq × T × D; keys is Hkv × S × D and values Hkv × S × Dv. Query head h uses KV head floor(h/(Hq/Hkv)).

Queries represent the last T positions of S keys. Key j is allowed for query i only if key_mask[j] is true and j ≤ S−T+i. Use scale 1/sqrt(D). A row with no allowed keys returns Dv zeros. Return Hq × T × Dv.

def grouped_query_attention(queries, keys, values, key_mask):

Constraints

• Hq is a positive multiple of Hkv; 1 ≤ Hq,Hkv,T,S,D,Dv ≤ 8 and T ≤ S.
• key_mask is S booleans; all tensor entries are finite.
• Inputs satisfy these constraints; invalid-input handling is not graded. Use only Python standard library and NumPy.

Examples

Arguments: [
  [
    [
      [
        1
      ]
    ]
  ],
  [
    [
      [
        1
      ]
    ]
  ],
  [
    [
      [
        3
      ]
    ]
  ],
  [
    true
  ]
]
Expected result: [
  [
    [
      3
    ]
  ]
]

Arguments: [
  [
    [
      [
        0
      ]
    ],
    [
      [
        1
      ]
    ]
  ],
  [
    [
      [
        0
      ],
      [
        1
      ]
    ]
  ],
  [
    [
      [
        2
      ],
      [
        4
      ]
    ]
  ],
  [
    true,
    true
  ]
]
Expected result: [
  [
    [
      3
    ]
  ],
  [
    [
      3.4621171572600096
    ]
  ]
]