feat(models): add native GPT-OSS implementation

[S1ro1]

Summary

Replace the Transformers/hub-kernel GPT-OSS runtime with a Prime-owned model, FA4 attention with learned sinks, and canonical grouped experts. BF16 GPT-OSS checkpoints load through lossless expert-layout conversion; inference policy updates route prepared expert and sink tensors through vLLM's registered parameter loaders.

This branch targets main and supports runtime QKV and gate/up fusions, canonical checkpointing, and the CP capability interface. Main's dependency pins are retained; the obsolete gpt-oss hub-kernel extra is removed.

Supported configuration

• model.impl = "auto" selects native GPT-OSS. Validated checkpoint: unsloth/gpt-oss-20b-BF16. Original MXFP4 checkpoint loading is outside this change.
• Direct FA4 calls preserve learned sinks, sliding/full attention, and packed-document boundaries. FA4 explicitly returns (output, lse); FA2/FA3 default calls return a tensor.
• Ring and Ulysses CP preserve sink gradients; common expert parallelism and runtime fusions are supported.
• Hopper is the validation target. The vLLM loader patch handles prepared expert and sink tensors with the pinned vLLM.

GPT-OSS and shared FA4 attention no longer use torch._dynamo.disable. There is no forwarding wrapper around the native FA4 function. Native configuration/input guards and HF GPT-OSS execution-comparison tests are removed; checkpoint conversion and CP checks remain. There are no example or skill changes relative to main.

Context-parallel attention

Every CP helper replaces an attention compute method. GPT-OSS's ring adapter is separate because the shared ring custom operators have no learned-sink input or sink-gradient output. It uses differentiable K/V gathers and FA4 to preserve those gradients. Its Ulysses adapter passes the model's FA4 function, scale, and sinks into the shared sink-aware Ulysses implementation. That thin adapter is an organizational choice; directly assigning the existing generic closure would omit these arguments.

Validation

• H200 full unit suite after cleanup: 701 passed, 11 skipped. After the final direct-FA4 assignment, the targeted checkpoint roundtrip and compiled attention/CP checks passed again.
• Cold torch.compile(fullgraph=False) forward/backward and eager-output agreement passed for GPT-OSS and shared FA2/FA3/FA4 attention.
• Two-GPU eager and compiled GPT-OSS ring/Ulysses output and Q/K/V/sink-gradient checks passed. Compiled shared FA4 Ulysses output and Q/K/V gradients passed on both ranks.
• The CP tests reuse one process group across cases, avoiding teardown/reinitialization reconnect failures.
• Earlier integration validation covered eight checkpoint/resume combinations and a real 20B export matching all 411 checkpoint keys/shapes and sampled tensor values.
• Ruff lint/format and git whitespace checks passed.

Final 20-step H200 RL run

After all code changes, uv run rl @ config.toml completed through the standard Slurm launcher on eight H200 GPUs: four trainer and four inference GPUs, unsloth/gpt-oss-20b-BF16, reverse-text, batch 64, group 8, sequence length 1024, completion budget 512, EP=4, QKV/gate-up fusions, compilation, selective activation checkpointing, and NCCL policy reloads. Rollout concurrency starts at 32 and is capped at 64.

Slurm job 326 completed with exit code 0, in 24m42s. All 20 optimizer steps completed and inference acknowledged 21 successful reloads (startup plus 20 updates). Rollout errors were 0%. The short completion budget truncates most reasoning; this validates numerical consistency and pipeline operation.

The table reports mismatch_kl/all/mean on commit c3fe30c29.

Completed steps: 20
Mean of step means: 0.0014423477
Min: 0.0007825278
Max: 0.0027635854

Step	Mismatch KL
1	0.0007825278
2	0.0008648058
3	0.0010610037
4	0.0015761884
5	0.0012348450
6	0.0023151361
7	0.0016181732
8	0.0012400812
9	0.0010534789
10	0.0017238446
11	0.0027635854
12	0.0023905332
13	0.0015134896
14	0.0014993388
15	0.0011815884
16	0.0012801469
17	0.0012795344
18	0.0012338012
19	0.0011749482
20	0.0010599031

---

Note

Medium Risk
Touches core training attention/MoE and vLLM weight loading for a new architecture; mistakes could affect numerics or checkpoint load, though conversion and CP tests mitigate this.

Overview
Replaces the Transformers / hub-kernel GPT-OSS path with a Prime-owned gpt_oss stack: custom config, packed-batch forward (seq_lens), grouped MoE, and FlashAttention 4 attention with learnable sink parameters wired through ring and Ulysses context parallelism (including sink gradients).

Training/runtime: Drops Hopper-only checks and explicit hub-kernel enablement; hub kernels stay disabled globally. Shared FA4 call sites now handle (output, lse) returns and drop torch._dynamo.disable wrappers. Ulysses varlen attention gains optional learnable_sink sharding.

Inference: Registers a vLLM startup patch to load prepared BF16 MoE expert and .sinks tensors via weight loaders until upstream vLLM catches up.

Packaging/docs: Removes the gpt-oss optional extra (kernels dependency); docs point at BF16 checkpoints (e.g. unsloth/gpt-oss-20b-BF16) and SM90/SM100+ for FA4 sinks. Tests add 2-GPU ring/Ulysses CP checks and drop HF MoE parity.

^{Reviewed by Cursor Bugbot for commit c3fe30c. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

Want reviews to match your repository better? Bugbot Learning can learn team-specific rules from PR activity. A team admin can enable Learning in the Cursor dashboard.

^{Reviewed by Cursor Bugbot for commit 4097a00. Configure here.}

[garrett361]

What exactly was the problem here? The change in dtype was causing issues?

[garrett361]

why is this early return and the one below correct? Not obvious, to me.

[garrett361]

LGTM. Checked the CP sink handling for ulysses, in particular.

Just had some minor questions.