From 3a1ae6e8e4f965a17e6b4eed368f98143fe069e9 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 16 Sep 2026 22:41:54 +0000 Subject: [PATCH 1/6] Initial plan From 508e25acc566e3110d1b37f8f9609bc0d4cf6936 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 16 Sep 2026 23:07:22 +0000 Subject: [PATCH 2/6] Add opt-in packed Qwen hybrid serving ABI Emit native SEPARATE PagedAttention and varlen DeltaNet state, with distinct Engine and workflow state groups. Preserve dense inference and record runtime pins. Add structural, CPU projection parity, and opt-in CUDA state-carry/Engine tests; full CUDA qualification remains outstanding. Signed-off-by: GitHub Co-authored-by: titaiwangms <18010845+titaiwangms@users.noreply.github.com> --- README.md | 37 +- src/mobius/__main__.py | 2 + src/mobius/_builder.py | 24 +- src/mobius/components/__init__.py | 3 + src/mobius/components/_attention.py | 59 +++ src/mobius/components/_gated_deltanet.py | 69 ++++ src/mobius/components/_paged_attention.py | 129 ++++++ src/mobius/components/_rotary_embedding.py | 19 +- src/mobius/integrations/_paged_hybrid.py | 149 +++++++ .../_schema/inference_metadata.schema.json | 38 ++ .../onnx_genai/paged_hybrid_metadata_test.py | 252 ++++++++++++ .../onnx_genai/workflow_metadata.py | 135 +++++++ .../integrations/ort_genai/auto_export.py | 70 +++- .../integrations/ort_genai/genai_config.py | 17 +- .../ort_genai/genai_config_test.py | 32 ++ .../integrations/transformers/_builder.py | 71 +++- src/mobius/models/paged_qwen35_export_test.py | 118 ++++++ src/mobius/models/qwen35.py | 66 +++- src/mobius/tasks/__init__.py | 3 + src/mobius/tasks/_causal_lm.py | 198 ++++++++++ tests/integration/paged_hybrid_test.py | 374 ++++++++++++++++++ 21 files changed, 1826 insertions(+), 39 deletions(-) create mode 100644 src/mobius/components/_paged_attention.py create mode 100644 src/mobius/integrations/_paged_hybrid.py create mode 100644 src/mobius/integrations/onnx_genai/paged_hybrid_metadata_test.py create mode 100644 src/mobius/models/paged_qwen35_export_test.py create mode 100644 tests/integration/paged_hybrid_test.py diff --git a/README.md b/README.md index 0a83d15cd..03311e459 100644 --- a/README.md +++ b/README.md @@ -117,7 +117,7 @@ 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 @@ -125,6 +125,41 @@ 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 +``` + +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: diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index f75387b9c..81ae23437 100644 --- a/src/mobius/__main__.py +++ b/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 diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py index 1abf62d62..dba24cb12 100644 --- a/src/mobius/_builder.py +++ b/src/mobius/_builder.py @@ -76,6 +76,7 @@ def _enable_prefill_prefix_pruning_task(task: str | ModelTask) -> str | ModelTas Gemma4Task, Gemma4TextCausalLMTask, HybridCausalLMTask, + PagedHybridCausalLMTask, ) if task == "text-generation": @@ -94,6 +95,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 +114,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 +169,21 @@ 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 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) diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index 7393061fa..3d1312407 100644 --- a/src/mobius/components/__init__.py +++ b/src/mobius/components/__init__.py @@ -101,6 +101,8 @@ "MobileNetV5Encoder", "MoELayer", "OffsetRMSNorm", + "PagedAttentionState", + "PagedHybridContext", "PatchEmbed", "PatchEmbedding", "ParakeetFastConformerEncoder", @@ -238,6 +240,7 @@ create_decoder_layer, ) from mobius.components._deepseek_mla import DeepSeekMLA as DeepSeekMLA +from mobius.components._paged_attention import PagedAttentionState, PagedHybridContext from mobius.components._diffusion import ( AdaLayerNormOutput, AdaLayerNormZero, diff --git a/src/mobius/components/_attention.py b/src/mobius/components/_attention.py index 2fdd1ee40..d10df9bdc 100644 --- a/src/mobius/components/_attention.py +++ b/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 @@ -593,6 +594,8 @@ def forward( past_key_value: tuple | None = None, static_cache: StaticCacheState | None = None, ): + if isinstance(past_key_value, PagedAttentionState): + return self.forward_paged(op, hidden_states, position_embeddings, 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 +653,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) diff --git a/src/mobius/components/_gated_deltanet.py b/src/mobius/components/_gated_deltanet.py index 44be2185d..505b6ae75 100644 --- a/src/mobius/components/_gated_deltanet.py +++ b/src/mobius/components/_gated_deltanet.py @@ -34,6 +34,10 @@ from mobius._configs import ArchitectureConfig from mobius.components._common import Linear +from mobius.components._paged_attention import ( + gated_delta_net, + varlen_causal_conv_with_state, +) from mobius.components._rms_norm import PostGatedRMSNorm @@ -64,6 +68,7 @@ def forward( op: OpBuilder, input_val: ir.Value, conv_state: ir.Value, + cumulative_sequence_lengths: ir.Value | None = None, ): """Run CausalConvWithState function op. @@ -82,6 +87,10 @@ def forward( op.CastLike(op.Constant(value_float=0.0), self.weight), op.Constant(value_ints=[self._channels]), ) + if cumulative_sequence_lengths is not None: + return varlen_causal_conv_with_state( + op, input_val, self.weight, cumulative_sequence_lengths, conv_bias, conv_state + ) return op.CausalConvWithState( input_val, self.weight, @@ -153,6 +162,10 @@ def __init__(self, config: ArchitectureConfig, linear_class: type | None = None) # Learnable parameters for decay computation self.dt_bias = nn.Parameter([self.num_v_heads]) self.A_log = nn.Parameter([self.num_v_heads]) + if config.export_paged_attention: + # Preserve raw FP32 native gates without changing dense parameter dtypes. + self.dt_bias._keep_float32 = True + self.A_log._keep_float32 = True # Gated output normalization self.norm = PostGatedRMSNorm(self.head_v_dim, eps=config.rms_norm_eps) @@ -166,6 +179,7 @@ def forward( hidden_states: ir.Value, conv_state: ir.Value, recurrent_state: ir.Value, + cumulative_sequence_lengths: ir.Value | None = None, ): """Forward pass for the Gated DeltaNet layer. @@ -183,6 +197,15 @@ def forward( new_conv_state: (batch, conv_dim, kernel_size-1) new_recurrent_state: (batch, num_v_heads, k_dim, v_dim) """ + if cumulative_sequence_lengths is not None: + return self.forward_paged( + op, + hidden_states, + conv_state, + recurrent_state, + cumulative_sequence_lengths, + ) + batch_dim = op.Shape(hidden_states, start=0, end=1) # === Projections === @@ -311,3 +334,49 @@ def forward( output = self.out_proj(op, output) return output, new_conv_state, new_recurrent_state + + def forward_paged( + self, + op: OpBuilder, + hidden_states: ir.Value, + conv_state: ir.Value, + recurrent_state: ir.Value, + cumulative_sequence_lengths: ir.Value, + ): + """Packed native path: ``[N,H]`` activations and V-major recurrent state.""" + mixed_qkv = self.in_proj_qkv(op, hidden_states) # (N, 2*K + V) + z = self.in_proj_z(op, hidden_states) # (N, Hv*Dv) + raw_b = op.Cast(self.in_proj_b(op, hidden_states), to=ir.DataType.FLOAT) + raw_a = op.Cast(self.in_proj_a(op, hidden_states), to=ir.DataType.FLOAT) + + # Invoke the child module so nn realizes its qualified conv1d.weight. + conv_out, new_conv_state = self.conv1d( + op, mixed_qkv, conv_state, cumulative_sequence_lengths + ) + query, key, value = op.Split( + conv_out, + op.Constant(value_ints=[self.key_dim, self.key_dim, self.value_dim]), + axis=-1, + _outputs=3, + ) + query = op.Reshape(query, [-1, self.num_k_heads, self.head_k_dim]) + key = op.Reshape(key, [-1, self.num_k_heads, self.head_k_dim]) + value = op.Reshape(value, [-1, self.num_v_heads, self.head_v_dim]) + + output, new_recurrent_state = gated_delta_net( + op, + query, + key, + value, + cumulative_sequence_lengths, + raw_a, + raw_b, + recurrent_state, + self.A_log, + self.dt_bias, + ) + output = op.Reshape(output, [-1, self.head_v_dim]) + z = op.Reshape(z, [-1, self.head_v_dim]) + output = self.norm(op, output, z) + output = op.Reshape(output, [-1, self.value_dim]) + return self.out_proj(op, output), new_conv_state, new_recurrent_state diff --git a/src/mobius/components/_paged_attention.py b/src/mobius/components/_paged_attention.py new file mode 100644 index 000000000..0ef702b8f --- /dev/null +++ b/src/mobius/components/_paged_attention.py @@ -0,0 +1,129 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Model-agnostic adapters for ORT's packed hybrid serving operators.""" + +from __future__ import annotations + +from typing import NamedTuple + +import onnx_ir as ir +from onnxscript import OpBuilder + + +DOMAIN = "com.microsoft" +ORT_REVISION = "f38538cd5a4b5945a4c839565a8eebc65e1e2ef8" +GENAI_REVISION = "d5b40851ba80ffa8e95b6b01f921dbb9008fac80" + + +class PagedAttentionState(NamedTuple): + """Per-layer SEPARATE paged KV state plus request-level packed metadata.""" + + key_cache: ir.Value + value_cache: ir.Value + cumulative_sequence_lengths: ir.Value + past_sequence_lengths: ir.Value + block_table: ir.Value + attention_metadata: ir.Value + + +class PagedHybridContext(NamedTuple): + """Shared packed-request context passed through a hybrid decoder stack.""" + + cumulative_sequence_lengths: ir.Value + past_sequence_lengths: ir.Value + block_table: ir.Value + attention_metadata: ir.Value + last_token_indices: ir.Value | None + + +def paged_attention( + op: OpBuilder, + query: ir.Value, + key: ir.Value, + value: ir.Value, + state: PagedAttentionState, + *, + num_heads: int, + kv_num_heads: int, +) -> tuple[ir.Value, ir.Value, ir.Value]: + """Emit the pinned 17-input SEPARATE PagedAttention ABI.""" + return op.PagedAttention( + query, + key, + value, + state.key_cache, + state.value_cache, + state.cumulative_sequence_lengths, + state.past_sequence_lengths, + state.block_table, + None, # cos cache: RoPE is applied externally after Q/K normalization + None, # sin cache + None, # slot mapping: the GenAI page manager derives packed writes + None, # attention sinks + None, # q norm: external OffsetRMSNorm must not be repeated + None, # k norm + None, # k scale + None, # v scale + state.attention_metadata, + num_heads=num_heads, + kv_num_heads=kv_num_heads, + kv_cache_layout="SEPARATE", + do_rotary=0, + _domain=DOMAIN, + _outputs=3, + ) + + +def varlen_causal_conv_with_state( + op: OpBuilder, + packed: ir.Value, + weight: ir.Value, + cumulative_sequence_lengths: ir.Value, + bias: ir.Value, + past_conv: ir.Value, +) -> tuple[ir.Value, ir.Value]: + """Emit packed depthwise convolution with fixed per-sequence carry state.""" + return op.VarlenCausalConvWithState( + packed, + weight, + cumulative_sequence_lengths, + bias, + past_conv, + activation="silu", + _domain=DOMAIN, + _outputs=2, + ) + + +def gated_delta_net( + op: OpBuilder, + query: ir.Value, + key: ir.Value, + value: ir.Value, + cumulative_sequence_lengths: ir.Value, + raw_a: ir.Value, + raw_b: ir.Value, + past_recurrent: ir.Value, + a_log: ir.Value, + dt_bias: ir.Value, +) -> tuple[ir.Value, ir.Value]: + """Emit the native packed Qwen GatedDeltaNet recurrence.""" + return op.GatedDeltaNet( + query, + key, + value, + cumulative_sequence_lengths, + raw_a, + raw_b, + past_recurrent, + a_log, + dt_bias, + gate_activation="qwen", + beta_activation="sigmoid", + qk_l2_norm=1, + update_rule="gated_delta", + scale=0.0, + _domain=DOMAIN, + _outputs=2, + ) diff --git a/src/mobius/components/_rotary_embedding.py b/src/mobius/components/_rotary_embedding.py index 4e286d40b..7dd781552 100644 --- a/src/mobius/components/_rotary_embedding.py +++ b/src/mobius/components/_rotary_embedding.py @@ -478,7 +478,9 @@ def __init__( data=ir.tensor(w_mask), ) - def forward(self, op: OpBuilder, position_ids: ir.Value): + def forward( + self, op: OpBuilder, position_ids: ir.Value, *, packed: bool = False + ): """Compute MRoPE cos/sin embeddings. Args: @@ -490,6 +492,21 @@ def forward(self, op: OpBuilder, position_ids: ir.Value): Returns: Tuple of ``(cos, sin)`` each with shape ``(batch, seq, rotary_dim)``. """ + if packed: + # Packed serving supplies exactly (3,N), with no batch padding. + pos_t = op.Squeeze(op.Gather(position_ids, [0], axis=0), [0]) + pos_h = op.Squeeze(op.Gather(position_ids, [1], axis=0), [0]) + pos_w = op.Squeeze(op.Gather(position_ids, [2], axis=0), [0]) + cos_t, sin_t = ( + op.Gather(self.cos_cache, pos_t), + op.Gather(self.sin_cache, pos_t), + ) + cos = op.Where(self.h_mask, op.Gather(self.cos_cache, pos_h), cos_t) + cos = op.Where(self.w_mask, op.Gather(self.cos_cache, pos_w), cos) + sin = op.Where(self.h_mask, op.Gather(self.sin_cache, pos_h), sin_t) + sin = op.Where(self.w_mask, op.Gather(self.sin_cache, pos_w), sin) + return self._cast_embeddings(op, cos, sin) + # For 2D text-only position_ids (batch, seq), expand to (3, batch, seq) # by stacking the same positions for T, H, W dimensions. # For 3D position_ids, this is a no-op reshape. diff --git a/src/mobius/integrations/_paged_hybrid.py b/src/mobius/integrations/_paged_hybrid.py new file mode 100644 index 000000000..5fb9117f1 --- /dev/null +++ b/src/mobius/integrations/_paged_hybrid.py @@ -0,0 +1,149 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Graph-verified packed hybrid ABI shared by runtime metadata exporters.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass + +import onnx_ir as ir + + +@dataclass(frozen=True) +class PagedHybridAbi: + block_size: int + full_layers: tuple[int, ...] + linear_layers: tuple[int, ...] + + def state_groups(self) -> list[dict]: + return [ + {"kind": "paged_kv", "layer_ids": list(self.full_layers)}, + {"kind": "fixed_conv", "layer_ids": list(self.linear_layers)}, + {"kind": "fixed_recurrent", "layer_ids": list(self.linear_layers)}, + ] + + +def inspect_paged_hybrid(model: ir.Model) -> PagedHybridAbi | None: + """Require all three state disciplines, not merely a block-table input.""" + if "mobius.paged_hybrid" not in model.metadata_props: + return None + if model.metadata_props["mobius.paged_hybrid"] != "qwen3_5_text": + raise ValueError("Unknown packed hybrid ABI") + inputs = {value.name: value for value in model.graph.inputs} + outputs = {value.name: value for value in model.graph.outputs} + required = { + "input_ids": (ir.DataType.INT64, 1), + "position_ids": (ir.DataType.INT64, 2), + "block_table": (ir.DataType.INT32, 2), + "cumulative_sequence_lengths": (ir.DataType.INT32, 1), + "past_sequence_lengths": (ir.DataType.INT32, 1), + "attention_metadata": (ir.DataType.INT32, 1), + } + for name, (dtype, rank) in required.items(): + value = inputs.get(name) + if ( + value is None + or value.dtype != dtype + or value.shape is None + or len(value.shape) != rank + ): + raise ValueError(f"Packed hybrid input {name!r} must have {dtype} rank {rank}") + if inputs["position_ids"].shape[0] != 3 or inputs["attention_metadata"].shape != ir.Shape( + [3] + ): + raise ValueError("Packed hybrid requires three position planes and three CPU bounds") + if "attention_mask" in inputs or "slot_mapping" in inputs: + raise ValueError("Packed hybrid does not expose attention_mask or slot_mapping") + logits = outputs.get("logits") + if ( + logits is None + or logits.dtype != ir.DataType.FLOAT + or logits.shape is None + or len(logits.shape) != 2 + ): + raise ValueError("Packed hybrid logits must be rank-2 FLOAT") + block_size = int(model.metadata_props.get("mobius.paged_block_size", "0")) + if block_size <= 0 or block_size % 256: + raise ValueError("Packed hybrid block size must be a positive multiple of 256") + layers: dict[str, set[int]] = { + kind: set() for kind in ("key", "value", "conv_state", "recurrent_state") + } + for name, past in inputs.items(): + match = re.fullmatch( + r"past_key_values\.(\d+)\.(key|value|conv_state|recurrent_state)", name or "" + ) + if match is None: + if name not in required: + raise ValueError(f"Unknown packed hybrid input {name!r}") + continue + layer, kind = int(match[1]), match[2] + layers[kind].add(layer) + present = outputs.get(f"present.{layer}.{kind}") + rank = 3 if kind == "conv_state" else 4 + dtype = ir.DataType.FLOAT if kind == "recurrent_state" else past.dtype + if ( + past.shape is None + or len(past.shape) != rank + or dtype not in {ir.DataType.FLOAT16, ir.DataType.BFLOAT16, ir.DataType.FLOAT} + or (kind != "recurrent_state" and dtype == ir.DataType.FLOAT) + or past.dtype != dtype + or present is None + or present.dtype != dtype + or present.shape != past.shape + ): + raise ValueError(f"Invalid packed hybrid state pair {name!r}") + if kind in {"key", "value"} and past.shape[1] != block_size: + raise ValueError(f"Page extent disagrees with block size for {name!r}") + full, linear = layers["key"], layers["conv_state"] + if ( + not full + or not linear + or full != layers["value"] + or linear != layers["recurrent_state"] + or full & linear + or full | linear != set(range(max(full | linear) + 1)) + ): + raise ValueError( + "Packed hybrid requires disjoint, complete paged and fixed layer groups" + ) + native = {} + state_operands = {"PagedAttention": 3, "VarlenCausalConvWithState": 4, "GatedDeltaNet": 6} + for node in model.graph: + if node.domain != "com.microsoft" or node.op_type not in state_operands: + continue + operand = state_operands[node.op_type] + if len(node.inputs) <= operand or node.inputs[operand] is None: + raise ValueError(f"Missing native {node.op_type} state operand") + key = (node.op_type, node.inputs[operand].name) + if key in native: + raise ValueError(f"Duplicate native state consumer {key}") + native[key] = node + for kind, ids, op_type in ( + ("key", full, "PagedAttention"), + ("conv_state", linear, "VarlenCausalConvWithState"), + ("recurrent_state", linear, "GatedDeltaNet"), + ): + for layer in ids: + node = native.get((op_type, f"past_key_values.{layer}.{kind}")) + if node is None: + raise ValueError(f"Missing native {op_type} for layer {layer}") + cu_index = ( + 5 + if op_type == "PagedAttention" + else (2 if op_type == "VarlenCausalConvWithState" else 3) + ) + if node.inputs[cu_index] is not inputs["cumulative_sequence_lengths"]: + raise ValueError(f"Missing packed sequence boundaries at layer {layer}") + if op_type == "PagedAttention" and ( + node.attributes.get_string("kv_cache_layout", "") != "SEPARATE" + or len(node.inputs) != 17 + or any(value is not None for value in node.inputs[8:16]) + or node.inputs[16] is not inputs["attention_metadata"] + or node.inputs[4] is not inputs[f"past_key_values.{layer}.value"] + or node.inputs[6] is not inputs["past_sequence_lengths"] + or node.inputs[7] is not inputs["block_table"] + ): + raise ValueError(f"Invalid SEPARATE PagedAttention contract at layer {layer}") + return PagedHybridAbi(block_size, tuple(sorted(full)), tuple(sorted(linear))) diff --git a/src/mobius/integrations/onnx_genai/_schema/inference_metadata.schema.json b/src/mobius/integrations/onnx_genai/_schema/inference_metadata.schema.json index 48fbb616b..486be7916 100644 --- a/src/mobius/integrations/onnx_genai/_schema/inference_metadata.schema.json +++ b/src/mobius/integrations/onnx_genai/_schema/inference_metadata.schema.json @@ -3654,6 +3654,44 @@ "StateUpdate": { "description": "How a state group's buffers absorb each step's new positions.\n\nBoth variants describe what the GRAPH does. Neither selects a storage\nstrategy, a slot allocator, or a device: a runtime is free to back an\n`append` group with a fixed arena or an `indexed_scatter` group with paged\nstorage, so long as the graph sees what it declared.", "oneOf": [ + { + "additionalProperties": false, + "description": "Packed tokens update caller-owned page buffers in place. Logical positions are mapped through the block table; neither page axis is a request batch axis. Request boundaries and committed lengths identify writes without exposing a slot_mapping input.", + "properties": { + "kind": { + "const": "paged_scatter", + "type": "string" + }, + "block_size": { + "type": "integer", + "minimum": 1, + "description": "Graph-visible number of token slots in each page." + }, + "block_table_ports": { + "type": "object", + "additionalProperties": {"type": "string"}, + "description": "Per-component logical-to-physical page table input." + }, + "past_sequence_length_ports": { + "type": "object", + "additionalProperties": {"type": "string"}, + "description": "Per-component input carrying committed lengths before the invocation." + }, + "cumulative_sequence_length_ports": { + "type": "object", + "additionalProperties": {"type": "string"}, + "description": "Per-component exclusive prefix sums delimiting packed query rows." + } + }, + "required": [ + "kind", + "block_size", + "block_table_ports", + "past_sequence_length_ports", + "cumulative_sequence_length_ports" + ], + "type": "object" + }, { "additionalProperties": false, "description": "Each step's positions extend the buffer along `sequence_axis`.\n\nThe valid region is the whole tensor, so no write cursor is graph-visible\nand the buffer's shape carries the length.", diff --git a/src/mobius/integrations/onnx_genai/paged_hybrid_metadata_test.py b/src/mobius/integrations/onnx_genai/paged_hybrid_metadata_test.py new file mode 100644 index 000000000..acbbcf116 --- /dev/null +++ b/src/mobius/integrations/onnx_genai/paged_hybrid_metadata_test.py @@ -0,0 +1,252 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Packed hybrid metadata must describe pages and fixed state independently.""" + +from __future__ import annotations + +import dataclasses +import json +from pathlib import Path + +import jsonschema +import numpy as np +import onnx_ir as ir +import onnx_ir.passes.common as common_passes +import onnxruntime as ort +import pytest + +from mobius import build_from_module +from mobius._testing import make_config +from mobius.integrations._paged_hybrid import inspect_paged_hybrid +from mobius.integrations.onnx_genai.workflow_metadata import build_decoder_workflow_metadata +from mobius.integrations.ort_genai.auto_export import _inspect_decoder_abi, _write_genai_config +from mobius.models.qwen35 import Qwen35CausalLMModel +from mobius.tasks import HybridCausalLMTask, PagedHybridCausalLMTask + + +def tiny_config(dtype=ir.DataType.FLOAT16): + return make_config( + model_type="qwen3_5_text", + dtype=dtype, + hidden_size=128, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=64, + intermediate_size=192, + vocab_size=128, + max_position_embeddings=2048, + num_hidden_layers=4, + layer_types=["linear_attention"] * 3 + ["full_attention"], + partial_rotary_factor=0.5, + mrope_section=[8, 4, 4], + mrope_interleaved=True, + linear_num_value_heads=2, + linear_num_key_heads=1, + linear_key_head_dim=64, + linear_value_head_dim=32, + linear_conv_kernel_dim=4, + ) + + +@pytest.fixture +def package(): + config = tiny_config() + return build_from_module( + Qwen35CausalLMModel(config), + config, + task=PagedHybridCausalLMTask(), + execution_provider="cuda", + ) + + +def test_exact_state_disciplines_and_application_scheduling(package): + metadata = build_decoder_workflow_metadata(package, package.config) + schema = json.loads( + (Path(__file__).with_name("_schema") / "inference_metadata.schema.json").read_text() + ) + jsonschema.validate(metadata, schema) + workflow = metadata["pipeline"]["workflow"] + assert workflow["steps"][0]["kind"] == "invoke" + assert all(step["kind"] == "emit" for step in workflow["steps"][1:]) + assert all( + value["source"]["kind"] == "application" for value in workflow["inputs"].values() + ) + groups = workflow["serving"]["state_service"]["groups"] + assert set(groups) == {"paged_kv", "fixed_conv", "fixed_recurrent"} + assert groups["paged_kv"]["update"]["kind"] == "paged_scatter" + assert groups["paged_kv"]["aliasing"] == "required" + for name in ("fixed_conv", "fixed_recurrent"): + assert groups[name]["update"] == {"kind": "replace"} + assert "sequence_axis" not in groups[name] + aliases = groups["paged_kv"]["ports"]["decoder"] + assert {alias["layer"] for alias in aliases.values()} == {3} + assert {alias["role"] for alias in aliases.values()} == {"key", "value"} + for cell in workflow["state"].values(): + if cell["service_group"] == "paged_kv": + assert "batch_layout" not in cell["contract"] + else: + assert cell["contract"]["batch_layout"]["kind"] == "request_aligned" + + +def test_genai_export_uses_verified_groups_and_no_capture(package, tmp_path): + _write_genai_config( + package.config, + str(tmp_path), + pkg=package, + ort_model_type="decoder", + context_length=2048, + ep="cuda", + bos_token_id=1, + eos_token_id=127, + pad_token_id=0, + is_vlm=False, + has_speech=False, + ) + config = json.loads((tmp_path / "genai_config.json").read_text()) + decoder = config["model"]["decoder"] + assert decoder["state_groups"] == [ + {"kind": "paged_kv", "layer_ids": [3]}, + {"kind": "fixed_conv", "layer_ids": [0, 1, 2]}, + {"kind": "fixed_recurrent", "layer_ids": [0, 1, 2]}, + ] + assert decoder["inputs"]["past_conv_names"] == "past_key_values.%d.conv_state" + assert decoder["inputs"]["past_recurrent_names"] == "past_key_values.%d.recurrent_state" + assert config["engine"]["dynamic_batching"]["block_size"] == 256 + assert config["search"]["past_present_share_buffer"] is True + assert ( + decoder["session_options"]["provider_options"][0]["cuda"].get("enable_cuda_graph", "0") + != "1" + ) + + +def test_packed_export_precision_and_rotary_shapes(package): + graph = package["model"].graph + assert graph.outputs[0].dtype == ir.DataType.FLOAT + for layer in range(3): + for suffix in ("A_log", "dt_bias"): + assert ( + graph.initializers[f"model.layers.{layer}.linear_attn.{suffix}"].dtype + == ir.DataType.FLOAT + ) + assert f"model.layers.{layer}.linear_attn.conv1d.weight" in graph.initializers + assert not any( + key[1] + in { + "GatedDeltaNet", + "LinearAttention", + "VarlenCausalConvWithState", + "CausalConvWithState", + } + for key in package["model"].functions + ) + assert not any( + node.op_type in {"Attention", "LinearAttention", "CausalConvWithState"} + for node in graph + ) + + +@pytest.mark.parametrize("mutation", ["missing_state", "bad_dtype", "bad_pages", "bad_node"]) +def test_incomplete_hybrid_cannot_claim_engine_compatibility(package, mutation): + model = package["model"] + if mutation == "missing_state": + model.graph.inputs.remove( + next(v for v in model.graph.inputs if v.name.endswith("0.conv_state")) + ) + elif mutation == "bad_dtype": + model.graph.outputs[0].type = ir.TensorType(ir.DataType.FLOAT16) + elif mutation == "bad_pages": + model.metadata_props["mobius.paged_block_size"] = "512" + else: + node = next(node for node in model.graph if node.op_type == "PagedAttention") + model.graph.remove(node, safe=False) + with pytest.raises(ValueError): + _inspect_decoder_abi(model, model_type="decoder") + + +def test_unmarked_block_table_is_not_hybrid_compatibility(package): + del package["model"].metadata_props["mobius.paged_hybrid"] + assert inspect_paged_hybrid(package["model"]) is None + + +@pytest.mark.parametrize("ep", ["default", "cpu", "dml"]) +def test_direct_task_rejects_non_cuda(ep): + config = tiny_config() + with pytest.raises(ValueError, match="requires execution_provider"): + build_from_module( + Qwen35CausalLMModel(config), + config, + PagedHybridCausalLMTask(), + execution_provider=ep, + ) + + +def test_dense_gate_parameters_still_follow_compute_dtype(): + from mobius._builder import _cast_module_dtype + + config = dataclasses.replace(tiny_config(), export_paged_attention=False) + module = Qwen35CausalLMModel(config) + _cast_module_dtype(module, config.dtype) + parameters = dict(module.named_parameters()) + assert parameters["model.layers.0.linear_attn.A_log"].dtype == config.dtype + assert parameters["model.layers.0.linear_attn.dt_bias"].dtype == config.dtype + + +def test_packed_qk_norm_and_interleaved_mrope_match_dense_on_cpu(tmp_path): + """Execute the real projection/RoPE subgraphs without requiring paged kernels.""" + config = dataclasses.replace( + tiny_config(), + num_hidden_layers=2, + layer_types=["full_attention", "linear_attention"], + ) + sessions = [] + for packed in (False, True): + module = Qwen35CausalLMModel(config) + rng = np.random.default_rng(51) + for _, parameter in module.named_parameters(): + if parameter.const_value is None: + parameter.const_value = ir.tensor( + rng.normal(0, 0.1, tuple(parameter.shape)).astype(np.float32) + ) + package = build_from_module( + module, + config, + PagedHybridCausalLMTask() if packed else HybridCausalLMTask(), + execution_provider="cuda" if packed else "cpu", + ) + model = package["model"] + attention = next( + node + for node in model.graph + if node.op_type == ("PagedAttention" if packed else "Attention") + ) + model.graph.outputs.clear() + for index, value in enumerate(attention.inputs[:3]): + value.name = f"projected_{index}" + model.graph.outputs.append(value) + common_passes.RemoveUnusedNodesPass()(model) + common_passes.RemoveUnusedFunctionsPass()(model) + # Production cleanup retains cache ports as an ABI promise. This + # projection-only test has deliberately removed that ABI. + for value in list(model.graph.inputs): + if not value.uses(): + model.graph.inputs.remove(value) + path = tmp_path / f"projections_{packed}.onnx" + ir.save(model, path) + options = ort.SessionOptions() + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL + sessions.append( + ort.InferenceSession(str(path), options, providers=["CPUExecutionProvider"]) + ) + tokens = np.asarray([[4, 7, 9], [12, 5, 3]], np.int64) + positions = np.asarray([[17, 18, 19], [41, 42, 43]], np.int64) + dense = sessions[0].run(None, {"input_ids": tokens, "position_ids": positions}) + packed = sessions[1].run( + None, {"input_ids": tokens.ravel(), "position_ids": np.tile(positions.ravel(), (3, 1))} + ) + for reference, actual in zip(dense, packed): + # Attention accepts both (B,H,S,D) and (B,S,H*D). + if reference.ndim == 4: + reference = reference.transpose(0, 2, 1, 3) + reference = reference.reshape(actual.shape) + np.testing.assert_allclose(actual, reference, atol=2e-3, rtol=2e-3) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index ce88cd537..7a0c07629 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -7945,6 +7945,8 @@ def build_decoder_workflow_metadata( """Build the exact workflow-policy contract for an autoregressive decoder.""" if len(pkg) != 1: raise ValueError("decoder workflow requires exactly one neural component") + if "mobius.paged_hybrid" in next(iter(pkg.values())).metadata_props: + return build_paged_hybrid_workflow_metadata(pkg) metadata = _build_autoregressive_workflow_metadata( pkg, config, sampler=sampler, source=source ) @@ -7952,6 +7954,139 @@ def build_decoder_workflow_metadata( return metadata +def build_paged_hybrid_workflow_metadata(pkg: Any) -> dict[str, Any]: + """Describe one externally scheduled packed invocation, not a dense loop. + + Engine owns admission, token packing, page allocation and request identity. + All scheduling tensors and current state are explicit application inputs; + every successor is returned. In particular, page buffers must not receive + the batch-row permutation used for the two fixed-state groups. + """ + from mobius.integrations._paged_hybrid import inspect_paged_hybrid + + if set(pkg) != {"model"}: + raise ValueError("Packed hybrid workflow requires one model component") + model = pkg["model"] + abi = inspect_paged_hybrid(model) + if abi is None: + raise ValueError("Packed hybrid workflow requires a verified packed hybrid ABI") + workflow_inputs = { + f"request.{value.name}": { + "contract": _contract(value), + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": f"request.{value.name}"}, + "required": True, + } + for value in model.graph.inputs + } + workflow_outputs = { + value.name: {"contract": _contract(value), "role": "tensor", "stage": "post_adapter"} + for value in model.graph.outputs + } + # The external scheduler, not this single neural invocation, owns request + # completion. These are workflow lifecycle inputs, not extra ONNX ports. + for name in ("active", "done"): + workflow_inputs[f"request.{name}"] = { + "contract": {"dtype": "bool", "rank": 1, "shape": ["batch"]}, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": f"request.{name}"}, + "required": True, + } + groups = {} + state = {} + for spec in abi.state_groups(): + name = spec["kind"] + paged = name == "paged_kv" + suffixes = ( + ("key", "value") + if paged + else ("conv_state" if name == "fixed_conv" else "recurrent_state",) + ) + aliases = {} + for layer in spec["layer_ids"]: + for suffix in suffixes: + port = f"past_key_values.{layer}.{suffix}" + cell = f"layer_{layer}_{suffix}" + aliases[cell] = { + "input": port, + "output": f"present.{layer}.{suffix}", + "layer": layer, + "access": "read_write", + **({"role": suffix} if paged else {}), + } + state[cell] = { + "contract": workflow_inputs[f"request.{port}"]["contract"], + "scope": "invocation", + "initializer": f"request.{port}", + "recurrence": {"kind": "invariant"}, + "management": "external", + "release_boundary": "invocation", + "service_group": name, + } + groups[name] = { + "kind": "full_attention" if paged else "recurrent", + "layout": "pkhd" if paged else ("bcw" if name == "fixed_conv" else "bhvk"), + "aliasing": "required" if paged else "permitted", + "update": ( + { + "kind": "paged_scatter", + "block_size": abi.block_size, + "block_table_ports": {"decoder": "block_table"}, + "past_sequence_length_ports": {"decoder": "past_sequence_lengths"}, + "cumulative_sequence_length_ports": { + "decoder": "cumulative_sequence_lengths" + }, + } + if paged + else {"kind": "replace"} + ), + "capabilities": { + "snapshot": False, + "fork": False, + "cascade": [ + other["kind"] for other in abi.state_groups() if other["kind"] != name + ], + }, + "ports": {"decoder": aliases}, + } + workflow = { + "manifest": {"capabilities": ["workflow_ssa", "typed_emit"]}, + "inputs": workflow_inputs, + "outputs": workflow_outputs, + "state": state, + "components": {"decoder": _component(model, "model.onnx")}, + "serving": { + "active": "request.active", + "done": "request.done", + "state_service": {"groups": groups}, + }, + "graph": { + "kind": "sequence", + "nodes": [ + _invoke( + "decoder", + {value.name: f"request.{value.name}" for value in model.graph.inputs}, + {value.name: f"decoder.{value.name}" for value in model.graph.outputs}, + ), + *[ + { + "kind": "emit", + "value": f"decoder.{value.name}", + "output": value.name, + "mode": "replace", + } + for value in model.graph.outputs + ], + ], + }, + } + return { + "schema_version": "v1.1", + "required_capabilities": ["packed_hybrid", "paged_scatter"], + "pipeline": {"workflow": _publish_workflow_v1(workflow)}, + } + + def _build_autoregressive_workflow_metadata( pkg: Any, config: Any, diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index ef906199c..68ec40cbd 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -187,6 +187,8 @@ class _DecoderAbi: outputs: dict[str, str] cache_slots: int has_recurrent_state: bool + state_groups: list[dict[str, Any]] | None = None + paged_block_size: int | None = None _GEMMA4_MODEL_TYPES = frozenset( @@ -386,6 +388,10 @@ def _inspect_decoder_abi(model: ir.Model, *, model_type: str) -> _DecoderAbi: recurrent_indices = set(input_cache.get("conv_state", {})) has_recurrent_state = bool(input_cache.get("recurrent_state")) + from mobius.integrations._paged_hybrid import inspect_paged_hybrid + + paged_hybrid = inspect_paged_hybrid(model) + is_paged_hybrid = paged_hybrid is not None if model_type == "lfm2": if recurrent_indices != set(output_cache.get("conv_state", {})): raise ValueError("LFM2 conv_state outputs must match its conv_state inputs") @@ -424,30 +430,39 @@ def _inspect_decoder_abi(model: ir.Model, *, model_type: str) -> _DecoderAbi: output_cache["conv_state"], label="present-convolution" ) elif recurrent_indices: - expected_input_prefix = decoder_inputs["past_key_names"].rsplit(".", 1)[0] - expected_output_prefix = decoder_outputs["present_key_names"].rsplit(".", 1)[0] - recurrent_templates = { - _name_template(input_cache["conv_state"], label="past-convolution"), - _name_template(input_cache["recurrent_state"], label="past-recurrent"), - } - present_templates = { - _name_template(output_cache["conv_state"], label="present-convolution"), - _name_template(output_cache["recurrent_state"], label="present-recurrent"), - } - # Preserve graph-derived names even when a runtime derives different - # names from the key-cache templates. - del ( - expected_input_prefix, - expected_output_prefix, - recurrent_templates, - present_templates, - ) + if is_paged_hybrid: + decoder_inputs["past_conv_names"] = _name_template( + input_cache["conv_state"], label="past-convolution" + ) + decoder_inputs["past_recurrent_names"] = _name_template( + input_cache["recurrent_state"], label="past-recurrent" + ) + decoder_outputs["present_conv_names"] = _name_template( + output_cache["conv_state"], label="present-convolution" + ) + decoder_outputs["present_recurrent_names"] = _name_template( + output_cache["recurrent_state"], label="present-recurrent" + ) + else: + # Preserve graph-derived names for diagnostics even when the + # released generic runtime has no fields for these templates. + _name_template(input_cache["conv_state"], label="past-convolution") + _name_template(input_cache["recurrent_state"], label="past-recurrent") + _name_template(output_cache["conv_state"], label="present-convolution") + _name_template(output_cache["recurrent_state"], label="present-recurrent") all_indices = set().union(*(set(indices) for indices in input_cache.values())) + state_groups = None + paged_block_size = None + if paged_hybrid is not None: + state_groups = paged_hybrid.state_groups() + paged_block_size = paged_hybrid.block_size return _DecoderAbi( inputs=decoder_inputs, outputs=decoder_outputs, cache_slots=max(all_indices) + 1, has_recurrent_state=has_recurrent_state, + state_groups=state_groups, + paged_block_size=paged_block_size, ) @@ -1543,6 +1558,10 @@ def _write_genai_config( # config, raise a clear error so the caller picks an EP/dtype combination # (e.g. fp32 on CPU) that lowers full attention to GQA. supports_in_place_kv_cache: bool | None = None + is_paged_hybrid = ( + decoder_model is not None + and decoder_model.metadata_props.get("mobius.paged_hybrid") == "qwen3_5_text" + ) if decoder_model is not None: has_gqa = any( node.op_type == "GroupQueryAttention" and node.domain == "com.microsoft" @@ -1556,7 +1575,9 @@ def _write_genai_config( node.op_type == "Attention" and node.domain in ("", "ai.onnx") for node in decoder_model.graph ) - if has_recurrent_state and has_standard_attention: + if is_paged_hybrid: + supports_in_place_kv_cache = True + elif has_recurrent_state and has_standard_attention: supports_in_place_kv_cache = True else: supports_in_place_kv_cache = has_gqa or has_recurrent_state @@ -1607,6 +1628,17 @@ def _write_genai_config( ), sliding_window=sliding_window, has_specialized_topology=not _is_single_model_decoder_package(pkg), + decoder_graph_capture=False if is_paged_hybrid else None, + state_groups=decoder_abi.state_groups if decoder_abi is not None else None, + dynamic_batching=( + { + "block_size": decoder_abi.paged_block_size, + "max_batch_size": 100, + "gpu_utilization_factor": 0.6, + } + if decoder_abi is not None and decoder_abi.paged_block_size is not None + else None + ), ) generator.with_special_tokens( **_special_token_ids_from_tokenizer_config(output_dir, config.vocab_size) diff --git a/src/mobius/integrations/ort_genai/genai_config.py b/src/mobius/integrations/ort_genai/genai_config.py index 70b47e4a2..4836db6fa 100644 --- a/src/mobius/integrations/ort_genai/genai_config.py +++ b/src/mobius/integrations/ort_genai/genai_config.py @@ -221,6 +221,8 @@ def __init__( sliding_window: dict[str, Any] | None = None, uses_longrope: bool = False, has_specialized_topology: bool = False, + state_groups: list[dict[str, Any]] | None = None, + dynamic_batching: dict[str, Any] | None = None, ): self.model_type = model_type self.vocab_size = vocab_size @@ -252,6 +254,8 @@ def __init__( self._sliding_window = sliding_window self._uses_longrope = uses_longrope self._has_specialized_topology = has_specialized_topology + self._state_groups = state_groups + self._dynamic_batching = dynamic_batching # Optional VLM fields (set via with_vision()) self._vision: dict[str, Any] | None = None @@ -283,6 +287,9 @@ def from_config( num_cache_layer_slots: int | None = None, sliding_window: dict[str, Any] | None = None, has_specialized_topology: bool = False, + decoder_graph_capture: bool | None = None, + state_groups: list[dict[str, Any]] | None = None, + dynamic_batching: dict[str, Any] | None = None, ) -> GenaiConfigGenerator: """Create a generator from a BaseModelConfig-like dataclass. @@ -331,6 +338,7 @@ def from_config( decoder_outputs=decoder_outputs, decoder_filename=decoder_filename, supports_in_place_kv_cache=supports_in_place_kv_cache, + decoder_graph_capture=decoder_graph_capture, layer_types=getattr(config, "layer_types", None), conv_cache_size=( getattr(config, "short_conv_kernel", 1) - 1 @@ -343,6 +351,8 @@ def from_config( and getattr(config, "rope_type", None) == "longrope" ), has_specialized_topology=has_specialized_topology, + state_groups=state_groups, + dynamic_batching=dynamic_batching, ) def with_vision( @@ -627,6 +637,8 @@ def generate(self) -> dict[str, Any]: decoder["outputs"].setdefault("present_conv_names", "present.%d.conv_state") if self._sliding_window is not None: decoder["sliding_window"] = self._sliding_window + if self._state_groups is not None: + decoder["state_groups"] = self._state_groups # Model section model: dict[str, Any] = { @@ -658,10 +670,13 @@ def generate(self) -> dict[str, Any]: model["speech"] = self._audio model.update(self._vlm_token_ids) - return { + result = { "model": model, "search": search, } + if self._dynamic_batching is not None: + result["engine"] = {"dynamic_batching": self._dynamic_batching} + return result def write(self, output_dir: str) -> str: """Write genai_config.json to the output directory. diff --git a/src/mobius/integrations/ort_genai/genai_config_test.py b/src/mobius/integrations/ort_genai/genai_config_test.py index 5a22044e1..7e9140279 100644 --- a/src/mobius/integrations/ort_genai/genai_config_test.py +++ b/src/mobius/integrations/ort_genai/genai_config_test.py @@ -16,6 +16,38 @@ ) +def test_packed_state_groups_and_dynamic_batching_are_explicit(): + config = GenaiConfigGenerator( + "qwen3_5_text", + vocab_size=100, + hidden_size=64, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + ep="cuda", + supports_in_place_kv_cache=True, + decoder_graph_capture=False, + state_groups=[ + {"kind": "paged_kv", "layer_ids": [3]}, + {"kind": "fixed_conv", "layer_ids": [0, 1, 2]}, + {"kind": "fixed_recurrent", "layer_ids": [0, 1, 2]}, + ], + dynamic_batching={ + "block_size": 256, + "max_batch_size": 100, + "gpu_utilization_factor": 0.6, + }, + ).generate() + assert config["model"]["type"] == "decoder" + assert config["model"]["decoder"]["state_groups"][0] == { + "kind": "paged_kv", + "layer_ids": [3], + } + assert config["engine"]["dynamic_batching"]["block_size"] == 256 + assert config["search"]["past_present_share_buffer"] is True + + class TestGenaiConfigGeneratorLLM: """Test genai_config generation for decoder-only LLMs.""" diff --git a/src/mobius/integrations/transformers/_builder.py b/src/mobius/integrations/transformers/_builder.py index 1059ecd37..5fda99afb 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -609,15 +609,70 @@ def build_transformers_model( ) config = dataclasses.replace(config, use_dsa=False) if export_paged_attention: - from mobius.components._paged_mla import paged_attention_rejection - config = dataclasses.replace(config, export_paged_attention=True) - reason = paged_attention_rejection(config) - if reason is not None: - raise ValueError( - "export_paged_attention=True (--features paged-attention) is not " - f"supported for model_type '{model_type}': {reason}" - ) + from mobius.tasks import CausalLMTask + + paged_task_placeholder = isinstance(task, CausalLMTask) and getattr( + task, "_paged_cache", False + ) + if model_type == "qwen3_5_text": + from mobius.tasks import PagedHybridCausalLMTask + + unsupported_options = [ + name + for name, enabled in { + "fp8_kv_cache": fp8_kv_cache, + "kv_cache_scales": kv_cache_scales is not None, + "output_layer_indices": output_layer_indices is not None, + "glm_full_attention": glm_full_attention, + }.items() + if enabled + ] + if unsupported_options: + raise ValueError( + "Qwen packed paged-attention does not support options: " + + ", ".join(unsupported_options) + ) + if execution_provider != "cuda": + raise ValueError( + "Qwen packed paged-attention requires execution_provider='cuda'; " + f"got {execution_provider!r}" + ) + if ( + task is not None + and not paged_task_placeholder + and not isinstance(task, PagedHybridCausalLMTask) + ): + raise ValueError( + "Qwen packed paged-attention owns its task and cannot be combined " + f"with task={task!r}" + ) + if config.dtype not in {ir.DataType.FLOAT16, ir.DataType.BFLOAT16}: + raise ValueError( + "Qwen packed paged-attention requires dtype f16 or bf16; " + f"got {config.dtype!r}" + ) + if config.quantization is not None and config.quantization.quant_method != "none": + raise ValueError( + "Qwen packed paged-attention does not support quantized checkpoints" + ) + if not isinstance(task, PagedHybridCausalLMTask): + task = PagedHybridCausalLMTask() + else: + from mobius.components._paged_mla import paged_attention_rejection + + reason = paged_attention_rejection(config) + if reason is not None: + raise ValueError( + "export_paged_attention=True (--features paged-attention) is not " + f"supported for model_type '{model_type}': {reason}" + ) + if task is not None and not paged_task_placeholder: + raise ValueError( + f"PagedAttention owns its task and cannot be combined with task={task!r}" + ) + if task is None: + task = CausalLMTask(paged_cache=True) if task is None: task = _default_task_for_model(model_type) diff --git a/src/mobius/models/paged_qwen35_export_test.py b/src/mobius/models/paged_qwen35_export_test.py new file mode 100644 index 000000000..973b730bf --- /dev/null +++ b/src/mobius/models/paged_qwen35_export_test.py @@ -0,0 +1,118 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Structural tests for the packed Qwen3.5 hybrid serving graph.""" + +from __future__ import annotations + +import onnx_ir as ir +import pytest + +from mobius._testing import make_config +from mobius.models.qwen35 import Qwen35CausalLMModel +from mobius.tasks import PagedHybridCausalLMTask + + +def _config(*, layers: int = 2, layer_types: list[str] | None = None): + return make_config( + model_type="qwen3_5_text", + dtype=ir.DataType.FLOAT16, + num_hidden_layers=layers, + layer_types=layer_types or ["full_attention", "linear_attention"], + partial_rotary_factor=0.5, + mrope_section=[2, 1, 1], + mrope_interleaved=True, + linear_num_value_heads=4, + linear_num_key_heads=2, + linear_key_head_dim=16, + linear_value_head_dim=16, + linear_conv_kernel_dim=4, + ) + + +def _build(config, *, prune: bool = False): + return PagedHybridCausalLMTask(prune_prefill_prefix=prune).build( + Qwen35CausalLMModel(config), config + )["model"] + + +def _nodes(model, op_type: str): + return [node for node in model.graph if node.op_type == op_type] + + +def test_packed_hybrid_io_and_native_operands(): + model = _build(_config()) + inputs = {value.name: value for value in model.graph.inputs} + assert {"input_ids", "position_ids", "block_table"} <= inputs.keys() + assert "attention_mask" not in inputs + assert "slot_mapping" not in inputs + assert inputs["input_ids"].shape == ir.Shape(["num_tokens"]) + assert inputs["position_ids"].shape == ir.Shape([3, "num_tokens"]) + assert inputs["attention_metadata"].dtype == ir.DataType.INT32 + assert inputs["past_key_values.0.key"].shape[1:] == ir.Shape([256, 2, 16]) + assert inputs["past_key_values.1.recurrent_state"].shape[1:] == ir.Shape( + [4, 16, 16] + ) + assert inputs["past_key_values.1.recurrent_state"].dtype == ir.DataType.FLOAT + + paged = _nodes(model, "PagedAttention") + conv = _nodes(model, "VarlenCausalConvWithState") + delta = _nodes(model, "GatedDeltaNet") + assert len(paged) == len(conv) == len(delta) == 1 + assert paged[0].domain == conv[0].domain == delta[0].domain == "com.microsoft" + assert len(paged[0].inputs) == 17 + assert all(value is None for value in paged[0].inputs[8:16]) + assert paged[0].inputs[16].name == "attention_metadata" + attrs = {name: attr.value for name, attr in paged[0].attributes.items()} + assert attrs["kv_cache_layout"] == "SEPARATE" + assert attrs["do_rotary"] == 0 + delta_attrs = {name: attr.value for name, attr in delta[0].attributes.items()} + assert delta_attrs["gate_activation"] == "qwen" + assert delta_attrs["qk_l2_norm"] == 1 + assert delta[0].inputs[4].dtype == ir.DataType.FLOAT + assert delta[0].inputs[5].dtype == ir.DataType.FLOAT + + assert model.metadata_props["mobius.paged_hybrid"] == "qwen3_5_text" + assert model.metadata_props["mobius.paged_block_size"] == "256" + assert "mobius.ort_revision" in model.metadata_props + assert "mobius.genai_revision" in model.metadata_props + + +def test_fp32_native_decay_parameters_and_qualified_weight_names(): + model = _build(_config()) + initializers = model.graph.initializers + for name in ( + "model.layers.1.linear_attn.A_log", + "model.layers.1.linear_attn.dt_bias", + ): + assert initializers[name].dtype == ir.DataType.FLOAT + assert "model.layers.0.self_attn.q_norm.weight" in initializers + assert "model.layers.1.linear_attn.conv1d.weight" in initializers + + +def test_pruning_gathers_each_packed_row_end_before_lm_head(): + model = _build(_config(), prune=True) + assert model.graph.outputs[0].shape == ir.Shape(["batch", 100]) + assert any( + node.op_type == "Gather" + and node.inputs[1].producer() is not None + and node.inputs[1].producer().op_type == "Sub" + for node in model.graph + ) + + +def test_production_topology_emits_16_paged_and_48_native_layers(): + schedule = [ + "full_attention" if index % 4 == 3 else "linear_attention" + for index in range(64) + ] + model = _build(_config(layers=64, layer_types=schedule)) + assert len(_nodes(model, "PagedAttention")) == 16 + assert len(_nodes(model, "VarlenCausalConvWithState")) == 48 + assert len(_nodes(model, "GatedDeltaNet")) == 48 + + +@pytest.mark.parametrize("block_size", [0, 1, 255, 257]) +def test_block_size_must_be_positive_multiple_of_256(block_size: int): + with pytest.raises(ValueError, match="positive multiple of 256"): + PagedHybridCausalLMTask(paged_block_size=block_size) diff --git a/src/mobius/models/qwen35.py b/src/mobius/models/qwen35.py index 62647cbb8..04d5f93f0 100644 --- a/src/mobius/models/qwen35.py +++ b/src/mobius/models/qwen35.py @@ -22,6 +22,10 @@ ) from mobius.components._gated_deltanet import GatedDeltaNet from mobius.components._mlp import MLP +from mobius.components._paged_attention import ( + PagedAttentionState, + PagedHybridContext, +) from mobius.components._quantized_linear import make_quantized_linear_factory from mobius.components._rms_norm import OffsetRMSNorm from mobius.components._rotary_embedding import initialize_rope @@ -188,7 +192,15 @@ def forward( conv_state, recurrent_state = past_key_value attn_output, new_conv_state, new_recurrent_state = self.linear_attn( - op, hidden_states, conv_state, recurrent_state + op, + hidden_states, + conv_state, + recurrent_state, + cumulative_sequence_lengths=( + attention_bias.cumulative_sequence_lengths + if isinstance(attention_bias, PagedHybridContext) + else None + ), ) present_key_value = (new_conv_state, new_recurrent_state) else: @@ -246,20 +258,28 @@ def forward( inputs_embeds: ir.Value | None = None, deepstack_embeds: list | None = None, ): - # Embed tokens: (batch, seq_len) → (batch, seq_len, hidden_size) + packed = isinstance(attention_mask, PagedHybridContext) + # Embed tokens: dense (B,S,H), or packed (N,H). if inputs_embeds is not None: hidden_states = inputs_embeds else: hidden_states = self.embed_tokens(op, input_ids) # Compute (cos, sin) for RoPE: each (batch, seq_len, rotary_dim) - position_embeddings = self.rotary_emb(op, position_ids) + if packed: + position_embeddings = self.rotary_emb(op, position_ids, packed=True) + else: + position_embeddings = self.rotary_emb(op, position_ids) # Causal attention mask: (batch, 1, seq_len, total_seq_len) - attention_bias = create_attention_bias( - op, - input_ids=hidden_states if input_ids is None else input_ids, - attention_mask=attention_mask, - dtype=self._dtype, + attention_bias = ( + attention_mask + if packed + else create_attention_bias( + op, + input_ids=hidden_states if input_ids is None else input_ids, + attention_mask=attention_mask, + dtype=self._dtype, + ) ) present_key_values: list = [] @@ -306,6 +326,36 @@ def __init__(self, config: ArchitectureConfig): super().__init__(config) self._replace_text_model(Qwen35TextModel(config)) + def forward( + self, + op: OpBuilder, + input_ids: ir.Value, + attention_mask: ir.Value | PagedHybridContext | None, + position_ids: ir.Value, + past_key_values: list | None = None, + ): + """Preserve the dense ABI while supporting the dedicated packed task.""" + if not isinstance(attention_mask, PagedHybridContext): + return super().forward( + op, + input_ids, + attention_mask, + position_ids, + past_key_values, + ) + hidden_states, present_key_values = self.model( + op, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + ) + if attention_mask.last_token_indices is not None: + hidden_states = op.Gather( + hidden_states, attention_mask.last_token_indices, axis=0 + ) + return self.lm_head(op, hidden_states), present_key_values + def preprocess_weights( self, state_dict: dict[str, torch.Tensor] ) -> dict[str, torch.Tensor]: diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py index 07d5ad13a..339057e80 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -61,6 +61,7 @@ "GlmMoeDsaTask", "GlmOcrVLTask", "HybridCausalLMTask", + "PagedHybridCausalLMTask", "HyV3MtpTask", "FalconH1CausalLMTask", "Cosmos3EdgeVLTask", @@ -132,6 +133,7 @@ from mobius.tasks._causal_lm import ( CausalLMTask, HybridCausalLMTask, + PagedHybridCausalLMTask, SmallThinkerGGUFCausalLMTask, ) from mobius.tasks._codec import CodecTask @@ -265,6 +267,7 @@ "t5-text-encoding": T5TextEncoderTask, "deepseek-v4": DeepSeekV4Task, "hybrid-text-generation": HybridCausalLMTask, + "paged-hybrid-text-generation": PagedHybridCausalLMTask, "hy-v3-mtp": HyV3MtpTask, "kimi-k3-text-generation": KimiK3CausalLMTask, "kimi-linear-text-generation": KimiLinearCausalLMTask, diff --git a/src/mobius/tasks/_causal_lm.py b/src/mobius/tasks/_causal_lm.py index 68f0d0836..d6ce73319 100644 --- a/src/mobius/tasks/_causal_lm.py +++ b/src/mobius/tasks/_causal_lm.py @@ -16,6 +16,12 @@ ) from mobius._model_package import ModelPackage from mobius.components._attention import StaticCacheState +from mobius.components._paged_attention import ( + GENAI_REVISION, + ORT_REVISION, + PagedAttentionState, + PagedHybridContext, +) from mobius.tasks._base import ( ModelTask, _make_graph, @@ -27,6 +33,7 @@ _register_hybrid_cache_outputs, _register_kv_cache_outputs, _register_linear_attention_functions, + linear_attention_dims, ) @@ -429,6 +436,197 @@ def build( return ModelPackage({"model": model}, config=config) +class PagedHybridCausalLMTask(ModelTask): + """Packed Qwen hybrid decoder using native paged attention and DeltaNet ops. + + Tokens from all request rows are concatenated into one ``[N]`` input. Full + attention layers update caller-owned page buffers while linear layers carry + fixed convolution and V-major recurrent state per request row. + """ + + def __init__( + self, + *, + paged_block_size: int = 256, + prune_prefill_prefix: bool = False, + ): + if ( + isinstance(paged_block_size, bool) + or not isinstance(paged_block_size, int) + or paged_block_size <= 0 + or paged_block_size % 256 + ): + raise ValueError("paged_block_size must be a positive multiple of 256") + self._paged_block_size = paged_block_size + self._prune_prefill_prefix = prune_prefill_prefix + + def build(self, module: nn.Module, config: ArchitectureConfig) -> ModelPackage: + if config.model_type != "qwen3_5_text": + raise ValueError( + "PagedHybridCausalLMTask supports only text-only model_type " + f"'qwen3_5_text'; got {config.model_type!r}" + ) + if config.dtype not in {ir.DataType.FLOAT16, ir.DataType.BFLOAT16}: + raise ValueError( + "PagedHybridCausalLMTask native CUDA ops support only float16/bfloat16; " + f"got {config.dtype!r}" + ) + quantization = config.quantization + if quantization is not None and quantization.quant_method != "none": + raise ValueError("PagedHybridCausalLMTask does not support quantized checkpoints") + if config.component_quantization is not None: + raise ValueError("PagedHybridCausalLMTask does not support component quantization") + if config.output_layer_indices or config.output_final_hidden_state: + raise ValueError( + "PagedHybridCausalLMTask does not support auxiliary hidden-state outputs" + ) + layer_types = config.layer_types or [] + if len(layer_types) != config.num_hidden_layers or set(layer_types) != { + "full_attention", + "linear_attention", + }: + raise ValueError( + "PagedHybridCausalLMTask requires a complete Qwen hybrid " + "full_attention/linear_attention layer_types schedule" + ) + + num_tokens = ir.SymbolicDim("num_tokens") + batch = ir.SymbolicDim("batch") + num_pages = ir.SymbolicDim("num_pages") + max_blocks = ir.SymbolicDim("max_blocks_per_sequence") + graph, builder = _make_graph() + op = builder.op + + input_ids = builder.input("input_ids", dtype=ir.DataType.INT64, shape=[num_tokens]) + position_ids = builder.input( + "position_ids", dtype=ir.DataType.INT64, shape=[3, num_tokens] + ) + block_table = builder.input( + "block_table", dtype=ir.DataType.INT32, shape=[batch, max_blocks] + ) + cumulative_sequence_lengths = builder.input( + "cumulative_sequence_lengths", + dtype=ir.DataType.INT32, + shape=["batch + 1"], + ) + past_sequence_lengths = builder.input( + "past_sequence_lengths", dtype=ir.DataType.INT32, shape=[batch] + ) + attention_metadata = builder.input( + "attention_metadata", dtype=ir.DataType.INT32, shape=[3] + ) + last_token_indices = None + if self._prune_prefill_prefix: + ends = op.Slice( + cumulative_sequence_lengths, + op.Constant(value_ints=[1]), + op.Constant(value_ints=[2**63 - 1]), + op.Constant(value_ints=[0]), + ) + last_token_indices = op.Sub( + op.Cast(ends, to=ir.DataType.INT64), op.Constant(value_int=1) + ) + context = PagedHybridContext( + cumulative_sequence_lengths, + past_sequence_lengths, + block_table, + attention_metadata, + last_token_indices, + ) + + dims = linear_attention_dims(config) if "linear_attention" in layer_types else None + states = [] + for index, layer_type in enumerate(layer_types): + if layer_type == "full_attention": + key_cache = builder.input( + f"past_key_values.{index}.key", + dtype=config.dtype, + shape=[ + num_pages, + self._paged_block_size, + config.num_key_value_heads, + config.head_dim, + ], + ) + value_cache = builder.input( + f"past_key_values.{index}.value", + dtype=config.dtype, + shape=[ + num_pages, + self._paged_block_size, + config.num_key_value_heads, + config.head_dim, + ], + ) + states.append( + PagedAttentionState( + key_cache, + value_cache, + cumulative_sequence_lengths, + past_sequence_lengths, + block_table, + attention_metadata, + ) + ) + else: + assert dims is not None + conv_state = builder.input( + f"past_key_values.{index}.conv_state", + dtype=config.dtype, + shape=[batch, dims.conv_dim, dims.conv_kernel - 1], + ) + # The native GDN ABI is V-major (Dv,Dk), unlike the dense + # LinearAttention function's K-major (Dk,Dv) state. + recurrent_state = builder.input( + f"past_key_values.{index}.recurrent_state", + dtype=ir.DataType.FLOAT, + shape=[batch, dims.num_v_heads, dims.head_v_dim, dims.head_k_dim], + ) + states.append((conv_state, recurrent_state)) + + logits, present = module( + op, + input_ids=input_ids, + attention_mask=context, + position_ids=position_ids, + past_key_values=states, + ) + logits = op.Cast(logits, to=ir.DataType.FLOAT) + logits.shape = ir.Shape( + [batch if self._prune_prefill_prefix else num_tokens, config.vocab_size] + ) + builder.add_output(logits, "logits") + for layer_type, layer_present in zip(layer_types, present): + if layer_type == "full_attention": + for cache in layer_present: + cache.type = ir.TensorType(config.dtype) + cache.shape = ir.Shape( + [ + num_pages, + self._paged_block_size, + config.num_key_value_heads, + config.head_dim, + ] + ) + else: + assert dims is not None + conv, recurrent = layer_present + conv.type = ir.TensorType(config.dtype) + conv.shape = ir.Shape([batch, dims.conv_dim, dims.conv_kernel - 1]) + recurrent.type = ir.TensorType(ir.DataType.FLOAT) + recurrent.shape = ir.Shape( + [batch, dims.num_v_heads, dims.head_v_dim, dims.head_k_dim] + ) + _register_hybrid_cache_outputs(builder, present, layer_types) + + model = _make_model(graph) + model.metadata_props["mobius.paged_hybrid"] = "qwen3_5_text" + model.metadata_props["mobius.paged_block_size"] = str(self._paged_block_size) + model.metadata_props["mobius.ort_revision"] = ORT_REVISION + model.metadata_props["mobius.genai_revision"] = GENAI_REVISION + return ModelPackage({"model": model}, config=config) + + def _validate_pruned_logits( logits: ir.Value, prune_prefill_prefix: bool, diff --git a/tests/integration/paged_hybrid_test.py b/tests/integration/paged_hybrid_test.py new file mode 100644 index 000000000..778b2d2a5 --- /dev/null +++ b/tests/integration/paged_hybrid_test.py @@ -0,0 +1,374 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Native packed-state parity and pinned Engine probes (CUDA, no Hub downloads). + +Run with pytest -m integration tests/integration/paged_hybrid_test.py. +The Engine probe additionally requires MOBIUS_GENAI_REVISION to identify the +source build being qualified. These tests are not a full 27B checkpoint run. +""" + +from __future__ import annotations + +import os + +import numpy as np +import onnx_ir as ir +import onnxruntime as ort +import pytest +import torch + +from mobius import build_from_module +from mobius.components._paged_attention import GENAI_REVISION, ORT_REVISION +from mobius.integrations.onnx_genai.paged_hybrid_metadata_test import tiny_config +from mobius.integrations.ort_genai.auto_export import _write_genai_config +from mobius.models.qwen35 import Qwen35CausalLMModel +from mobius.tasks import HybridCausalLMTask, PagedHybridCausalLMTask + +pytestmark = pytest.mark.integration + + +@pytest.fixture(autouse=True) +def cuda_runtime(): + if ( + not torch.cuda.is_available() + or "CUDAExecutionProvider" not in ort.get_available_providers() + ): + pytest.skip("Native packed hybrid operators require a CUDA ORT build") + if os.environ.get("MOBIUS_ORT_REVISION") != ORT_REVISION: + pytest.skip(f"Set MOBIUS_ORT_REVISION={ORT_REVISION} for the pinned CUDA build") + + +def _package(dtype, *, packed=True, prune=False): + config = tiny_config(dtype) + module = Qwen35CausalLMModel(config) + rng = np.random.default_rng(731) + for name, parameter in module.named_parameters(): + if parameter.const_value is not None: + continue + data = rng.normal(0, 0.04, tuple(parameter.shape)).astype(np.float32) + if name.endswith("linear_attn.norm.weight"): + data += 1 + parameter.const_value = ir.tensor(data) + task = ( + PagedHybridCausalLMTask(prune_prefill_prefix=prune) if packed else HybridCausalLMTask() + ) + return build_from_module( + module, config, task, execution_provider="cuda" if packed else "cpu" + ) + + +def _session(package, path, provider): + ir.save(package["model"], path) + options = ort.SessionOptions() + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL + return ort.InferenceSession(str(path), options, providers=[provider]) + + +def _torch_type(dtype): + return torch.float16 if dtype == ir.DataType.FLOAT16 else torch.bfloat16 + + +def _run_packed(session, feed, *, pruned, vocab_size): + """Bind page outputs onto inputs, as the native kernel requires.""" + binding = session.io_binding() + element_types = { + torch.float16: 10, + torch.bfloat16: 16, + torch.float32: 1, + torch.int64: 7, + torch.int32: 6, + } + tensors = {} + for name, value in feed.items(): + tensor = value.contiguous() + tensors[name] = tensor + binding.bind_input( + name, + tensor.device.type, + 0, + element_types[tensor.dtype], + tuple(tensor.shape), + tensor.data_ptr(), + ) + outputs = {} + for output in session.get_outputs(): + name = output.name + if name == "logits": + rows = len(feed["past_sequence_lengths"]) if pruned else len(feed["input_ids"]) + tensor = torch.empty((rows, vocab_size), dtype=torch.float32, device="cuda") + else: + past = feed[name.replace("present.", "past_key_values.", 1)] + tensor = past if name.endswith((".key", ".value")) else torch.empty_like(past) + binding.bind_output( + name, + "cuda", + 0, + element_types[tensor.dtype], + tuple(tensor.shape), + tensor.data_ptr(), + ) + outputs[name] = tensor + session.run_with_iobinding(binding) + torch.cuda.synchronize() + return outputs + + +def _initial_fixed(config, rng=None): + result = {} + conv_dim = ( + 2 * config.linear_num_key_heads * config.linear_key_head_dim + + config.linear_num_value_heads * config.linear_value_head_dim + ) + for layer in range(3): + for suffix, shape in ( + ("conv_state", (1, conv_dim, config.linear_conv_kernel_dim - 1)), + ( + "recurrent_state", + ( + 1, + config.linear_num_value_heads, + config.linear_key_head_dim, + config.linear_value_head_dim, + ), + ), + ): + result[f"past_key_values.{layer}.{suffix}"] = ( + np.zeros(shape, np.float32) + if rng is None + else rng.normal(0, 0.02, shape).astype(np.float32) + ) + return result + + +@pytest.mark.parametrize("dtype", [ir.DataType.FLOAT16, ir.DataType.BFLOAT16]) +def test_packed_prefill_continuation_reorder_and_page_reuse(dtype, tmp_path): + if dtype == ir.DataType.BFLOAT16 and not torch.cuda.is_bf16_supported(): + pytest.skip("BF16 requires a supported CUDA device") + dense_pkg = _package(ir.DataType.FLOAT, packed=False) + packed_pkg = _package(dtype) + pruned_pkg = _package(dtype, prune=True) + dense = _session(dense_pkg, tmp_path / "dense.onnx", "CPUExecutionProvider") + packed = _session(packed_pkg, tmp_path / "packed.onnx", "CUDAExecutionProvider") + pruned = _session(pruned_pkg, tmp_path / "pruned.onnx", "CUDAExecutionProvider") + config = packed_pkg.config + compute = _torch_type(dtype) + rng = np.random.default_rng(42) + # A uses nonconsecutive physical pages. B's page is later reused for C + # without clearing it, so masking must prevent stale-token contamination. + pages = {"A": [2, 0], "B": [1], "C": [1]} + page_shape = (4, 256, config.num_key_value_heads, config.head_dim) + caches = [ + { + f"past_key_values.3.{kind}": torch.zeros(page_shape, device="cuda", dtype=compute) + for kind in ("key", "value") + } + for _ in range(2) + ] + fixed = [{}, {}] + dense_state = {} + lengths = {} + tolerance = 0.025 if dtype == ir.DataType.FLOAT16 else 0.09 + schedule = [ + [("A", 255), ("B", 3)], + [("B", 2), ("A", 3)], + [("A", 1), ("C", 7)], + [("C", 3), ("A", 2)], + ] + for step, requests in enumerate(schedule): + ids, expected = [], [] + past_lengths = [] + for request, count in requests: + if request not in lengths: + lengths[request] = 0 + initial = _initial_fixed(config, rng if request != "C" else None) + dense_state[request] = { + **initial, + **{ + f"past_key_values.3.{kind}": np.zeros( + (1, config.num_key_value_heads, 0, config.head_dim), np.float32 + ) + for kind in ("key", "value") + }, + } + for state in fixed: + state[request] = { + name: torch.as_tensor( + value.transpose(0, 1, 3, 2).copy() + if name.endswith("recurrent_state") + else value, + device="cuda", + dtype=torch.float32 + if name.endswith("recurrent_state") + else compute, + ) + for name, value in initial.items() + } + past = lengths[request] + past_lengths.append(past) + tokens = rng.integers(1, config.vocab_size - 1, count, dtype=np.int64) + ids.append(tokens) + feed = { + **dense_state[request], + "input_ids": tokens[None], + "position_ids": np.arange(past, past + count, dtype=np.int64)[None], + "attention_mask": np.ones((1, past + count), np.int64), + } + result = dict(zip([v.name for v in dense.get_outputs()], dense.run(None, feed))) + expected.append(result) + dense_state[request] = { + name.replace("present.", "past_key_values.", 1): value + for name, value in result.items() + if name != "logits" + } + counts = [count for _, count in requests] + cu = np.asarray([0, *np.cumsum(counts)], np.int32) + positions = np.concatenate( + [np.arange(past, past + count) for past, count in zip(past_lengths, counts)] + ) + block_table = np.zeros((len(requests), 2), np.int32) + for row, (request, _) in enumerate(requests): + block_table[row, : len(pages[request])] = pages[request] + shared = { + "input_ids": torch.tensor(np.concatenate(ids), device="cuda"), + "position_ids": torch.tensor(np.tile(positions, (3, 1)), device="cuda"), + "block_table": torch.tensor(block_table, device="cuda"), + "cumulative_sequence_lengths": torch.tensor(cu, device="cuda"), + "past_sequence_lengths": torch.tensor( + past_lengths, device="cuda", dtype=torch.int32 + ), + "attention_metadata": torch.tensor( + [max(counts), max(p + c for p, c in zip(past_lengths, counts)), 0], + dtype=torch.int32, + ), + } + results = [] + for index, session in enumerate((packed, pruned)): + feed = { + **shared, + **caches[index], + **{ + name: torch.cat([fixed[index][request][name] for request, _ in requests]) + for name in fixed[index][requests[0][0]] + }, + } + result = _run_packed( + session, feed, pruned=bool(index), vocab_size=config.vocab_size + ) + results.append(result) + for row, (request, count) in enumerate(requests): + for name, value in result.items(): + if name.endswith(("conv_state", "recurrent_state")): + fixed[index][request][ + name.replace("present.", "past_key_values.", 1) + ] = value[row : row + 1].clone() + reference = expected[row][name] + if name.endswith("recurrent_state"): + reference = reference.transpose(0, 1, 3, 2) + np.testing.assert_allclose( + value[row : row + 1].float().cpu().numpy(), + reference, + atol=tolerance, + rtol=tolerance, + err_msg=f"{step=} {request=} {name=}", + ) + length = past_lengths[row] + count + for kind in ("key", "value"): + logical = result[f"present.3.{kind}"][pages[request]].reshape( + -1, config.num_key_value_heads, config.head_dim + )[:length] + reference = expected[row][f"present.3.{kind}"][0].transpose(1, 0, 2) + np.testing.assert_allclose( + logical.float().cpu().numpy(), + reference, + atol=tolerance, + rtol=tolerance, + ) + full_logits = results[0]["logits"].cpu().numpy() + np.testing.assert_allclose( + full_logits, + np.concatenate([result["logits"][0] for result in expected]), + atol=tolerance, + rtol=tolerance, + ) + np.testing.assert_allclose( + results[1]["logits"].cpu().numpy(), full_logits[cu[1:] - 1], atol=1e-5, rtol=1e-5 + ) + for request, count in requests: + lengths[request] += count + if step == 1: + for state in fixed: + del state["B"] + del dense_state["B"] + + +def test_pinned_engine_overlap_and_late_admission(tmp_path): + if os.environ.get("MOBIUS_GENAI_REVISION") != GENAI_REVISION: + pytest.skip(f"Set MOBIUS_GENAI_REVISION={GENAI_REVISION} for the pinned Engine build") + import onnxruntime_genai as og + + package = _package(ir.DataType.FLOAT16) + ir.save(package["model"], tmp_path / "model.onnx") + _write_genai_config( + package.config, + str(tmp_path), + pkg=package, + ort_model_type="decoder", + ep="cuda", + context_length=2048, + bos_token_id=1, + eos_token_id=127, + pad_token_id=0, + is_vlm=False, + has_speech=False, + ) + model = og.Model(str(tmp_path)) + prompts = [ + np.arange(2, 19, dtype=np.int32), + np.arange(21, 25, dtype=np.int32), + np.arange(30, 38, dtype=np.int32), + ] + + def run(indices, *, overlap=False): + engine = og.Engine(model) + requests, generated = {}, {} + + def admit(index): + options = og.RequestOptions() + options.set_max_session_tokens(64) + request = engine.create_request(options=options) + turn = og.TurnOptions(request) + turn.set_do_sample(False) + turn.set_max_generated_tokens(6) + request.begin_turn(prompts[index], turn) + requests[request] = index + generated[index] = [] + + for index in indices[:2] if overlap else indices: + admit(index) + late = overlap + finished = set() + event_buffer = engine.create_event_buffer(1) + for _ in range(1000): + events = engine.run(event_buffer) + for event in events: + assert not event.flags & og.EngineEventFlags.FAILED + if event.request is None: + continue + index = requests[event.request] + if event.flags & og.EngineEventFlags.TOKEN: + generated[index].append(event.token) + if event.flags & og.EngineEventFlags.TURN_FINISHED: + finished.add(index) + event.request.close() + if late and any(generated.values()): + admit(indices[2]) + late = False + if not late and len(finished) == len(indices): + break + assert not late and len(finished) == len(indices) + assert all(generated.values()) + return generated + + isolated = {index: run([index])[index] for index in range(3)} + assert run([0, 1, 2], overlap=True) == isolated From f07b251c0e8b2bb7c6e54f93b01f676cca3f4d60 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 16 Sep 2026 23:13:42 +0000 Subject: [PATCH 3/6] Harden packed ABI typing and task validation Declare native output shapes before downstream construction, reject incompatible direct tasks and RoPE configurations, and verify the 64-layer state manifest. Exercise both CLI runtime packagers and strengthen page-reuse coverage. Focused checks: 20 passed, 3 CUDA-dependent skips. Signed-off-by: GitHub Co-authored-by: titaiwangms <18010845+titaiwangms@users.noreply.github.com> --- README.md | 2 +- src/mobius/_builder.py | 6 +++ src/mobius/components/__init__.py | 2 +- src/mobius/components/_paged_attention.py | 18 +++++-- .../onnx_genai/paged_hybrid_metadata_test.py | 53 +++++++++++++++++++ src/mobius/tasks/_causal_lm.py | 6 +++ tests/integration/paged_hybrid_test.py | 7 ++- 7 files changed, 86 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 03311e459..0af82c117 100644 --- a/README.md +++ b/README.md @@ -132,7 +132,7 @@ 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 + --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 diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py index dba24cb12..3476c6a6c 100644 --- a/src/mobius/_builder.py +++ b/src/mobius/_builder.py @@ -171,6 +171,12 @@ def build_from_module( 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'") diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index 3d1312407..c805f00b2 100644 --- a/src/mobius/components/__init__.py +++ b/src/mobius/components/__init__.py @@ -240,7 +240,6 @@ create_decoder_layer, ) from mobius.components._deepseek_mla import DeepSeekMLA as DeepSeekMLA -from mobius.components._paged_attention import PagedAttentionState, PagedHybridContext from mobius.components._diffusion import ( AdaLayerNormOutput, AdaLayerNormZero, @@ -349,6 +348,7 @@ PaddleOCRProjector, YouTuVLProjector, ) +from mobius.components._paged_attention import PagedAttentionState, PagedHybridContext from mobius.components._paged_mla import ( PagedCacheState as PagedCacheState, ) diff --git a/src/mobius/components/_paged_attention.py b/src/mobius/components/_paged_attention.py index 0ef702b8f..0e0ca37f5 100644 --- a/src/mobius/components/_paged_attention.py +++ b/src/mobius/components/_paged_attention.py @@ -10,7 +10,6 @@ import onnx_ir as ir from onnxscript import OpBuilder - DOMAIN = "com.microsoft" ORT_REVISION = "f38538cd5a4b5945a4c839565a8eebc65e1e2ef8" GENAI_REVISION = "d5b40851ba80ffa8e95b6b01f921dbb9008fac80" @@ -48,7 +47,7 @@ def paged_attention( kv_num_heads: int, ) -> tuple[ir.Value, ir.Value, ir.Value]: """Emit the pinned 17-input SEPARATE PagedAttention ABI.""" - return op.PagedAttention( + output, key_cache, value_cache = op.PagedAttention( query, key, value, @@ -73,6 +72,11 @@ def paged_attention( _domain=DOMAIN, _outputs=3, ) + # Generic ONNX shape inference does not know these pinned contrib schemas. + output.type, output.shape = query.type, query.shape + key_cache.type, key_cache.shape = state.key_cache.type, state.key_cache.shape + value_cache.type, value_cache.shape = state.value_cache.type, state.value_cache.shape + return output, key_cache, value_cache def varlen_causal_conv_with_state( @@ -84,7 +88,7 @@ def varlen_causal_conv_with_state( past_conv: ir.Value, ) -> tuple[ir.Value, ir.Value]: """Emit packed depthwise convolution with fixed per-sequence carry state.""" - return op.VarlenCausalConvWithState( + output, present = op.VarlenCausalConvWithState( packed, weight, cumulative_sequence_lengths, @@ -94,6 +98,9 @@ def varlen_causal_conv_with_state( _domain=DOMAIN, _outputs=2, ) + output.type, output.shape = packed.type, packed.shape + present.type, present.shape = past_conv.type, past_conv.shape + return output, present def gated_delta_net( @@ -109,7 +116,7 @@ def gated_delta_net( dt_bias: ir.Value, ) -> tuple[ir.Value, ir.Value]: """Emit the native packed Qwen GatedDeltaNet recurrence.""" - return op.GatedDeltaNet( + output, present = op.GatedDeltaNet( query, key, value, @@ -127,3 +134,6 @@ def gated_delta_net( _domain=DOMAIN, _outputs=2, ) + output.type, output.shape = value.type, value.shape + present.type, present.shape = past_recurrent.type, past_recurrent.shape + return output, present diff --git a/src/mobius/integrations/onnx_genai/paged_hybrid_metadata_test.py b/src/mobius/integrations/onnx_genai/paged_hybrid_metadata_test.py index acbbcf116..f5adcb43d 100644 --- a/src/mobius/integrations/onnx_genai/paged_hybrid_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/paged_hybrid_metadata_test.py @@ -181,6 +181,59 @@ def test_direct_task_rejects_non_cuda(ep): ) +@pytest.mark.parametrize( + "overrides,match", + [ + ({"dtype": ir.DataType.FLOAT}, "float16/bfloat16"), + ({"model_type": "llama"}, "text-only"), + ({"mrope_section": None}, "interleaved MRoPE"), + ({"mrope_interleaved": False}, "interleaved MRoPE"), + ({"layer_types": ["full_attention"] * 4}, "layer_types"), + ], +) +def test_direct_task_rejects_incompatible_contracts(overrides, match): + config = dataclasses.replace(tiny_config(), **overrides) + with pytest.raises(ValueError, match=match): + build_from_module( + Qwen35CausalLMModel(config), + config, + PagedHybridCausalLMTask(), + execution_provider="cuda", + ) + + +def test_paged_flag_cannot_be_combined_with_a_dense_task(): + config = dataclasses.replace(tiny_config(), export_paged_attention=True) + with pytest.raises(ValueError, match="requires PagedHybridCausalLMTask"): + build_from_module( + Qwen35CausalLMModel(config), + config, + HybridCausalLMTask(), + execution_provider="cuda", + ) + + +def test_64_layer_manifest_matches_native_operator_topology(): + config = dataclasses.replace( + tiny_config(), + num_hidden_layers=64, + layer_types=tiny_config().layer_types * 16, + ) + package = build_from_module( + Qwen35CausalLMModel(config), + config, + PagedHybridCausalLMTask(), + execution_provider="cuda", + ) + abi = inspect_paged_hybrid(package["model"]) + assert abi.full_layers == tuple(range(3, 64, 4)) + assert len(abi.linear_layers) == 48 + nodes = list(package["model"].graph) + assert sum(node.op_type == "PagedAttention" for node in nodes) == 16 + assert sum(node.op_type == "VarlenCausalConvWithState" for node in nodes) == 48 + assert sum(node.op_type == "GatedDeltaNet" for node in nodes) == 48 + + def test_dense_gate_parameters_still_follow_compute_dtype(): from mobius._builder import _cast_module_dtype diff --git a/src/mobius/tasks/_causal_lm.py b/src/mobius/tasks/_causal_lm.py index d6ce73319..31a207078 100644 --- a/src/mobius/tasks/_causal_lm.py +++ b/src/mobius/tasks/_causal_lm.py @@ -461,6 +461,10 @@ def __init__( self._prune_prefill_prefix = prune_prefill_prefix def build(self, module: nn.Module, config: ArchitectureConfig) -> ModelPackage: + from mobius.models.qwen35 import Qwen35CausalLMModel + + if not isinstance(module, Qwen35CausalLMModel): + raise TypeError("PagedHybridCausalLMTask requires Qwen35CausalLMModel") if config.model_type != "qwen3_5_text": raise ValueError( "PagedHybridCausalLMTask supports only text-only model_type " @@ -471,6 +475,8 @@ def build(self, module: nn.Module, config: ArchitectureConfig) -> ModelPackage: "PagedHybridCausalLMTask native CUDA ops support only float16/bfloat16; " f"got {config.dtype!r}" ) + if not config.mrope_section or not config.mrope_interleaved: + raise ValueError("PagedHybridCausalLMTask requires interleaved MRoPE sections") quantization = config.quantization if quantization is not None and quantization.quant_method != "none": raise ValueError("PagedHybridCausalLMTask does not support quantized checkpoints") diff --git a/tests/integration/paged_hybrid_test.py b/tests/integration/paged_hybrid_test.py index 778b2d2a5..5f7a6419c 100644 --- a/tests/integration/paged_hybrid_test.py +++ b/tests/integration/paged_hybrid_test.py @@ -172,7 +172,7 @@ def test_packed_prefill_continuation_reorder_and_page_reuse(dtype, tmp_path): schedule = [ [("A", 255), ("B", 3)], [("B", 2), ("A", 3)], - [("A", 1), ("C", 7)], + [("A", 1), ("C", 2)], [("C", 3), ("A", 2)], ] for step, requests in enumerate(schedule): @@ -292,7 +292,10 @@ def test_packed_prefill_continuation_reorder_and_page_reuse(dtype, tmp_path): rtol=tolerance, ) np.testing.assert_allclose( - results[1]["logits"].cpu().numpy(), full_logits[cu[1:] - 1], atol=1e-5, rtol=1e-5 + results[1]["logits"].cpu().numpy(), + full_logits[cu[1:] - 1], + atol=2e-3 if dtype == ir.DataType.FLOAT16 else 1e-2, + rtol=2e-3 if dtype == ir.DataType.FLOAT16 else 1e-2, ) for request, count in requests: lengths[request] += count From 370655b7009ecaa22a9cf2510798cfbf8a59475f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 16 Sep 2026 23:32:33 +0000 Subject: [PATCH 4/6] Fix paged pruning review feedback Enable prefix pruning when the paged hybrid task is selected by its registered name and cover that dispatch path. Correct the packed MRoPE contract documentation and apply the CI lint fixes. Signed-off-by: GitHub Co-authored-by: titaiwangms <18010845+titaiwangms@users.noreply.github.com> --- src/mobius/_builder.py | 2 ++ src/mobius/_builder_test.py | 11 ++++++++++- src/mobius/components/_rotary_embedding.py | 11 ++++++----- src/mobius/models/paged_qwen35_export_test.py | 7 ++----- src/mobius/models/qwen35.py | 5 +---- 5 files changed, 21 insertions(+), 15 deletions(-) diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py index 3476c6a6c..355d40018 100644 --- a/src/mobius/_builder.py +++ b/src/mobius/_builder.py @@ -83,6 +83,8 @@ def _enable_prefill_prefix_pruning_task(task: str | ModelTask) -> str | ModelTas 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": diff --git a/src/mobius/_builder_test.py b/src/mobius/_builder_test.py index 354308ed1..ed6a05574 100644 --- a/src/mobius/_builder_test.py +++ b/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( diff --git a/src/mobius/components/_rotary_embedding.py b/src/mobius/components/_rotary_embedding.py index 7dd781552..ac3c97b63 100644 --- a/src/mobius/components/_rotary_embedding.py +++ b/src/mobius/components/_rotary_embedding.py @@ -478,19 +478,20 @@ def __init__( data=ir.tensor(w_mask), ) - def forward( - self, op: OpBuilder, position_ids: ir.Value, *, packed: bool = False - ): + def forward(self, op: OpBuilder, position_ids: ir.Value, *, packed: bool = False): """Compute MRoPE cos/sin embeddings. Args: op: ONNX op builder. position_ids: Either ``(batch, seq)`` for text-only or ``(3, batch, seq)`` for multimodal (T, H, W dimensions). - For 2D input, the same positions are used for all 3 dims. + With ``packed=True``, accepts ``(3, num_tokens)`` instead. + For dense 2D input, the same positions are used for all 3 dims. + packed: Whether ``position_ids`` uses the packed serving layout. Returns: - Tuple of ``(cos, sin)`` each with shape ``(batch, seq, rotary_dim)``. + Tuple of ``(cos, sin)`` each with shape ``(batch, seq, rotary_dim)``, + or ``(num_tokens, rotary_dim)`` when ``packed=True``. """ if packed: # Packed serving supplies exactly (3,N), with no batch padding. diff --git a/src/mobius/models/paged_qwen35_export_test.py b/src/mobius/models/paged_qwen35_export_test.py index 973b730bf..ebc87cd16 100644 --- a/src/mobius/models/paged_qwen35_export_test.py +++ b/src/mobius/models/paged_qwen35_export_test.py @@ -50,9 +50,7 @@ def test_packed_hybrid_io_and_native_operands(): assert inputs["position_ids"].shape == ir.Shape([3, "num_tokens"]) assert inputs["attention_metadata"].dtype == ir.DataType.INT32 assert inputs["past_key_values.0.key"].shape[1:] == ir.Shape([256, 2, 16]) - assert inputs["past_key_values.1.recurrent_state"].shape[1:] == ir.Shape( - [4, 16, 16] - ) + assert inputs["past_key_values.1.recurrent_state"].shape[1:] == ir.Shape([4, 16, 16]) assert inputs["past_key_values.1.recurrent_state"].dtype == ir.DataType.FLOAT paged = _nodes(model, "PagedAttention") @@ -103,8 +101,7 @@ def test_pruning_gathers_each_packed_row_end_before_lm_head(): def test_production_topology_emits_16_paged_and_48_native_layers(): schedule = [ - "full_attention" if index % 4 == 3 else "linear_attention" - for index in range(64) + "full_attention" if index % 4 == 3 else "linear_attention" for index in range(64) ] model = _build(_config(layers=64, layer_types=schedule)) assert len(_nodes(model, "PagedAttention")) == 16 diff --git a/src/mobius/models/qwen35.py b/src/mobius/models/qwen35.py index 04d5f93f0..2d556a80c 100644 --- a/src/mobius/models/qwen35.py +++ b/src/mobius/models/qwen35.py @@ -22,10 +22,7 @@ ) from mobius.components._gated_deltanet import GatedDeltaNet from mobius.components._mlp import MLP -from mobius.components._paged_attention import ( - PagedAttentionState, - PagedHybridContext, -) +from mobius.components._paged_attention import PagedHybridContext from mobius.components._quantized_linear import make_quantized_linear_factory from mobius.components._rms_norm import OffsetRMSNorm from mobius.components._rotary_embedding import initialize_rope From cdff31aafaf671b750f236c4449ba66c0c8abc8c Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 16 Sep 2026 23:36:29 +0000 Subject: [PATCH 5/6] Apply remaining Qwen formatter fix Match Ruff's formatting for the packed prefix-pruning gather after removing the unused paged-attention import. Signed-off-by: GitHub Co-authored-by: titaiwangms <18010845+titaiwangms@users.noreply.github.com> --- src/mobius/models/qwen35.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/mobius/models/qwen35.py b/src/mobius/models/qwen35.py index 2d556a80c..bd42b3205 100644 --- a/src/mobius/models/qwen35.py +++ b/src/mobius/models/qwen35.py @@ -348,9 +348,7 @@ def forward( past_key_values=past_key_values, ) if attention_mask.last_token_indices is not None: - hidden_states = op.Gather( - hidden_states, attention_mask.last_token_indices, axis=0 - ) + hidden_states = op.Gather(hidden_states, attention_mask.last_token_indices, axis=0) return self.lm_head(op, hidden_states), present_key_values def preprocess_weights( From feceaa434109029ffb3827899a706ef8dde372bb Mon Sep 17 00:00:00 2001 From: titaiwang Date: Thu, 17 Sep 2026 16:50:07 +0000 Subject: [PATCH 6/6] Harden Qwen packed hybrid runtime contract Verify native packed-state dataflow, operand layouts, and Phase 1 attributes before emitting runtime metadata. Route paged context explicitly, expand malformed-graph coverage, and fix the CI dependency and stale VibeVoice streaming denylist regressions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang --- src/mobius/components/_attention.py | 11 +- src/mobius/components/_attention_test.py | 17 + src/mobius/components/_gated_deltanet.py | 5 + src/mobius/integrations/_paged_hybrid.py | 465 ++++++++++++++++-- .../onnx_genai/paged_hybrid_metadata_test.py | 299 +++++++++++ src/mobius/models/paged_qwen35_export_test.py | 20 +- src/mobius/models/qwen35.py | 46 +- src/mobius/models/vibevoice.py | 8 - src/mobius/tasks/_causal_lm.py | 3 +- tests/integration/paged_hybrid_test.py | 28 +- 10 files changed, 843 insertions(+), 59 deletions(-) diff --git a/src/mobius/components/_attention.py b/src/mobius/components/_attention.py index d10df9bdc..ed3440410 100644 --- a/src/mobius/components/_attention.py +++ b/src/mobius/components/_attention.py @@ -593,9 +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): - return self.forward_paged(op, hidden_states, position_embeddings, past_key_value) + 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 diff --git a/src/mobius/components/_attention_test.py b/src/mobius/components/_attention_test.py index 4023c0dc5..61e4ccb18 100644 --- a/src/mobius/components/_attention_test.py +++ b/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.""" diff --git a/src/mobius/components/_gated_deltanet.py b/src/mobius/components/_gated_deltanet.py index 505b6ae75..d57b47c34 100644 --- a/src/mobius/components/_gated_deltanet.py +++ b/src/mobius/components/_gated_deltanet.py @@ -348,6 +348,11 @@ def forward_paged( z = self.in_proj_z(op, hidden_states) # (N, Hv*Dv) raw_b = op.Cast(self.in_proj_b(op, hidden_states), to=ir.DataType.FLOAT) raw_a = op.Cast(self.in_proj_a(op, hidden_states), to=ir.DataType.FLOAT) + # The generic builder does not infer Cast output types until the + # production optimization pipeline runs. Keep the native ABI explicit + # even for directly-built graphs. + raw_b.type = ir.TensorType(ir.DataType.FLOAT) + raw_a.type = ir.TensorType(ir.DataType.FLOAT) # Invoke the child module so nn realizes its qualified conv1d.weight. conv_out, new_conv_state = self.conv1d( diff --git a/src/mobius/integrations/_paged_hybrid.py b/src/mobius/integrations/_paged_hybrid.py index 5fb9117f1..a3b22a7a2 100644 --- a/src/mobius/integrations/_paged_hybrid.py +++ b/src/mobius/integrations/_paged_hybrid.py @@ -10,6 +10,48 @@ import onnx_ir as ir +_COMPUTE_DTYPES = {ir.DataType.FLOAT16, ir.DataType.BFLOAT16} + +_OPERANDS = { + "PagedAttention": { + "query": 0, + "key": 1, + "value": 2, + "key_cache": 3, + "value_cache": 4, + "cumulative_sequence_lengths": 5, + "past_sequence_lengths": 6, + "block_table": 7, + "attention_metadata": 16, + }, + "VarlenCausalConvWithState": { + "packed": 0, + "weight": 1, + "cumulative_sequence_lengths": 2, + "bias": 3, + "state": 4, + }, + "GatedDeltaNet": { + "query": 0, + "key": 1, + "value": 2, + "cumulative_sequence_lengths": 3, + "raw_a": 4, + "raw_b": 5, + "state": 6, + "a_log": 7, + "dt_bias": 8, + }, +} +_STATE_OPERANDS = { + op_type: operands[state_name] + for op_type, operands, state_name in ( + ("PagedAttention", _OPERANDS["PagedAttention"], "key_cache"), + ("VarlenCausalConvWithState", _OPERANDS["VarlenCausalConvWithState"], "state"), + ("GatedDeltaNet", _OPERANDS["GatedDeltaNet"], "state"), + ) +} + @dataclass(frozen=True) class PagedHybridAbi: @@ -25,12 +67,359 @@ def state_groups(self) -> list[dict]: ] +def _require_value( + value: ir.Value | None, + *, + dtype: ir.DataType, + rank: int, + description: str, +) -> ir.Value: + if ( + value is None + or value.dtype != dtype + or value.shape is None + or len(value.shape) != rank + ): + raise ValueError(f"{description} must have {dtype} rank {rank}") + return value + + +def _attribute_value(node: ir.Node, name: str): + attribute = node.attributes.get(name) + return None if attribute is None else attribute.value + + +def _has_default_attribute(node: ir.Node, name: str, default) -> bool: + value = _attribute_value(node, name) + return value is None or value == default + + +def _operand(node: ir.Node, name: str) -> ir.Value | None: + return node.inputs[_OPERANDS[node.op_type][name]] + + +def _require_native_outputs( + node: ir.Node, + count: int, + state: ir.Value, + present: ir.Value, + layer: int, +) -> None: + if len(node.outputs) != count or node.outputs[-1] is not present: + raise ValueError( + f"Native {node.op_type} state output is disconnected at layer {layer}" + ) + native_state = node.outputs[-1] + if native_state.dtype != state.dtype or not _shapes_compatible( + native_state.shape, state.shape + ): + raise ValueError(f"Invalid native {node.op_type} state output at layer {layer}") + + +def _known_dimension_disagrees(left, right) -> bool: + """Compare dimensions only when both are concrete graph facts.""" + return isinstance(left, int) and isinstance(right, int) and left != right + + +def _shapes_compatible(left: ir.Shape | None, right: ir.Shape | None) -> bool: + """Reject only rank or concrete-dimension disagreements.""" + return ( + left is not None + and right is not None + and len(left) == len(right) + and not any(_known_dimension_disagrees(a, b) for a, b in zip(left, right)) + ) + + +def _validate_request_dimensions( + inputs: dict[str, ir.Value], layers: dict[str, set[int]] +) -> None: + batch_dimensions = { + "block_table": inputs["block_table"].shape[0], + "past_sequence_lengths": inputs["past_sequence_lengths"].shape[0], + } + for layer in layers["conv_state"]: + batch_dimensions[f"conv_state at layer {layer}"] = inputs[ + f"past_key_values.{layer}.conv_state" + ].shape[0] + batch_dimensions[f"recurrent_state at layer {layer}"] = inputs[ + f"past_key_values.{layer}.recurrent_state" + ].shape[0] + concrete_batches = { + dimension for dimension in batch_dimensions.values() if isinstance(dimension, int) + } + if len(concrete_batches) > 1: + raise ValueError("Packed hybrid request-aligned batch dimensions disagree") + cumulative_extent = inputs["cumulative_sequence_lengths"].shape[0] + if concrete_batches and isinstance(cumulative_extent, int): + batch = next(iter(concrete_batches)) + if cumulative_extent != batch + 1: + raise ValueError( + "Packed hybrid cumulative_sequence_lengths extent must equal batch + 1" + ) + + +def _validate_paged_attention( + node: ir.Node, + *, + layer: int, + inputs: dict[str, ir.Value], + outputs: dict[str, ir.Value], +) -> None: + key_cache = inputs[f"past_key_values.{layer}.key"] + value_cache = inputs[f"past_key_values.{layer}.value"] + if key_cache.dtype != value_cache.dtype or not _shapes_compatible( + key_cache.shape, value_cache.shape + ): + raise ValueError(f"Paged key/value cache layouts disagree at layer {layer}") + num_heads = _attribute_value(node, "num_heads") + kv_num_heads = _attribute_value(node, "kv_num_heads") + if ( + len(node.inputs) != 17 + or _operand(node, "key_cache") is not key_cache + or _operand(node, "value_cache") is not value_cache + or _operand(node, "cumulative_sequence_lengths") + is not inputs["cumulative_sequence_lengths"] + or _operand(node, "past_sequence_lengths") is not inputs["past_sequence_lengths"] + or _operand(node, "block_table") is not inputs["block_table"] + or any(operand is not None for operand in node.inputs[8:16]) + or _operand(node, "attention_metadata") is not inputs["attention_metadata"] + or not _has_default_attribute(node, "kv_cache_layout", "SEPARATE") + or not _has_default_attribute(node, "do_rotary", 0) + or not _has_default_attribute(node, "is_causal", 1) + or not _has_default_attribute(node, "local_window_size", -1) + or not _has_default_attribute(node, "softcap", 0.0) + or not _has_default_attribute(node, "scale", 0.0) + or not isinstance(num_heads, int) + or num_heads <= 0 + or not isinstance(kv_num_heads, int) + or kv_num_heads <= 0 + or num_heads % kv_num_heads != 0 + or _known_dimension_disagrees(kv_num_heads, key_cache.shape[2]) + or len(node.outputs) != 3 + or node.outputs[1] is not outputs[f"present.{layer}.key"] + or node.outputs[2] is not outputs[f"present.{layer}.value"] + ): + raise ValueError(f"Invalid SEPARATE PagedAttention contract at layer {layer}") + + query, key, value = (_operand(node, name) for name in ("query", "key", "value")) + if any( + operand is None + or operand.dtype not in _COMPUTE_DTYPES + or operand.shape is None + or len(operand.shape) != 2 + for operand in (query, key, value) + ): + raise ValueError(f"Invalid PagedAttention QKV operands at layer {layer}") + assert query is not None and query.shape is not None + assert key is not None and key.shape is not None + assert value is not None and value.shape is not None + if ( + key.dtype != query.dtype + or value.dtype != query.dtype + or _known_dimension_disagrees(query.shape[0], key.shape[0]) + or _known_dimension_disagrees(query.shape[0], value.shape[0]) + or _known_dimension_disagrees(key.shape[0], value.shape[0]) + ): + raise ValueError(f"Invalid PagedAttention QKV layout at layer {layer}") + head_dim = key_cache.shape[3] + if isinstance(head_dim, int) and ( + _known_dimension_disagrees(query.shape[1], num_heads * head_dim) + or _known_dimension_disagrees(key.shape[1], kv_num_heads * head_dim) + or _known_dimension_disagrees(value.shape[1], kv_num_heads * head_dim) + ): + raise ValueError(f"Invalid PagedAttention QKV width at layer {layer}") + + data_output = node.outputs[0] + if ( + data_output.dtype not in _COMPUTE_DTYPES + or data_output.dtype != query.dtype + or data_output.shape is None + or len(data_output.shape) != 2 + or _known_dimension_disagrees(data_output.shape[0], query.shape[0]) + or _known_dimension_disagrees(data_output.shape[1], query.shape[1]) + or not data_output.uses() + ): + raise ValueError( + f"Invalid or disconnected PagedAttention data output at layer {layer}" + ) + if ( + node.outputs[1].dtype != key_cache.dtype + or not _shapes_compatible(node.outputs[1].shape, key_cache.shape) + or node.outputs[2].dtype != value_cache.dtype + or not _shapes_compatible(node.outputs[2].shape, value_cache.shape) + ): + raise ValueError(f"Invalid PagedAttention state outputs at layer {layer}") + + +def _validate_varlen_conv( + node: ir.Node, + *, + layer: int, + inputs: dict[str, ir.Value], + outputs: dict[str, ir.Value], +) -> None: + state = inputs[f"past_key_values.{layer}.conv_state"] + if ( + len(node.inputs) != 5 + or _operand(node, "cumulative_sequence_lengths") + is not inputs["cumulative_sequence_lengths"] + or _operand(node, "state") is not state + or _attribute_value(node, "activation") != "silu" + or not _has_default_attribute(node, "dilation", 1) + or not _has_default_attribute(node, "state_update_capacity", 0) + ): + raise ValueError(f"Invalid VarlenCausalConvWithState contract at layer {layer}") + _require_native_outputs(node, 2, state, outputs[f"present.{layer}.conv_state"], layer) + packed = _operand(node, "packed") + weight = _operand(node, "weight") + bias = _operand(node, "bias") + if ( + packed is None + or packed.dtype not in _COMPUTE_DTYPES + or packed.shape is None + or len(packed.shape) != 2 + or weight is None + or weight.dtype != packed.dtype + or weight.shape is None + or len(weight.shape) != 3 + or _known_dimension_disagrees(weight.shape[1], 1) + or bias is None + or bias.dtype != packed.dtype + or bias.shape is None + or len(bias.shape) != 1 + or _known_dimension_disagrees(packed.shape[1], state.shape[1]) + or _known_dimension_disagrees(weight.shape[0], state.shape[1]) + or _known_dimension_disagrees(bias.shape[0], state.shape[1]) + ): + raise ValueError(f"Invalid VarlenCausalConvWithState operands at layer {layer}") + if ( + isinstance(weight.shape[2], int) + and isinstance(state.shape[2], int) + and weight.shape[2] != state.shape[2] + 1 + ): + raise ValueError(f"Invalid VarlenCausalConvWithState layout at layer {layer}") + + +def _validate_gated_delta_net( + node: ir.Node, + *, + layer: int, + inputs: dict[str, ir.Value], + outputs: dict[str, ir.Value], +) -> None: + state = inputs[f"past_key_values.{layer}.recurrent_state"] + required_attributes = { + "gate_activation": "qwen", + "beta_activation": "sigmoid", + "qk_l2_norm": 1, + "update_rule": "gated_delta", + } + if ( + len(node.inputs) != 9 + or _operand(node, "cumulative_sequence_lengths") + is not inputs["cumulative_sequence_lengths"] + or _operand(node, "state") is not state + or any( + ( + not _has_default_attribute(node, name, expected) + if name == "update_rule" + else _attribute_value(node, name) != expected + ) + for name, expected in required_attributes.items() + ) + or not _has_default_attribute(node, "scale", 0.0) + or not _has_default_attribute(node, "state_update_capacity", 0) + ): + raise ValueError(f"Invalid GatedDeltaNet contract at layer {layer}") + _require_native_outputs(node, 2, state, outputs[f"present.{layer}.recurrent_state"], layer) + query, key, value = (_operand(node, name) for name in ("query", "key", "value")) + if any( + operand is None + or operand.dtype not in _COMPUTE_DTYPES + or operand.shape is None + or len(operand.shape) != 3 + for operand in (query, key, value) + ): + raise ValueError(f"Invalid GatedDeltaNet QKV operands at layer {layer}") + assert query is not None and query.shape is not None + assert key is not None and key.shape is not None + assert value is not None and value.shape is not None + if ( + key.dtype != query.dtype + or value.dtype != query.dtype + or any( + _known_dimension_disagrees(query.shape[axis], key.shape[axis]) for axis in range(3) + ) + or _known_dimension_disagrees(query.shape[0], value.shape[0]) + ): + raise ValueError(f"Invalid GatedDeltaNet QKV layout at layer {layer}") + raw_a = _operand(node, "raw_a") + raw_b = _operand(node, "raw_b") + a_log = _operand(node, "a_log") + dt_bias = _operand(node, "dt_bias") + for name, operand, rank in ( + ("raw_a", raw_a, 2), + ("raw_b", raw_b, 2), + ("A_log", a_log, 1), + ("dt_bias", dt_bias, 1), + ): + _require_value( + operand, + dtype=ir.DataType.FLOAT, + rank=rank, + description=f"GatedDeltaNet {name} at layer {layer}", + ) + assert raw_a is not None and raw_a.shape is not None + assert raw_b is not None and raw_b.shape is not None + assert a_log is not None and a_log.shape is not None + assert dt_bias is not None and dt_bias.shape is not None + if ( + any( + _known_dimension_disagrees(raw_a.shape[axis], raw_b.shape[axis]) + for axis in range(2) + ) + or _known_dimension_disagrees(query.shape[0], raw_a.shape[0]) + or _known_dimension_disagrees(query.shape[0], raw_b.shape[0]) + or _known_dimension_disagrees(value.shape[1], state.shape[1]) + or _known_dimension_disagrees(value.shape[2], state.shape[2]) + or _known_dimension_disagrees(query.shape[2], state.shape[3]) + or _known_dimension_disagrees(raw_a.shape[1], state.shape[1]) + or _known_dimension_disagrees(a_log.shape[0], state.shape[1]) + or _known_dimension_disagrees(dt_bias.shape[0], state.shape[1]) + or ( + isinstance(query.shape[1], int) + and ( + query.shape[1] <= 0 + or (isinstance(state.shape[1], int) and state.shape[1] % query.shape[1] != 0) + ) + ) + ): + raise ValueError(f"Invalid V-major GatedDeltaNet layout at layer {layer}") + data_output = node.outputs[0] + if ( + data_output.dtype != value.dtype + or data_output.shape is None + or len(data_output.shape) != 3 + or _known_dimension_disagrees(data_output.shape[0], value.shape[0]) + or _known_dimension_disagrees(data_output.shape[1], state.shape[1]) + or _known_dimension_disagrees(data_output.shape[2], state.shape[2]) + ): + raise ValueError(f"Invalid GatedDeltaNet data output at layer {layer}") + + def inspect_paged_hybrid(model: ir.Model) -> PagedHybridAbi | None: """Require all three state disciplines, not merely a block-table input.""" if "mobius.paged_hybrid" not in model.metadata_props: return None if model.metadata_props["mobius.paged_hybrid"] != "qwen3_5_text": raise ValueError("Unknown packed hybrid ABI") + input_names = [value.name for value in model.graph.inputs] + output_names = [value.name for value in model.graph.outputs] + if len(input_names) != len(set(input_names)) or len(output_names) != len( + set(output_names) + ): + raise ValueError("Packed hybrid graph ports must have unique names") inputs = {value.name: value for value in model.graph.inputs} outputs = {value.name: value for value in model.graph.outputs} required = { @@ -86,12 +475,12 @@ def inspect_paged_hybrid(model: ir.Model) -> PagedHybridAbi | None: if ( past.shape is None or len(past.shape) != rank - or dtype not in {ir.DataType.FLOAT16, ir.DataType.BFLOAT16, ir.DataType.FLOAT} - or (kind != "recurrent_state" and dtype == ir.DataType.FLOAT) + or dtype not in _COMPUTE_DTYPES | {ir.DataType.FLOAT} + or (kind != "recurrent_state" and dtype not in _COMPUTE_DTYPES) or past.dtype != dtype or present is None or present.dtype != dtype - or present.shape != past.shape + or not _shapes_compatible(present.shape, past.shape) ): raise ValueError(f"Invalid packed hybrid state pair {name!r}") if kind in {"key", "value"} and past.shape[1] != block_size: @@ -108,42 +497,52 @@ def inspect_paged_hybrid(model: ir.Model) -> PagedHybridAbi | None: raise ValueError( "Packed hybrid requires disjoint, complete paged and fixed layer groups" ) - native = {} - state_operands = {"PagedAttention": 3, "VarlenCausalConvWithState": 4, "GatedDeltaNet": 6} + expected_outputs = {"logits"} | { + f"present.{layer}.{kind}" for kind, ids in layers.items() for layer in ids + } + if set(outputs) != expected_outputs: + raise ValueError("Packed hybrid graph has unknown or missing outputs") + _validate_request_dimensions(inputs, layers) + + native: dict[tuple[str, str], ir.Node] = {} for node in model.graph: - if node.domain != "com.microsoft" or node.op_type not in state_operands: + if node.domain != "com.microsoft" or node.op_type not in _STATE_OPERANDS: continue - operand = state_operands[node.op_type] + operand = _STATE_OPERANDS[node.op_type] if len(node.inputs) <= operand or node.inputs[operand] is None: raise ValueError(f"Missing native {node.op_type} state operand") key = (node.op_type, node.inputs[operand].name) if key in native: raise ValueError(f"Duplicate native state consumer {key}") native[key] = node - for kind, ids, op_type in ( - ("key", full, "PagedAttention"), - ("conv_state", linear, "VarlenCausalConvWithState"), - ("recurrent_state", linear, "GatedDeltaNet"), - ): - for layer in ids: - node = native.get((op_type, f"past_key_values.{layer}.{kind}")) - if node is None: - raise ValueError(f"Missing native {op_type} for layer {layer}") - cu_index = ( - 5 - if op_type == "PagedAttention" - else (2 if op_type == "VarlenCausalConvWithState" else 3) - ) - if node.inputs[cu_index] is not inputs["cumulative_sequence_lengths"]: - raise ValueError(f"Missing packed sequence boundaries at layer {layer}") - if op_type == "PagedAttention" and ( - node.attributes.get_string("kv_cache_layout", "") != "SEPARATE" - or len(node.inputs) != 17 - or any(value is not None for value in node.inputs[8:16]) - or node.inputs[16] is not inputs["attention_metadata"] - or node.inputs[4] is not inputs[f"past_key_values.{layer}.value"] - or node.inputs[6] is not inputs["past_sequence_lengths"] - or node.inputs[7] is not inputs["block_table"] - ): - raise ValueError(f"Invalid SEPARATE PagedAttention contract at layer {layer}") + + expected_native_keys: set[tuple[str, str]] = set() + for layer in full: + key = inputs[f"past_key_values.{layer}.key"] + node_key = ("PagedAttention", key.name) + expected_native_keys.add(node_key) + node = native.get(node_key) + if node is None: + raise ValueError(f"Missing native PagedAttention for layer {layer}") + _validate_paged_attention(node, layer=layer, inputs=inputs, outputs=outputs) + + for layer in linear: + conv_state = inputs[f"past_key_values.{layer}.conv_state"] + conv_key = ("VarlenCausalConvWithState", conv_state.name) + expected_native_keys.add(conv_key) + conv = native.get(conv_key) + if conv is None: + raise ValueError(f"Missing native VarlenCausalConvWithState for layer {layer}") + _validate_varlen_conv(conv, layer=layer, inputs=inputs, outputs=outputs) + + recurrent_state = inputs[f"past_key_values.{layer}.recurrent_state"] + delta_key = ("GatedDeltaNet", recurrent_state.name) + expected_native_keys.add(delta_key) + delta = native.get(delta_key) + if delta is None: + raise ValueError(f"Missing native GatedDeltaNet for layer {layer}") + _validate_gated_delta_net(delta, layer=layer, inputs=inputs, outputs=outputs) + + if set(native) != expected_native_keys: + raise ValueError("Packed hybrid graph has extra or cross-bound native state bindings") return PagedHybridAbi(block_size, tuple(sorted(full)), tuple(sorted(linear))) diff --git a/src/mobius/integrations/onnx_genai/paged_hybrid_metadata_test.py b/src/mobius/integrations/onnx_genai/paged_hybrid_metadata_test.py index f5adcb43d..2bb3ae46c 100644 --- a/src/mobius/integrations/onnx_genai/paged_hybrid_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/paged_hybrid_metadata_test.py @@ -61,6 +61,11 @@ def package(): def test_exact_state_disciplines_and_application_scheduling(package): + abi = inspect_paged_hybrid(package["model"]) + assert abi is not None + assert abi.full_layers == (3,) + assert abi.linear_layers == (0, 1, 2) + metadata = build_decoder_workflow_metadata(package, package.config) schema = json.loads( (Path(__file__).with_name("_schema") / "inference_metadata.schema.json").read_text() @@ -169,6 +174,300 @@ def test_unmarked_block_table_is_not_hybrid_compatibility(package): assert inspect_paged_hybrid(package["model"]) is None +def test_verifier_accepts_normal_saved_export(package, tmp_path): + package.save(str(tmp_path), check_weights=False, progress_bar=False) + + assert inspect_paged_hybrid(ir.load(tmp_path / "model.onnx")) is not None + + +def test_verifier_accepts_unrelated_symbolic_request_dimensions(package): + graph = package["model"].graph + request_inputs = { + value.name: value + for value in graph.inputs + if value.name + in {"block_table", "past_sequence_lengths", "cumulative_sequence_lengths"} + or value.name.endswith((".conv_state", ".recurrent_state")) + } + request_inputs["block_table"].shape = ir.Shape(["blocks_batch", "max_blocks"]) + request_inputs["past_sequence_lengths"].shape = ir.Shape(["lengths_batch"]) + request_inputs["cumulative_sequence_lengths"].shape = ir.Shape(["boundaries"]) + for name, state in request_inputs.items(): + if not name.endswith((".conv_state", ".recurrent_state")): + continue + state.shape = ir.Shape([f"{name}.batch", *state.shape[1:]]) + layer, kind = name.split(".")[1:3] + present = next( + value for value in graph.outputs if value.name == f"present.{layer}.{kind}" + ) + present.shape = ir.Shape([f"present.{name}.batch", *state.shape[1:]]) + + key = next(value for value in graph.inputs if value.name == "past_key_values.3.key") + value = next(value for value in graph.inputs if value.name == "past_key_values.3.value") + key.shape = ir.Shape(["key_pages", 256, "key_heads", "key_dim"]) + value.shape = ir.Shape(["value_pages", 256, "value_heads", "value_dim"]) + next(value for value in graph.outputs if value.name == "present.3.key").shape = ir.Shape( + ["present_key_pages", 256, "present_key_heads", "present_key_dim"] + ) + next(value for value in graph.outputs if value.name == "present.3.value").shape = ir.Shape( + ["present_value_pages", 256, "present_value_heads", "present_value_dim"] + ) + + paged = next(node for node in graph if node.op_type == "PagedAttention") + for index, name in enumerate(("query", "key", "value")): + operand = paged.inputs[index] + operand.shape = ir.Shape([f"{name}_tokens", operand.shape[1]]) + paged.outputs[0].shape = ir.Shape(["output_tokens", paged.outputs[0].shape[1]]) + + assert inspect_paged_hybrid(package["model"]) is not None + + +@pytest.mark.parametrize( + ("operand_index", "mutation"), + [ + (0, "dtype"), + (1, "dtype"), + (2, "dtype"), + (0, "rank"), + (1, "rank"), + (2, "rank"), + (0, "width"), + (1, "width"), + (2, "width"), + ], +) +def test_verifier_rejects_invalid_paged_attention_qkv(package, operand_index, mutation): + node = next(node for node in package["model"].graph if node.op_type == "PagedAttention") + operand = node.inputs[operand_index] + assert operand is not None and operand.shape is not None + if mutation == "dtype": + operand.type = ir.TensorType(ir.DataType.FLOAT) + elif mutation == "rank": + operand.shape = ir.Shape([1, *operand.shape]) + else: + operand.shape = ir.Shape([operand.shape[0], 7]) + + with pytest.raises(ValueError, match="PagedAttention"): + inspect_paged_hybrid(package["model"]) + + +@pytest.mark.parametrize("mutation", ["dtype", "rank", "width", "disconnected"]) +def test_verifier_rejects_invalid_paged_attention_data_output(package, mutation): + node = next(node for node in package["model"].graph if node.op_type == "PagedAttention") + output = node.outputs[0] + assert output.shape is not None + if mutation == "dtype": + output.type = ir.TensorType(ir.DataType.FLOAT) + elif mutation == "rank": + output.shape = ir.Shape([1, *output.shape]) + elif mutation == "width": + output.shape = ir.Shape([output.shape[0], 7]) + else: + replacement = ir.Value(type=output.type, shape=output.shape) + output.replace_all_uses_with(replacement) + + with pytest.raises(ValueError, match="PagedAttention data output"): + inspect_paged_hybrid(package["model"]) + + +@pytest.mark.parametrize("mutation", ["disconnected_key", "swapped_key_value"]) +def test_verifier_rejects_paged_attention_output_misbinding(package, mutation): + graph = package["model"].graph + key = next(value for value in graph.outputs if value.name == "present.3.key") + value = next(value for value in graph.outputs if value.name == "present.3.value") + if mutation == "swapped_key_value": + key.name, value.name = value.name, key.name + else: + key.name = "orphaned.present.3.key" + graph.outputs.remove(key) + graph.outputs.append( + ir.Value( + name="present.3.key", + type=key.type, + shape=key.shape, + ) + ) + + with pytest.raises(ValueError, match="PagedAttention"): + inspect_paged_hybrid(package["model"]) + + +@pytest.mark.parametrize( + ("state_name", "native_op"), + [ + ("present.0.conv_state", "VarlenCausalConvWithState"), + ("present.0.recurrent_state", "GatedDeltaNet"), + ], +) +def test_verifier_rejects_disconnected_linear_state_output(package, state_name, native_op): + graph = package["model"].graph + state = next(value for value in graph.outputs if value.name == state_name) + state.name = f"orphaned.{state_name}" + graph.outputs.remove(state) + graph.outputs.append(ir.Value(name=state_name, type=state.type, shape=state.shape)) + + with pytest.raises(ValueError, match=native_op): + inspect_paged_hybrid(package["model"]) + + +@pytest.mark.parametrize( + ("op_type", "attribute", "replacement"), + [ + ( + "VarlenCausalConvWithState", + "activation", + ir.AttrString("activation", "relu"), + ), + ("GatedDeltaNet", "gate_activation", ir.AttrString("gate_activation", "silu")), + ("GatedDeltaNet", "scale", ir.AttrFloat32("scale", 1.0)), + ], +) +def test_verifier_rejects_malformed_native_attributes( + package, op_type, attribute, replacement +): + node = next(node for node in package["model"].graph if node.op_type == op_type) + node.attributes[attribute] = replacement + + with pytest.raises(ValueError, match=op_type): + inspect_paged_hybrid(package["model"]) + + +@pytest.mark.parametrize( + ("op_type", "attribute", "replacement"), + [ + ("PagedAttention", "scale", ir.AttrFloat32("scale", 0.5)), + ("PagedAttention", "is_causal", ir.AttrInt64("is_causal", 0)), + ( + "PagedAttention", + "local_window_size", + ir.AttrInt64("local_window_size", 128), + ), + ("PagedAttention", "softcap", ir.AttrFloat32("softcap", 30.0)), + ("VarlenCausalConvWithState", "dilation", ir.AttrInt64("dilation", 2)), + ( + "VarlenCausalConvWithState", + "state_update_capacity", + ir.AttrInt64("state_update_capacity", 1), + ), + ("GatedDeltaNet", "state_update_capacity", ir.AttrInt64("state_update_capacity", 1)), + ], +) +def test_verifier_rejects_non_phase1_native_options(package, op_type, attribute, replacement): + node = next(node for node in package["model"].graph if node.op_type == op_type) + node.attributes[attribute] = replacement + + with pytest.raises(ValueError, match=op_type): + inspect_paged_hybrid(package["model"]) + + +@pytest.mark.parametrize(("mutation", "size"), [("input", 10), ("output", 3)]) +def test_verifier_rejects_extended_gated_delta_contract(package, mutation, size): + node = next(node for node in package["model"].graph if node.op_type == "GatedDeltaNet") + if mutation == "input": + node.resize_inputs(size) + else: + node.resize_outputs(size) + + with pytest.raises(ValueError, match="GatedDeltaNet"): + inspect_paged_hybrid(package["model"]) + + +@pytest.mark.parametrize("operand_index", [4, 5, 7, 8]) +def test_verifier_rejects_non_float_gated_delta_operands(package, operand_index): + node = next(node for node in package["model"].graph if node.op_type == "GatedDeltaNet") + operand = node.inputs[operand_index] + assert operand is not None + operand.type = ir.TensorType(ir.DataType.FLOAT16) + + with pytest.raises(ValueError, match="GatedDeltaNet"): + inspect_paged_hybrid(package["model"]) + + +def test_verifier_requires_matching_gated_delta_qkv_dtypes(package): + node = next(node for node in package["model"].graph if node.op_type == "GatedDeltaNet") + key = node.inputs[1] + assert key is not None + key.type = ir.TensorType(ir.DataType.BFLOAT16) + + with pytest.raises(ValueError, match="GatedDeltaNet QKV layout"): + inspect_paged_hybrid(package["model"]) + + +@pytest.mark.parametrize("operand_index", [2, 4, 5]) +def test_verifier_requires_matching_gated_delta_token_extents(package, operand_index): + node = next(node for node in package["model"].graph if node.op_type == "GatedDeltaNet") + for index in (0, 1, 2, 4, 5): + value = node.inputs[index] + assert value is not None and value.shape is not None + value.shape = ir.Shape([8, *value.shape[1:]]) + node.outputs[0].shape = ir.Shape([8, *node.outputs[0].shape[1:]]) + operand = node.inputs[operand_index] + assert operand is not None and operand.shape is not None + operand.shape = ir.Shape([7, *operand.shape[1:]]) + + with pytest.raises(ValueError, match="GatedDeltaNet"): + inspect_paged_hybrid(package["model"]) + + +def test_verifier_requires_depthwise_varlen_conv_weight(package): + node = next( + node for node in package["model"].graph if node.op_type == "VarlenCausalConvWithState" + ) + weight = node.inputs[1] + assert weight is not None and weight.shape is not None + weight.shape = ir.Shape([weight.shape[0], 2, weight.shape[2]]) + + with pytest.raises(ValueError, match="VarlenCausalConvWithState"): + inspect_paged_hybrid(package["model"]) + + +@pytest.mark.parametrize( + "target", + [ + "past_sequence_lengths", + "conv_state", + "recurrent_state", + "cumulative_sequence_lengths", + ], +) +def test_verifier_rejects_concrete_request_dimension_mismatch(package, target): + graph = package["model"].graph + inputs = {value.name: value for value in graph.inputs} + inputs["block_table"].shape = ir.Shape([2, inputs["block_table"].shape[1]]) + inputs["past_sequence_lengths"].shape = ir.Shape([2]) + inputs["cumulative_sequence_lengths"].shape = ir.Shape([3]) + for layer in range(3): + for kind in ("conv_state", "recurrent_state"): + state = inputs[f"past_key_values.{layer}.{kind}"] + state.shape = ir.Shape([2, *state.shape[1:]]) + present = next( + value for value in graph.outputs if value.name == f"present.{layer}.{kind}" + ) + present.shape = state.shape + + if target in {"past_sequence_lengths", "cumulative_sequence_lengths"}: + inputs[target].shape = ir.Shape([4]) + else: + state = inputs[f"past_key_values.0.{target}"] + state.shape = ir.Shape([4, *state.shape[1:]]) + present = next(value for value in graph.outputs if value.name == f"present.0.{target}") + present.shape = state.shape + + with pytest.raises(ValueError, match=r"request-aligned|batch \+ 1"): + inspect_paged_hybrid(package["model"]) + + +def test_verifier_rejects_cross_layer_recurrent_state_binding(package): + nodes = [node for node in package["model"].graph if node.op_type == "GatedDeltaNet"] + assert len(nodes) >= 2 + first_state, second_state = nodes[0].inputs[6], nodes[1].inputs[6] + nodes[0].replace_input_with(6, second_state) + nodes[1].replace_input_with(6, first_state) + + with pytest.raises(ValueError, match=r"state output|cross-bound"): + inspect_paged_hybrid(package["model"]) + + @pytest.mark.parametrize("ep", ["default", "cpu", "dml"]) def test_direct_task_rejects_non_cuda(ep): config = tiny_config() diff --git a/src/mobius/models/paged_qwen35_export_test.py b/src/mobius/models/paged_qwen35_export_test.py index ebc87cd16..cf650e156 100644 --- a/src/mobius/models/paged_qwen35_export_test.py +++ b/src/mobius/models/paged_qwen35_export_test.py @@ -5,11 +5,17 @@ from __future__ import annotations +import inspect + import onnx_ir as ir import pytest from mobius._testing import make_config -from mobius.models.qwen35 import Qwen35CausalLMModel +from mobius.models.qwen35 import ( + Qwen35CausalLMModel, + Qwen35DecoderLayer, + Qwen35TextModel, +) from mobius.tasks import PagedHybridCausalLMTask @@ -76,6 +82,18 @@ def test_packed_hybrid_io_and_native_operands(): assert "mobius.genai_revision" in model.metadata_props +def test_packed_context_has_an_explicit_component_api(): + causal_parameters = inspect.signature(Qwen35CausalLMModel.forward).parameters + text_parameters = inspect.signature(Qwen35TextModel.forward).parameters + layer_parameters = inspect.signature(Qwen35DecoderLayer.forward).parameters + + assert causal_parameters["paged_context"].kind is inspect.Parameter.KEYWORD_ONLY + assert "paged_context" in text_parameters + assert "paged_context" in layer_parameters + assert "PagedHybridContext" not in str(causal_parameters["attention_mask"].annotation) + assert "PagedHybridContext" not in str(layer_parameters["attention_bias"].annotation) + + def test_fp32_native_decay_parameters_and_qualified_weight_names(): model = _build(_config()) initializers = model.graph.initializers diff --git a/src/mobius/models/qwen35.py b/src/mobius/models/qwen35.py index bd42b3205..b717df733 100644 --- a/src/mobius/models/qwen35.py +++ b/src/mobius/models/qwen35.py @@ -22,7 +22,7 @@ ) from mobius.components._gated_deltanet import GatedDeltaNet from mobius.components._mlp import MLP -from mobius.components._paged_attention import PagedHybridContext +from mobius.components._paged_attention import PagedAttentionState, PagedHybridContext from mobius.components._quantized_linear import make_quantized_linear_factory from mobius.components._rms_norm import OffsetRMSNorm from mobius.components._rotary_embedding import initialize_rope @@ -176,9 +176,10 @@ def forward( self, op: OpBuilder, hidden_states: ir.Value, - attention_bias: ir.Value, + attention_bias: ir.Value | None, position_embeddings: tuple[ir.Value, ir.Value], - past_key_value: tuple[ir.Value, ir.Value] | None, + past_key_value: tuple[ir.Value, ir.Value] | PagedAttentionState | None, + paged_context: PagedHybridContext | None = None, ): residual = hidden_states hidden_states = self.input_layernorm(op, hidden_states) @@ -186,6 +187,8 @@ def forward( if self.layer_type == "linear_attention": # DeltaNet states are passed through past_key_value as # (conv_state, recurrent_state), same tuple pattern as KV cache + if isinstance(past_key_value, PagedAttentionState) or past_key_value is None: + raise TypeError("Linear attention requires conv/recurrent state") conv_state, recurrent_state = past_key_value attn_output, new_conv_state, new_recurrent_state = self.linear_attn( @@ -194,19 +197,29 @@ def forward( conv_state, recurrent_state, cumulative_sequence_lengths=( - attention_bias.cumulative_sequence_lengths - if isinstance(attention_bias, PagedHybridContext) + paged_context.cumulative_sequence_lengths + if paged_context is not None else None ), ) present_key_value = (new_conv_state, new_recurrent_state) else: + paged_state = None + dense_state = past_key_value + if paged_context is not None: + if not isinstance(past_key_value, PagedAttentionState): + raise TypeError("Packed full attention requires PagedAttentionState") + paged_state = past_key_value + dense_state = None + elif isinstance(past_key_value, PagedAttentionState): + raise TypeError("PagedAttentionState requires paged_context") attn_output, present_key_value = self.self_attn( op, hidden_states=hidden_states, attention_bias=attention_bias, position_embeddings=position_embeddings, - past_key_value=past_key_value, + past_key_value=dense_state, + paged_state=paged_state, ) hidden_states = op.Add(residual, attn_output) @@ -249,13 +262,16 @@ def forward( self, op: OpBuilder, input_ids: ir.Value | None, - attention_mask: ir.Value, + attention_mask: ir.Value | None, position_ids: ir.Value, past_key_values: list | None = None, inputs_embeds: ir.Value | None = None, deepstack_embeds: list | None = None, + paged_context: PagedHybridContext | None = None, ): - packed = isinstance(attention_mask, PagedHybridContext) + packed = paged_context is not None + if packed and attention_mask is not None: + raise ValueError("paged_context cannot be combined with attention_mask") # Embed tokens: dense (B,S,H), or packed (N,H). if inputs_embeds is not None: hidden_states = inputs_embeds @@ -269,7 +285,7 @@ def forward( # Causal attention mask: (batch, 1, seq_len, total_seq_len) attention_bias = ( - attention_mask + None if packed else create_attention_bias( op, @@ -290,6 +306,7 @@ def forward( attention_bias=attention_bias, position_embeddings=position_embeddings, past_key_value=past_kv, + paged_context=paged_context, ) present_key_values.append(present_kv) # DeepStack injection (see TextModel.forward for the rationale). @@ -327,12 +344,14 @@ def forward( self, op: OpBuilder, input_ids: ir.Value, - attention_mask: ir.Value | PagedHybridContext | None, + attention_mask: ir.Value | None, position_ids: ir.Value, past_key_values: list | None = None, + *, + paged_context: PagedHybridContext | None = None, ): """Preserve the dense ABI while supporting the dedicated packed task.""" - if not isinstance(attention_mask, PagedHybridContext): + if paged_context is None: return super().forward( op, input_ids, @@ -346,9 +365,10 @@ def forward( attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, + paged_context=paged_context, ) - if attention_mask.last_token_indices is not None: - hidden_states = op.Gather(hidden_states, attention_mask.last_token_indices, axis=0) + if paged_context.last_token_indices is not None: + hidden_states = op.Gather(hidden_states, paged_context.last_token_indices, axis=0) return self.lm_head(op, hidden_states), present_key_values def preprocess_weights( diff --git a/src/mobius/models/vibevoice.py b/src/mobius/models/vibevoice.py index 83fecf253..590a2ad0e 100644 --- a/src/mobius/models/vibevoice.py +++ b/src/mobius/models/vibevoice.py @@ -68,14 +68,6 @@ class VibeVoiceSources: _UNSUPPORTED_VIBEVOICE_MODELS = { - "microsoft/VibeVoice-ASR-Streaming-7B": ( - "VibeVoice ASR Streaming requires the VibeVoice-ASR streaming task, " - "which Mobius does not export yet." - ), - "microsoft/VibeVoice-ASR-Streaming-1.5B": ( - "VibeVoice ASR Streaming requires the VibeVoice-ASR streaming task, " - "which Mobius does not export yet." - ), "microsoft/VibeVoice-ASR-BitNet": ( "VibeVoice ASR BitNet requires the VibeVoice-ASR task and BitNet " "weight loader, which Mobius does not export yet." diff --git a/src/mobius/tasks/_causal_lm.py b/src/mobius/tasks/_causal_lm.py index 31a207078..5eded9e34 100644 --- a/src/mobius/tasks/_causal_lm.py +++ b/src/mobius/tasks/_causal_lm.py @@ -593,9 +593,10 @@ def build(self, module: nn.Module, config: ArchitectureConfig) -> ModelPackage: logits, present = module( op, input_ids=input_ids, - attention_mask=context, + attention_mask=None, position_ids=position_ids, past_key_values=states, + paged_context=context, ) logits = op.Cast(logits, to=ir.DataType.FLOAT) logits.shape = ir.Shape( diff --git a/tests/integration/paged_hybrid_test.py b/tests/integration/paged_hybrid_test.py index 5f7a6419c..b28e6699f 100644 --- a/tests/integration/paged_hybrid_test.py +++ b/tests/integration/paged_hybrid_test.py @@ -19,8 +19,8 @@ import torch from mobius import build_from_module +from mobius._testing import make_config from mobius.components._paged_attention import GENAI_REVISION, ORT_REVISION -from mobius.integrations.onnx_genai.paged_hybrid_metadata_test import tiny_config from mobius.integrations.ort_genai.auto_export import _write_genai_config from mobius.models.qwen35 import Qwen35CausalLMModel from mobius.tasks import HybridCausalLMTask, PagedHybridCausalLMTask @@ -39,8 +39,32 @@ def cuda_runtime(): pytest.skip(f"Set MOBIUS_ORT_REVISION={ORT_REVISION} for the pinned CUDA build") +def _tiny_config(dtype=ir.DataType.FLOAT16): + return make_config( + model_type="qwen3_5_text", + dtype=dtype, + hidden_size=128, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=64, + intermediate_size=192, + vocab_size=128, + max_position_embeddings=2048, + num_hidden_layers=4, + layer_types=["linear_attention"] * 3 + ["full_attention"], + partial_rotary_factor=0.5, + mrope_section=[8, 4, 4], + mrope_interleaved=True, + linear_num_value_heads=2, + linear_num_key_heads=1, + linear_key_head_dim=64, + linear_value_head_dim=32, + linear_conv_kernel_dim=4, + ) + + def _package(dtype, *, packed=True, prune=False): - config = tiny_config(dtype) + config = _tiny_config(dtype) module = Qwen35CausalLMModel(config) rng = np.random.default_rng(731) for name, parameter in module.named_parameters():