[Feature] Support gpt-oss-20b model

[sii-xinglong]

Checklist

• If the issue you raised is not a feature but a question, please raise a discussion at https://github.com/sgl-project/sgl-jax/discussions/new/choose Otherwise, it will be closed.
• Please use English, otherwise it will be closed.

Motivation

https://huggingface.co/openai/gpt-oss-20b

Related resources

[sii-xinglong]

@SII-WeinanLiu will focus on it

[weinanliu]

1

[weinanliu]

Hi all, preliminary support for gpt-oss 20b is completed. The model definition is based on:
https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_oss/modeling_gpt_oss.py
https://github.com/openai/gpt-oss/blob/main/gpt_oss/torch/model.py

Special features of gpt-oss 20b: sliding window, attention sinks, MXFP4 for MoE.

1. We have implemented support for sliding window and attention sinks in the native attention.
2. During weight loading, we dequantize the MXFP4 format weights to bf16.
3. Regarding precision alignment, we found significant absolute errors between the bf16 JAX implementation and the PyTorch implementation. Therefore, we decided to use the official PyTorch fp32 implementation as the reference for precision alignment. Ultimately, we found that using fp32 for attention calculations and bf16 for MoE calculations brings the layer-by-layer precision closest to the official PyTorch fp32 benchmark.
4. Currently, FlashAttention lacks the attention sinks feature. We have delegated the implementation of this feature to Google engineers. The lack of FlashAttention makes the speed very, very slow, so we have not conducted large-scale end-to-end precision tests.

Please review!