Merge branch 'main' into dependabot/pip/requirements/lintrunner/ruff-0.16.0

Author: titaiwangms

.github/workflows/publish.yml

@@ -27,11 +27,11 @@ jobs:
     name: Build distribution
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v6
+      - uses: actions/checkout@v7
         with:
           persist-credentials: false
       - name: Setup Python
-        uses: actions/setup-python@v6
+        uses: actions/setup-python@v7
         with:
           python-version: "3.12"
       - name: Install build dependencies

CODE_OF_CONDUCT.md

@@ -0,0 +1,10 @@
+# Microsoft Open Source Code of Conduct
+
+This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/).
+
+Resources:
+
+- [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/)
+- [Microsoft Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/)
+- Contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with questions or concerns
+- Employees can reach out at [aka.ms/opensource/moderation-support](https://aka.ms/opensource/moderation-support)

scripts/generate_golden.py

@@ -237,12 +237,23 @@ def _generate_causal_lm(case: TestCase, json_path: Path, device: str) -> None:
         _apply_nemotron_h_generate_patch(model)
 
         model_device = _get_model_device(model, device)
+        gen_ids = torch.from_numpy(input_ids).to(model_device)
+        max_new = case.generation_params.get("max_new_tokens", 20)
         with torch.no_grad():
-            gen_output = model.generate(
-                torch.from_numpy(input_ids).to(model_device),
-                max_new_tokens=case.generation_params.get("max_new_tokens", 20),
-                do_sample=False,
-            )
+            try:
+                gen_output = model.generate(gen_ids, max_new_tokens=max_new, do_sample=False)
+            except ValueError as e:
+                # All-attention GraniteMoeHybrid variants (e.g. granite-4.0-1b)
+                # trip transformers' hybrid Mamba/attention generation cache,
+                # which assumes at least one linear-attention (Mamba) layer:
+                # "`has_previous_state` can only be called on LinearAttention
+                # layers". Greedy output is cache-independent, so fall back to
+                # the (slower) cache-free path.
+                if "has_previous_state" not in str(e):
+                    raise
+                gen_output = model.generate(
+                    gen_ids, max_new_tokens=max_new, do_sample=False, use_cache=False
+                )
         generated_ids = gen_output[0, seq_len:].cpu().numpy()
 
     save_golden_ref(

src/mobius/_builder.py

@@ -417,7 +417,8 @@ def build(
             can use ``GroupQueryAttention`` on GQA-capable execution providers.
             Raises :class:`ValueError` if the resolved ``model_type`` has no
             text-only sibling. Currently supported for ``gemma4_unified``
-            (``google/gemma-4-12B``).
+            (``google/gemma-4-12B``) and ``qwen3_5_moe_vl``
+            (``Qwen/Qwen3.6-35B-A3B``).
 
     Returns:
         A :class:`ModelPackage` containing the built model(s).

src/mobius/_configs/_base.py

@@ -2226,10 +2226,39 @@ class GraniteMoeHybridConfig(BambaConfig):
 
     @classmethod
     def from_transformers(cls, config, parent_config=None) -> GraniteMoeHybridConfig:
-        # Reuse BambaConfig.from_transformers for mamba fields + layer_types conversion
-        # (converts HF "mamba"→"mamba2" and "attention"→"full_attention")
+        # Reuse BambaConfig.from_transformers for mamba fields, MoE/RoPE/multiplier
+        # extraction, then rebuild layer_types from GraniteMoeHybrid's own naming.
         bamba = BambaConfig.from_transformers(config, parent_config)
         bamba_fields = _shallow_fields(bamba)
+
+        # GraniteMoeHybrid names layers "full_attention" / "linear_attention"
+        # (linear_attention == Mamba2/SSD), unlike Bamba's "attention" / "mamba".
+        raw_layer_types = (
+            getattr(config, "layer_types", None)
+            or getattr(config, "layers_block_type", None)
+            or []
+        )
+        _attn = {"full_attention", "attention"}
+        _mamba = {"linear_attention", "mamba", "mamba2"}
+        layer_types: list[str] = []
+        for ltype in raw_layer_types:
+            if ltype in _attn:
+                layer_types.append("full_attention")
+            elif ltype in _mamba:
+                layer_types.append("mamba2")
+            else:
+                raise ValueError(f"Unknown GraniteMoeHybrid layer type: {ltype!r}")
+        if layer_types:
+            bamba_fields["layer_types"] = layer_types
+
+        # Respect position_embedding_type: GraniteMoeHybrid checkpoints ship
+        # default ``rope_parameters`` even for the NoPE variant
+        # (granite-4.0-tiny-preview: position_embedding_type='nope'). Only apply
+        # RoPE when explicitly requested; otherwise disable it so
+        # ``initialize_rope`` returns None and attention runs NoPE.
+        if getattr(config, "position_embedding_type", "rope") != "rope":
+            bamba_fields["rope_type"] = None
+
         return cls(
             **bamba_fields,
             shared_intermediate_size=getattr(config, "shared_intermediate_size", 1024),

src/mobius/_registry.py

@@ -604,6 +604,13 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None:
     "qwen2_vl_text": ModelRegistration(Qwen25VLTextModel),
     "qwen3_5": ModelRegistration(Qwen35VL3ModelCausalLMModel, task="hybrid-qwen-vl"),
     "qwen3_5_moe_vl": ModelRegistration(Qwen35MoEVL3ModelCausalLMModel, task="hybrid-qwen-vl"),
+    # Text-only sibling of ``qwen3_5_moe_vl`` (Qwen3.6-35B-A3B). The MoE
+    # backbone ``Qwen35MoECausalLMModel`` already strips ``language_model.``
+    # and drops ``visual.``/MTP keys, so it consumes the VL checkpoint's text
+    # weights directly; ``build(..., text_only=True)`` routes here via
+    # ``_TEXT_ONLY_MODEL_TYPE``. It also matches the VL ``text_config``'s own
+    # ``model_type=qwen3_5_moe_text`` so that config resolves cleanly.
+    "qwen3_5_moe_text": ModelRegistration(Qwen35MoECausalLMModel),
     "qwen3_5_vl": ModelRegistration(Qwen35VL3ModelCausalLMModel, task="hybrid-qwen-vl"),
     "qwen3_5_vl_text": ModelRegistration(Qwen35VLTextModel),
     "qwen3_vl": ModelRegistration(Qwen3VL3ModelCausalLMModel, task="qwen-vl"),
@@ -791,6 +798,12 @@ def _create_default_registry() -> ModelRegistry:
 _TEXT_ONLY_MODEL_TYPE: dict[str, str] = {
     "gemma4_unified": "gemma4_unified_text",
     "gemma4_unified_text": "gemma4_unified_text",
+    # Qwen3.5-MoE-VL (Qwen3.6-35B-A3B): export just the hybrid MoE text
+    # backbone as a standalone decoder-only LLM. The builder overrides
+    # ``qwen3_5_moe`` -> ``qwen3_5_moe_vl`` when a ``vision_config`` is present,
+    # so the text-only override keys off the VL type here.
+    "qwen3_5_moe_vl": "qwen3_5_moe_text",
+    "qwen3_5_moe_text": "qwen3_5_moe_text",
 }
 
 
@@ -894,6 +907,7 @@ def _create_default_registry() -> ModelRegistry:
     "qwen2_moe": "Qwen/Qwen1.5-MoE-A2.7B-Chat",
     "qwen3_moe": "Qwen/Qwen3-30B-A3B",
     "qwen3_5_moe": "Qwen/Qwen3.5-MoE-A3B-128K",
+    "qwen3_5_moe_text": "Qwen/Qwen3.6-35B-A3B",
     "qwen3_next": "Qwen/Qwen3-235B-A22B",
     "granitemoe": "ibm-granite/granite-3.0-1b-a400m-instruct",
     "olmoe": "allenai/OLMoE-1B-7B-0924",
@@ -1141,6 +1155,7 @@ def _create_default_registry() -> ModelRegistry:
     "qwen3_moe": "qwen",
     "qwen3_5_text": "qwen",
     "qwen3_5_moe": "qwen",
+    "qwen3_5_moe_text": "qwen",
     "qwen3_next": "qwen",
     "qwen2_vl": "qwen",
     "qwen2_vl_text": "qwen",

src/mobius/_testing/torch_reference.py

@@ -334,7 +334,18 @@ def torch_forward(
             )
         kwargs["past_key_values"] = cache
 
-    outputs = model(**kwargs)
+    try:
+        outputs = model(**kwargs)
+    except ValueError as e:
+        # All-attention hybrid models (e.g. GraniteMoeHybrid granite-4.0-1b)
+        # trip transformers' recurrent-mask builder, which assumes the hybrid
+        # cache contains a linear-attention (Mamba) layer: "`has_previous_state`
+        # can only be called on LinearAttention layers". A single-pass forward's
+        # logits don't depend on caching, so retry without a cache.
+        if "has_previous_state" not in str(e) or "past_key_values" in kwargs:
+            raise
+        kwargs["use_cache"] = False
+        outputs = model(**kwargs)
     logits = outputs.logits.cpu().numpy()
 
     # Extract KV cache if available (Mamba models don't have it)

src/mobius/functions/packed_multi_head_attention.py

@@ -17,16 +17,19 @@
 
     # Build block-diagonal bias (0 for same segment, -inf for different)
     same = Equal(segment_ids[:, None], segment_ids[None, :])
-    bias = Where(same, 0.0, -10000.0)
+    attn_bias = Where(same, 0.0, -10000.0)
 
     # Standard Attention
-    output = Attention(query, key, value, bias,
+    output = Attention(query, key, value, attn_bias,
                        q_num_heads=<num_heads>, kv_num_heads=<num_heads>,
                        scale=<scale>)
 
-The ``token_offset`` input is consumed by the native kernel but is
-unused by the fallback body (segment boundaries from
-``cumulative_sequence_length`` are sufficient).
+The optional ``bias`` (slot 4) and ``token_offset`` inputs are consumed by
+the native kernel but are unused by the fallback body (the ``attn_bias``
+computed above from ``cumulative_sequence_length`` is sufficient, and is
+unrelated to the formal ``bias`` input).  They are still declared
+as formal inputs to preserve the positional slot alignment expected by the
+ORT ``PackedMultiHeadAttention`` signature.
 
 Attributes:
     num_heads (int): Number of attention heads.
@@ -53,6 +56,7 @@ def packed_multi_head_attention() -> ir.Function:
         query:  (token_count, hidden_size)
         key:    (token_count, hidden_size)
         value:  (token_count, v_hidden_size)
+        bias:   (optional) — unused in fallback
         token_offset: (batch_size, sequence_length) — unused in fallback
         cumulative_sequence_length: (batch_size + 1,) INT32
 
@@ -69,9 +73,14 @@ def body(
         query_input,
         key_input,
         value_input,
+        bias_input,
         token_offset_input,
         cumulative_sequence_length_input,
     ):
+        # bias_input is the optional slot-4 input of the ORT
+        # PackedMultiHeadAttention signature.  It is unused by this fallback
+        # (the block-diagonal bias is reconstructed from cu_seqlens below) but
+        # must exist as a formal input to keep the positional slots aligned.
         # --- Compute sequence length from query shape ---
         # query: (token_count, hidden_size)
         token_count = op.Shape(query_input, start=0, end=1)
@@ -159,6 +168,7 @@ def body(
             ir.Value(name="query"),
             ir.Value(name="key"),
             ir.Value(name="value"),
+            ir.Value(name="bias"),
             ir.Value(name="token_offset"),
             ir.Value(name="cumulative_sequence_length"),
         ],

src/mobius/functions/packed_multi_head_attention_test.py

@@ -0,0 +1,62 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Unit tests for the PackedMultiHeadAttention fallback ``ir.Function``.
+
+These guard the formal-input arity/order of the standard-ONNX fallback for
+``com.microsoft::PackedMultiHeadAttention``.  ORT's op has a 6-7 input
+positional signature::
+
+    query, key(opt), value(opt), bias(opt),
+    token_offset, cumulative_sequence_length, attention_bias(opt)
+
+Because ``token_offset`` / ``cumulative_sequence_length`` occupy positional
+slots 5 and 6, the optional ``bias`` at slot 4 must exist as a formal input
+even though the fallback body ignores it.  Call sites emit 6 inputs
+``(q, k, v, "", token_offset, cu_seqlens)``; if the function declared only 5
+formals, onnx-genai's function-inline admission rejects the call with a
+FunctionArityMismatch (``call.input.len()=6 > func.input.len()=5``).
+"""
+
+from __future__ import annotations
+
+from mobius.functions.packed_multi_head_attention import (
+    packed_multi_head_attention,
+)
+
+EXPECTED_INPUT_ORDER = [
+    "query",
+    "key",
+    "value",
+    "bias",
+    "token_offset",
+    "cumulative_sequence_length",
+]
+
+
+def test_packed_mha_declares_six_positional_formal_inputs() -> None:
+    func = packed_multi_head_attention()
+
+    actual_order = [value.name for value in func.inputs]
+    # Assert the full positional order, not just the count: the ``bias`` slot
+    # must sit at index 3 (between ``value`` and ``token_offset``) to match the
+    # ORT PackedMultiHeadAttention signature.
+    assert actual_order == EXPECTED_INPUT_ORDER
+    assert func.inputs[3].name == "bias"
+
+
+def test_packed_mha_admits_six_input_call() -> None:
+    func = packed_multi_head_attention()
+
+    # A call site emits 6 inputs: (q, k, v, "", token_offset, cu_seqlens).
+    # onnx-genai admits a call when len(call inputs) <= len(func inputs), so
+    # the function must declare at least 6 formals for the call to be inlined.
+    call_input_count = 6
+    assert call_input_count <= len(func.inputs)
+
+
+def test_packed_mha_function_identity() -> None:
+    func = packed_multi_head_attention()
+
+    assert func.domain == "com.microsoft"
+    assert func.name == "PackedMultiHeadAttention"

src/mobius/integrations/ort_genai/auto_export.py

@@ -87,10 +87,10 @@
     # ORT GenAI (see onnxruntime-genai/src/models/model_type.h LLM list).
     "hunyuan_v1_dense": "decoder",
     "deepseek_v4": "decoder",
-    # Qwen VL models all use the same GenAI pipeline as qwen2_5_vl
+    # Qwen VL model families have separate ORT GenAI model types.
     "qwen2_vl": "qwen2_5_vl",
-    "qwen3_vl": "qwen2_5_vl",
-    "qwen3_vl_text": "qwen2_5_vl",
+    "qwen3_vl": "qwen3_vl",
+    "qwen3_vl_text": "qwen3_vl",
     "qwen3_5": "qwen2_5_vl",
     "qwen3_5_vl": "qwen2_5_vl",
 }
@@ -880,6 +880,16 @@ def _write_genai_config(
                 if sms is not None:
                     vision_kwargs["spatial_merge_size"] = sms
                 vision_kwargs["config_filename"] = "processor_config.json"
+                if model_type in {"qwen3_vl", "qwen3_vl_text"}:
+                    patch_size = getattr(vision_cfg, "patch_size", None)
+                    window_size = getattr(vision_cfg, "window_size", None)
+                    if patch_size is not None:
+                        vision_kwargs["patch_size"] = patch_size
+                    if window_size is not None:
+                        vision_kwargs["window_size"] = window_size
+                    vision_kwargs["tokens_per_second"] = float(
+                        getattr(config, "tokens_per_second", 2.0)
+                    )
 
             if vision_input_mapping is not None:
                 vision_kwargs["input_names"] = vision_input_mapping
@@ -889,6 +899,12 @@ def _write_genai_config(
             embedding_output_mapping = _introspect_outputs(pkg, "embedding")
             if embedding_output_mapping is not None:
                 vision_kwargs["embedding_output_names"] = embedding_output_mapping
+            vision_start_token_id = getattr(config, "vision_start_token_id", None)
+            video_token_id = getattr(config, "video_token_id", None)
+            if vision_start_token_id is not None:
+                vision_kwargs["vision_start_token_id"] = vision_start_token_id
+            if video_token_id is not None:
+                vision_kwargs["video_token_id"] = video_token_id
 
             generator.with_vision(image_token_id=image_token_id, **vision_kwargs)
 
@@ -946,8 +962,8 @@ def write_ort_genai_config(
         pkg: Already-built :class:`~mobius._model_package.ModelPackage` with
             weights applied and ``config`` set.
         directory: Output directory (created if needed).
-        hf_model_id: HuggingFace model ID. When provided, used to fetch token
-            IDs (``bos``/``eos``/``pad``) and download tokenizer files.
+        hf_model_id: HuggingFace model ID or local model directory. When provided,
+            used to fetch token IDs (``bos``/``eos``/``pad``) and copy tokenizer files.
             When ``None``, token IDs are read from ``pkg.config`` fields
             (``bos_token_id``, ``eos_token_id``, ``pad_token_id``) populated
             by :meth:`~mobius._configs.ArchitectureConfig.from_transformers`,
@@ -1112,11 +1128,15 @@ def write_ort_genai_config(
             f.write("\n")
         result["mtp_config"] = mtp_path
 
-    # Copy tokenizer files — HF Hub takes precedence; local dir is the fallback
-    # for --config mode where no HF model ID is available.
+    # Copy tokenizer files. A local hf_model_id is a local model directory, not a
+    # Hub repo id; copy directly instead of calling hf_hub_download.
     if hf_model_id is not None:
-        logger.info("Copying tokenizer files from %s", hf_model_id)
-        tokenizer_files = _copy_tokenizer_files(hf_model_id, directory)
+        if os.path.isdir(hf_model_id):
+            logger.info("Copying tokenizer files from local model directory %s", hf_model_id)
+            tokenizer_files = _copy_tokenizer_files_from_local(hf_model_id, directory)
+        else:
+            logger.info("Copying tokenizer files from %s", hf_model_id)
+            tokenizer_files = _copy_tokenizer_files(hf_model_id, directory)
         for tf in tokenizer_files:
             result[tf] = os.path.join(directory, tf)
     elif local_config_dir is not None: