Datadog interview questions

Write it once as GQA: plain MHA is the case where key-value heads equal query heads, and multi-query is one key-value head.

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

class GQA(nn.Module):
    def __init__(self, d, n_q, n_kv):
        super().__init__()
        self.hd, self.nq, self.nkv = d // n_q, n_q, n_kv
        self.q = nn.Linear(d, d, bias=False)
        self.kv = nn.Linear(d, 2 * n_kv * self.hd,
                            bias=False)
        self.o = nn.Linear(d, d, bias=False)

    def forward(self, x, cache=None):
        B, T, _ = x.shape
        q = self.q(x).view(B, T, self.nq, self.hd)
        k, v = self.kv(x).view(
            B, T, 2, self.nkv, self.hd).unbind(2)
        q, k, v = (t.transpose(1, 2) for t in (q, k, v))
        if cache is not None:
            k = torch.cat([cache[0], k], dim=2)
            v = torch.cat([cache[1], v], dim=2)
        cache = (k, v)  # n_kv heads: the memory saving
        r = self.nq // self.nkv
        k = k.repeat_interleave(r, dim=1)
        v = v.repeat_interleave(r, dim=1)
        S = k.shape[2]  # cached + new positions
        # new query i sits at position S - T + i
        m = torch.ones(T, S, dtype=torch.bool,
                       device=x.device).tril(S - T)
        a = q @ k.transpose(-2, -1) / self.hd ** 0.5
        a = a.masked_fill(~m, float("-inf"))
        y = F.softmax(a, dim=-1) @ v
        y = y.transpose(1, 2).reshape(B, T, -1)
        return self.o(y), cache

Interviewers check three things: the mask is offset by the cached length, the cache stores keys and values before repeating them, so it shrinks by the group factor, and cached decoding reproduces a full forward pass token for token.

A pipeline with an agent in the middle. Triage first: take only tickets with a clear reproduction and a scoped component. The agent gathers context with tools: code search, file reads, git history, and tests run in a sandbox without production credentials. It writes a failing test, edits code and loops until tests pass, under step and cost limits.

The output is a pull request with a summary and links, always reviewed by a human.

Focal loss scales cross-entropy by (1 − p_t) to the power gamma, so well-classified examples add little and training focuses on hard ones; alpha rebalances the classes. Work from logits throughout.

import torch

def focal_loss(z, y, alpha=0.25, gamma=2.0):
    # z: logits, y: float 0/1 targets, same shape
    # stable BCE: max(z,0) - z*y + log(1+exp(-|z|))
    ce = (z.clamp(min=0) - z * y
          + torch.log1p(torch.exp(-z.abs())))
    p_t = torch.exp(-ce)
    a_t = alpha * y + (1 - alpha) * (1 - y)
    return (a_t * (1 - p_t) ** gamma * ce).mean()

Interviewers check you never take the log of a sigmoid output, that p_t comes from exp of minus the loss, and that gamma 0 recovers alpha-weighted cross-entropy.