Add experimental Qwen3.8 packed hybrid export for ORT GenAI Engine

README.md

@@ -117,14 +117,49 @@ mobius build --model openai/whisper-tiny --output output_dir/
 ```
 
 Build-mode toggles use the cargo-style `--features` option. Available features
-are `static-cache`, `fp8-kv-cache`, `prune-prefill-prefix`, and `text-only`. Pass them
+include `static-cache`, `paged-attention`, `fp8-kv-cache`, `prune-prefill-prefix`, and `text-only`. Pass them
 as a comma-separated list or repeat the option:
 
 ```sh
 mobius build --model meta-llama/Llama-3.2-1B --output output_dir/ \
       --features static-cache,prune-prefill-prefix --max-seq-len 2048
 ```
 
+#### Experimental Qwen packed hybrid serving
+
+`qwen3_5_text` checkpoints, including `Qwen/Qwen3.8-27B`, have an opt-in
+CUDA FP16/BF16 export for ORT GenAI's continuous-batching Engine:
+
+```sh
+mobius build --model Qwen/Qwen3.8-27B --output qwen-paged/ \
+      --ep cuda --dtype f16 --features paged-attention --runtime ort-genai
+```
+
+This is a separate packed-token ABI, not a change to dense inference. It uses
+SEPARATE K/V pages only for full-attention layers, native varlen convolution,
+and FP32 V-major `[B, H_v, D_v, D_k]` recurrent state. The default page size is
+256; `PagedHybridCausalLMTask(paged_block_size=...)` accepts positive multiples
+of 256. `prune-prefill-prefix` selects each request's final packed row before
+the LM head. Static cache, quantized checkpoints/KV, multimodal serving, and
+CUDA graph capture are not supported in this path.
+
+The operator contract is pinned to ORT
+`f38538cd5a4b5945a4c839565a8eebc65e1e2ef8` and GenAI
+`d5b40851ba80ffa8e95b6b01f921dbb9008fac80`; released wheels are not assumed
+compatible. Generated GenAI configuration declares `paged_kv`, `fixed_conv`,
+and `fixed_recurrent` groups and dynamic batching. The workflow metadata
+describes one externally scheduled invocation, not a dense generation loop;
+the scheduler supplies packed boundaries, page tables, and request lifecycle
+inputs. `attention_metadata` is CPU INT32 containing query-length upper bound,
+KV-length upper bound, and KV-length lower bound. Paged cache outputs **must**
+alias their inputs; fixed states must follow the same request ordering.
+
+**Not production-qualified:** a full Qwen3.8-27B CUDA smoke run is still required.
+The network-free CUDA parity and Engine probes are in
+`tests/integration/paged_hybrid_test.py`; set `MOBIUS_ORT_REVISION` and
+`MOBIUS_GENAI_REVISION` to the above revisions when testing matching source
+builds. CPU-only runs skip these probes.
+
 Use `--release` with either `build` or `build-gguf` to potentially reduce saved model size
 by stripping build-time debug and provenance metadata. Functional metadata with keys prefixed by
 `mobius.` is preserved:

src/mobius/__main__.py

@@ -271,6 +271,8 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
                 "Error: --features paged-attention cannot be combined with --task. "
                 "Remove --task to use --features paged-attention."
             )
+        # The Transformers builder replaces this placeholder with Qwen's
+        # dedicated packed hybrid task when the effective config requires it.
         task = CausalLMTask(paged_cache=True)
     trust_remote_code = args.trust_remote_code
     revision = args.revision

src/mobius/_builder.py

@@ -76,12 +76,15 @@ def _enable_prefill_prefix_pruning_task(task: str | ModelTask) -> str | ModelTas
         Gemma4Task,
         Gemma4TextCausalLMTask,
         HybridCausalLMTask,
+        PagedHybridCausalLMTask,
     )
 
     if task == "text-generation":
         return CausalLMTask(prune_prefill_prefix=True)
     if task == "hybrid-text-generation":
         return HybridCausalLMTask(prune_prefill_prefix=True)
+    if task == "paged-hybrid-text-generation":
+        return PagedHybridCausalLMTask(prune_prefill_prefix=True)
     if task == "gemma4-text-generation":
         return Gemma4TextCausalLMTask(prune_prefill_prefix=True)
     if task == "gemma4":
@@ -94,6 +97,11 @@ def _enable_prefill_prefix_pruning_task(task: str | ModelTask) -> str | ModelTas
         )
     if isinstance(task, HybridCausalLMTask):
         return HybridCausalLMTask(prune_prefill_prefix=True)
+    if isinstance(task, PagedHybridCausalLMTask):
+        return PagedHybridCausalLMTask(
+            paged_block_size=task._paged_block_size,
+            prune_prefill_prefix=True,
+        )
     if isinstance(task, Gemma4TextCausalLMTask):
         return Gemma4TextCausalLMTask(
             static_cache=getattr(task, "_static_cache", False),
@@ -108,7 +116,8 @@ def _enable_prefill_prefix_pruning_task(task: str | ModelTask) -> str | ModelTas
         )
     raise ValueError(
         "prune_prefill_prefix=True is only supported for text-generation, "
-        "hybrid-text-generation, gemma4-text-generation, and gemma4 tasks."
+        "hybrid-text-generation, paged-hybrid-text-generation, "
+        "gemma4-text-generation, and gemma4 tasks."
     )
 
 
@@ -162,6 +171,27 @@ def build_from_module(
     if prune_prefill_prefix:
         task = _enable_prefill_prefix_pruning_task(task)
     resolved_task = get_task(task)
+    from mobius.tasks import PagedHybridCausalLMTask
+
+    if (
+        getattr(config, "export_paged_attention", False)
+        and getattr(config, "model_type", None) == "qwen3_5_text"
+        and not isinstance(resolved_task, PagedHybridCausalLMTask)
+    ):
+        raise ValueError("Qwen paged attention requires PagedHybridCausalLMTask")
+    if isinstance(resolved_task, PagedHybridCausalLMTask):
+        if execution_provider != "cuda":
+            raise ValueError("PagedHybridCausalLMTask requires execution_provider='cuda'")
+        if fp8_kv_cache or kv_cache_scales is not None:
+            raise ValueError("PagedHybridCausalLMTask supports only unquantized KV caches")
+        # Direct task selection must preserve raw native gate precision too,
+        # even when the model was constructed without the feature flag.
+        from mobius.components import GatedDeltaNet
+
+        for child in module.modules():
+            if isinstance(child, GatedDeltaNet):
+                child.A_log._keep_float32 = True
+                child.dt_bias._keep_float32 = True
     component_manifest = resolved_task.component_manifest()
     configure_component_quantization(module, config, resolved_task)
     _cast_module_dtype(module, dtype)

src/mobius/_builder_test.py

@@ -20,6 +20,7 @@
     flags,
 )
 from mobius._model_package import ModelPackage
+from mobius.tasks import PagedHybridCausalLMTask
 
 
 def _make_value(name: str) -> ir.Value:
@@ -66,12 +67,20 @@ def test_prefill_prefix_pruning_error_lists_supported_tasks() -> None:
     with pytest.raises(
         ValueError,
         match=(
-            "text-generation, hybrid-text-generation, gemma4-text-generation, and gemma4 tasks"
+            "text-generation, hybrid-text-generation, paged-hybrid-text-generation, "
+            "gemma4-text-generation, and gemma4 tasks"
         ),
     ):
         _enable_prefill_prefix_pruning_task("feature-extraction")
 
 
+def test_prefill_prefix_pruning_supports_registered_paged_hybrid_task() -> None:
+    task = _enable_prefill_prefix_pruning_task("paged-hybrid-text-generation")
+
+    assert isinstance(task, PagedHybridCausalLMTask)
+    assert task._prune_prefill_prefix
+
+
 def test_graph_requires_opset24_tensor_scatter() -> None:
     # A TensorScatter node (opset-24-only) must force opset 24 retention.
     node = ir.Node(

src/mobius/components/__init__.py

@@ -101,6 +101,8 @@
     "MobileNetV5Encoder",
     "MoELayer",
     "OffsetRMSNorm",
+    "PagedAttentionState",
+    "PagedHybridContext",
     "PatchEmbed",
     "PatchEmbedding",
     "ParakeetFastConformerEncoder",
@@ -346,6 +348,7 @@
     PaddleOCRProjector,
     YouTuVLProjector,
 )
+from mobius.components._paged_attention import PagedAttentionState, PagedHybridContext
 from mobius.components._paged_mla import (
     PagedCacheState as PagedCacheState,
 )

src/mobius/components/_attention.py

@@ -11,6 +11,7 @@
 
 from mobius._configs import ArchitectureConfig
 from mobius.components._common import Linear
+from mobius.components._paged_attention import PagedAttentionState, paged_attention
 from mobius.components._rms_norm import OffsetRMSNorm, RMSNorm
 from mobius.components._rotary_embedding import apply_rotary_pos_emb
 
@@ -592,7 +593,18 @@ def forward(
         position_embeddings: tuple,
         past_key_value: tuple | None = None,
         static_cache: StaticCacheState | None = None,
+        paged_state: PagedAttentionState | None = None,
     ):
+        if paged_state is not None:
+            if past_key_value is not None or static_cache is not None:
+                raise ValueError(
+                    "Paged attention state cannot be combined with dense cache state"
+                )
+            return self.forward_paged(op, hidden_states, position_embeddings, paged_state)
+        if isinstance(past_key_value, PagedAttentionState):
+            raise TypeError(
+                "PagedAttentionState must be passed through paged_state, not past_key_value"
+            )
         # Q projection (doubled) → split into Q and gate per head
         q_gate = self.q_proj(op, hidden_states)
         # Reshape to per-head view so split separates Q/gate within each head
@@ -650,3 +662,59 @@ def forward(
 
         attn_output = self.o_proj(op, attn_output)
         return attn_output, (present_key, present_value)
+
+    def forward_paged(
+        self,
+        op: OpBuilder,
+        hidden_states: ir.Value,
+        position_embeddings: tuple,
+        state: PagedAttentionState,
+    ):
+        """Packed Qwen attention with external Q/K norm and partial M-RoPE."""
+        q_gate = self.q_proj(op, hidden_states)
+        q_gate = op.Reshape(q_gate, [-1, self.num_attention_heads, self.head_dim * 2])
+        query_states, gate = op.Split(q_gate, num_outputs=2, axis=-1, _outputs=2)
+        key_states = op.Reshape(
+            self.k_proj(op, hidden_states), [-1, self.num_key_value_heads, self.head_dim]
+        )
+        value_states = self.v_proj(op, hidden_states)
+
+        query_states = self.q_norm(op, query_states)
+        key_states = self.k_norm(op, key_states)
+        # RotaryEmbedding's rank-3 ABI is (batch,sequence,hidden), not (N,H,D).
+        # A singleton batch is safe here: only RoPE, never attention, sees it.
+        query_states = op.Reshape(
+            query_states, [1, -1, self.num_attention_heads * self.head_dim]
+        )
+        key_states = op.Reshape(key_states, [1, -1, self.num_key_value_heads * self.head_dim])
+        position_embeddings = tuple(op.Unsqueeze(x, [0]) for x in position_embeddings)
+        query_states = apply_rotary_pos_emb(
+            op,
+            query_states,
+            position_embeddings,
+            self.num_attention_heads,
+            self.rotary_embedding_dim,
+            self._rope_interleave,
+        )
+        key_states = apply_rotary_pos_emb(
+            op,
+            key_states,
+            position_embeddings,
+            self.num_key_value_heads,
+            self.rotary_embedding_dim,
+            self._rope_interleave,
+        )
+        query_states = op.Reshape(query_states, [-1, self.num_attention_heads * self.head_dim])
+        key_states = op.Reshape(key_states, [-1, self.num_key_value_heads * self.head_dim])
+        output, key_cache, value_cache = paged_attention(
+            op,
+            query_states,
+            key_states,
+            value_states,
+            state,
+            num_heads=self.num_attention_heads,
+            kv_num_heads=self.num_key_value_heads,
+        )
+        q_width = self.num_attention_heads * self.head_dim
+        output = op.Mul(output, op.Sigmoid(op.Reshape(gate, [-1, q_width])))
+        return self.o_proj(op, output), (key_cache, value_cache)

src/mobius/components/_attention_test.py

@@ -15,6 +15,7 @@
     make_config,
 )
 from mobius.components._attention import Attention, FusedQKVAttention, Qwen35Attention
+from mobius.components._paged_attention import PagedAttentionState
 
 
 class TestAttention:
@@ -215,6 +216,22 @@ def test_forward_builds_graph(self):
         assert count_op_type(graph, "Attention") >= 1
         assert count_op_type(graph, "Sigmoid") >= 1
 
+    def test_legacy_paged_state_fails_fast(self):
+        config = make_config(partial_rotary_factor=0.5)
+        attn = Qwen35Attention(config)
+        builder, op, _ = create_test_builder()
+        hidden = create_test_input(builder, "hidden", [1, 8, 64])
+        paged_state = PagedAttentionState(*([None] * 6))
+
+        with pytest.raises(TypeError, match="paged_state"):
+            attn(
+                op,
+                hidden,
+                attention_bias=None,
+                position_embeddings=(),
+                past_key_value=paged_state,
+            )
+
 
 class TestGQAContextDispatch:
     """Tests for the GQAContext direct GroupQueryAttention emission path."""