#!/usr/bin/env python3
"""Deterministically patch vLLM's CMakeLists for a DGX Spark (GB10 / sm_121a)
build.

Two problems in the newest local-inference-lab fork vs. the stock cu130 torch:

1. CUDA_SUPPORTED_ARCHS for the CUDA>=13 branch omits the 12.1 (sm_121) family
   and 12.1a (GB10 sm_121a) targets, so a 12.1a request is silently dropped to
   sm_120 by cuda_archs_loose_intersection().

2. The GPU arch is derived from the torch wheel's arch list (sm_120) rather than
   honouring TORCH_CUDA_ARCH_LIST, so vLLM compiles for the wrong arch.

Both are fixed here idempotently (safe to run after the patches/ dir too).

`cvt.e2m1x2.f32` (FP4) requires sm_121a specifically — neither sm_120 nor the
sm_121 family accepts it — so we must keep the `a` suffix end-to-end.
"""
import pathlib

P = pathlib.Path("CMakeLists.txt")  # cwd is the vLLM source tree
s = P.read_text()
changed = False

# 1) Allow the sm_121 family + GB10 sm_121a in the CUDA>=13 supported list.
old = 'set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;11.0;12.0")'
new = 'set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;11.0;12.0;12.1;12.1a")'
if '12.0;12.1;12.1a"' not in s and old in s:
    s = s.replace(old, new, 1)
    changed = True

# 2) Honour TORCH_CUDA_ARCH_LIST by overriding CUDA_ARCHS (keep the 'a' suffix).
if "_forced_arch" not in s:
    anchor = 'extract_unique_cuda_archs_ascending(CUDA_ARCHS "${CUDA_ARCH_FLAGS}")'
    ovr = (
        anchor
        + "\n"
        + "  # Deterministically honour TORCH_CUDA_ARCH_LIST (e.g. 12.1a -> sm_121a),\n"
        + "  # else the stock torch wheel arch list (sm_120) is used and the FP4\n"
        + "  # cvt.e2m1x2.f32 path fails in ptxas.\n"
        + '  if(DEFINED ENV{TORCH_CUDA_ARCH_LIST} AND NOT "$ENV{TORCH_CUDA_ARCH_LIST}" STREQUAL "")\n'
        + '    string(REGEX MATCH "[0-9]+\\.[0-9]+[af]?" _forced_arch "$ENV{TORCH_CUDA_ARCH_LIST}")\n'
        + '    if(NOT _forced_arch STREQUAL "")\n'
        + '      set(CUDA_ARCHS "${_forced_arch}")\n'
        + "    endif()\n"
        + "  endif()"
    )
    if anchor in s:
        s = s.replace(anchor, ovr, 1)
        changed = True

P.write_text(s)
print("CMakeLists.txt arch-patched" if changed else "CMakeLists.txt already patched")