[ops] feat: Support GPT-OSS learnable attention sinks

examples/moe/run_moonlight_gsm8k.sh

@@ -93,6 +93,7 @@ RAY_DEDUP_LOGS=1 PYTHONUNBUFFERED=1 python3 -m verl.trainer.main_ppo \
     actor_rollout_ref.rollout.max_num_batched_tokens=4096 \
     ++actor_rollout_ref.rollout.engine_kwargs.vexact.max_cache_blocks=$vexact_max_cache_blocks \
     ++actor_rollout_ref.rollout.engine_kwargs.vexact.attn_impl=$INFER_FA_IMPL \
+    ++actor_rollout_ref.rollout.engine_kwargs.vexact.moe_implementation=$moe_implementation \
     actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \
     actor_rollout_ref.rollout.disable_log_stats=False \
     actor_rollout_ref.rollout.free_cache_engine=True \

tests/batch_invariant_ops/test_flash_attention.py

@@ -15,7 +15,11 @@
 import pytest
 import torch
 
-from vexact.batch_invariant_ops.flash_attention import _get_fa_window_size, flash_attention_forward
+from vexact.batch_invariant_ops.flash_attention import (
+    _get_fa_learnable_sink,
+    _get_fa_window_size,
+    flash_attention_forward,
+)
 from vexact.batch_invariant_ops.kv_cache_context import set_kv_cache_context
 
 
@@ -37,6 +41,21 @@ def test_get_fa_window_size_rejects_non_positive_values(sliding_window):
         _get_fa_window_size(sliding_window)
 
 
+def test_get_fa_learnable_sink_casts_to_query_dtype():
+    query = torch.empty(1, dtype=torch.bfloat16)
+    sink = torch.randn(2, dtype=torch.float32)
+
+    actual = _get_fa_learnable_sink(sink, query, use_cute=True)
+
+    assert actual.dtype == query.dtype
+    torch.testing.assert_close(actual.float(), sink.to(torch.bfloat16).float())
+
+
+def test_get_fa_learnable_sink_rejects_fa3():
+    with pytest.raises(NotImplementedError, match="require the FA4"):
+        _get_fa_learnable_sink(torch.zeros(2), torch.empty(1), use_cute=False)
+
+
 @pytest.mark.parametrize("use_cute", [False, True], ids=["fa3", "fa4"])
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 def test_flash_attention_sliding_window_matches_reference(use_cute):
@@ -97,3 +116,62 @@ def test_flash_attention_sliding_window_matches_reference(use_cute):
     expected = expected.transpose(1, 2).reshape_as(actual)
 
     torch.testing.assert_close(actual.float(), expected, rtol=2e-2, atol=2e-2)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_fa4_learnable_sink_matches_reference():
+    if torch.cuda.get_device_capability()[0] < 9:
+        pytest.skip("FA4 requires SM90+")
+    pytest.importorskip("flash_attn.cute")
+
+    torch.manual_seed(1)
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+    batch_size, seq_len, num_heads, head_dim = 1, 16, 2, 64
+    page_size = 256
+    sliding_window = 4
+    scaling = head_dim**-0.5
+
+    query = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype)
+    key = torch.randn_like(query)
+    value = torch.randn_like(query)
+    sinks = torch.randn(num_heads, device=device, dtype=dtype)
+    key_cache = torch.zeros(1, page_size, num_heads, head_dim, device=device, dtype=dtype)
+    value_cache = torch.zeros_like(key_cache)
+    set_kv_cache_context(
+        is_paged_attn=True,
+        key_cache={0: key_cache},
+        value_cache={0: value_cache},
+        block_tables=torch.tensor([[0]], device=device, dtype=torch.int32),
+        context_lens=torch.tensor([seq_len], device=device, dtype=torch.int32),
+        slot_mapping=torch.arange(seq_len, device=device, dtype=torch.int64),
+        query_start_loc=torch.tensor([0, seq_len], device=device, dtype=torch.int32),
+        max_seqlen_q=seq_len,
+    )
+
+    module = torch.nn.Module()
+    module.layer_idx = 0
+    actual, _ = flash_attention_forward(
+        module,
+        query,
+        key,
+        value,
+        attention_mask=None,
+        scaling=scaling,
+        sliding_window=sliding_window,
+        s_aux=sinks,
+        use_cute=True,
+    )
+
+    scores = torch.matmul(query.float(), key.float().transpose(-1, -2)) * scaling
+    positions = torch.arange(seq_len, device=device)
+    query_positions = positions[:, None]
+    key_positions = positions[None, :]
+    mask = (key_positions <= query_positions) & (key_positions > query_positions - sliding_window)
+    scores.masked_fill_(~mask, float("-inf"))
+    sink_scores = sinks.float().view(1, num_heads, 1, 1).expand(batch_size, -1, seq_len, -1)
+    attention_probs = torch.softmax(torch.cat((scores, sink_scores), dim=-1), dim=-1)[..., :-1]
+    expected = torch.matmul(attention_probs, value.float())
+    expected = expected.transpose(1, 2).reshape_as(actual)
+
+    torch.testing.assert_close(actual.float(), expected, rtol=2e-2, atol=2e-2)

tests/test_gpt_oss_model.py

@@ -0,0 +1,92 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from transformers import AutoModelForCausalLM, GptOssConfig
+from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts, GptOssForCausalLM
+
+from vexact.inferencer.model_loader import init_on_device_without_buffers
+from vexact.models.gpt_oss import modeling_gpt_oss
+
+
+def test_gpt_oss_experts_use_fused_quack(monkeypatch):
+    captured = {}
+
+    def fake_fused_quack(**kwargs):
+        captured.update(kwargs)
+        return torch.tensor([1.0])
+
+    monkeypatch.setattr(modeling_gpt_oss, "_run_fused_quack", fake_fused_quack)
+    experts = SimpleNamespace(
+        num_experts=4,
+        gate_up_proj=torch.tensor([1.0]),
+        gate_up_proj_bias=torch.tensor([2.0]),
+        down_proj=torch.tensor([3.0]),
+        down_proj_bias=torch.tensor([4.0]),
+        alpha=1.702,
+        limit=7.0,
+    )
+    hidden_states = torch.tensor([[5.0]])
+    router_indices = torch.tensor([[1]])
+    routing_weights = torch.tensor([[0.5]])
+
+    output = modeling_gpt_oss.gpt_oss_experts_forward(experts, hidden_states, router_indices, routing_weights)
+
+    assert torch.equal(output, torch.tensor([1.0]))
+    assert captured == {
+        "num_experts": 4,
+        "routing_weights": routing_weights,
+        "selected_experts": router_indices,
+        "hidden_states": hidden_states,
+        "gate_up_proj": experts.gate_up_proj,
+        "gate_up_proj_bias": experts.gate_up_proj_bias,
+        "down_proj": experts.down_proj,
+        "down_proj_bias": experts.down_proj_bias,
+        "alpha": 1.702,
+        "limit": 7.0,
+    }
+
+
+def test_gpt_oss_experts_require_routing_inputs():
+    with pytest.raises(ValueError, match="requires router_indices and routing_weights"):
+        modeling_gpt_oss.gpt_oss_experts_forward(SimpleNamespace(), torch.tensor([[1.0]]))
+
+
+def test_gpt_oss_patch_replaces_experts_forward():
+    modeling_gpt_oss.apply_gpt_oss_patches()
+
+    assert GptOssExperts.forward is modeling_gpt_oss.gpt_oss_experts_forward
+
+
+def test_gpt_oss_uses_common_transformers_model_path():
+    config = GptOssConfig(
+        vocab_size=128,
+        hidden_size=64,
+        intermediate_size=32,
+        num_hidden_layers=1,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        num_local_experts=4,
+        num_experts_per_tok=2,
+        head_dim=16,
+    )
+    with init_on_device_without_buffers("meta"):
+        model = AutoModelForCausalLM.from_config(config)
+
+    assert isinstance(model, GptOssForCausalLM)
+    assert type(model).__module__.startswith("transformers.models.gpt_oss")
+    assert model.model.layers[0].mlp.experts.forward.__func__ is modeling_gpt_oss.gpt_oss_experts_forward

tests/test_model_config.py

@@ -0,0 +1,56 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from types import SimpleNamespace
+
+import pytest
+
+from vexact.config import ModelConfig
+
+
+def _gpt_oss_config():
+    return SimpleNamespace(model_type="gpt_oss", max_position_embeddings=128)
+
+
+def test_gpt_oss_requires_fa4_attention():
+    with pytest.raises(ValueError, match="learnable attention sinks require"):
+        ModelConfig(model_path="unused", attn_impl="fa-invariant", hf_config=_gpt_oss_config())
+
+
+def test_gpt_oss_accepts_fa4_attention():
+    config = ModelConfig(model_path="unused", attn_impl="fa-invariant-cute", hf_config=_gpt_oss_config())
+
+    assert config.attn_impl == "fa-invariant-cute"
+
+
+@pytest.mark.parametrize("moe_implementation", ["fused_triton", "eager"])
+def test_gpt_oss_rejects_non_quack_moe(moe_implementation):
+    with pytest.raises(ValueError, match="requires moe_implementation='fused_quack'"):
+        ModelConfig(
+            model_path="unused",
+            attn_impl="fa-invariant-cute",
+            moe_implementation=moe_implementation,
+            hf_config=_gpt_oss_config(),
+        )
+
+
+def test_gpt_oss_accepts_fused_quack_moe():
+    config = ModelConfig(
+        model_path="unused",
+        attn_impl="fa-invariant-cute",
+        moe_implementation="fused_quack",
+        hf_config=_gpt_oss_config(),
+    )
+
+    assert config.moe_implementation == "fused_quack"

vexact/batch_invariant_ops/README.md

@@ -61,7 +61,7 @@ batch_invariant_ops/
 
 Two attention backends, both supporting paged KV cache:
 
-- **Flash Attention** — Uses `flash_attn_with_kvcache` (FA3, SM90) or `flash_attn.cute` (FA4, SM90+/SM100). Both paths support causal sliding-window attention. Requires `is_paged_attn=True` in the KV cache context. Forces `num_splits=1` for determinism.
+- **Flash Attention** — Uses `flash_attn_with_kvcache` (FA3, SM90) or `flash_attn.cute` (FA4, SM90+/SM100). Both paths support causal sliding-window attention; FA4 also supports per-head learnable attention sinks used by GPT-OSS. Requires `is_paged_attn=True` in the KV cache context. Forces `num_splits=1` for determinism.
 - **Flex Attention** — Uses `torch.nn.attention.flex_attention` with custom `mask_mod` for causal masking. Supports both packed (non-paged) and paged modes. Paged mode builds a physical-to-logical block mapping and caches the block mask across layers.
 
 Both expect the KV cache context to be set via `set_kv_cache_context()` before the forward pass.

vexact/batch_invariant_ops/flash_attention.py

@@ -35,6 +35,17 @@ def _get_fa_window_size(sliding_window: Optional[int]) -> tuple[int, int]:
     return (window_bound, window_bound)
 
 
+def _get_fa_learnable_sink(
+    s_aux: Optional[torch.Tensor], query: torch.Tensor, use_cute: bool
+) -> Optional[torch.Tensor]:
+    """Translate Transformers' attention sink argument to the supported FA backend."""
+    if s_aux is None:
+        return None
+    if not use_cute:
+        raise NotImplementedError("Learnable attention sinks require the FA4 (fa-invariant-cute) backend")
+    return s_aux.to(dtype=query.dtype)
+
+
 @proton.scope("flash_attention_forward")
 def flash_attention_forward(
     module: nn.Module,
@@ -45,6 +56,7 @@ def flash_attention_forward(
     scaling: float,
     dropout: float = 0.0,
     sliding_window: Optional[int] = None,
+    s_aux: Optional[torch.Tensor] = None,
     use_cute: bool = False,
     **kwargs,
 ):
@@ -70,6 +82,7 @@ def flash_attention_forward(
         dropout (float): The dropout rate. Not currently used.
         sliding_window (Optional[int]): The number of tokens visible in causal local-attention layers,
             including the current token. ``None`` selects full causal attention.
+        s_aux (Optional[torch.Tensor]): Per-head learnable attention sink logits. Supported by FA4 only.
 
     Returns:
         Tuple[torch.Tensor, None]: A tuple containing the attention output and None.
@@ -150,6 +163,7 @@ def flash_attention_forward(
     # max_seqlen_k = context_lens.max().item()
     from vexact.utils.device import DEVICE_MAJOR
 
+    learnable_sink = _get_fa_learnable_sink(s_aux, query, use_cute)
     if use_cute:
         assert DEVICE_MAJOR >= 9, f"FA4 (flash_attn.cute) requires SM90+, got SM{DEVICE_MAJOR}0"
         from flash_attn.cute import flash_attn_varlen_func
@@ -168,6 +182,7 @@ def flash_attention_forward(
             softmax_scale=scaling,
             causal=True,
             window_size=window_size,
+            learnable_sink=learnable_sink,
             num_splits=1,
         )
     else:

vexact/config.py

@@ -100,6 +100,10 @@ class ModelConfig:
         default="fa-invariant",
         metadata={"help": "Attention implementation to use (fa-invariant, fa-invariant-cute, triton-invariant, flex)."},
     )
+    moe_implementation: str = field(
+        default="fused_quack",
+        metadata={"help": "MoE experts implementation. GPT-OSS supports fused_quack only."},
+    )
     enable_batch_invariant: bool = field(
         default=True,
         metadata={"help": "Enable batch invariant operations for deterministic inference."},
@@ -143,6 +147,16 @@ def __post_init__(self):
         if self.hf_config is None:
             self._load_hf_config()
 
+        if getattr(self.hf_config, "model_type", None) == "gpt_oss":
+            if self.attn_impl != "fa-invariant-cute":
+                raise ValueError(
+                    f"GPT-OSS learnable attention sinks require attn_impl='fa-invariant-cute'; got {self.attn_impl!r}"
+                )
+            if self.moe_implementation != "fused_quack":
+                raise ValueError(
+                    f"GPT-OSS requires moe_implementation='fused_quack' (SM90+); got {self.moe_implementation!r}"
+                )
+
         # Initialize max_model_len
         self._initialize_max_model_len()
 

vexact/models/gpt_oss/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .modeling_gpt_oss import apply_gpt_oss_patches
+
+
+__all__ = ["apply_gpt_oss_patches"]