gpt-oss-20B is OpenAI's open-weight MoE: 24 layers, 32 experts per layer with 4 active per token, and an attention stack that strictly alternates between sliding-window layers and full-attention layers. This is a components study rather than a full training run: I reimplemented the two places where the architecture deviates from a vanilla decoder, the attention pattern and the expert layer, in plain PyTorch with manual attention.
The config is copied from the official release, and I kept even the fields this file never reads: 64 query heads over 8 KV heads at head dimension 64, hidden size 2880, vocabulary 201,088, sliding window 128, RoPE theta 150,000, a SwiGLU clamp at 7.0, and the YaRN long-context parameters. A config that mirrors the release is documentation: it tells you exactly what the production model has that the study version does not.
Sliding-window attention is one line of masking
Half the layers see the full causal past; the other half see only the last 128 positions. Both are expressed as masks over the same score matrix. The banded mask composes triu and tril: take the lower triangle for causality, then cut everything older than the window.
if self.layer_type == "sliding_attention":
mask = torch.triu(torch.tril(torch.ones(T, T, device=x.device)),
diagonal=-self.config.sliding_window)
else:
mask = torch.tril(torch.ones(T, T, device=x.device))
scores = q @ k.transpose(-1, -2) # B, n_kv, repeat, T, T
scores = scores / math.sqrt(self.config.head_dim)
scores = scores.masked_fill(mask == 0, float("-inf"))
scores = F.softmax(scores, dim=-1)The layer_types tuple in the config drives which mask each layer gets, so the alternation is data, not control flow. Grouped queries use the same broadcast trick as the rest of my reimplementations: queries reshaped to a 5-D tensor over 8 KV groups, keys and values left singleton, and matmul broadcasting does the grouping.
RoPE in the rotation-matrix convention
Positional encoding is rotary, written here in the rotation-matrix sign convention over interleaved even and odd channels. Frequencies and angles are recomputed per forward pass, which keeps the code stateless and readable.
@staticmethod
def apply_rope(x, cos, sin):
x_even = x[..., 0::2] # B, T, C//2
x_odd = x[..., 1::2] # B, T, C//2
u = x_even * cos + x_odd * sin
v = -x_even * sin + x_odd * cos
x_rotated = torch.stack([u, v], dim=-1)
return x_rotated.flatten(-2) # B, T, CThe expert layer: top-4 routing and a clamped SwiGLU
The router keeps the top 4 of 32 expert logits per token and softmaxes over only those 4, matching the release. gpt-oss bounds its SwiGLU activations; this version applies the clamp to the gated product with the config's limit of 7.0. Dispatch loops over experts and scatters weighted outputs back with index_add_, clarity over throughput.
u = self.fc1[expert_id](x_selected) # tokens routed to this expert
v = self.fc2[expert_id](x_selected)
act = F.silu(u) * v
if self.config.swiglu_limit is not None:
act = torch.clamp(act, -self.config.swiglu_limit, self.config.swiglu_limit)
h = self.fc3[expert_id](act)
y_flat.index_add_(0, expert_tokens, expert_weights.unsqueeze(1) * h)Condensed from the MoE forward
What the release adds on top
Three gpt-oss features stay config-only in this study: attention sinks (learned logits that keep early positions attendable in long contexts), YaRN frequency scaling for the 131k context (the rope_factor and beta fields are all present, the RoPE here is vanilla), and mxfp4 quantization of the expert weights (the quant_method and skip-list fields document which modules the release leaves in higher precision). Each is a good exercise on its own; the point of this file is that the skeleton of the model, mask plus router plus clamp, fits in an afternoon of PyTorch.