Amazon EM 计算题：MLP Batch 1 与 256 的 Roofline 下界

题意与假设

这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。

• 训练解释X/H/Y为B×4096，W为4096×4096，全为2-byte BF16，H与X为不同输入；题目dimension未展开shape，故先核此合同。
• 采用十进制TB=10^12、MB=10^6，FMA计2 FLOP；冷权重一次从HBM加载，融合bias/GELU/residual且不写中间结果。
• 只计算理想下界，GELU额外计算未计入主matmul FLOPs；不把峰值指标等同实测吞吐或精确端到端latency。

解题路线

1. 主matmul FLOPs=2Bd²，最低冷HBM traffic=2(d²+3Bd+d)字节，分别含W、X/H/Y与bias。与硬件单位一致，不能把bits当bytes或把FMA只算1却仍用双FLOP峰。
2. B=1：33,554,432 FLOP、33,587,200 byte；compute≥0.33554432µs，HBM≥33.5872µs，理想max下界约33.59µs。算术强度约0.999 FLOP/byte，低于峰/带宽=100，模型为memory bound。
3. B=256：8,589,934,592 FLOP、39,854,080 byte；compute≥85.89934592µs，HBM≥39.85408µs，理想max下界约85.90µs。强度约215.535 FLOP/byte超过100，理想模型为compute bound。
4. W有33,554,432 byte=32MiB，不能全装25MB SRAM；但不等于完全不能复用。X/H/Y合约6.29MB，理想tile可在剩余SRAM内分块W并在batch内复用，实际bank/layout与工作区会限制可达性。
5. 冷权重假设必须说明：跨调用缓存驻留会降低HBM读，未融合中间张量或spill会增加读写。独立compute与memory下界取max仅为可重叠理想界，无重叠粗估可取sum，真实runtime还含launch、同步与非矩阵算子。
6. 证明实际瓶颈需profile有效FLOP/s、bytes、cache和stall，不能仅靠理论数宣称每种实现都如此。batch1 GPU利用率可能远低峰，batch256也可能因差tile/带宽争用偏离模型。

手动走一遍

训练同一W被256行复用，主FLOPs变256倍而冷HBM总bytes只从33.59MB增到39.85MB，因此强度增长。上述数字是指定shape/融合/冷cache假设的算术估算，不是候选人实测。

规模与取舍

主计算随B·d²增长，冷IO随d²+B·d增长；SRAM限制影响tile与数据重读，不能把25MB与HBM带宽当同一种资源。算式已用独立整数计数核对。

易错点与追问

容易失分的边界

• 把两个下界当精确实测latency。
• 忘记读取residual或写Y。
• W不全fit就认定无tiling复用。
• 缓存/融合假设不写。

继续深挖

• residual与X同一buffer如何改traffic？
• GELU分kernel怎样多计HBM？
• 有效BW降一半时瓶颈分界？