Interview prep — LLM coding drills

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

def attention(Q, K, V, causal=False, padding_mask=None):
    """
    Q, K, V: (B, h, T, d_h)
    padding_mask: (B, T) bool, True for valid keys.
    """
    d_h = Q.size(-1)
    scores = (Q @ K.transpose(-2, -1)) / math.sqrt(d_h)            # (B, h, T_q, T_k)

    if causal:
        T_q, T_k = scores.size(-2), scores.size(-1)
        cm = torch.ones(T_q, T_k, dtype=torch.bool, device=scores.device).tril()
        scores = scores.masked_fill(~cm, float("-inf"))

    if padding_mask is not None:
        scores = scores.masked_fill(~padding_mask[:, None, None, :], float("-inf"))

    weights = scores.softmax(dim=-1)
    return weights @ V, weights

class MHA(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        assert d % h == 0
        self.h, self.dh = h, d // h
        self.qkv = nn.Linear(d, 3 * d, bias=False)
        self.o   = nn.Linear(d, d, bias=False)

    def forward(self, x, kv_cache=None):
        B, T, d = x.shape
        qkv = self.qkv(x).view(B, T, 3, self.h, self.dh).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]               # (B, h, T, dh)

        if kv_cache is not None:
            k = torch.cat([kv_cache["k"], k], dim=2)
            v = torch.cat([kv_cache["v"], v], dim=2)
            kv_cache["k"], kv_cache["v"] = k, v

        scores = q @ k.transpose(-2, -1) / (self.dh ** 0.5)
        T_q, T_k = scores.size(-2), scores.size(-1)
        cm = torch.ones(T_q, T_k, dtype=torch.bool, device=x.device).tril(T_k - T_q)
        scores = scores.masked_fill(~cm, float("-inf"))
        out = scores.softmax(-1) @ v                   # (B, h, T, dh)
        return self.o(out.transpose(1, 2).reshape(B, T, d))

def precompute_rope(d_head, max_T, theta=10000.0, device="cpu"):
    inv = 1.0 / (theta ** (torch.arange(0, d_head, 2, device=device).float() / d_head))
    t = torch.arange(max_T, device=device).float()
    freqs = torch.outer(t, inv)                   # (T, d_head/2)
    return freqs.cos(), freqs.sin()

def apply_rope(x, cos, sin):
    # x: (..., T, d_head), with last dim split into pairs
    x1, x2 = x[..., 0::2], x[..., 1::2]
    rotated = torch.stack([x1 * cos - x2 * sin,
                           x1 * sin + x2 * cos], dim=-1)
    return rotated.flatten(-2)

import torch
import torch.nn.functional as F

def top_p_sample(logits, p=0.95, T=1.0):
    """logits: (V,)"""
    logits = logits / max(T, 1e-6)
    sorted_logits, idx = torch.sort(logits, descending=True)
    probs = F.softmax(sorted_logits, dim=-1)
    cum = probs.cumsum(0)
    cutoff = (cum > p).nonzero()[0].item() + 1
    sorted_logits[cutoff:] = -float("inf")
    final = torch.full_like(logits, -float("inf"))
    final.scatter_(0, idx, sorted_logits)
    return torch.multinomial(F.softmax(final, -1), 1).item()

from collections import Counter

def get_pairs(corpus):
    pairs = Counter()
    for word in corpus:                           # word is a list of token ids
        for a, b in zip(word, word[1:]):
            pairs[(a, b)] += 1
    return pairs

def merge(corpus, pair, new_id):
    out = []
    a, b = pair
    for word in corpus:
        new = []
        i = 0
        while i < len(word):
            if i < len(word)-1 and word[i] == a and word[i+1] == b:
                new.append(new_id); i += 2
            else:
                new.append(word[i]); i += 1
        out.append(new)
    return out

def train_bpe(corpus_bytes, vocab_size=512):
    # corpus_bytes: list[list[int]] of byte ids per pre-token
    next_id = 256
    merges = []
    while next_id < vocab_size:
        pairs = get_pairs(corpus_bytes)
        if not pairs: break
        pair, _ = pairs.most_common(1)[0]
        corpus_bytes = merge(corpus_bytes, pair, next_id)
        merges.append((pair, next_id))
        next_id += 1
    return merges

@torch.no_grad()
def speculative_step(target, draft, prompt, k=4, T=1.0):
    cur = prompt
    draft_ids, draft_logits = [], []
    for _ in range(k):
        out = draft(cur)
        l = out[:, -1] / T
        nxt = torch.multinomial(F.softmax(l, -1), 1)
        draft_ids.append(int(nxt))
        draft_logits.append(l)
        cur = torch.cat([cur, nxt], 1)

    # one big target forward over prompt + drafts
    full_logits = target(cur)[:, -k-1:] / T
    accepted = []
    for i in range(k):
        p_t = F.softmax(full_logits[:, i], -1)[0, draft_ids[i]]
        p_d = F.softmax(draft_logits[i], -1)[0, draft_ids[i]]
        if torch.rand(1).item() < min(1.0, (p_t / max(p_d, 1e-9)).item()):
            accepted.append(draft_ids[i])
        else:
            diff = F.softmax(full_logits[:, i], -1) - F.softmax(draft_logits[i], -1)
            diff = torch.clamp(diff, min=0); diff = diff / diff.sum()
            accepted.append(int(torch.multinomial(diff, 1)))
            return prompt[0].tolist() + accepted
    bonus = int(torch.multinomial(F.softmax(full_logits[:, -1], -1), 1))
    accepted.append(bonus)
    return prompt[0].tolist() + accepted

def dpo_loss(logp_pi_chosen, logp_pi_rejected,
             logp_ref_chosen, logp_ref_rejected, beta=0.1):
    pi_logratios  = logp_pi_chosen - logp_pi_rejected
    ref_logratios = logp_ref_chosen - logp_ref_rejected
    logits = beta * (pi_logratios - ref_logratios)
    return -F.logsigmoid(logits).mean()

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
import asyncio, json

app = FastAPI()

@app.post("/v1/completions")
async def completions(req: dict):
    async def gen():
        async for tok in stream_tokens(req["prompt"], req.get("max_tokens", 200)):
            payload = {"choices":[{"delta":{"content": tok}}]}
            yield f"data: {json.dumps(payload)}\n\n"
        yield "data: [DONE]\n\n"
    return StreamingResponse(gen(), media_type="text/event-stream")

class Scheduler:
    def __init__(self, engine, max_batch=32):
        self.engine = engine
        self.max_batch = max_batch
        self.pending: list[Request] = []
        self.running: list[Request] = []

    def add(self, req): self.pending.append(req)

    def step(self):
        # 1. promote pending into running while batch has slots
        while self.pending and len(self.running) < self.max_batch:
            r = self.pending.pop(0)
            self.engine.prefill(r)              # fills KV cache
            self.running.append(r)
        # 2. one forward step
        next_tokens = self.engine.decode_step(self.running)
        finished = []
        for r, tok in zip(self.running, next_tokens):
            r.append(tok)
            if r.is_done(): finished.append(r)
        # 3. clean up
        for r in finished:
            self.running.remove(r)
            r.set_result()

LLAMA3 = (
    "<|begin_of_text|>"
    "{turns}"
    "<|start_header_id|>assistant<|end_header_id|>\n\n"   # if add_generation_prompt
)
TURN = "<|start_header_id|>{role}<|end_header_id|>\n\n{content}<|eot_id|>"

def apply_llama3_chat_template(messages, add_generation_prompt=True):
    turns = "".join(TURN.format(**m) for m in messages)
    if add_generation_prompt:
        return "<|begin_of_text|>" + turns + "<|start_header_id|>assistant<|end_header_id|>\n\n"
    return "<|begin_of_text|>" + turns