Update preset code to support new versions of transformers

roundpipe/models/gpt_oss.py

@@ -20,9 +20,30 @@
 from .function import CompileForCausalLMLoss, ChunkedCompileLinearForCausalLMLoss
 
 
+def _resolve_attention_type(layer: GptOssDecoderLayer) -> str:
+    attention_type = getattr(layer.self_attn, "layer_type", None)
+    if attention_type is None:
+        attention_type = getattr(layer.self_attn, "attention_type", None)
+    if attention_type is None:
+        attention_type = getattr(layer, "attention_type", None)
+    if attention_type is None:
+        raise AttributeError("GPT-OSS decoder layer does not expose its attention type")
+    return cast(str, attention_type)
+
+
 class GptOssOptExperts(nn.Module):
     def __init__(self, mod: GptOssExperts) -> None:
         super().__init__()
+        expert_weights = (mod.gate_up_proj, mod.down_proj)
+        if any(
+            weight.ndim != 3 or not torch.is_floating_point(weight)
+            for weight in expert_weights
+        ):
+            raise TypeError(
+                "RoundPipe's GPT-OSS preset requires unpacked floating-point "
+                "expert weights; load a BF16/FP16 checkpoint instead of an "
+                "MXFP4/packed checkpoint."
+            )
         self.num_experts = mod.num_experts
         self.hidden_size = mod.hidden_size
         self.alpha = mod.alpha
@@ -39,13 +60,20 @@ def forward(
         router_indices: torch.Tensor,
         routing_weights: torch.Tensor,
     ) -> torch.Tensor:
-        batch_size, sequence_length, hidden_dim = hidden_states.shape
-        hidden_states = hidden_states.view(-1, hidden_dim)
+        input_shape = hidden_states.shape
+        hidden_dim = input_shape[-1]
+        hidden_states = hidden_states.reshape(-1, hidden_dim)
 
         top_k = router_indices.shape[-1]
-        selected_experts = router_indices.view(-1)
-        routing_weights = torch.gather(routing_weights, 1, router_indices)
-        routing_weights = routing_weights.view(-1)
+        selected_experts = router_indices.reshape(-1)
+        if routing_weights.shape != router_indices.shape:
+            if routing_weights.shape[-1] != self.num_experts:
+                raise ValueError(
+                    "GPT-OSS routing weights must contain either top-k weights "
+                    "or one weight per expert"
+                )
+            routing_weights = torch.gather(routing_weights, 1, router_indices)
+        routing_weights = routing_weights.reshape(-1)
 
         _, sort_idx = torch.sort(selected_experts)
         permute_weight = routing_weights[sort_idx]
@@ -66,7 +94,7 @@ def forward(
             save_for_recompute(token_per_expert_cpu)
 
         final_hidden_states = torch.zeros(
-            (batch_size * sequence_length, hidden_dim),
+            (hidden_states.shape[0], hidden_dim),
             dtype=hidden_states.dtype,
             device=hidden_states.device,
         )
@@ -97,7 +125,7 @@ def forward(
             final_hidden_states.index_add_(0, expert_tokens, expert_output)
             start_idx += num_tokens
 
-        return final_hidden_states.view(batch_size, sequence_length, hidden_dim)
+        return final_hidden_states.reshape(input_shape)
 
 
 class GptOssForCausalLMPrefix(nn.Module):
@@ -203,7 +231,7 @@ def __init__(self, layer: GptOssDecoderLayer) -> None:
         self.mlp = layer.mlp
         self.input_layernorm = layer.input_layernorm
         self.post_attention_layernorm = layer.post_attention_layernorm
-        self.attention_type = layer.self_attn.attention_type
+        self.attention_type = _resolve_attention_type(layer)
 
     def forward(self, input):
         (

roundpipe/models/qwen3.py

@@ -1,4 +1,5 @@
 from typing_extensions import *
+import inspect
 import warnings
 
 import torch
@@ -15,6 +16,45 @@
 from .function import CompileForCausalLMLoss, ChunkedCompileLinearForCausalLMLoss
 
 
+def _resolve_attention_type(layer: nn.Module) -> str:
+    self_attn = layer.self_attn
+    attention_type = getattr(self_attn, "layer_type", None)
+    if attention_type is None:
+        attention_type = getattr(self_attn, "attention_type", None)
+    if attention_type is None:
+        attention_type = getattr(layer, "attention_type", None)
+    if attention_type is None:
+        raise AttributeError("Qwen3 decoder layer does not expose its attention type")
+    return cast(str, attention_type)
+
+
+def _create_mask_compat(
+    mask_function: Callable[..., torch.Tensor],
+    *,
+    config: Any,
+    inputs_embeds: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    past_key_values: Optional[Any],
+    position_ids: torch.Tensor,
+    cache_position: Optional[torch.Tensor],
+) -> torch.Tensor:
+    parameters = inspect.signature(mask_function).parameters
+    mask_kwargs: Dict[str, Any] = {
+        "config": config,
+        "attention_mask": attention_mask,
+        "past_key_values": past_key_values,
+    }
+    embeds_parameter = (
+        "inputs_embeds" if "inputs_embeds" in parameters else "input_embeds"
+    )
+    mask_kwargs[embeds_parameter] = inputs_embeds
+    if "position_ids" in parameters:
+        mask_kwargs["position_ids"] = position_ids
+    if "cache_position" in parameters:
+        mask_kwargs["cache_position"] = cache_position
+    return mask_function(**mask_kwargs)
+
+
 class Qwen3ForCausalLMPrefix(nn.Module):
     def __init__(self, model: Qwen3ForCausalLM) -> None:
         super().__init__()
@@ -70,29 +110,39 @@ def forward(
             )
         past_key_values = None
 
-        if position_ids is None:
-            position_ids = torch.arange(
+        if cache_position is None:
+            cache_position = torch.arange(
                 inputs_embeds.shape[1], device=inputs_embeds.device
-            ).unsqueeze(0)
+            )
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
 
         # It may already have been prepared by e.g. `generate`
         if not isinstance(causal_mask_mapping := attention_mask, dict):
-            # Prepare mask arguments
-            mask_kwargs = {
-                "config": self.config,
-                "inputs_embeds": inputs_embeds,
-                "attention_mask": attention_mask,
-                "past_key_values": past_key_values,
-                "position_ids": position_ids,
-            }
             # Create the masks
             causal_mask_mapping = {
-                "full_attention": create_causal_mask(**mask_kwargs),
+                "full_attention": _create_mask_compat(
+                    create_causal_mask,
+                    config=self.config,
+                    inputs_embeds=inputs_embeds,
+                    attention_mask=attention_mask,
+                    past_key_values=past_key_values,
+                    position_ids=position_ids,
+                    cache_position=cache_position,
+                ),
             }
             # The sliding window alternating layers are not always activated depending on the config
             if self.has_sliding_layers:
                 causal_mask_mapping["sliding_attention"] = (
-                    create_sliding_window_causal_mask(**mask_kwargs)
+                    _create_mask_compat(
+                        create_sliding_window_causal_mask,
+                        config=self.config,
+                        inputs_embeds=inputs_embeds,
+                        attention_mask=attention_mask,
+                        past_key_values=past_key_values,
+                        position_ids=position_ids,
+                        cache_position=cache_position,
+                    )
                 )
 
         hidden_states = inputs_embeds
@@ -114,6 +164,7 @@ class Qwen3ForCausalLMWrappedLayer(nn.Module):
     def __init__(self, layer: nn.Module) -> None:
         super().__init__()
         self.layer = layer
+        self.attention_type = _resolve_attention_type(layer)
 
     def forward(self, input):
         (
@@ -127,7 +178,7 @@ def forward(self, input):
         ) = input
         hidden_states = self.layer(
             hidden_states,
-            attention_mask=causal_mask_mapping[self.layer.self_attn.layer_type],
+            attention_mask=causal_mask_mapping[self.attention_type],
             position_ids=position_ids,
             past_key_values=None,
             use_cache=False,

tests/models/test_preset_compatibility.py

@@ -0,0 +1,189 @@
+from types import SimpleNamespace
+
+import pytest
+import torch
+from transformers.models.gpt_oss.configuration_gpt_oss import GptOssConfig
+from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts
+
+from roundpipe.models import gpt_oss, qwen3
+
+
+def test_qwen3_resolves_current_and_legacy_attention_type_locations():
+    current_layer = SimpleNamespace(
+        self_attn=SimpleNamespace(layer_type="sliding_attention")
+    )
+    legacy_layer = SimpleNamespace(
+        self_attn=SimpleNamespace(), attention_type="full_attention"
+    )
+
+    assert qwen3._resolve_attention_type(current_layer) == "sliding_attention"
+    assert qwen3._resolve_attention_type(legacy_layer) == "full_attention"
+
+
+def test_qwen3_calls_current_mask_contract():
+    inputs_embeds = torch.randn(1, 3, 4)
+    position_ids = torch.arange(3).unsqueeze(0)
+
+    def current_mask(
+        *,
+        config,
+        inputs_embeds,
+        attention_mask,
+        past_key_values,
+        position_ids,
+    ):
+        return (
+            config,
+            inputs_embeds,
+            attention_mask,
+            past_key_values,
+            position_ids,
+        )
+
+    result = qwen3._create_mask_compat(
+        current_mask,
+        config="config",
+        inputs_embeds=inputs_embeds,
+        attention_mask="mask",
+        past_key_values=None,
+        position_ids=position_ids,
+        cache_position=None,
+    )
+
+    assert result == ("config", inputs_embeds, "mask", None, position_ids)
+
+
+def test_qwen3_calls_legacy_mask_contract():
+    inputs_embeds = torch.randn(1, 3, 4)
+    position_ids = torch.arange(3).unsqueeze(0)
+    cache_position = position_ids.squeeze(0)
+
+    def legacy_mask(
+        *,
+        config,
+        input_embeds,
+        attention_mask,
+        cache_position,
+        past_key_values,
+        position_ids,
+    ):
+        return (
+            config,
+            input_embeds,
+            attention_mask,
+            cache_position,
+            past_key_values,
+            position_ids,
+        )
+
+    result = qwen3._create_mask_compat(
+        legacy_mask,
+        config="config",
+        inputs_embeds=inputs_embeds,
+        attention_mask="mask",
+        past_key_values=None,
+        position_ids=position_ids,
+        cache_position=cache_position,
+    )
+
+    assert result == (
+        "config",
+        inputs_embeds,
+        "mask",
+        cache_position,
+        None,
+        position_ids,
+    )
+
+
+def _make_gpt_oss_experts() -> GptOssExperts:
+    config = GptOssConfig(
+        hidden_size=8,
+        intermediate_size=4,
+        num_local_experts=3,
+        num_experts_per_tok=2,
+        num_hidden_layers=1,
+        num_attention_heads=2,
+        num_key_value_heads=1,
+        head_dim=4,
+        vocab_size=32,
+    )
+    torch.manual_seed(0)
+    experts = GptOssExperts(config)
+    for parameter in experts.parameters():
+        torch.nn.init.uniform_(parameter, -0.1, 0.1)
+    return experts
+
+
+def test_gpt_oss_optimized_experts_support_current_2d_topk_contract():
+    experts = _make_gpt_oss_experts()
+    optimized = gpt_oss.GptOssOptExperts(experts)
+    router_indices = torch.tensor([[0, 2], [1, 0], [2, 1]])
+    routing_weights = torch.tensor([[0.7, 0.3], [0.4, 0.6], [0.8, 0.2]])
+    reference_input = torch.randn(3, 8, requires_grad=True)
+    optimized_input = reference_input.detach().clone().requires_grad_(True)
+
+    expected = experts(reference_input, router_indices, routing_weights)
+    actual = optimized(optimized_input, router_indices, routing_weights)
+
+    torch.testing.assert_close(actual, expected)
+    expected.sum().backward()
+    actual.sum().backward()
+    torch.testing.assert_close(optimized_input.grad, reference_input.grad)
+
+
+def test_gpt_oss_optimized_experts_support_legacy_3d_full_router_contract():
+    experts = _make_gpt_oss_experts()
+    optimized = gpt_oss.GptOssOptExperts(experts)
+    hidden_states = torch.randn(2, 2, 8)
+    router_indices = torch.tensor([[0, 2], [1, 0], [2, 1], [1, 2]])
+    full_routing_weights = torch.tensor(
+        [
+            [0.7, 0.0, 0.3],
+            [0.6, 0.4, 0.0],
+            [0.0, 0.2, 0.8],
+            [0.0, 0.4, 0.6],
+        ]
+    )
+    topk_routing_weights = torch.gather(
+        full_routing_weights, 1, router_indices
+    )
+
+    expected = experts(
+        hidden_states.reshape(-1, 8), router_indices, topk_routing_weights
+    ).reshape_as(hidden_states)
+    actual = optimized(hidden_states, router_indices, full_routing_weights)
+
+    torch.testing.assert_close(actual, expected)
+
+
+def test_gpt_oss_resolves_current_and_legacy_attention_type_locations():
+    current_layer = SimpleNamespace(
+        self_attn=SimpleNamespace(layer_type="sliding_attention")
+    )
+    intermediary_layer = SimpleNamespace(
+        self_attn=SimpleNamespace(attention_type="full_attention")
+    )
+    legacy_layer = SimpleNamespace(
+        self_attn=SimpleNamespace(), attention_type="sliding_attention"
+    )
+
+    assert gpt_oss._resolve_attention_type(current_layer) == "sliding_attention"
+    assert gpt_oss._resolve_attention_type(intermediary_layer) == "full_attention"
+    assert gpt_oss._resolve_attention_type(legacy_layer) == "sliding_attention"
+
+
+def test_gpt_oss_optimized_experts_reject_packed_weights_at_wrap_time():
+    packed_experts = SimpleNamespace(
+        num_experts=3,
+        hidden_size=8,
+        alpha=1.702,
+        limit=7.0,
+        gate_up_proj=torch.empty(3, 8, 2, dtype=torch.uint8),
+        gate_up_proj_bias=torch.empty(3, 8),
+        down_proj=torch.empty(3, 4, 2, dtype=torch.uint8),
+        down_proj_bias=torch.empty(3, 8),
+    )
+
+    with pytest.raises(TypeError, match="floating-point expert weights"):
+        gpt_oss.GptOssOptExperts(packed_experts)