Implement grouped-query attention (GQA) in PyTorch by repeating/broadcasting KV heads acro

Medium

Attention Efficiency and Long Context3 / 5

Chapter 033 / 5

Attention Efficiency and Long Context

Implement grouped-query attention (GQA) in PyTorch by repeating/broadcasting KV heads acro

Medium

Implement grouped-query attention (GQA) in PyTorch by repeating/broadcasting KV heads across query-head groups, and verify it matches full multi-head attention when num_kv_heads equals num_query_heads.

Target interface

class GQA

Complete the starter code in the editor. Any correct approach is accepted.

Hints0 / 2

Reference solution

Your own code stays in the editor.

solution.pyPython

Local draft

import torch
import torch.nn as nn
import torch.nn.functional as F

class GQA(nn.Module):

    def __init__(self, d_model, n_heads, n_kv_heads):
        raise NotImplementedError

    def forward(self, x):
        raise NotImplementedError

def mha(x):
    raise NotImplementedError

Local draft

⌘/Ctrl↵to submit

AI review

Ready when you are

Submit your solution for a score and feedback on correctness, edge cases, and clarity. Any correct approach is accepted.

Chapter 033 / 5

Attention Efficiency and Long Context

Implement grouped-query attention (GQA) in PyTorch by repeating/broadcasting KV heads acro

Medium

Target interface

class GQA

Complete the starter code in the editor. Any correct approach is accepted.

Hints0 / 2

Reference solution

Your own code stays in the editor.