Module utils 

Apply penalty and repeat_kv

Functions

apply_repeat_penalty

repeat_kv

Repeats a key or value tensor for grouped query attention The input tensor should have a shape (batch, num_kv_heads, seq_len, head_dim),