The Faster Transformers Stack Behind GPT-OSS

Part 2 of 4 in the LLM fine-tuning series.

The Hugging Face post Tricks from OpenAI gpt-oss YOU can use with transformers is easy to read as release excitement.

That would undersell it.

The more useful reading is that transformers absorbed several practical runtime ideas that used to require separate serving stacks, custom forks, or hand-built CUDA/Triton work. The features were accelerated by GPT-OSS, but many of them point beyond GPT-OSS: downloadable kernels, MXFP4 quantization, tensor parallelism, expert parallelism, dynamic sliding-window KV cache, continuous batching, and faster model loading.

That list sounds like infrastructure noise until you try to run a large model and hit the wall.

The wall may be memory. It may be cold-start time. It may be KV cache growth on long prompts. It may be static batches wasting GPU time. It may be a model that needs multiple GPUs for compute, not just for storage. It may be an MoE model where experts should be sharded differently from dense layers.

The runtime stack decides whether the model is usable.

The Runtime Map#

The features fit into one mental model:

The failure mode I see in teams is treating these as independent toggles. Someone enables a kernel because it sounds faster. Someone else uses quantization because the model barely fits. A third person turns on tensor parallelism because there are multiple GPUs. Nobody checks whether the choices are compatible, measurable, or relevant to the workload.

The real tradeoff is not "fast versus slow." It is memory footprint, implementation trust, hardware support, communication overhead, batch shape, and production maturity.

Downloadable Kernels From the Hub#

Hugging Face added opt-in downloadable kernels through the Hub. The idea is direct: instead of forcing users to compile every low-level optimization locally, the kernels package can fetch prebuilt CUDA or Triton kernels when the runtime supports them.

For GPT-OSS, the post calls out kernels such as Liger RMSNorm, MegaBlocks MoE MLP, Flash Attention 3 with attention sinks, and MXFP4 Triton kernels. That matters because modern LLM runtime cost is often dominated by repeated hot paths: normalization, attention, and MoE feed-forward layers.

The code shape is small:

from transformers import AutoModelForCausalLM, AutoTokenizer
 
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,
)

The caveat is not small. Downloadable kernels are an execution trust boundary. They are opt-in and logged, which is good. In regulated or locked-down environments, the team should audit what is pulled, pin versions, and treat kernel source/provenance as part of the deployment artifact.

There is also a compatibility caveat from the Hugging Face post: the custom Hub kernels they discuss are not compatible with MXFP4 in that path, so inference falls back to bfloat16 if those are combined. That means use_kernels=True is not automatically better than MXFP4. Benchmark both for the actual batch and sequence profile.

MXFP4 Is a Memory Decision#

MXFP4 is a 4-bit floating format with blockwise scaling. The Hugging Face post explains it as E2M1 values grouped into blocks of 32 with a shared scale. The practical result for GPT-OSS is dramatic on supported hardware: GPT-OSS 20B can fit in roughly 16 GB of VRAM and GPT-OSS 120B in roughly 80 GB when the MXFP4 path is active.

The model card for openai/gpt-oss-20b also shows the operational implication: the model can be used with Transformers, vLLM, Ollama, LM Studio, and other paths, but the runtime path determines memory and compatibility.

You can inspect the config:

from transformers import GptOssConfig
 
cfg = GptOssConfig.from_pretrained("openai/gpt-oss-20b")
print(cfg.quantization_config)

The important check is whether the model declares an MXFP4 quantization method and whether the machine can actually run that path. Hugging Face's post names the requirements: accelerate, kernels, triton>=3.4, and an NVIDIA GPU with compute capability at least 7.5. If the constraints are not met, the runtime falls back to a higher-precision path, with much larger memory usage.

That fallback is exactly the kind of thing I want in an ops checklist:

python - <<'PY'
import torch
print("cuda:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("device:", torch.cuda.get_device_name(0))
    print("capability:", torch.cuda.get_device_capability(0))
PY
 
hf cache scan

If the model silently loads in bfloat16, the memory story changes. The team should notice before the benchmark.

Tensor Parallelism Is Not Device Map#

device_map="auto" places model components across devices to fit memory.

Tensor parallelism splits tensors inside layers across GPUs and coordinates the collective operations needed to compute each layer. Those are different things.

The new Transformers path lets supported models use tensor parallelism directly from from_pretrained:

from transformers import AutoModelForCausalLM, AutoTokenizer
 
model_id = "openai/gpt-oss-120b"
tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side="left")
 
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    tp_plan="auto",
    dtype="auto",
).eval()
 
print(model._tp_plan)

Run this with torchrun on the target number of processes.

Tensor parallelism helps when the model is too large for one GPU or when the team needs parallel compute on the same layer. It is communication-heavy. It works best on a single node with fast interconnects. Across slow links, the collectives can eat the gain.

The failure mode is using TP as a default multi-GPU badge. If a smaller model fits cleanly and the batch is small, TP can add overhead without enough payoff. If the workload is long-context or large-batch, TP is more likely to earn its complexity.

Expert Parallelism Fits MoE Models#

GPT-OSS is a mixture-of-experts model. MoE changes the runtime question because not every token uses every expert. Expert parallelism shards experts across ranks and routes tokens to the experts they need.

Transformers exposes this through the distributed config:

from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.distributed import DistributedConfig
 
model_id = "openai/gpt-oss-120b"
tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side="left")
 
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    distributed_config=DistributedConfig(enable_expert_parallel=True),
    dtype="auto",
).eval()

The Hugging Face post notes that enabling expert parallelism also enables tensor parallelism for GPT-OSS. That makes sense operationally: MoE experts and dense layer shards both need distributed execution.

The tradeoff is that MoE serving is harder to reason about than dense-model serving. Routing, load balance, communication, and batch shape matter. If you are only experimenting, the built-in path is a gift. If you are serving production traffic, you still need to compare against purpose-built serving stacks.

Dynamic Sliding-Window KV Cache#

Long-context models often fail by memory, not by weights.

During generation, the model stores key/value tensors so it does not recompute the full past on every token. That KV cache grows with sequence length unless the model architecture allows something smarter. Sliding-window attention only needs a recent window for some layers, so cache memory should not grow forever for those layers.

Transformers added a DynamicSlidingWindowLayer and config-aware DynamicCache. If the model config declares sliding or hybrid attention, the cache can stop growing beyond the window for sliding layers.

The explicit version looks like this:

from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache
 
model_id = "openai/gpt-oss-20b"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    dtype="auto",
    device_map="auto",
).eval()
 
messages = [{"role": "user", "content": "Explain KV caching in one paragraph."}]
inputs = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    return_tensors="pt",
    return_dict=True,
).to(model.device)
 
cache = DynamicCache(config=model.config)
outputs = model.generate(
    **inputs,
    past_key_values=cache,
    max_new_tokens=256,
)

For many users, this works by default when the config declares the pattern. The explicit code is useful because it makes the cache a visible runtime artifact.

The failure mode is benchmarking only short prompts. A KV-cache improvement may barely show up in tiny examples and matter a lot when prompts or generations are long.

Continuous Batching Is for Utilization#

Static generation batches waste GPU time because requests finish at different lengths. The slowest sequence holds the batch hostage.

Continuous batching fills newly freed slots with new requests. Hugging Face added a generate_batch API for this kind of experimentation. Their own caveat is important: use it for evaluation and experiments, not as a production-grade online serving replacement.

That is the right boundary.

For offline evals, continuous batching can make throughput tests more realistic. For production serving, frameworks such as vLLM and SGLang are still designed around the larger serving problem: scheduling, paged attention, multi-tenant traffic, monitoring, request cancellation, backpressure, and deployment shape.

The practical rule:

offline eval throughput -> try continuous batching in Transformers
online production serving -> evaluate vLLM, SGLang, TGI, NIM, or a managed endpoint

Do not confuse a useful library API with a complete serving platform.

Faster Model Loading Is a Cold-Start Feature#

Large models can spend surprising time allocating memory during load. Hugging Face describes a pre-allocation change that reserves larger GPU memory blocks per device map before copying weights, avoiding thousands of tiny allocations.

There is no new user-facing code:

model = AutoModelForCausalLM.from_pretrained(
    "openai/gpt-oss-20b",
    dtype="auto",
    device_map="auto",
)

The point is operational. Cold-start time affects development loops, autoscaling behavior, ephemeral jobs, and eval throughput. It does not make a bad serving architecture good, but it reduces friction in the normal from_pretrained path.

How I Would Use the Stack#

For a single NVIDIA GPU, I would start with MXFP4 if the hardware and package versions support it. If custom kernels look attractive, benchmark bf16 + use_kernels=True against MXFP4 instead of assuming they compose.

For a same-node multi-GPU box, I would test tp_plan="auto" when the model needs parallel compute. For GPT-OSS-style MoE models, I would evaluate expert parallelism when the model and workload justify it.

For long prompts and generations, I would check KV cache behavior explicitly. If memory climbs linearly when the model config says it should not, something is wrong.

For offline evaluations, I would try continuous batching to avoid measuring an artificially wasteful static batch path.

For online serving, I would still ask a separate question: what stack owns traffic, scheduling, observability, and rollback?

The Hardware Boundary#

The biggest mistake is importing the NVIDIA mental model onto every machine.

MXFP4 Triton kernels, Flash Attention 3, and many downloadable CUDA kernels are not the Mac path. They belong to the NVIDIA/ROCm side of the story. Apple Silicon has its own stack: PyTorch MPS, MLX, Metal, GGUF runners, and different quantization tradeoffs.

That boundary becomes the next post.

Related Notes#

• NVFP4 and the Infrastructure Meaning of Precision
• Fine-Tuning Open Source LLMs With NVIDIA NeMo

Part 2 of 4. Previous: Fine-Tuning LLMs Is an Operating Loop, Not a Training Command. Next: Fine-Tuning LLMs on a MacBook Pro With MPS and MLX.