【Bug已解决】Model trained with Flash Attention 2.0 raises RuntimeError： query and key must have the same

【Bug已解决】Model trained with Flash Attention 2.0 raises RuntimeError query and key must have the same dtype when generating 解决方案一、现象长什么样你用 Flash Attention 2.0attn_implementationflash_attention_2训练好一个模型推理单条forward正常但一调用model.generate()就报RuntimeError: query and key must have the same dtype File .../flash_attention_2.py, line 120, in forward assert q.dtype k.dtype, query and key must have the same dtype # 更具辨识度的线索只在 generate 时报forward 不报 # 或者第一轮prefill正常从第二步decode cache开始报最诡异的是同样的输入直接model(inputs)没问题套上generate()就炸而且往往发生在生成的第 2 个 token 之后即开始使用past_key_values缓存时。二、背景Flash Attention 2 的实现有一条硬约束query 和 key 的张量必须完全相同的数据类型不能一个 bf16 一个 fp32。这是 Flash Attention 2 内核的要求普通 PyTorch SDPA 的 math 后端则宽容得多会自动 promote。问题出在generate()的流程生成是自回归的第 1 步prefill把整段 prompt 喂进去算一遍并把每一层的 key/value 存进past_key_valuesKV 缓存第 2 步起只喂新生成的 1 个 token和缓存里的历史 key/value 拼接后做注意力。如果历史缓存的 dtype 与新 query 的 dtype 不一致Flash Attention 2 在拼接后的注意力计算里就抛query and key must have the same dtype。为什么会出现不一致常见三种来源模型/缓存在generate前被部分.float()或.half()导致某一层权重是 fp32、另一层是 bf16。past_key_values被显式转成 fp32 保存为省显存或精度而新 query 是 bf16。用了torch.compile或autocastprefill 走了一条 dtype 路径、decode 走另一条。三、根因根因落在KV 缓存与新 query 的 dtype 在 generate 路径上分叉prefill 与 decode 的 dtype 路径不同。 prefill 时 q/k 都来自同一份输入dtype 一致但 decode 时q 来自当前步的隐状态可能经过某层.float()或被 autocast 影响k 来自历史缓存之前保存时的 dtype。二者分叉 → 不一致。KV 缓存被手动转了 dtype。 用户为了省显存把past_key_values转 fp32或为了精度把某层 cache 转 fp32但模型权重仍是 bf16新算出的 q 是 bf16 → 拼接后 q(bf16) vs k(fp32)。Flash Attention 2 不自动 promote。 普通注意力后端遇到 dtype 不同会偷偷k k.to(q.dtype)但 Flash Attention 2 选择直接报错因为内核严格于是把潜在问题暴露成RuntimeError而 math 后端会默默吞掉——这就是forward 不报、generate 报且换后端不报的原因。四、最小可运行复现下面用纯 Python 模拟prefill 的 k 缓存是 fp32、decode 的 q 是 bf16拼接后 Flash 报错from dataclasses import dataclass from typing import List dataclass class FakeTensor: dtype: str # fp32 / bf16 def flash_attn_2(q: FakeTensor, k: FakeTensor) - str: Flash Attention 2 的硬约束q、k 必须同 dtype。 if q.dtype ! k.dtype: raise RuntimeError(query and key must have the same dtype) return attn_ok # 模拟 generate 流程 # prefillq/k 都是 bf16一致 q_prefill FakeTensor(bf16) k_prefill FakeTensor(bf16) print(prefill:, flash_attn_2(q_prefill, k_prefill)) # attn_ok # 用户为省显存把 KV 缓存转成 fp32 k_cache FakeTensor(fp32) # decode 第 2 步新 q 来自 bf16 权重 q_decode FakeTensor(bf16) try: flash_attn_2(q_decode, k_cache) # 复现bf16 vs fp32 print(复现失败) except RuntimeError as e: print(复现成功:, e) # 修正decode 前把 q、k 对齐到同一 dtype def align_dtype(q, k, targetbf16): return FakeTensor(target), FakeTensor(target) q2, k2 align_dtype(q_decode, k_cache, targetbf16) print(修正后:, flash_attn_2(q2, k2)) # attn_ok运行后decode 步q(bf16)与k_cache(fp32)触发RuntimeError对齐 dtype 后通过正好复现根因。五、解决方案第一层最小直接修复最快的止血在generate前/中强制模型、输入、KV 缓存都用同一 dtype尤其是不要让缓存被转成不同精度import torch # 1) 模型整体统一 dtype训练是 bf16就全程 bf16 model model.to(torch.bfloat16) model.config.torch_dtype bfloat16 # 2) 输入也转同一 dtype inputs inputs.to(torch.bfloat16) # 3) 关键不要手动把 past_key_values 转 fp32若必须转generate 时统一回 bf16 def generate_stable(model, inputs, max_new20): out model.generate( **inputs, max_new_tokensmax_new, # 让 HF 用模型自身的 dtype 维护 KV 缓存不要额外 cast do_sampleFalse, ) return out # 4) 若你自定义了 past_key_values 的 dtype在送入注意力前对齐 def align_kv_dtype(hidden_states, past_key_value, target_dtype): q hidden_states.to(target_dtype) if past_key_value is not None: k, v past_key_value k, v k.to(target_dtype), v.to(target_dtype) past_key_value (k, v) return q, past_key_value第一层让用户立刻消除query and key must have the same dtypegenerate正常。六、解决方案第二层结构性改进用FlashDtypeGuard把q/k 同 dtype的约束收进注意力调用处无论 prefill 还是 decode 都自动对齐from dataclasses import dataclass from typing import Optional, Tuple import torch dataclass class FlashDtypeGuard: 保证送入 Flash Attention 2 的 q、k、v 同 dtype吸收 generate 路径的分叉。 policy: str bf16 # 统一目标 dtype def _target(self, *tensors) - torch.dtype: # 默认取模型权重的 dtype若指定 policy 则用它 if self.policy bf16: return torch.bfloat16 if self.policy fp16: return torch.float16 return torch.float32 def align(self, q, k, v, past_key_valueNone): tgt self._target(q, k, v) q, k, v q.to(tgt), k.to(tgt), v.to(tgt) if past_key_value is not None: pk, pv past_key_value pk, pv pk.to(tgt), pv.to(tgt) past_key_value (pk, pv) return q, k, v, past_key_value # 在模型注意力 forward 里使用 guard FlashDtypeGuard(policybf16) def attention_with_guard(module, hidden_states, past_key_valueNone): q module.q_proj(hidden_states) k module.k_proj(hidden_states) v module.v_proj(hidden_states) # 统一 dtypeprefill 和 decode 都走这里保证一致 q, k, v, past_key_value guard.align(q, k, v, past_key_value) out flash_attention_2(q, k, v, past_key_valuepast_key_value) return outFlashDtypeGuard的语义是Flash Attention 2 的 dtype 约束不该依赖调用方记得对齐而应由注意力层自身在入口强制对齐从而根治 generate 路径的分叉。七、解决方案第三层断言 / CI 守护用 pytest 固化generate 路径下 q/k 永远同 dtypeimport pytest import torch def test_flash_requires_same_dtype(): from flash_guard import flash_attention_2 with pytest.raises(RuntimeError): flash_attention_2(torch.randn(1,1,8,16, dtypetorch.bfloat16), torch.randn(1,1,8,16, dtypetorch.float32)) def test_guard_aligns_dtype(): from flash_guard import FlashDtypeGuard g FlashDtypeGuard(policybf16) q torch.randn(1,1,8,16, dtypetorch.bfloat16) k torch.randn(1,1,8,16, dtypetorch.float32) # 缓存是 fp32 q2, k2, v2, _ g.align(q, k, k) assert q2.dtype k2.dtype torch.bfloat16 def test_generate_does_not_raise_dtype_error(): # 端到端开启 flash_attention_2 的模型 generate 不应报 dtype 错 from transformers import AutoModelForCausalLM, AutoTokenizer # 用一个小模型本地构造无网络验证 dtype 一致性逻辑 # 这里以 guard 覆盖为例 g FlashDtypeGuard(policybf16) q torch.randn(2,4,8, dtypetorch.bfloat16) k torch.randn(2,4,8, dtypetorch.bfloat16) q2, k2, _, _ g.align(q, k, k) assert q2.dtype k2.dtypeCI 跑pytest tests/test_flash_dtype.py以后只要有人又手动把 KV 缓存转成别的 dtype或 prefill/decode 路径分叉测试立刻红灯。八、排查清单当 Flash Attention 2 训练模型在generate时报 dtype 错按顺序查报错含query and key must have the same dtype→ KV 缓存与新 query dtype 不一致先用第一层统一模型/输入/cache dtype。只在 generate 第 2 步后报、forward 不报 → 必然是 KV 缓存历史 k与新 q 分叉重点检查过去_key_values 是否被.float()/.half()。换attn_implementationeager(math) 不报 → 确认是 Flash 2 的硬约束暴露了潜在 dtype 问题应修 dtype 而非躲到 math 后端math 只是掩盖。用了torch.compile/autocast → 确认 prefill 与 decode 走相同 dtype 路径。长期方案用FlashDtypeGuard在注意力入口强制对齐结构性杜绝分叉。九、小结Flash Attention 2.0 generate 报 query and key must have the same dtype 的根因是Flash Attention 2 要求 q/k 严格同 dtype而generate的自回归路径prefill 存 KV 缓存、decode 取缓存拼接容易让历史缓存 dtype 与新 query 分叉math 后端会偷偷 promote 掩盖问题Flash 2 直接报错。第一层统一模型/输入/KV 缓存 dtype绝不手动把缓存转成不同精度立刻消除报错。第二层用FlashDtypeGuard在注意力入口强制对齐 q/k/v含 past_key_value无论 prefill/decode 都一致。第三层pytest 断言Flash 要求同 dtype、guard 能对齐、generate 路径不报 dtype 错防止回归。记住Flash Attention 2 不会替你 promote dtype凡是走 KV 缓存的生成路径都要保证历史缓存与新 query 同 dtype否则它就把潜在问题暴露成 RuntimeError。