Tricks from OpenAI gpt-oss YOU 🫵 can use with transformers

transformers , we have upgraded the library considerably. The updates make it very efficient to load, run, and fine-tunethe models. In this blog post, we talk about all the upgrades in-depth, and how they become part of the transformers toolkit so other models (current and future) can benefit from them. Providing clean implementations of new methods in transformers also allows the community to quickly understand and adopt them. Frameworks such as MLX , llama.cpp or vLLM can use the transformers code as a reference to build their own implementations. For this release, we worked on: - Zero-build Kernels, downloadable from the Hub - MXFP4 Quantization - Tensor Parallelism - Expert Parallelism - Dynamic Sliding Window Layer & Cache - Continuous Batching & Paged Attention - Load larger models faster Best part: Most of these features should work across all major models within transformers ! A kernel is a * specialized*, compact program that runs on accelerators to execute tasks like matrix multiplications, activations, or normalizations. In eager PyTorch, operations trigger individual kernels sequentially, which is straightforward but can incur extra memory transfers and launch overheads. PyTorch 2.0's torch.compile with backends like TorchInductor addresses this by automatically fusing and optimizing kernels, delivering 2–10× performance gains.In addition, the community has created custom kernels for frequent combinations of operations, not just individual PyTorch ops like matmul. For example, Flash Attention was created to optimize the critical attention block that defines the transformers architecture, and is present in many models including most LLMs. By carefully combining all the attention operations inside a single kernel, memory transfers are minimized, memory use is reduced, and speedups can be achieved. The problem is that all these various kernels are available in separate libraries, which creates a dependency bloat if they were to be added to the transformers library. Furthermore, these kernels are not just Python code, they consist of low-level cuda code, glued together with C++ and exposed through a Python layer. This means they have to be compiled in the target system, which in turn requires whatever build system is required by each kernel library. The kernels package solves this problem by downloading pre-built binaries of supported kernels from the Hub. You just indicate the kernel you want to use, and kernels will look for a version compatible with your system and download it on first use. GPT-OSS, a Mixture of Experts (MoE) model, is a big user of Kernels from the Hub. It leverages several custom kernels: - Liger RMSNorm, used as @use_kernel_forward_from_hub("RMSNorm") - Megablocks MoE kernels: @use_kernel_forward_from_hub("MegaBlocksMoeMLP") - Flash Attention 3 with support for attention sinks. - MXFP4 triton kernels (covered later) Let's take a look at the first two ones. Behind the scenes, the decorators (1 and 2) simply point to community-contributed kernels. For example, RMSNorm comes from liger_kernels , while the MegaBlocksMoeMLP kernel comes from megablocks . Depending on your device (CUDA or ROCm) and whether you’re training or running inference, the right kernel is pulled in automatically. This design is both specific and general: the RMSNorm liger kernels are already being reused across multiple models, and the MoE kernel could be applied to future MoEs as well. Because kernels pulls code from the Hub, you have to opt-in to this feature by passing use_kernels=True in your model instantiation, as shown below. We enable INFO logging in the example so you can easily verify that downloadable kernels are in use. These kernels are not compatible with mxfp4 , so inference will happen inbfloat16 if you use them. Please, benchmark your system for the best combination in memory and throughput that suits your project! from transformers import AutoTokenizer, AutoModelForCausalLM import logging logging.basicConfig(level=logging.INFO) model_id = "openai/gpt-oss-20b" tokenizer = AutoTokenizer.from_pretrained(model_id) model = AutoModelForCausalLM.from_pretrained( model_id, dtype="auto", device_map="auto", use_kernels=True, ) Running a quick generation yields log messages like INFO:root:Using layer `LigerRMSNorm` from repo `kernels-community/liger_kernels` INFO:root:Using layer `MegaBlocksMoeMLP` from repo `kernels-community/megablocks` Figure 1 shows that, in the system we tested, these kernels work best for larger batch sizes. We always recommend to benchmark any performance-related changes as closely to your production conditions as possible. You can explore and play with the benchmarking script here OpenAI gpt-oss models use attention sinks, which improves quality and facilitates the use of longer contexts. The vLLM team added this feature to the latest version of Flash Attention (Flash Attention 3), and the resulting custom kernel is available on the Hub. Currently, this kernel is compatible with the Hopper architecture. If you have one, this is the way to enable it: model = AutoModelForCausalLM.from_pretrained( model_id, dtype="auto", device_map="auto", + # Flash Attention with Sinks + attn_implementation="kernels-community/vllm-flash-attn3", ) Large language models are memory-hungry. Quantization reduces memory footprint by storing weights (and sometimes activations) in lower-precision formats. For reference, FP32 uses 32 bits per number and BF16 uses 16. By reducing bit width, we trade some precision for smaller models and faster memory movement. If you want a visual primer on quantization trade-offs, Maarten Grootendorst’s article is excellent: A Visual Guide to Quantization. MXFP4 is a 4-bit floating format with E2M1 layout: 1 sign bit, 2 exponent bits, and 1 mantissa bit, as shown in Figure 2. On its own, E2M1 is very coarse. MXFP4 compensates with blockwise scaling: - Vectors are grouped into blocks of 32 elements. - Each block stores a shared scale that restores dynamic range when dequantizing. - Inside each block, 4-bit values represent numbers relative to that scale. This blockwise scheme lets MXFP4 keep range while using very few bits. In practice, GPT-OSS 20B fits in roughly 16 GB of VRAM and GPT-OSS 120B fits in roughly 80 GB when MXFP4 is active, which is the difference between “cannot load” and “can run on a single GPU.” The catch is that matrix multiplies now have to respect block scales. Doing this efficiently at scale requires dedicated kernels. transformers now includes native support for MXFP4, leveraging optimized triton (MXFP4) kernels for enhanced performance. This builds on the community-driven kernel distribution discussed earlier, utilizing pre-compiled kernels from the Hub to simplify deployment. Key implementation details: - Quantizer logic: Found in the MXFP4 quantizer file, this handles the core quantization process for MXFP4. - Integration hooks: The MXFP4 integration file enables seamless use of MXFP4 within the transformers framework. To check if a model supports MXFP4 , inspect its configuration: from transformers import GptOssConfig model_id = "openai/gpt-oss-120b" cfg = GptOssConfig.from_pretrained(model_id) print(cfg.quantization_config) # Example output: # { # 'modules_to_not_convert': [ # 'model.layers.*.self_attn', # 'model.layers.*.mlp.router', # 'model.embed_tokens', # 'lm_head' # ], # 'quant_method': 'mxfp4' # } If 'quant_method': 'mxfp4' is present, the model will automatically use the MXFP4 pathway with Triton kernels when supported. Thanks to this pull request, you can fine-tune gpt-oss models and save them directly to the Hub in MXFP4 format, streamlining deployment with optimized performance. To run MXFP4 on GPU you need: accelerate ,kernels , andtriton>=3.4 installed.