Qwen3.6-27B iluvatar bi-v100 adaptation
This commit is contained in:
186
qwen3_6_scripts/patch_xformers_sdpa_seq.py
Normal file
186
qwen3_6_scripts/patch_xformers_sdpa_seq.py
Normal file
@@ -0,0 +1,186 @@
|
||||
"""
|
||||
策略：顺序（per-sequence）fallback — 纯 PyTorch 数学实现
|
||||
==========================================================
|
||||
逐条序列用 matmul + softmax 手写 attention，完全绕开所有硬件
|
||||
flash attention kernel（ixformer / cudnnFlashAttnForward）。
|
||||
|
||||
背景：
|
||||
Iluvatar cudnnFlashAttnForward 存在两个已知问题：
|
||||
1. 不支持 is_causal=True（报错）
|
||||
2. 使用 attn_mask 路径时数值结果不正确（静默错误，输出全为"!"）
|
||||
与华为昇腾 910B4 上 llama.cpp --flash-attn off 修复同类问题的原理相同。
|
||||
纯数学路径（matmul + softmax）在任何 PyTorch 后端上结果都正确。
|
||||
|
||||
优点：
|
||||
数值正确，不依赖任何硬件特定 attention kernel。
|
||||
峰值显存 = max(seq_len)² × H × dtype_size，由 --max-model-len 控制。
|
||||
|
||||
缺点：
|
||||
并发请求的 prefill attention 串行执行。
|
||||
O(L²) 显存（无 flash attention 的 O(L) 优化）。
|
||||
|
||||
内存参考（fp16，H_local=6）：
|
||||
max-model-len=4096 → 峰值 ~200 MB
|
||||
max-model-len=8192 → 峰值 ~800 MB
|
||||
max-model-len=16384 → 峰值 ~3.2 GB
|
||||
|
||||
Deploy:
|
||||
python3 modified_scripts/patch_xformers_sdpa_seq.py
|
||||
"""
|
||||
|
||||
XFORMERS_PATH = (
|
||||
"/usr/local/corex/lib64/python3/dist-packages/"
|
||||
"vllm/attention/backends/xformers.py"
|
||||
)
|
||||
|
||||
FALLBACK_METHOD = '''
|
||||
def _run_sdpa_fallback(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_metadata: "XFormersMetadata",
|
||||
) -> torch.Tensor:
|
||||
"""顺序纯数学 attention fallback。
|
||||
|
||||
完全绕开 ixformer / cudnnFlashAttnForward，用 matmul + softmax
|
||||
手写 attention。Iluvatar cudnnFlashAttnForward 的 attn_mask 路径
|
||||
存在静默数值错误（输出全为"!"），纯数学路径结果正确。
|
||||
|
||||
softmax 在 float32 下计算以防止 float16 溢出，结果转回原始 dtype。
|
||||
|
||||
Args:
|
||||
query : [1, total_prefill_tokens, num_heads, head_dim]
|
||||
key : [1, total_prefill_tokens, num_kv_heads, head_dim]
|
||||
value : [1, total_prefill_tokens, num_kv_heads, head_dim]
|
||||
Returns:
|
||||
[1, total_prefill_tokens, num_heads, head_dim]
|
||||
"""
|
||||
assert attn_metadata.seq_lens is not None
|
||||
orig_dtype = query.dtype
|
||||
|
||||
q_flat = query.squeeze(0) # [T, H, D]
|
||||
k_flat = key.squeeze(0) # [T, Hkv, D]
|
||||
v_flat = value.squeeze(0)
|
||||
|
||||
output = torch.empty_like(q_flat)
|
||||
start = 0
|
||||
for seq_len in attn_metadata.seq_lens:
|
||||
end = start + seq_len
|
||||
# [1, H, L, D]
|
||||
q_s = q_flat[start:end].permute(1, 0, 2).contiguous().unsqueeze(0)
|
||||
k_s = k_flat[start:end].permute(1, 0, 2).contiguous().unsqueeze(0)
|
||||
v_s = v_flat[start:end].permute(1, 0, 2).contiguous().unsqueeze(0)
|
||||
|
||||
# GQA：展开 KV heads 至与 query heads 一致
|
||||
if k_s.shape[1] != q_s.shape[1]:
|
||||
n = q_s.shape[1] // k_s.shape[1]
|
||||
k_s = k_s.repeat_interleave(n, dim=1).contiguous()
|
||||
v_s = v_s.repeat_interleave(n, dim=1).contiguous()
|
||||
|
||||
# 纯数学 attention：完全绕开硬件 flash attention kernel
|
||||
# [1, H, L, L]
|
||||
attn_w = torch.matmul(q_s.float(), k_s.float().transpose(-2, -1))
|
||||
attn_w = attn_w * self.scale
|
||||
|
||||
# 上三角填 -inf（future tokens）
|
||||
causal_mask = torch.triu(
|
||||
torch.ones(seq_len, seq_len, dtype=torch.bool, device=attn_w.device),
|
||||
diagonal=1,
|
||||
)
|
||||
attn_w = attn_w.masked_fill(causal_mask, float("-inf"))
|
||||
|
||||
# float32 softmax 防止 float16 溢出
|
||||
attn_w = torch.softmax(attn_w, dim=-1)
|
||||
|
||||
out_s = torch.matmul(attn_w, v_s.float()).to(orig_dtype)
|
||||
# [1, H, L, D] → [L, H, D]
|
||||
output[start:end] = out_s.squeeze(0).permute(1, 0, 2)
|
||||
start = end
|
||||
|
||||
return output.unsqueeze(0) # [1, T, H, D]
|
||||
|
||||
'''
|
||||
|
||||
OLD_XFORMER_BLOCK = """\
|
||||
self.attn_op = xops.fmha.flash.FwOp()
|
||||
if self.alibi_slopes is None:
|
||||
# Add the batch dimension.
|
||||
query = query.unsqueeze(0)
|
||||
key = key.unsqueeze(0)
|
||||
value = value.unsqueeze(0)
|
||||
out = xops.memory_efficient_attention_forward(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
attn_bias=attn_bias[0],
|
||||
p=0.0,
|
||||
scale=self.scale,
|
||||
op = self.attn_op
|
||||
)
|
||||
return out.view_as(original_query)\
|
||||
"""
|
||||
|
||||
NEW_XFORMER_BLOCK = """\
|
||||
self.attn_op = xops.fmha.flash.FwOp()
|
||||
if self.alibi_slopes is None:
|
||||
# Add the batch dimension.
|
||||
query = query.unsqueeze(0)
|
||||
key = key.unsqueeze(0)
|
||||
value = value.unsqueeze(0)
|
||||
if self.head_size > 128:
|
||||
out = self._run_sdpa_fallback(query, key, value, attn_metadata)
|
||||
else:
|
||||
out = xops.memory_efficient_attention_forward(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
attn_bias=attn_bias[0],
|
||||
p=0.0,
|
||||
scale=self.scale,
|
||||
op=self.attn_op,
|
||||
)
|
||||
return out.view_as(original_query)\
|
||||
"""
|
||||
|
||||
INJECT_ANCHOR = " def _run_memory_efficient_xformers_forward("
|
||||
|
||||
|
||||
def patch_file(path):
|
||||
with open(path, "r") as f:
|
||||
content = f.read()
|
||||
changed = False
|
||||
|
||||
if "_run_sdpa_fallback" in content:
|
||||
print(" [skip] _run_sdpa_fallback already present")
|
||||
elif INJECT_ANCHOR not in content:
|
||||
print(" [warn] inject anchor not found")
|
||||
else:
|
||||
content = content.replace(INJECT_ANCHOR, FALLBACK_METHOD + INJECT_ANCHOR, 1)
|
||||
print(" [ok] injected _run_sdpa_fallback (sequential, pure-math)")
|
||||
changed = True
|
||||
|
||||
if NEW_XFORMER_BLOCK in content:
|
||||
print(" [skip] dispatch block already patched")
|
||||
elif OLD_XFORMER_BLOCK in content:
|
||||
content = content.replace(OLD_XFORMER_BLOCK, NEW_XFORMER_BLOCK, 1)
|
||||
print(" [ok] patched dispatch block")
|
||||
changed = True
|
||||
else:
|
||||
print(" [warn] dispatch block anchor not found")
|
||||
|
||||
if changed:
|
||||
with open(path, "w") as f:
|
||||
f.write(content)
|
||||
print(f" Written: {path}")
|
||||
|
||||
|
||||
def main():
|
||||
print("=== patch_xformers_sdpa_seq (sequential, pure-math) ===")
|
||||
print(f"Target: {XFORMERS_PATH}")
|
||||
patch_file(XFORMERS_PATH)
|
||||
print("\nDone.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user