手写题和项目故事

v0 到 v7 你都能讲清楚，也能对到 Megatron 的代码。NVIDIA 的面试还要现场交付两样东西，限时手写代码，以及用数字讲清你在项目里做了什么。这一课先练五道手写题（按 NVIDIA 面经里出现的次数排），再把你的 TikTok SFT 整理成一个 infra 故事。

面试会这样问Implement a decoder KV cache. · Implement MoE and show the dimensions. · What were your exact contributions?

考过NVIDIA MLE 电面 2026（ML coding 实现 KV cache，1p3a 1184818）· NVIDIA 一亩三分地题库 OJ "Implement a Decoder KV Cache" · NVIDIA SWE（"Implement MoE and show the dimensions"，Glassdoor RVW105422785，年份未核实）· NVIDIA 框架岗 2023（电话面写 Column / Row Parallel Linear 的 pseudocode；take-home 和 onsite 从零实现 gradient，知乎）· NVIDIA 一亩三分地题库 hot 题（"FP32 Tensor to Int8 Quantization"，汇总帖 1190285）

手写题的节奏

1. 先复述题目，问清输入输出的 shape 和 dtype
2. 写最短的能跑的版本，边写边说
3. 用能手算的小例子或笨办法的参考实现对一遍
4. 再答追问

下面的代码都是 NumPy，完整练习放进 ml-coding 里写。

手写 1 · causal MHA + KV cache

原题：“Implement a Decoder KV Cache”。写完要能证明逐 token decode 的输出和整段重算完全一样。

关键步骤：只为新 token 算 q、k、v → 把 k、v 接到 cache 后面 → 新的 q 对 cache 里全部的 K、V 做 attention → causal mask 只在一次进多个 token（prefill）时起作用。

def attn(q, K, V, mask=None):              # q (h,n,dk)  K,V (h,t,dk)
    s = q @ K.transpose(0, 2, 1) / np.sqrt(q.shape[-1])
    if mask is not None:
        s = np.where(mask, -1e9, s)        # -1e9 比 -inf 安全
    p = np.exp(s - s.max(-1, keepdims=True))
    return (p / p.sum(-1, keepdims=True)) @ V

class CachedMHA:
    def __init__(self, Wq, Wk, Wv, Wo, h):
        self.W, self.h, self.K, self.V = (Wq, Wk, Wv, Wo), h, None, None
    def split(self, X):                    # (n, d) -> (h, n, dk)
        n, d = X.shape
        return X.reshape(n, self.h, d // self.h).transpose(1, 0, 2)
    def __call__(self, X):                 # X (n, d)：整段 prompt 或 1 个新 token
        Wq, Wk, Wv, Wo = self.W
        q, k, v = (self.split(X @ W) for W in (Wq, Wk, Wv))
        self.K = k if self.K is None else np.concatenate([self.K, k], 1)
        self.V = v if self.V is None else np.concatenate([self.V, v], 1)
        n, t = X.shape[0], self.K.shape[1]
        mask = np.triu(np.ones((n, t), bool), k=t - n + 1)  # 不看未来
        o = attn(q, self.K, self.V, mask)  # (h, n, dk)
        return o.transpose(1, 0, 2).reshape(n, -1) @ Wo

怎么验证：同一组 weight，一次喂整段，和先 prefill 再一个个 decode，两种输出用 np.allclose 比。

• 每个 token 占多少？2 × L × n_kv × d_head × 2 bytes。Qwen3-8B：2 × 36 × 8 × 128 × 2 = 144 KiB，8K tokens 约 1.1 GiB
• np.concatenate 的问题？每步拷贝整个 cache，O(t)。实际做法是预分配 max_len，或者像 vLLM 按 block 分页
• 为什么不缓存 Q、GQA 怎么改、RoPE 存什么，见 3.11

手写 2 · MoE forward，写出每一步的 shape

原题：NVIDIA SWE “Implement MoE and show the dimensions”，同一场还追问 kernel fusion 和 launch overhead。

关键步骤（数字和第 7 课一样，expert 中间维 f = 12288）：x [T, h] → router logits [T, E] → top-k 的下标和 gate [T, k] → 每个 expert 最多收 C = cf·kT/E = 1280 个 token，[≤C, h] → [≤C, f] → [≤C, h] → 按 gate 加权加回 y [T, h]。

def softmax(z):
    e = np.exp(z - z.max(-1, keepdims=True)); return e / e.sum(-1, keepdims=True)

def moe(x, Wr, W1, W2, k=2, cf=1.25):
    # x [T,h]  Wr [h,E]  W1 [E,h,f]  W2 [E,f,h]
    T, E = x.shape[0], Wr.shape[1]
    p = softmax(x @ Wr)                          # router probs  [T,E]
    idx = np.argsort(-p, axis=1)[:, :k]          # top-k experts [T,k]
    w = np.take_along_axis(p, idx, axis=1)       # gate values   [T,k]
    w = w / w.sum(1, keepdims=True)
    C = int(cf * k * T / E)                      # capacity per expert
    y = np.zeros_like(x)                         # output        [T,h]
    for e in range(E):
        t, j = np.nonzero(idx == e)              # tokens sent to e
        t, j = t[:C], j[:C]                      # overflow dropped
        hid = np.maximum(x[t] @ W1[e], 0)        # [n,h] -> [n,f]
        y[t] += w[t, j, None] * (hid @ W2[e])    # [n,f] -> [n,h], combine
    f = np.bincount(idx[:, 0], minlength=E) / T  # top-1 fraction [E]
    aux = E * np.sum(f * p.mean(0))              # Switch aux loss
    return y, aux

怎么验证：写一个逐 token 循环的笨版本（按 token 顺序占 capacity），两者输出对比。

• 为什么要 aux loss？防 expert collapse。上面的 aux 均匀分配时等于 1，全挤到一个 expert 时等于 E，实际再乘系数 α
• kernel 太碎、launch overhead 大？（NVIDIA 原追问）按 expert 循环要发 E 个小 GEMM。先按 expert 排序 token（permute），用一个 grouped GEMM 算完，再排回去（unpermute），即 Megatron 的 --moe-grouped-gemm、--moe-permute-fusion

手写 3 · Column / Row Parallel Linear

原题：“Column Parallel和Row Parallel Linear Layer的实现和pros & cons”，电话面写少量 pseudocode。

关键步骤：MLP \(Z = \mathrm{ReLU}(XA)\,B\)。A 按列切（column），B 按行切（row）。forward 只在最后 all-reduce 一次（g），backward 只有 dX 要 all-reduce 一次（f），dA、dB 都在本卡算完。

# NumPy 模拟 t = 2；dZ 是从后面传回来的 gradient
A1, A2 = np.split(A, 2, axis=1)          # column：每卡 [h, 4h/2]
B1, B2 = np.split(B, 2, axis=0)          # row：   每卡 [4h/2, h]
Y = [relu(X @ Ai) for Ai in (A1, A2)]    # 每卡本地算，不通信
Z = Y[0] @ B1 + Y[1] @ B2                # g：forward all-reduce
dY = [dZ @ Bi.T * (Yi > 0) for Bi, Yi in zip((B1, B2), Y)]
dB = [Yi.T @ dZ for Yi in Y]             # 本地，不通信
dA = [X.T @ dYi for dYi in dY]           # 本地，不通信
dX = dY[0] @ A1.T + dY[1] @ A2.T         # f：backward all-reduce

# PyTorch 伪代码（省略 @staticmethod），对应 Megatron 的 mappings.py
class f(torch.autograd.Function):        # _CopyToModelParallelRegion
    def forward(ctx, x):  return x
    def backward(ctx, g): return all_reduce(g)
class g(torch.autograd.Function):        # _ReduceFromModelParallelRegion
    def forward(ctx, x):  return all_reduce(x)
    def backward(ctx, g): return g
# 每个 rank：Z = g.apply(relu(f.apply(X) @ A_i) @ B_i)

怎么验证：Z、dX 和单卡结果一样；两卡的 dA 按列拼、dB 按行拼，也和单卡一样。

• pros 和 cons？pros：weight、optimizer state、activation 都变成 1/t；cons：all-reduce 在 critical path 上，只能放在 NVLink 内，t 大了 GEMM 变小、效率下降
• 为什么先 column 后 row、bias 怎么切、多卡怎么测见 6.1，attention 的切法在第 9 课

手写 4 · FP32 → INT8 quantization

原题：实现 quantize_fp32_to_int8(x, scale)，任意 shape 的 FP32 tensor 每个值除以 scale、round、限制到 [−128, 127]，返回同 shape 的 int8 tensor。之后有 5 个 follow-up。

关键步骤：\(q = \mathrm{clamp}(\mathrm{round}(x/s) + z,\ -128,\ 127)\)，还原 \(\hat x = s\,(q - z)\)。先 clamp 再 cast。全 0 的 tensor 要防除以 0。注意 np.round 是 round half to even（np.round(2.5) 是 2），题目要“四舍五入”时先问清楚。

def quantize_fp32_to_int8(x, scale, zero_point=0, qmin=-128, qmax=127):
    q = np.round(x / scale) + zero_point           # 任意 shape，逐元素
    return np.clip(q, qmin, qmax).astype(np.int8)  # 先 clamp 再 cast

def dequantize(q, scale, zero_point=0):
    return scale * (q.astype(np.float32) - zero_point)

def qparams(x, symmetric=True):
    if symmetric:                                  # z = 0，配 qmin=-127
        return max(np.abs(x).max(), 1e-8) / 127, 0
    lo, hi = min(x.min(), 0.0), max(x.max(), 0.0)  # 区间要包含 0
    scale = max(hi - lo, 1e-8) / 255
    return scale, int(np.round(-lo / scale)) - 128

自检：quantize_fp32_to_int8(np.array([-1., 0., 1.]), 0.1) → [−10, 0, 10]。x = [−1, 0, 0.5, 3] 用 asymmetric：s = 4/255 ≈ 0.0157，z = −64，q = [−128, −64, −32, 127]。

1. 加 zero_point 做 asymmetric：s = (max − min)/255，z = round(−min/s) − 128
2. symmetric 什么时候不够：分布偏向一边时（ReLU 之后、softmax 输出）一半格子用不上。上面的 x 用 symmetric 最大误差 0.0079，asymmetric 是 0.0039
3. scale 和 zero_point 怎么配合：s 定一格多宽，z 定实数 0 落在哪个整数
4. 超出 [−128, 127]：clamp 到边界（saturation），scale = 0.1 时 20 → 127 → 还原成 12.7
5. skewed 或有 outlier：max 被 outlier 撑大，小值全变 0，[0.1, 0.2, 0.3, 100] → [0, 0, 0, 127]。办法有 per-channel scale、percentile calibration、LLM.int8()、SmoothQuant（6.16）

手写 5 · mini autograd

原题：take-home 在一个能 trace graph 的框架里从零实现 gradient（weight、bias、activation、softmax），onsite 又在 colab 写 backprop，追问 gradient accumulation。

关键步骤：每个节点存 data、grad、parents 和一个 backward closure → backward 时先用 DFS 排出 topological order → loss 的 grad 设为 1 → 逆序调用每个 closure → gradient 一律用 +=。

class Value:
    def __init__(self, data, parents=()):
        self.data, self.grad, self.parents = data, 0.0, parents
        self._backward = lambda: None      # 把 out.grad 传给 parents
    def __add__(self, o):
        out = Value(self.data + o.data, (self, o))
        def bw(): self.grad += out.grad; o.grad += out.grad
        out._backward = bw; return out
    def __mul__(self, o):
        out = Value(self.data * o.data, (self, o))
        def bw(): self.grad += o.data * out.grad; o.grad += self.data * out.grad
        out._backward = bw; return out
    def backward(self):
        topo, seen = [], set()
        def build(v):                      # DFS 后序 = topological order
            if v not in seen:
                seen.add(v)
                for p in v.parents: build(p)
                topo.append(v)
        build(self)
        self.grad = 1.0                    # dL/dL = 1
        for v in reversed(topo): v._backward()   # 从 loss 往输入走

自检：L = xy + x，x = 2、y = 3 → x.grad = 4，y.grad = 2；再用 central difference 对一遍。

• 为什么用 +=？一个变量用了几次，gradient 就是几条路径之和。c * c（c = 3）用 += 得 6，写成 = 只剩 3
• 为什么要 topological order？保证一个节点收齐所有下游的 gradient 才往上传，每个 closure 只跑一次
• gradient accumulation？（onsite 原追问）k 个 micro-batch 之间不清零，loss 除以 k，step 之后再 zero_grad
• 扩展到 tensor 的 backward 见 6.19

项目故事：TikTok 的 32 卡 SFT

NVIDIA 会问 “What were your exact contributions?”，然后顺着数字往下追。已有的事实都填进了下面的模板（42.8% = 136 ÷ 318），空着的要你去查。

填空模板（英文，面试时直接说）：

1. Context：At TikTok I fine-tuned Qwen3-VL at 2B, 4B and 8B for ____, on an account-level dataset of 74K accounts and 1.33M videos.
2. Setup：Full-parameter SFT on 32 H100s, 4 nodes × 8, with DeepSpeed ZeRO-__, bf16 ____, gradient checkpointing ____, ____ frames or ____ tokens per sample.
3. Batch and speed：micro-batch __ × accumulation __ × 32 GPUs = global batch __; a step took __ s, about __ tokens/s per GPU, peak memory __ GB, so MFU was roughly __%.
4. Problem：The hardest issue was ____. I found it by ____, fixed it with ____, and it went from __ to __.
5. My part：I owned ____; the team owned ____.
6. Result：I evaluated 53 checkpoints across domains and calibrated thresholds; the final 4B model cut false positives by 42.8%, from 318 to 182, versus the 8B baseline at about 70% recall.

先去查的数

提前算好三笔账

• 为什么非用 ZeRO 不可：8B × 16 bytes = 128 GB，超过 80 GB，DDP 放不下
• 你那一级每卡多少（N = 32，只算 model state）：

实测峰值减去这个数，剩下的主要是 activation 和临时 buffer。按第 1 课的估算，一条 8K tokens 的样本（假设的数）开 FlashAttention、不开 checkpointing 约 41 GB，比 ZeRO-2 的 19.5 GB 还大，这就是 gradient checkpointing 的理由。

• 估一个 MFU：

参照点：Ψ = 8B 时每个 token 6Ψ = 48B FLOPs，MFU 40% 等于每卡约 8,200 tokens/s。粗估只算 LLM 部分，token 数用进入 LLM 的文本加vision token；s = 8192 时 attention 还要再加约 s/(6h) ≈ 33%，6ND 没算它。

去查一下模板里的空和上面那张表，都要从你当时的 ds_config、训练 log 和 wandb 里查出来，不要凭印象说。查完把模板填满，用英文讲两遍，控制在 2 分钟内。

准备一个问题和解决过程

结构：症状（带数字）→ 怎么定位 → 根因 → 改了什么 → 改完的数字。四类里选一个你真遇到过的：

• OOM：根因多是长视频样本的 activation，办法有 gradient checkpointing、减帧数或分辨率、按长度分桶、ZeRO-3 或 offload
• 慢：视频解码拖住 dataloader（多开 worker、预先抽帧）、ZeRO-3 每层 all-gather 走跨机 IB、padding 太多
• NaN：查 learning rate 和 warmup、gradient clipping、坏样本
• hang：各 rank 的 collective 次数对不上，比如某些 rank 的 batch 里没有图像，跳过了 vision encoder，ZeRO-3 的 parameter all-gather 就对不齐（查法见第 8 课）

可能的追问

• “Why not Megatron?”：8B 用 ZeRO 就放得下（ZeRO-2 每卡 19.5 GB），DP 的通信大多能藏在计算后面，模型代码和 DeepSpeed 又都直接接 HF。TP、PP 要到 70B 以上、超长序列，或者卡多到 global batch 不够分时才划算
• “What limited throughput?”：先报实测的 MFU，再说时间花在数据、通信、recompute（多约 1/3 的 compute）还是 padding。没测过就说会用 torch profiler 抓一个 step 的 trace
• “How would you scale to 70B?”：70B × 16 = 1,120 GB，ZeRO-3 切到 32 卡每卡 35 GB 放得下，但每层都要跨机 all-gather，activation 也不切。所以换 Megatron，TP 8 放机内、PP 4 跨机正好 32 卡，长序列再加 SP、CP（第 8 课的 70B 配置）

30 秒英文回答

自测

想看细一点 3.11 KV cache · 3.3 MHA · 6.8 MoE 与 Expert Parallel · 6.1 TP Linear · 6.16 Quantization · 6.19 autograd · 6.20 gradient accumulation · 6.6 DP、ZeRO、FSDP · 6.11 FLOPs、6ND、MFU · 6.18 训练慢、hang 怎么查