File size: 2,789 Bytes
03afa3a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | #!/usr/bin/env python3
"""LNA-LAB: DSV4 attention.py reads compressor.fused_wkv_wgate.weight directly (torch.mm). When the
module was built by vllm-exl3 (no .weight, has _exl3_linears), run the EXL3 shards instead.
Idempotent, exact anchors, backup attention.py.orig-lna2."""
import shutil, sys
from pathlib import Path
HELPER = '''
def _lna_kv_score(module, hidden_states):
"""LNA-LAB: fused_wkv_wgate as bf16 weight (torch.mm) or as EXL3 shards (vllm-exl3)."""
linears = getattr(module, "_exl3_linears", None)
if linears:
x = hidden_states.reshape(-1, hidden_states.shape[-1]).to(torch.float16).contiguous()
outs = [lin.forward(x, {}, out_dtype=torch.float32) for lin in linears]
out = outs[0] if len(outs) == 1 else torch.cat(outs, dim=-1)
return out.reshape(*hidden_states.shape[:-1], out.shape[-1])
return torch.mm(hidden_states, module.fused_wkv_wgate.weight.T, out_dtype=torch.float32)
'''
EDITS = [
(''' return torch.mm(
hidden_states,
compressor.fused_wkv_wgate.weight.T,
out_dtype=torch.float32,
)
''', ''' return _lna_kv_score(compressor.fused_wkv_wgate, hidden_states) # LNA-LAB
'''),
(''' return torch.mm(
hidden_states,
indexer.compressor.fused_wkv_wgate.weight.T,
out_dtype=torch.float32,
)
''', ''' return _lna_kv_score(indexer.compressor.fused_wkv_wgate, hidden_states) # LNA-LAB
'''),
]
def main():
root = Path(sys.argv[1]) if len(sys.argv) > 1 else None
if root is None:
import vllm
root = Path(vllm.__file__).resolve().parent
f = root / "models" / "deepseek_v4" / "attention.py"
text = f.read_text()
if "_lna_kv_score" in text:
print(f"{f}: already patched"); return 0
for old, new in EDITS:
if text.count(old) != 1:
raise SystemExit(f"{f}: expected 1 anchor, found {text.count(old)}: {old[:60]!r}")
text = text.replace(old, new)
# helper: insert before the first class definition
idx = text.index("\nclass ")
text = text[:idx] + "\n" + HELPER + text[idx:]
# the helper only reads .weight in the fallback path; keep it after fused_wkv_wgate module refs
text = text.replace("return torch.mm(hidden_states, module.fused_wkv_wgate.weight.T, out_dtype=torch.float32)",
"return torch.mm(hidden_states, module.weight.T, out_dtype=torch.float32)")
bak = f.with_suffix(f.suffix + ".orig-lna2")
if not bak.exists(): shutil.copy2(f, bak)
f.write_text(text); compile(text, str(f), "exec"); print(f"{f}: patched")
return 0
if __name__ == "__main__":
sys.exit(main())
|