| #!/usr/bin/env python3 | |
| """LNA-LAB: let LNA_DSV4_AUX_STREAMS=0 disable DSV4's attention aux CUDA streams (EXL3 dense GEMMs are | |
| cooperative kernels sharing one device lock buffer; running them concurrently on aux streams deadlocks). | |
| Idempotent exact-anchor patch of nvidia/model.py, backup model.py.orig-lna3.""" | |
| import shutil, sys | |
| from pathlib import Path | |
| OLD = " aux_stream_list = [torch.cuda.Stream() for _ in range(3)]\n" | |
| NEW = (" import os as _lna_os # LNA-LAB\n" | |
| " aux_stream_list = (None if _lna_os.environ.get(\"LNA_DSV4_AUX_STREAMS\", \"1\") == \"0\"\n" | |
| " else [torch.cuda.Stream() for _ in range(3)]) # LNA-LAB: EXL3 dense needs sequential GEMMs\n") | |
| def main(): | |
| import vllm | |
| f = Path(vllm.__file__).resolve().parent / "models" / "deepseek_v4" / "nvidia" / "model.py" | |
| t = f.read_text() | |
| if "LNA_DSV4_AUX_STREAMS" in t: print(f"{f}: already patched"); return 0 | |
| if t.count(OLD) != 1: raise SystemExit(f"{f}: expected 1 anchor, found {t.count(OLD)}") | |
| b = f.with_suffix(f.suffix + ".orig-lna3"); b.exists() or shutil.copy2(f, b) | |
| t = t.replace(OLD, NEW); compile(t, str(f), "exec"); f.write_text(t); print(f"{f}: patched"); return 0 | |
| if __name__ == "__main__": sys.exit(main()) | |