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

Medium

Attention Efficiency & Long Context3 / 5

Chapter 033 / 5

Attention Efficiency & 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

Implement the function/class skeleton 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

⌘/Ctrl↵to submit

AI review

No review yet

Submit your solution for a verdict, a score out of 100 and concrete feedback. Any correct approach passes.

Chapter 033 / 5

Attention Efficiency & 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

Implement the function/class skeleton in the editor. Any correct approach is accepted.

Hints0 / 2

Reference solution

Your own code stays in the editor.