[RDY] feat: add basic MoE support

packages/llm/tests/test_gpt_oss_handler.py

@@ -0,0 +1,158 @@
+from types import SimpleNamespace
+
+import torch
+from torch import nn
+
+from torch_to_nnef_llm.models.handlers import GptOssArchitectureHandler
+from torch_to_nnef_llm.models.handlers.registry import get_handler
+
+
+class FakeGptOssExpert(nn.Module):
+    def __init__(self, experts_implementation):
+        super().__init__()
+        self.config = SimpleNamespace(
+            _experts_implementation=experts_implementation
+        )
+
+
+class FakeGptOssModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.config = SimpleNamespace(_experts_implementation="grouped_mm")
+        self.grouped = FakeGptOssExpert("grouped_mm")
+        self.eager = FakeGptOssExpert("eager")
+
+
+def test_gpt_oss_handler_registered():
+    assert get_handler("gpt_oss") is GptOssArchitectureHandler
+
+
+def test_gpt_oss_handler_uses_traceable_experts_implementation():
+    model = FakeGptOssModel()
+
+    GptOssArchitectureHandler().prepare_model_for_export(model)
+
+    assert model.config._experts_implementation == "batched_mm"
+    assert model.grouped.config._experts_implementation == "batched_mm"
+    assert model.eager.config._experts_implementation == "eager"
+
+
+NEG = torch.finfo(torch.float32).min
+
+
+def _wrapper(*, sliding_window=4, layer_types=None, force_causal_mask=True):
+    if layer_types is None:
+        layer_types = ["sliding_attention", "full_attention"]
+    return SimpleNamespace(
+        model=SimpleNamespace(
+            config=SimpleNamespace(
+                sliding_window=sliding_window,
+                layer_types=layer_types,
+            )
+        ),
+        force_causal_mask=force_causal_mask,
+        with_dyn_cache=False,
+    )
+
+
+def _inputs(*, seq_length, past_length, n_layers=2):
+    input_ids = torch.zeros(1, seq_length, dtype=torch.long)
+    kv = [torch.zeros(1, 2, past_length, 8) for _ in range(2 * n_layers)]
+    return (input_ids, *kv)
+
+
+def _visible(mask):
+    """Boolean [S, K] view of an additive mask: True where attention is kept."""
+    return mask[0, 0] == 0.0
+
+
+def test_full_and_sliding_masks_are_distinct():
+    handler = GptOssArchitectureHandler()
+    ctx = handler.build_forward_inputs(
+        inputs=_inputs(seq_length=8, past_length=0),
+        wrapper=_wrapper(sliding_window=4),
+    )
+    mapping = ctx.model_inputs["attention_mask"]
+
+    assert isinstance(mapping, dict)
+    assert set(mapping) == {"full_attention", "sliding_attention"}
+    # The whole point of the handler: these must not be the same tensor.
+    assert not torch.equal(
+        mapping["full_attention"], mapping["sliding_attention"]
+    )
+
+
+def test_full_mask_is_plain_causal():
+    handler = GptOssArchitectureHandler()
+    ctx = handler.build_forward_inputs(
+        inputs=_inputs(seq_length=4, past_length=0),
+        wrapper=_wrapper(sliding_window=2),
+    )
+    visible = _visible(ctx.model_inputs["attention_mask"]["full_attention"])
+    expected = torch.tril(torch.ones(4, 4, dtype=torch.bool))
+    assert torch.equal(visible, expected)
+
+
+def test_sliding_mask_keeps_exactly_window_keys():
+    handler = GptOssArchitectureHandler()
+    window = 3
+    ctx = handler.build_forward_inputs(
+        inputs=_inputs(seq_length=6, past_length=0),
+        wrapper=_wrapper(sliding_window=window),
+    )
+    visible = _visible(ctx.model_inputs["attention_mask"]["sliding_attention"])
+
+    for q in range(6):
+        kept = visible[q].nonzero().flatten().tolist()
+        # Query q sees keys in (q - window, q], clipped at 0.
+        assert kept == list(range(max(0, q - window + 1), q + 1))
+        assert len(kept) <= window
+
+
+def test_sliding_mask_accounts_for_past_length():
+    handler = GptOssArchitectureHandler()
+    window, past = 4, 10
+    ctx = handler.build_forward_inputs(
+        inputs=_inputs(seq_length=2, past_length=past),
+        wrapper=_wrapper(sliding_window=window),
+    )
+    visible = _visible(ctx.model_inputs["attention_mask"]["sliding_attention"])
+
+    assert visible.shape == (2, past + 2)
+    # First new token sits at absolute position 10 and sees keys 7..10.
+    assert visible[0].nonzero().flatten().tolist() == [7, 8, 9, 10]
+    assert visible[1].nonzero().flatten().tolist() == [8, 9, 10, 11]
+
+
+def test_masks_are_additive_with_neg_inf_where_hidden():
+    handler = GptOssArchitectureHandler()
+    ctx = handler.build_forward_inputs(
+        inputs=_inputs(seq_length=3, past_length=0),
+        wrapper=_wrapper(sliding_window=2),
+    )
+    mask = ctx.model_inputs["attention_mask"]["sliding_attention"]
+    assert mask.shape == (1, 1, 3, 3)
+    assert mask.dtype == torch.float32
+    # Row 2 with window 2 hides key 0 only.
+    assert mask[0, 0, 2, 0] == NEG
+    assert mask[0, 0, 2, 1] == 0.0
+    assert mask[0, 0, 2, 2] == 0.0
+
+
+def test_full_attention_only_model_keeps_base_single_mask():
+    """No sliding layers means the base handler's single mask is correct."""
+    handler = GptOssArchitectureHandler()
+    ctx = handler.build_forward_inputs(
+        inputs=_inputs(seq_length=4, past_length=0),
+        wrapper=_wrapper(sliding_window=0, layer_types=["full_attention"] * 2),
+    )
+    assert isinstance(ctx.model_inputs["attention_mask"], torch.Tensor)
+
+
+def test_no_mask_mapping_when_causal_mask_not_forced():
+    handler = GptOssArchitectureHandler()
+    ctx = handler.build_forward_inputs(
+        inputs=_inputs(seq_length=4, past_length=0),
+        wrapper=_wrapper(force_causal_mask=False),
+    )
+    assert ctx.model_inputs["attention_mask"] is None

packages/llm/tests/test_sdpa_attention_export.py

@@ -162,6 +162,16 @@ def test_reify_sdpa_operator_rejects_eager_attention():
         exporter._resolve_attn_implementation("eager", True)
 
 
+def test_existing_export_test_dir_is_allowed_for_existing_export_root(
+    tmp_path,
+):
+    export_dir = tmp_path / "export"
+    test_dir = export_dir / "tests"
+    test_dir.mkdir(parents=True)
+
+    assert exporter._ensure_export_test_dir(export_dir, True) == test_dir
+
+
 def test_dump_llm_routes_reified_sdpa_to_loader(monkeypatch, tmp_path):
     captured_load = {}
     captured_dump = {}
@@ -214,6 +224,42 @@ def fake_load(*args, **kwargs):
     assert captured_load["experts_implementation"] == "batched_mm"
 
 
+def test_dump_llm_routes_reified_sdpa_input_upcast_policy(
+    monkeypatch, tmp_path
+):
+    captured_inference_target = {}
+
+    class _Exporter:
+        # non-multimodal model_type so dump_llm's multimodal guard passes
+        model_infos = SimpleNamespace(conf=SimpleNamespace(model_type="fake"))
+
+        def build_inference_target(self, **kwargs):
+            captured_inference_target.update(kwargs)
+            return object()
+
+        def dump_with_inference_target(self, **_kwargs):
+            return None
+
+        def dump(self, **kwargs):
+            return exporter.LLMExporter.dump(self, **kwargs)
+
+    def fake_load(*_args, **_kwargs):
+        return _Exporter()
+
+    monkeypatch.setattr(exporter.LLMExporter, "load", staticmethod(fake_load))
+    exporter.dump_llm(
+        "fake/model",
+        export_dirpath=tmp_path / "export",
+        reify_sdpa_operator=True,
+        upcast_reified_sdpa_inputs_to_f32=False,
+    )
+
+    assert captured_inference_target["reify_sdpa_operator"] is True
+    assert (
+        captured_inference_target["upcast_reified_sdpa_inputs_to_f32"] is False
+    )
+
+
 def test_cli_reify_sdpa_operator_implies_sdpa_attention(monkeypatch, tmp_path):
     captured = {}
 

packages/llm/torch_to_nnef_llm/exporter.py

@@ -139,6 +139,14 @@ def _resolve_attn_implementation(
     return attn_implementation
 
 
+def _ensure_export_test_dir(
+    export_dirpath: Path, ignore_already_exist_dir: bool
+) -> Path:
+    test_dir = export_dirpath / "tests"
+    test_dir.mkdir(parents=True, exist_ok=ignore_already_exist_dir)
+    return test_dir
+
+
 #: Default number of retries for transient Hugging Face download failures.
 DEFAULT_HF_DOWNLOAD_N_RETRIES = 5
 
@@ -791,8 +799,9 @@ def export_model(
             inference_target.dynamic_axes = dynamic_axes
 
             # Add io.npz test in exproted dir for dbg purpose
-            test_dir = export_dirpath / "tests"
-            test_dir.mkdir(parents=True)
+            test_dir = _ensure_export_test_dir(
+                export_dirpath, ignore_already_exist_dir
+            )
 
             if check_inference_modes:
                 self._dump_modes_json(
@@ -863,6 +872,7 @@ def dump(self, **kwargs):
                     "force_f32_linear_accumulator",
                     "force_f32_normalization",
                     "reify_sdpa_operator",
+                    "upcast_reified_sdpa_inputs_to_f32",
                     "tract_check_io_tolerance",
                 ]
                 if key in kwargs
@@ -884,6 +894,7 @@ def build_inference_target(
         force_f32_linear_accumulator: T.Optional[bool] = None,
         force_f32_normalization: T.Optional[bool] = None,
         reify_sdpa_operator: T.Optional[bool] = None,
+        upcast_reified_sdpa_inputs_to_f32: T.Optional[bool] = None,
         tract_check_io_tolerance: TractCheckTolerance = LM_CHECK_TOLERANCE,
         compression_method: T.Optional[str] = None,
         compression_registry: T.Optional[str] = None,
@@ -930,6 +941,10 @@ def build_inference_target(
 
         if reify_sdpa_operator is not None:
             inference_target.reify_sdpa_operator = reify_sdpa_operator
+        if upcast_reified_sdpa_inputs_to_f32 is not None:
+            inference_target.upcast_reified_sdpa_inputs_to_f32 = (
+                upcast_reified_sdpa_inputs_to_f32
+            )
 
         if (
             self.is_half_precision_model

packages/llm/torch_to_nnef_llm/loader.py

@@ -655,6 +655,7 @@ def _from_pretrained(
 
 
 @require_extra_decorator(extra=T2NExtra.LLM_TRACT, module="transformers")
+# pylint: disable-next=too-many-branches
 def load_model(
     hf_model_slug: T.Optional[str] = None,
     local_dir: T.Optional[Path] = None,

packages/llm/torch_to_nnef_llm/models/handlers/__init__.py

@@ -16,6 +16,7 @@
     Gemma4VideoEncoderHandler,
     Gemma4VisionEncoderHandler,
 )
+from .gpt_oss import GptOssArchitectureHandler
 from .idefics3_vl import (
     Idefics3ArchitectureHandler,
     Idefics3VisionEncoderHandler,
@@ -53,6 +54,7 @@
     "Gemma4ArchitectureHandler",
     "Gemma4VideoEncoderHandler",
     "Gemma4VisionEncoderHandler",
+    "GptOssArchitectureHandler",
     "IOSpec",
     "Idefics3ArchitectureHandler",
     "Idefics3VisionEncoderHandler",