【暑期实习】蚂蚁 AI Infra 一面复盘

算子优化、量化、图优化都问到了

最近面了蚂蚁的 AI Infra 一面，整体感受是，这一轮比较偏工程理解，不是只问八股，更想看你是不是真的做过优化，能不能把问题、分析过程和思路讲清楚。

一开始是自我介绍，后面主要围绕我现在实习里做过的事情展开。面试官先让我讲了实习中解决过的一些问题，重点问到一个算子要不要优化，该怎么判断，以及如果决定优化，具体会怎么做。我这里主要是从热点分析、瓶颈判断和优化路径选择几个角度来回答。比如先看它是不是耗时热点，再看问题到底是算力瓶颈、访存瓶颈，还是调度和启动开销，然后再决定是做 kernel 优化、算子融合、图优化，还是数据布局调整。

后面问得比较多的是量化。先考察了我对量化工作的理解，然后又给了一个开放题：如果给我一个还没有量化的模型，要怎么推进量化。我当时主要按自己的理解回答，大概是先分析模型结构和部署目标，再去选量化粒度、量化方案和校准方式，最后结合精度和性能去反复迭代。这一块我能答出思路，但说实话还不算特别系统，因为我现在接触更多还是偏推理侧和算子层优化，对完整模型量化流程的工程经验还不够多。

接着他追问了一个我觉得挺典型的问题，就是像 AWQ 这种量化方式，看起来其实给模型增加了很多步骤，比如把 BF16 压成更低比特，计算的时候又要做一些反量化，那为什么最后整体还是会有加速。我当时给了两个理解。一个是量化之后权重更小了，带宽压力会下降，模型读取和访存效率会提升，很多推理场景本来就是 memory bound，所以即使多了一些量化和反量化步骤，总体吞吐还是可能更高。另一个是量化往往不只是改数据类型，很多时候会和算子融合、数据排布优化一起做，这样虽然表面上步骤变多了，但中间 tensor 的搬运、读写和 kernel 切换次数减少了，所以总时间反而下降。

然后他又问了我在实习里写的一个 timing 插件，为什么要写，它底层是怎么实现的，解决了什么问题，以及为什么不用 Nsight、NCU 这种 profiling 工具直接做。我这里主要讲的是，这个插件更适合日常开发过程里的快速定位和粗粒度统计，接入成本低，迭代快，不需要每次都起完整 profiler，比较适合高频排查问题。

另外还问到了不同模型在本地部署和量化时会不会有明显差异，也问了图优化具体是怎么选的，开发流程是什么。还有静态图和动态图相关的问题，比如为什么很多时候要先把动态图转成静态图再去做融合优化，以及动态图和静态图转换本身的开销怎么处理。这部分我答得一般，因为我目前的经验还是零散的，没能把这块讲得特别体系化。

中间也问到了 vLLM。我这里其实有点可惜，因为源码我下过，但还没系统看完，所以只能回答一些比较浅的理解，比如它在推理调度、KV Cache 管理上的一些核心思路，以及我最近看过的一些 KV Cache 量化相关论文。这个地方也能明显感觉到，如果面的是 AI Infra，光知道概念还是不够，最好还是要往源码和实现细节里走。

后面还问了我在实习过程中有没有解决过比较困难的问题。我提到的是一个 OOM 和随机崩溃的问题，出现在长轮次推理或者训练中，会有显存异常增长以及随机退出的情况。我就讲了当时大概是怎么排查链路、定位原因，再一点点修复的。

最后还做了一道 LeetCode 题，是复制带随机指针的链表。这个题本身不算特别陌生，我给的思路就是先处理普通 next 链，再去处理 random 指针，整体思路还算顺下来了。

反问环节我主要问了两件事。一个是后续流程，面试官说大概还有三轮，应该是两轮技术面加一轮 HR 面。另一个是进去之后主要做什么，他说整体偏推理优化方向，他们团队本身是做强化学习相关的，但具体工作不只是单点算子，也会涉及上层框架、推理引擎、训练引擎、传输引擎和后面的算子优化。

整体下来，我觉得这一面不是那种很纯算法的面试，更像是在看你有没有真实做过 AI Infra 相关的事情，能不能把性能问题、定位过程和优化思路讲清楚。对我自己来说，这一面也暴露了几个短板，一个是模型量化流程理解还不够系统，一个是动态图、静态图和图优化这块还不够扎实，还有一个就是像 vLLM 这种推理系统，后面确实要继续补源码。