在计算attention score的时候如何对padding做mask操作

1️⃣ 考察意图

面试官想看你是否真正理解Transformer底层实现，而非只调过库。考察类型是工程实现+原理debug。刁钻点在于：很多人知道要mask，但说不清为什么用-inf而非0、训练和推理的差异、以及如何与因果mask共存。答好了能展示你对attention机制的底层掌控力、处理变长序列的工程经验，以及避免梯度泄漏的敏感度。

2️⃣ 标准答

核心原理：padding mask的目的是让模型忽略无效的padding token。在计算attention score后、softmax之前，将padding位置的score设为-inf（或极大负数如-1e9），这样softmax后这些位置的权重趋近于0，不会参与后续的加权求和。

具体实现步骤：

生成mask矩阵：根据batch内每个序列的实际长度，生成一个[batch_size, seq_len]的布尔矩阵，padding位置为True。例如，用(token_ids == pad_token_id)得到。
扩展维度：attention score的形状是[batch_size, num_heads, seq_len_q, seq_len_k]。mask需要广播到相同维度，通常扩展为[batch_size, 1, 1, seq_len_k]（对decoder自注意力）或[batch_size, 1, seq_len_q, seq_len_k]（cross-attention）。
应用mask：用masked_fill将True位置填充为-inf。为什么不用0？ 因为0在softmax后会产生非零权重，导致padding token泄漏信息到有效token，破坏语义。-inf确保softmax后权重为0。
softmax：对mask后的score做softmax(dim=-1)，得到注意力权重。

训练与推理的差异：

训练：必须用mask，因为batch内序列长度不同，padding是常态。
推理：如果是单条序列（无padding），不需要mask；如果是批量推理（如服务端并发），仍需mask。

实际落地的坑与解法：

坑1：数值稳定性。直接用float('-inf')可能导致softmax输出NaN（尤其在fp16下）。解法：用-1e9或-65504（fp16最大负值），或使用torch.where配合torch.finfo动态获取最小值。
坑2：与因果mask合并。在decoder中，需要同时应用padding mask和因果mask（上三角）。解法：先生成因果mask（torch.triu(..., diagonal=1)），再用torch.logical_or与padding mask合并，统一填充-inf。
坑3：FlashAttention兼容。FlashAttention内部自动处理padding mask（通过cu_seqlens），手动mask会破坏其分块计算。解法：使用FlashAttention时，传入cu_seqlens参数，不要额外做masked_fill。

工程取舍：

位置：在score后、softmax前做mask，而非在权重后。因为权重后mask无法消除padding token对softmax分母的贡献。
数据类型：mask矩阵用布尔型而非整型，因为PyTorch的masked_fill对布尔型有优化，且节省显存。

3️⃣ 答题模板（30 秒电梯版）

“这个问题我从原理、实现、工程坑三个层面回答。原理上，padding mask将padding位置的attention score置为-inf，使softmax后权重归零。实现上，先生成布尔mask，扩展维度后与score相加或masked_fill。工程上要注意：用-1e9而非-inf避免NaN，与因果mask用logical_or合并，FlashAttention下用cu_seqlens替代手动mask。总结一句：mask的本质是信息隔离，核心是-inf而非0。”

4️⃣ 高频追问 & 应对

追问 1：为什么不用0而用-inf？用0会怎样？

用0会导致softmax后padding位置仍有非零权重，因为softmax的分母包含所有位置的指数和。假设有效token的score是[10, 9]，padding是0，softmax后padding权重约0.0001，虽然小但非零。这会导致梯度回传到padding token，浪费计算且可能干扰有效token的梯度。-inf确保指数为0，权重精确为0，梯度不回流。

追问 2：在推理时，如果batch内序列长度不同，如何处理？

推理时仍需mask。常见做法是动态padding：将batch内序列padding到当前batch的最大长度。但更高效的是连续批处理（continuous batching），如vLLM的实现，用page attention管理KV cache，无需padding。如果必须用padding，mask逻辑与训练一致，但注意推理时通常不需要因果mask（因为自回归生成已通过KV cache控制）。

追问 3：FlashAttention如何避免手动mask？原理是什么？

FlashAttention通过分块计算和重计算避免显式mask。它接收cu_seqlens（每个序列的起始位置索引），在分块时跳过padding块。例如，序列长度[3,5]的batch，cu_seqlens为[0,3,8]。在计算第2个序列时，只处理索引3-7的块，padding块（索引8+）不参与计算。这避免了-inf填充，且节省显存和计算量。

5️⃣ 避坑 · 常见错误答法

❌ 说“在softmax之后用0 mask掉padding位置” → ✅ 正确做法是在softmax之前用-inf，因为softmax后mask无法消除padding对分母的贡献。
❌ 说“训练和推理都用相同mask逻辑” → ✅ 推理时如果使用KV cache，只需mask query对key的padding，且因果mask已由cache控制，需区分场景。
❌ 说“用float('-inf')没问题” → ✅ 在fp16下-inf可能导致NaN，应使用-1e9或torch.finfo(dtype).min。

6️⃣ 简历呼应

如果你有LLM训练项目：从“训练时处理变长batch的mask实现”切入，提到你对比过masked_fill和FlashAttention的cu_seqlens方式，并记录过显存节省比例（如20%）。
如果你只做过传统NLP：用“RNN中padding mask通过pack_padded_sequence实现”类比，迁移到Transformer的-inf机制，强调序列建模中信息隔离的通用性。
如果你是校招无项目：聚焦“复现GPT-2时实现padding+因果mask合并”的demo，展示你读过《Attention Is All You Need》和HuggingFace源码，能手写mask逻辑。
《Attention Is All You Need》——原始论文，Section 3.2.3对mask的描述
HuggingFace Transformers源码：modeling_gpt2.py中_attn函数的mask实现
FlashAttention论文（Dao et al., 2022）——Section 3.3对padding mask的处理
PyTorch官方文档：torch.nn.Transformer的src_key_padding_mask参数说明
《The Annotated Transformer》——Harvard NLP的逐行实现，含mask代码

—— 本场面试完 ——