Transformer serve gpt oss 20b cannot support multi-turn chat
Author: Taimin
Pull request: no
Labels: name: bug

### System Info

I installed the environment to run gpt-oss-20b using the following command:

> uv pip install --pre vllm==0.10.1+gptoss \
    --extra-index-url https://wheels.vllm.ai/gpt-oss/ \
    --extra-index-url https://download.pytorch.org/whl/nightly/cu128 \
    --index-strategy unsafe-best-match

My GPU is 5090 and vllm still has difficulty in support it so I tried transformers serve openai/gpt-oss-20b. I cna use the transformer chat command to chat with the model for one turn but as long as I entered the second turn then the following error will occur:

> Traceback (most recent call last):
  File "/home/yang/.local/share/uv/python/cpython-3.12.10-linux-x86_64-gnu/lib/python3.12/threading.py", line 1075, in _bootstrap_inner
    self.run()
  File "/home/yang/.local/share/uv/python/cpython-3.12.10-linux-x86_64-gnu/lib/python3.12/threading.py", line 1012, in run
    self._target(*self._args, **self._kwargs)
  File "/home/yang/.gptoss/lib/python3.12/site-packages/transformers/commands/serving.py", line 947, in generate_with_cache
    generate_output = model.generate(**kwargs)
                      ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yang/.gptoss/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/yang/.gptoss/lib/python3.12/site-packages/transformers/generation/utils.py", line 2629, in generate
    result = self._sample(
             ^^^^^^^^^^^^^
  File "/home/yang/.gptoss/lib/python3.12/site-packages/transformers/generation/utils.py", line 3610, in _sample
    outputs = self(**model_inputs, return_dict=True)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yang/.gptoss/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yang/.gptoss/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yang/.gptoss/lib/python3.12/site-packages/transformers/utils/generic.py", line 959, in wrapper
    output = func(self, *args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yang/.gptoss/lib/python3.12/site-packages/transformers/models/gpt_oss/modeling_gpt_oss.py", line 658, in forward
    outputs: MoeModelOutputWithPast = self.model(
                                      ^^^^^^^^^^^
  File "/home/yang/.gptoss/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yang/.gptoss/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yang/.gptoss/lib/python3.12/site-packages/transformers/utils/generic.py", line 1083, in wrapper
    outputs = func(self, *args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yang/.gptoss/lib/python3.12/site-packages/transformers/models/gpt_oss/modeling_gpt_oss.py", line 491, in forward
    hidden_states = decoder_layer(
                    ^^^^^^^^^^^^^^
  File "/home/yang/.gptoss/lib/python3.12/site-packages/transformers/modeling_layers.py", line 94, in __call__
    return super().__call__(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yang/.gptoss/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yang/.gptoss/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yang/.gptoss/lib/python3.12/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 "/home/yang/.gptoss/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yang/.gptoss/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yang/.gptoss/lib/python3.12/site-packages/transformers/integrations/mxfp4.py", line 300, in mlp_forward
    routed_out = self.experts(hidden_states, routing_data, gather_idx, scatter_idx)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yang/.gptoss/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yang/.gptoss/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yang/.gptoss/lib/python3.12/site-packages/transformers/integrations/mxfp4.py", line 190, in forward
    intermediate_cache1 = matmul_ogs(
                          ^^^^^^^^^^^
  File "/home/yang/.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 "/home/yang/.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 "/home/yang/.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



### 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

uv venv
uv pip install --pre vllm==0.10.1+gptoss
--extra-index-url https://wheels.vllm.ai/gpt-oss/
--extra-index-url https://download.pytorch.org/whl/nightly/cu128
--index-strategy unsafe-best-match

transformers serve 
transformers chat localhost:8000 --model-name-or-path openai/gpt-oss-20b

### Expected behavior

Should support multiturn chat. 

Comments:
Looks like something wacky is going on in the kernels - cc @gante for chat and @sunmarc for the triton kernel

yes, at first glance seems like a kernel issue 👀 (cc @SunMarc )

Hi,

I've been looking into the reported AssertionError: assert num_stages >= 1 that happens during multi-turn chat generation with the GPT-OSS-20B model using MXFP4 quantization.

The error occurs during the second turn and later because the model input tensor shape changes. In the first turn, it uses a longer 3D tensor (batch_size, seq_len, hidden_dim). In subsequent turns, with key-value caching, it becomes a shorter tensor with seq_len=1 to speed up generation. This shape change causes the Triton kernel’s internal autotuner to create an invalid configuration, leading to the assertion failure in kernel optimization flags.

I suggest changing the mlp_forward function in mxfp4.py to explicitly check if the input tensor is 3D (first turn) or effectively 2D (subsequent turns). We can then reshape the input tensor to a consistent shape before sending it to the Triton MXFP4 kernel and reshape the output back to the expected dimensions.

This method prevents the kernel from receiving unexpected input shapes. It makes the MXFP4 integration strong against different input shapes during multi-turn generation, avoiding the invalid num_stages kernel optimization configuration.

<h3>
I am taking this issue and will create PR. This will be my first PR in this repo. Please tell me if you have anything or do you think I need to change my approach or you don't want PR for this. I feel that this issue is not confirmed by you yet to solve, because I am not seeing any discussion on this. </h3>

@PrathmeshAdsod thank you for having a look at the issue 🙏 I've tagged the appropriate reviewer there :)

Hey @Taimin, can you retry on your side ? On my side, I'm getting an error on the cache instead https://github.com/huggingface/transformers/pull/40358#issuecomment-3228885751

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.