[Fix] Make DeepGEMM triton fallback more robust by remi-or · Pull Request #47126 · huggingface/transformers

remi-or · 2026-07-07T06:57:34Z

This PR makes the triton fallback to DeepGEMM more robust:

in case CUDA_HOME is not set, the current behavior is to fail with an AssertionError which is not caught by the try block because it specializes on ImportError. To avoid this, we call get_mk_alignment() before creating the DeepGemm object in a try block and pass the missing CUDA_HOME error as an ImportError -> triton fallback works
in case CUDA_HOME is set but does not match the CUDA toolkit, the JIT compilation will fail but this is not caught at loading time, hence the problem manifests at Runtime and we hard-error instead of falling back to triton. To avoid this, we run a tiny JIT compiled kernel at load time (transform_sf_into_required_layout) and fallback to triton if the JIT compilation fails.

Also, I think the functools.cache decorator never worked as intended because it does not catch error raised in the function, so any fail during _load_deepgemm_kernel was bound to happen all over again. To avoid this, instead of raising the expected ImportError in the cached function, we wrap it into a function that will raise by itself.

Result: on a machine with no CUDA_HOME or a mismatching CUDA_HOME, we fall back to triton instead of erroring out. I have not tested on a machine with DeepGEMM but it should not change the behavior or a succesful load. Perhaps @IlyasMoutawwakil can confirm please? Thx!

HuggingFaceDocBuilderDev · 2026-07-07T07:07:39Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

IlyasMoutawwakil · 2026-07-07T08:30:09Z

+            return f"DeepGEMM requires {arch}; current device is SM{major}{minor}."

        # Per the DeepGEMM README: SM90 needs CUDA 12.3+, SM100 needs CUDA 12.9+.
        cuda_major, cuda_minor = get_cuda_runtime_version()


do you know why this doesn't catch the cuda version mismatch ? maybe we can make get_cuda_runtime_version respect CUDA_HOME if it's not the case already

CUDA_HOME is user set / can not match between PATH and the env, this is an env-only check (torch does not use CUDA_HOME to check the version)
In any case it's nore about needing the CUDA toolkit at runtime: the best way to verify that is to actually know if it can be ran. As long as the result of load_deepgemm.... is cached it will only happen once

robust toolkit checks in #47154

IlyasMoutawwakil · 2026-07-07T08:31:37Z

+                f"DeepGEMM failed to JIT-compile a probe kernel, which usually means a broken or version-mismatched "
+                "CUDA toolchain. Check that `CUDA_HOME` points to a complete and version-consistent CUDA toolkit. You "
+                f"can re-run with `DG_JIT_DEBUG=1` for the compiler output. Original error: {type(e).__name__}: {e}"
+            )


this however i'm very reluctant about, we should in theory be able to know whether a cuda toolkit is reachable and its version programmatically without having to allocate or run anything, wdyt ?

Not sure about this really -- it depends on a few variables, and if the goal of the function is to load callables that will be JIT compiled, then it seems logical to check if the JIT compiling works. WDYT? Otherwise pretty hard to tell if a toolkit is working

If this became a load time only error for you cuz tookits where not available,we should indeed do it!

with the checks merged from #47154 the kernel will only be loaded when the env has the toolkit with nvcc, and nvcc claims to be of compatible with the kernels.

the real problem imo is that we used the wrong way to check the cuda versions (whichever libcudart loaded first), now we do the same as what deepgemm does internally (check cuda home, check nvcc, check their versions) so it should never be a problem unless the user manually created an environment with broken (cuda home + nvcc), in which case imo they should be told that their env is broken loudly

not against the final compilation check in itself, but rather against downloading the full kernel and only realizing that its not compatible with your system after a couple minutes of rate limiting.

IlyasMoutawwakil

Thanks for catching this ! we do have a flag to disable the deepgemm fp8 linear, we should probably mention it in the error message in finegrained_fp8's linear dispatch.

remi-or · 2026-07-08T02:28:52Z

Thanks for catching this ! we do have a flag to disable the deepgemm fp8 linear, we should probably mention it in the error message in finegrained_fp8's linear dispatch.

Not sure what you mean here, can you add a permalink?

I feel like this is ready for a second round of review / approval, thanks!

IlyasMoutawwakil · 2026-07-08T05:28:22Z

@remi-or

transformers/src/transformers/integrations/finegrained_fp8.py

Line 239 in d49affb

    
           # ``TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1`` forces the Triton fallback for this single

I have opened #47154 based and into this branch, tell me what you think.

ArthurZucker · 2026-07-08T12:40:16Z

+                f"DeepGEMM failed to JIT-compile a probe kernel, which usually means a broken or version-mismatched "
+                "CUDA toolchain. Check that `CUDA_HOME` points to a complete and version-consistent CUDA toolkit. You "
+                f"can re-run with `DG_JIT_DEBUG=1` for the compiler output. Original error: {type(e).__name__}: {e}"
+            )


If this became a load time only error for you cuz tookits where not available,we should indeed do it!

IlyasMoutawwakil

LGTM ! thanks for making it more robust 🙏

* Gate DeepGEMM JIT on the CUDA toolkit (nvcc/NVRTC), not the runtime DeepGEMM JIT-compiles its kernels with the system CUDA toolkit, so gating on the CUDA runtime version was the wrong signal (and on multi-toolkit boxes could read a stray/mismatched libcudart). Resolve the toolkit the way DeepGEMM does (respecting CUDA_HOME) and validate it statically before loading the kernel. - import_utils: add get_cuda_home + get_nvcc_version (read the toolkit version from version.json/version.txt/cuda.h, no subprocess); deprecate get_cuda_runtime_version. - deepgemm: gate on the resolved nvcc toolkit; when nvcc is missing/too old but torch's bundled libnvrtc is new enough, fall back to the NVRTC backend instead of failing. - test_modeling_common: mirror the gate (Hopper -> CUDA 12.3, Blackwell -> 12.9). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR in worst case scenario * remove nvrtc path as not usable with kernel hub build * trim comment * address comments * update comment --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

github-actions · 2026-07-09T02:14:22Z

CI recap

Dashboard: View test results in Grafana
Latest run: 28988683001:1
Result: success | Jobs: 15 | Tests: 170,701 | Failures: 0 | Duration: 19h 43m

remi-or requested review from ArthurZucker and IlyasMoutawwakil July 7, 2026 06:57

remi-or force-pushed the fix-deepgemm-integration branch from a608a3f to f8633dc Compare July 7, 2026 07:46

IlyasMoutawwakil reviewed Jul 7, 2026

View reviewed changes

Comment thread src/transformers/integrations/deepgemm.py Outdated

IlyasMoutawwakil reviewed Jul 7, 2026

View reviewed changes

Comment thread src/transformers/integrations/deepgemm.py Outdated

IlyasMoutawwakil reviewed Jul 7, 2026

View reviewed changes

Comment thread src/transformers/integrations/deepgemm.py Outdated

IlyasMoutawwakil reviewed Jul 7, 2026

View reviewed changes

remi-or force-pushed the fix-deepgemm-integration branch from f8633dc to 7d9d23b Compare July 8, 2026 02:25

ArthurZucker approved these changes Jul 8, 2026

View reviewed changes

IlyasMoutawwakil approved these changes Jul 8, 2026

View reviewed changes

remi-or and others added 4 commits July 8, 2026 18:00

Catch DeepGEMM errors on call

3d5ad18

Add real call to fallback to triton for broken homes

a7656f5

Remove extra wrapper

d1b5e44

remi-or force-pushed the fix-deepgemm-integration branch from deb5b95 to 3c650e6 Compare July 9, 2026 01:02

Nit

1de9036

remi-or enabled auto-merge July 9, 2026 02:06

remi-or added this pull request to the merge queue Jul 9, 2026

Merged via the queue into main with commit 17a09ee Jul 9, 2026
104 checks passed

remi-or deleted the fix-deepgemm-integration branch July 9, 2026 02:35

Uh oh!

Conversation

remi-or commented Jul 7, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

HuggingFaceDocBuilderDev commented Jul 7, 2026

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

IlyasMoutawwakil Jul 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

IlyasMoutawwakil left a comment

Choose a reason for hiding this comment

Uh oh!

remi-or commented Jul 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

IlyasMoutawwakil commented Jul 8, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

IlyasMoutawwakil left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Jul 9, 2026

CI recap

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

remi-or commented Jul 7, 2026 •

edited by github-actions Bot

Loading

IlyasMoutawwakil Jul 8, 2026 •

edited

Loading

remi-or commented Jul 8, 2026 •

edited

Loading