Amazon interview questions

Train byte-level BPE on a sample of the real mix of languages and code, so nothing is unknown. Vocabulary size trades shorter sequences against bigger embedding and output tables. Add digit and whitespace rules and chat tokens.

KL-regularised SFT adds β times the KL divergence from a frozen base model's next-token distributions to the cross-entropy loss. It limits drift and forgetting; β sets how far it may move.

import math

def sigmoid(x):
    if x >= 0:
        return 1 / (1 + math.exp(-x))
    e = math.exp(x)  # x is negative: no overflow
    return e / (1 + e)

def softmax(xs):
    m = max(xs)  # shift so the largest exp is 1
    es = [math.exp(x - m) for x in xs]
    s = sum(es)
    return [e / s for e in es]

def attention(Q, K, V):
    d = len(Q[0])
    out = []
    for q in Q:
        scores = []
        for k in K:
            dot = 0.0
            for a, b in zip(q, k):
                dot += a * b
            scores.append(dot / math.sqrt(d))
        w = softmax(scores)
        row = [0.0] * len(V[0])
        for wi, v in zip(w, V):
            for j in range(len(v)):
                row[j] += wi * v[j]
        out.append(row)
    return out

The trap is overflow: exp of a large number blows up. Branch sigmoid on the sign, and subtract the max before softmax, which leaves the result unchanged. Attention is then dot products, scaled by √d, a softmax per query, and a weighted sum of values.

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 decoder-only model runs prefill, then one decode step per token, and decode is memory-bound, so the win is batching many requests.

Continuous batching adds and removes requests at every step instead of waiting for a whole batch to finish. PagedAttention stores each request's KV cache in fixed-size blocks, like memory pages, so nothing is reserved for the maximum length and wasted: more requests fit, and shared prefixes can share blocks.

First find which cost hurts: attention compute grows quadratically with length and dominates long prefills, while the KV cache grows linearly and caps decode batch size.

Exact fixes first. FlashAttention tiles the work in on-chip memory, never storing the n × n score matrix: same result, memory linear in n. For the cache: grouped-query attention, quantisation, paged allocation.

Then approximations that change the model: sliding-window or sparse patterns like Longformer's local plus global attention, or linear-attention and state-space layers. Past one GPU, ring attention splits the sequence across devices. Re-check long-context quality.

Two loops. Offline: a versioned eval set scored by cheap deterministic checks first, like schema, required facts and banned content, then an LLM judge calibrated against human labels, gating every change. Online: the same checks on live outputs, sampled judge scores, and user signals such as retries, sliced by task. Failures flow back into the eval set, and humans review a sample so the judge doesn't drift.

Offline: parse documents, chunk on their structure, embed, and index chunks with metadata and permissions in a vector and keyword index, updated incrementally.

Online: rewrite the query, run hybrid search filtered by the user's access, rerank, and pack the best passages into a prompt that says to cite them and abstain when the answer isn't there.

Around it: an eval set scoring retrieval and faithfulness separately, logged passages, caching, and a per-stage latency budget.

GRPO drops PPO's value network. It samples a group of answers per prompt and uses each answer's reward minus the group mean, over the group standard deviation, as its advantage, inside PPO's clipped objective with a KL penalty to a reference model. No critic saves memory.

Rollout workers run data-parallel inference; the trainer shards with FSDP or tensor and pipeline parallelism, syncing weights each step. Rewards should be verifiable, like exact answer match, unit tests or a format check, not a learned reward model that's easier to game.

It's a loop: observe, decide, act, check. Perception pairs a screenshot with the accessibility tree when one exists; the tree gives element names and ids, which are far more reliable to target than raw coordinates. Without one, mark candidate elements on the image and let the model pick a mark.

Actions are typed tools: click, type, scroll, key press. After each, a fresh screenshot verifies the effect. Run in a sandboxed VM, treat on-screen text as data, not instructions, and ask the user before irreversible steps. Measure task success on benchmarks like OSWorld.

An MoE layer swaps one feed-forward block for many experts and a small router that sends each token to the top few, often 2.

Parameters grow, but each token uses a slice: Mixtral 8×7B has about 47B parameters and uses about 13B per token. Expert parallelism places experts on different GPUs and ships tokens to them in an all-to-all exchange, so the costs become communication and load balance.

Tensor parallelism splits each weight matrix across GPUs, so every layer needs a collective step: fast only over very fast links, usually inside one machine. Pipeline parallelism gives each GPU a run of consecutive layers and passes activations along: little communication, but GPUs sit idle in the pipeline bubble unless the batch is split into micro-batches. Large runs combine both with data parallelism.

All-reduce sums gradients across data-parallel replicas, and partial results in tensor parallelism. All-gather and reduce-scatter are its two halves: sharded training such as ZeRO or FSDP gathers each layer's weights just before use and reduce-scatters its gradients after. All-to-all routes tokens to experts in MoE. Each costs bytes over links, so a good layout puts the chattiest collectives on the fastest links.

A transformer stacks blocks of attention, which mixes tokens, and an MLP, each with a residual. MoE swaps the MLP for many experts and a router picking a few per token.

Training runs whole sequences in parallel under the causal mask; MoE adds a load-balancing loss so experts don't collapse. Inference prefills the prompt, then decodes token by token with a KV cache; every expert stays in memory though each token uses few.