Implement Grouped Query Attention Forward Pass

Problem Statement You are given three tensors—Query, Key, and Value—and need to implement the forward pass of Grouped Query Attention without any masking. The tensor shapes are: Here: B is the batch size. Hq is the number of query heads. Hkv is the number of key/value heads. Hq is always divisible by Hkv. Each group of Hq / Hkv query heads shares a single key/value head. For a given query head h, the key/value head it maps to is: For every batch element, query head, and…