Migrate QuTLASS kernels to torch Stable ABI

[cleonard530]

This PR is to continue the libtorch stable ABI migration (see vllm-project/vllm#26946) for vLLM.

Migrates the qutlass kernels used by vLLM from unstable libtorch C++ APIs to PyTorch’s stable ABI (torch::stable::, STABLE_TORCH_LIBRARY_, AOTI shims). This enables qutlass to ship as an abi3 extension (_qutlass_C.so) alongside vLLM’s other stable targets (e.g. _C_stable_libtorch), so the extension can be built against a narrower libtorch surface and remain compatible across PyTorch minor versions.

Migration progress using the Audit Python extension torch-abi-audit:

 -- extensions --
    [STABLE  ] [abi3-ok               ] _C_stable_libtorch.abi3.so  (stable_shim=88, unstable=0)
    [UNSTABLE] [abi3-ok               ] _flashmla_C.abi3.so  (stable_shim=0, unstable=72)
    [UNSTABLE] [abi3-ok               ] _flashmla_extension_C.abi3.so  (stable_shim=0, unstable=68)
    [STABLE  ] [abi3-ok               ] _moe_C_stable_libtorch.abi3.so  (stable_shim=72, unstable=0)
--> [STABLE  ] [abi3-ok               ] _qutlass_C.abi3.so  (stable_shim=60, unstable=0)
    [NO-TORCH] [abi3-ok               ] cumem_allocator.abi3.so
    [NO-TORCH] [abi3-ok               ] spinloop.abi3.so
    [UNSTABLE] [uses-private-api      ] third_party/deep_gemm/_C.cpython-312-x86_64-linux-gnu.so  (stable_shim=0, unstable=57)
    [UNSTABLE] [abi3-ok               ] vllm_flash_attn/_vllm_fa2_C.abi3.so  (stable_shim=0, unstable=84)
    [UNSTABLE] [abi3-ok               ] vllm_flash_attn/_vllm_fa3_C.abi3.so  (stable_shim=0, unstable=80)

I have not gained access to Blackwell GPUs yet, so I have not been able to run unit test locally to test the migration, I have only tested that the build work and is stable. I opened the vLLM PR vllm-project/vllm#47879 to test the changes against the vLLM CI.

cc @Harry-Chen @janeyx99 @LopezCastroRoberto

Note @LopezCastroRoberto, this migrated the kernels used by vLLM to the torch Stable ABI, but I think there are more that are still unstable. If you would like me to do a full migration, I can follow up this PR with one that migrates the rest of the kernels.

Test Plan

In vllm:
pytest --noconftest tests/kernels/quantization/test_nvfp4_qutlass.py tests/kernels/quantization/test_mxfp4_qutlass.py

In qutlass:
pytest tests/
python tests/quartet_test.py

Test Results

vllm test

qutlass test

[LopezCastroRoberto]

Why is this and many other methods being removed from the library?

[LopezCastroRoberto]

I understand these aren't used in vLLM, so we can either migrate them to the stable ABI or leave them as they are, but they can't be removed

[cleonard530]

Sorry, this was originally just for vLLM and I forgot to add them back when I pushed upstream. Adding them back now

[LopezCastroRoberto]

Why do we need all these methods in C++ if the library offers a Python/PyTorch interface where this kind of checks are much easier?

[cleonard530]

Do you mean why not have the checks in somewhere like qutlass/__init__.py? There is no reason they can't be in there. I was just trying to keep consistency with the bindings before: e.g. instead of the unstable torch method torch::checkAllContiguous, we call the stable check_all_contiguous from bindings_utils.h instead, and there's a 1-1 map that is easy to see in the git diff. I can definitely add these checks to the python interface instead though.

[cleonard530]

Actually on second thought, adding these checks to the Python interface won't carry over to the shared C++ library (e.g. _qutlass_C.so) used by other projects, like vLLM. So the check would have to be duplicated (or just not there) for every project that uses the qutlass C++ library.

Let me know if I have this wrong and am missing something, or if you still want it moved over to the python interface (which does have much simpler syntax).

[cleonard530]

Another option is to update vllm to use qutlass's Python package, which is what it does for every other third-party library. This would actually be my recommendation but it would just require more work on the vllm side. Flagging @Harry-Chen and @janeyx99 to see what they think.

[Harry-Chen]

@cleonard530 Yeah using the wheel is definitely a more elegant solution, but we often need to build our own wheel instead of the upstream version due to various reasons (e.g. frequent updates but infrequent release, etc.), and sometimes even to maintain a fork. I tend to keep the status quo as-is for the moment to reduce the complexity. But it we are using it in a not-so-tightly coupled way, maybe switching to wheels is a good idea.

[cleonard530]

@LopezCastroRoberto, I added the removed functions back and converted them to use torch ABI Stable APIs. Instead of the QUTLASS_DISABLE_PYBIND guard I replaced it with a QUTLASS_MINIMAL_BUILD guard, and moved the methods that are not built in vLLM behind it. This is really just to keep the library as small as possible for vLLM and to change as little as possible in it's build, but if you don't want the guard there then I doubt it would matter too much and I can remove it.

[cleonard530]

@LopezCastroRoberto I was able to run the test on a blackwell GPU and they are all passing (after a few updates). I ran unit tests in both vllm and in this repo (see description). I did not update vllm to use the qutlass wheel and it is still being built in the vllm repo, and so I still have the bindings_utils.h file for checking the CUDA checks.

Let me know if there is anything else I need to do before getting this merged. Thanks!

[LopezCastroRoberto]

Left some minor comments, but overall LGTM now.

[LopezCastroRoberto]

nit: this contradicts the new ABI target. I think we should require PyTorch >=2.11 or derive TORCH_TARGET_VERSION consistently.

[cleonard530]

I added the constant variable TORCH_TARGET_MAJOR = 2 and TORCH_TARGET_MINOR = 11 and use this for the check and to create the TORCH_TARGET_VERSION. It add's a few more lines of code (see commit c886612) but ensures they don't drift from each other. If you would rather just hard code "2.11" into the check, I can remove this commit.

[LopezCastroRoberto]

This should work across Python versions, but don't we need to specify a minimum CPython version if we want package installers to not reject the wheel on other versions?

[cleonard530]

Done! Added the minimum python version (requires-python = ">=3.9") to pyproject.toml‎, and added options={"bdist_wheel": {"py_limited_api": "cp39"}}, to setup.py‎. Let me know if you want to change the minumum python version.

[LopezCastroRoberto]

Also, seems like quartet_test.py didn't run based on your screenshots? Can you please check that @cleonard530? Thanks!

[cleonard530]

Also, seems like quartet_test.py didn't run based on your screenshots? Can you please check that @cleonard530? Thanks!

Ran the test and added the results to the PR description!