Reuse ABI-compatible JIT extension caches across Python/torch envs

[Qubitium]

Summary

JIT extension cache fingerprints hash the exact Python version (python=3.12.9) and exact torch wheel version unconditionally, so every Python or torch env change forces a full recompile — even though all TorchOpsJitExtension binaries are loaded via torch.ops.load_library (is_python_module=False) and most have no CPython ABI dependence. This PR makes fingerprints reflect actual ABI requirements so compiled extensions are reused wherever the binary is genuinely compatible.

What Changed

• gptqmodel/utils/cpp.py:
  • TorchOpsJitExtension gains python_abi_dependent / torch_stable_abi_target params. Python-ABI detection (default None = auto-detect by scanning sources + recursively resolved local includes) is deliberately conservative: Python.h / pybind11/* / torch/extension.h / torch/python.h includes and any Py[A-Z]* / py:: / PYBIND11_* usage flag the extension as python-dependent; audited torch.ops-only loaders (pack_block, floatx, awq, qqq, exllamav2 gptq+awq, exllamav3, machete, marlin fp16+bf16, paroquant, hadamard) set explicit python_abi_dependent=False (only unused registration macros / incidental includes, no runtime Python usage in compiled TUs).
  • Torch stable-ABI cross-version reuse requires an explicit torch_stable_abi_target=(major, minor). When set, the class injects -DTORCH_TARGET_VERSION into both host cflags and NVCC flags (so the two floors can't drift), fingerprints as torch_abi=stable-{major}.{minor}, validates the scanned sources (only torch/csrc/stable/ / torch/headeronly/ torch headers are stable; ATen/, c10/, caffe2/ and other torch/ headers raise a RuntimeError naming the offending include), and enforces the runtime floor itself: load() refuses to reuse a cached binary or compile when the running torch is below the target. Sources that auto-detect as stable but have no explicit target stay keyed to the current torch build.
  • Include recursion (both ABI scanning and source fingerprint hashing) now follows angle-bracket includes too, when they resolve under the including file's directory or an explicit extra_include_paths entry; unresolvable angle includes are treated as system headers and skipped.
  • Fingerprint payload pseudo-diff:

    - python=3.12.9
    + python_abi=agnostic            (or the interpreter SOABI, e.g. cpython-313t-x86_64-linux-gnu,
    +                                 when python-ABI dependent — free-threaded builds get distinct keys)
    - torch=2.8.0+cu128
    + torch_abi=stable-2.10          (extensions with an explicit stable-ABI target)
    + torch=2.8.0                    (non-stable CPU extensions: accelerator wheel tag stripped)
    + torch=2.8.0+cu128              (non-stable CUDA extensions: full wheel version kept)
    - torch_cuda=12.8                (now emitted only when requires_cuda)

  • Non-stable CPU normalization is allowlist-only (_torch_cpu_cache_version()): only +cpu, +cuNNN, +rocmX.Y, +xpu are stripped; source/vendor local identifiers (e.g. +gitabc1234) are preserved so different libtorch builds never collide in the cache.
  • New torch_stable_abi_target_define(major, minor) → -DTORCH_TARGET_VERSION=0x020A000000000000 (libtorch encoding: major<<56 | minor<<48).
• gptqmodel/utils/swordfish.py: swordfish (written against the torch stable ABI) is pinned via torch_stable_abi_target=(2, 10). 2.10 is the audited floor: torch_get_current_cuda_blas_handle (used by libtorch_stable/torch_utils.h) is gated at TORCH_VERSION_2_10_0 in torch/csrc/stable/c/shim.h; no swordfish API needs 2.11+. _swordfish_static_runtime_error() also gates torch < 2.10 with an explicit error so expected incompatibility never enters JIT compilation.
• Behavior: CPU extensions no longer recompile when switching between +cpu/+cu12x/+rocm*/+xpu wheels of the same torch version; python-agnostic extensions share caches across Python versions (including free-threaded); swordfish shares caches across torch versions >= 2.10 and refuses cached-binary reuse below the floor.
• Out of scope: CUDA extensions built against non-stable libtorch APIs still key on the full torch wheel version (required for correctness); no migration of other kernels to the stable ABI.

Tests

• I added a new simple/fast unit test for this change, or documented why that is not applicable.
• I ran the new targeted test locally before opening this PR.
• I ran any other directly relevant local tests.

New tests in tests/test_torch_ops_jit_extension.py cover: python-agnostic fingerprint stability across Python versions, SOABI keying (incl. free-threaded fallback), conservative python-ABI marker detection, wheel-switch normalization (+cpu ↔ +cu128/+rocm6.2/+xpu reuse; +gitabc1234 preserved), stable-target fingerprinting + define injection into both flag sets, explicit-stable-target rejection of non-stable headers, runtime floor refusal on torch 2.9 with target 2.10 (no cache lookup), angle-bracket include resolution (ABI detection + fingerprint invalidation + system-include skip), swordfish torch<2.10 gate, and the TORCH_TARGET_VERSION encoding.

$ python -m pytest -q tests/test_torch_ops_jit_extension.py
46 passed in 4.90s
$ ruff check gptqmodel/utils/cpp.py gptqmodel/utils/swordfish.py gptqmodel/utils/awq.py gptqmodel/utils/qqq.py gptqmodel/utils/exllamav2.py gptqmodel/exllamav3/ext.py gptqmodel/utils/hadamard.py gptqmodel/utils/machete.py gptqmodel/utils/marlin.py gptqmodel/utils/paroquant.py tests/test_torch_ops_jit_extension.py
All checks passed!

Review Requirements

• I personally reviewed every file in this diff.
• I checked that the code matches existing project structure, APIs, and conventions.
• I avoided unnecessary monkeypatching and used the project's normal extension points where possible.

Notes

• Full swordfish CUDA JIT compilation/runtime was not exercised here (no CUDA toolkit/GPU on this box); a stable-shim syntax check with TORCH_TARGET_VERSION=0x020A000000000000 confirms the required declarations are visible at the 2.10 floor. A real build/load on torch >= 2.10 + Blackwell is recommended before relying on cross-torch swordfish reuse.
• Cache-key changes mean existing JIT caches will be rebuilt once under the new fingerprints.

Link to Devin session: https://app.devin.ai/sessions/b1d3f556711a4711b7b82bd732c66de5
Open in Devin Desktop: https://app.devin.ai/desktop/session/b1d3f556711a4711b7b82bd732c66de5?variant=devin
Requested by: @Qubitium

[devin-ai-integration]

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

• Address comments on this PR. Add '(aside)' to your comment to have me ignore it.
• Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

• Disable automatic comment, CI, and merge conflict monitoring

[Qubitium]

Review follow-up (H100-only test environment):

• Found and fixed two edge cases in the ABI cache path. The SOABI fallback now works on supported platforms where sys.abiflags is absent, and stable-ABI loaders now fail closed when torch.__version__ cannot be parsed instead of compiling/reusing an unverifiable binary.
• Updated the stale Machete JIT assertion to match the current SM90/Blackwell diagnostic emitted by main (the mismatch predated this PR).

Validation:

• tests/test_torch_ops_jit_extension.py: 49 passed.
• Extension integration tests (AWQ include paths, ExLlama v2 GPTQ/AWQ, ExLlama v3, Machete deterministic tests, Marlin, QQQ, compile progress; CUDA smoke builds excluded because they invoke the multi-minute CUTLASS build): 76 passed, 47 deselected.
• ExLlama v2 GPTQ and AWQ JIT extensions also built and loaded successfully on the H100.
• Post-push GitHub Analyze, CodeQL, and Ruff checks are passing.