feat: support scaled dot-product attention

Author: RayZhao1998

cs336_basics/scaled_dot_product_attention.py

@@ -0,0 +1,21 @@
+import torch
+from jaxtyping import Float
+from cs336_basics.softmax import softmax
+
+def scaled_dot_product_attention(
+    Q: Float[torch.Tensor, "... queries d_k"],
+    K: Float[torch.Tensor, "... keys d_k"],
+    V: Float[torch.Tensor, "... values d_v"],
+    mask: Float[torch.Tensor, "... queries keys"] | None = None
+) -> Float[torch.Tensor, "... queries d_v"]:
+    d_k = Q.shape[-1]
+    sqrt_d_k: int = d_k ** 0.5
+    scores = Q @ K.transpose(-2, -1) / sqrt_d_k # (..., queries, keys)
+
+    if mask is not None:
+        scores = scores.masked_fill(~mask, float("-inf"))
+
+    attention = softmax(scores, dim=-1) # (..., queries, keys)
+    # (..., queries, keys) @ (..., values, d_v) -> (..., queries, d_v)
+    result = attention @ V 
+    return result

tests/adapters.py

@@ -17,6 +17,7 @@
 from cs336_basics.swiglu import SwiGLU
 from cs336_basics.rope import RoPE
 from cs336_basics.softmax import softmax
+from cs336_basics.scaled_dot_product_attention import scaled_dot_product_attention
 
 def run_linear(
     d_in: int,
@@ -116,7 +117,7 @@ def run_scaled_dot_product_attention(
     Returns:
         Float[Tensor, " ... queries d_v"]: Output of SDPA
     """
-    raise NotImplementedError
+    return scaled_dot_product_attention(Q, K, V, mask)
 
 
 def run_multihead_self_attention(