5090 - Random triton kernel crashes depending on batch size on MXFP4 (gpt-oss)
Author: cmp-nct
Pull request: no
Labels: name: bug

### System Info

I do not have a POC, it's happening in a project I am working on.
```
    return self.uncompiled_model.generate(
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/transformers/generation/utils.py", line 2629, in generate
    result = self._sample(
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/transformers/generation/utils.py", line 3610, in _sample
    outputs = self(**model_inputs, return_dict=True)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/transformers/utils/generic.py", line 959, in wrapper
    output = func(self, *args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/transformers/models/gpt_oss/modeling_gpt_oss.py", line 658, in forward
    outputs: MoeModelOutputWithPast = self.model(
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/transformers/utils/generic.py", line 1083, in wrapper
    outputs = func(self, *args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/transformers/models/gpt_oss/modeling_gpt_oss.py", line 491, in forward
    hidden_states = decoder_layer(
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/transformers/modeling_layers.py", line 94, in __call__
    return super().__call__(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/transformers/models/gpt_oss/modeling_gpt_oss.py", line 370, in forward
    hidden_states, _ = self.mlp(hidden_states)  # diff with llama: router scores
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/transformers/integrations/mxfp4.py", line 300, in mlp_forward
    routed_out = self.experts(hidden_states, routing_data, gather_idx, scatter_idx)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/transformers/integrations/mxfp4.py", line 190, in forward
    intermediate_cache1 = matmul_ogs(
  File "/root/.cache/huggingface/hub/models--kernels-community--triton_kernels/snapshots/1d2e9557ac0d4c651055a209055748d4db0fe65b/build/torch-universal/triton_kernels/matmul_ogs.py", line 444, in matmul_ogs
    opt_flags = make_opt_flags(out_dtype, x.dtype, w.dtype, precision_config,
  File "/root/.cache/huggingface/hub/models--kernels-community--triton_kernels/snapshots/1d2e9557ac0d4c651055a209055748d4db0fe65b/build/torch-universal/triton_kernels/matmul_ogs_details/opt_flags.py", line 297, in make_opt_flags
    return make_default_opt_flags_nvidia(*args)
  File "/root/.cache/huggingface/hub/models--kernels-community--triton_kernels/snapshots/1d2e9557ac0d4c651055a209055748d4db0fe65b/build/torch-universal/triton_kernels/matmul_ogs_details/opt_flags.py", line 214, in make_default_opt_flags_nvidia
    assert num_stages >= 1
AssertionError
```

I've had this happen based on prompt length when doing batched processing as well as depending on batch size.
So switching batch size can make it go away.
However, I also had it work randomly without changing any parameters strangely.

This is with triton 3.4.0 using CU 12.9 and torch 2.8.0, transformers 4.55.2
I've also had this on 4.56.dev0

I tried reinstalling triton kernels,  various combinations of versions.

Lastly - Very strangely this is not happening on my 4090 card.
I was under the impression that the 4090 does not even support the new 4 bit quantization, however it works (not defaulting to bf16) and I did not see that error on it.
When including my 5090 or single 5090 inference I get this error randomly - so it's probably related to the hopper kernel.

Sampling is turned off, flash attention not being used.

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Don't have a POC available.
I thought the stacktrace maybe is already helpful.

I tested it with python 3.10 and 3.12
Here my requirements.txt for 3.12:
```
transformers>=4.55.0
kernels>=0.9.0
datasets==4.0.0
bitsandbytes==0.46.1
trl>=0.20.0
peft>=0.17.0
lighteval==0.10.0
hf-transfer==0.1.8
hf_xet
tensorboard
liger-kernel==0.6.1
deepspeed==0.17.4
lm-eval[api]==0.4.9
Pillow
mlflow
sagemaker-mlflow==0.1.0
triton
git+https://github.com/triton-lang/triton.git@main#subdirectory=python/triton_kernels
pynvml
```



### Expected behavior

no crash

Comments:
Hmm, this is a pretty mysterious bug! I'm going to leave it up for now to see if anyone else is reporting it - if it turns out to be affecting multiple people then we can explore it in depth

> Hmm, this is a pretty mysterious bug! I'm going to leave it up for now to see if anyone else is reporting it - if it turns out to be affecting multiple people then we can explore it in depth

Thanks. I've literally spent two entire nights on this issue and my current solution is to hack the `opt_flags.py` and remove that assertion - that appears to solve it without (yet visible) sideeffects. Inference works though I've seen random sudden gibberish output in one single prompt generation within a batch (possibly not related though)

The platform is WSL(2) - that might matter.

I went through quite some steps in the meantime:
1) I updated the nvidia host driver and the host toolkit in case there is any interaction with that (no changes)
2) I changed toolkit to 12.8 12.9 and now 13.0 (no changes)
3) I tested it with a fresh python installation and different python versions (same behavior)
4) Tried to manually install kernels, manually compile triton and transformers (same behavior)

Removing the assertion removes the issue, so maybe the underlying calculation of num_stages is bugged.
Setting batch size to 1 appears to remove the issue - possibly it would still bug out at some point though. 
Or running it on the 4090, also fixes it.


There's an open issue for the MXFP4 kernels here - if nothing else helps, you may want to keep an eye on it and try any updated kernels https://github.com/huggingface/transformers/issues/40170

cc @MekkCyber and @danieldk if you have time to investigate!

Hi @cmp-nct, The kernels should run even on T4 GPUs: if the compute capability is insufficient, the weights are dequantized within the kernel. On 5090 GPUs, however, the kernels are expected to run without dequantizing the weights, which likely explains why they work reliably on 4090 but only inconsistently on 5090 since the behaviour is different. I will try to take a deeper look, but feel free to open an  issue on the `triton` github because all the kernels used are copied from : https://github.com/triton-lang/triton/tree/main/python/triton_kernels, especially the `matmul_ogs` which is the cause of the issue

I'm currently facing the same issue with the RTX 5090TI after ~520 tokens i get this assertion error

> I'm currently facing the same issue with the RTX 5090TI after ~520 tokens i get this assertion error

Try install this PR manually:
https://github.com/huggingface/transformers/pull/40358

It would be helpful if you can confirm if it solves or doesn't solve the issue.
It appears to have solved mine but I had very irregular behavior with those assertions

> > I'm currently facing the same issue with the RTX 5090TI after ~520 tokens i get this assertion error
> 
> Try install this PR manually: [#40358](https://github.com/huggingface/transformers/pull/40358)
> 
> It would be helpful if you can confirm if it solves or doesn't solve the issue. It appears to have solved mine but I had very irregular behavior with those assertions

After upgrading to Transformers v4.56.0, the issue was resolved. However, function calling is no longer being triggered. I'm not sure whether this is due to Transformers or my code, but I’ll look into it. Thanks for your help!

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.