diff --git a/README.md b/README.md index 0a83d15cd..0af82c117 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 --runtime ort-genai +``` + +This is a separate packed-token ABI, not a change to dense inference. It uses +SEPARATE K/V pages only for full-attention layers, native varlen convolution, +and FP32 V-major `[B, H_v, D_v, D_k]` recurrent state. The default page size is +256; `PagedHybridCausalLMTask(paged_block_size=...)` accepts positive multiples +of 256. `prune-prefill-prefix` selects each request's final packed row before +the LM head. Static cache, quantized checkpoints/KV, multimodal serving, and +CUDA graph capture are not supported in this path. + +The operator contract is pinned to ORT +`f38538cd5a4b5945a4c839565a8eebc65e1e2ef8` and GenAI +`d5b40851ba80ffa8e95b6b01f921dbb9008fac80`; released wheels are not assumed +compatible. Generated GenAI configuration declares `paged_kv`, `fixed_conv`, +and `fixed_recurrent` groups and dynamic batching. The workflow metadata +describes one externally scheduled invocation, not a dense generation loop; +the scheduler supplies packed boundaries, page tables, and request lifecycle +inputs. `attention_metadata` is CPU INT32 containing query-length upper bound, +KV-length upper bound, and KV-length lower bound. Paged cache outputs **must** +alias their inputs; fixed states must follow the same request ordering. + +**Not production-qualified:** a full Qwen3.8-27B CUDA smoke run is still required. +The network-free CUDA parity and Engine probes are in +`tests/integration/paged_hybrid_test.py`; set `MOBIUS_ORT_REVISION` and +`MOBIUS_GENAI_REVISION` to the above revisions when testing matching source +builds. CPU-only runs skip these probes. + Use `--release` with either `build` or `build-gguf` to potentially reduce saved model size by stripping build-time debug and provenance metadata. Functional metadata with keys prefixed by `mobius.` is preserved: diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index f9e4d4834..ff2c96183 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..355d40018 100644 --- a/src/mobius/_builder.py +++ b/src/mobius/_builder.py @@ -76,12 +76,15 @@ def _enable_prefill_prefix_pruning_task(task: str | ModelTask) -> str | ModelTas Gemma4Task, Gemma4TextCausalLMTask, HybridCausalLMTask, + PagedHybridCausalLMTask, ) if task == "text-generation": return CausalLMTask(prune_prefill_prefix=True) if task == "hybrid-text-generation": return HybridCausalLMTask(prune_prefill_prefix=True) + if task == "paged-hybrid-text-generation": + return PagedHybridCausalLMTask(prune_prefill_prefix=True) if task == "gemma4-text-generation": return Gemma4TextCausalLMTask(prune_prefill_prefix=True) if task == "gemma4": @@ -94,6 +97,11 @@ def _enable_prefill_prefix_pruning_task(task: str | ModelTask) -> str | ModelTas ) if isinstance(task, HybridCausalLMTask): return HybridCausalLMTask(prune_prefill_prefix=True) + if isinstance(task, PagedHybridCausalLMTask): + return PagedHybridCausalLMTask( + paged_block_size=task._paged_block_size, + prune_prefill_prefix=True, + ) if isinstance(task, Gemma4TextCausalLMTask): return Gemma4TextCausalLMTask( static_cache=getattr(task, "_static_cache", False), @@ -108,7 +116,8 @@ def _enable_prefill_prefix_pruning_task(task: str | ModelTask) -> str | ModelTas ) raise ValueError( "prune_prefill_prefix=True is only supported for text-generation, " - "hybrid-text-generation, gemma4-text-generation, and gemma4 tasks." + "hybrid-text-generation, paged-hybrid-text-generation, " + "gemma4-text-generation, and gemma4 tasks." ) @@ -162,6 +171,27 @@ def build_from_module( if prune_prefill_prefix: task = _enable_prefill_prefix_pruning_task(task) resolved_task = get_task(task) + from mobius.tasks import PagedHybridCausalLMTask + + if ( + getattr(config, "export_paged_attention", False) + and getattr(config, "model_type", None) == "qwen3_5_text" + and not isinstance(resolved_task, PagedHybridCausalLMTask) + ): + raise ValueError("Qwen paged attention requires PagedHybridCausalLMTask") + if isinstance(resolved_task, PagedHybridCausalLMTask): + if execution_provider != "cuda": + raise ValueError("PagedHybridCausalLMTask requires execution_provider='cuda'") + if fp8_kv_cache or kv_cache_scales is not None: + raise ValueError("PagedHybridCausalLMTask supports only unquantized KV caches") + # Direct task selection must preserve raw native gate precision too, + # even when the model was constructed without the feature flag. + from mobius.components import GatedDeltaNet + + for child in module.modules(): + if isinstance(child, GatedDeltaNet): + child.A_log._keep_float32 = True + child.dt_bias._keep_float32 = True component_manifest = resolved_task.component_manifest() configure_component_quantization(module, config, resolved_task) _cast_module_dtype(module, dtype) 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/__init__.py b/src/mobius/components/__init__.py index 7393061fa..c805f00b2 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", @@ -346,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/_attention.py b/src/mobius/components/_attention.py index 2fdd1ee40..ed3440410 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 @@ -592,7 +593,18 @@ def forward( position_embeddings: tuple, past_key_value: tuple | None = None, static_cache: StaticCacheState | None = None, + paged_state: PagedAttentionState | None = None, ): + if paged_state is not None: + if past_key_value is not None or static_cache is not None: + raise ValueError( + "Paged attention state cannot be combined with dense cache state" + ) + return self.forward_paged(op, hidden_states, position_embeddings, paged_state) + if isinstance(past_key_value, PagedAttentionState): + raise TypeError( + "PagedAttentionState must be passed through paged_state, not past_key_value" + ) # Q projection (doubled) → split into Q and gate per head q_gate = self.q_proj(op, hidden_states) # Reshape to per-head view so split separates Q/gate within each head @@ -650,3 +662,59 @@ def forward( attn_output = self.o_proj(op, attn_output) return attn_output, (present_key, present_value) + + def forward_paged( + self, + op: OpBuilder, + hidden_states: ir.Value, + position_embeddings: tuple, + state: PagedAttentionState, + ): + """Packed Qwen attention with external Q/K norm and partial M-RoPE.""" + q_gate = self.q_proj(op, hidden_states) + q_gate = op.Reshape(q_gate, [-1, self.num_attention_heads, self.head_dim * 2]) + query_states, gate = op.Split(q_gate, num_outputs=2, axis=-1, _outputs=2) + key_states = op.Reshape( + self.k_proj(op, hidden_states), [-1, self.num_key_value_heads, self.head_dim] + ) + value_states = self.v_proj(op, hidden_states) + + query_states = self.q_norm(op, query_states) + key_states = self.k_norm(op, key_states) + # RotaryEmbedding's rank-3 ABI is (batch,sequence,hidden), not (N,H,D). + # A singleton batch is safe here: only RoPE, never attention, sees it. + query_states = op.Reshape( + query_states, [1, -1, self.num_attention_heads * self.head_dim] + ) + key_states = op.Reshape(key_states, [1, -1, self.num_key_value_heads * self.head_dim]) + position_embeddings = tuple(op.Unsqueeze(x, [0]) for x in position_embeddings) + query_states = apply_rotary_pos_emb( + op, + query_states, + position_embeddings, + self.num_attention_heads, + self.rotary_embedding_dim, + self._rope_interleave, + ) + key_states = apply_rotary_pos_emb( + op, + key_states, + position_embeddings, + self.num_key_value_heads, + self.rotary_embedding_dim, + self._rope_interleave, + ) + query_states = op.Reshape(query_states, [-1, self.num_attention_heads * self.head_dim]) + key_states = op.Reshape(key_states, [-1, self.num_key_value_heads * self.head_dim]) + output, key_cache, value_cache = paged_attention( + op, + query_states, + key_states, + value_states, + state, + num_heads=self.num_attention_heads, + kv_num_heads=self.num_key_value_heads, + ) + q_width = self.num_attention_heads * self.head_dim + output = op.Mul(output, op.Sigmoid(op.Reshape(gate, [-1, q_width]))) + return self.o_proj(op, output), (key_cache, value_cache) 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 44be2185d..d57b47c34 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,54 @@ 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) + # 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( + 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..0e0ca37f5 --- /dev/null +++ b/src/mobius/components/_paged_attention.py @@ -0,0 +1,139 @@ +# 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.""" + output, key_cache, value_cache = 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, + ) + # 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( + 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.""" + output, present = op.VarlenCausalConvWithState( + packed, + weight, + cumulative_sequence_lengths, + bias, + past_conv, + activation="silu", + _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( + 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.""" + output, present = 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, + ) + 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/components/_rotary_embedding.py b/src/mobius/components/_rotary_embedding.py index 4e286d40b..ac3c97b63 100644 --- a/src/mobius/components/_rotary_embedding.py +++ b/src/mobius/components/_rotary_embedding.py @@ -478,18 +478,36 @@ 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: 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. + 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..a3b22a7a2 --- /dev/null +++ b/src/mobius/integrations/_paged_hybrid.py @@ -0,0 +1,548 @@ +# 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 + +_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: + 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 _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 = { + "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 _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 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: + 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" + ) + 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: + 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 + + 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/_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..2bb3ae46c --- /dev/null +++ b/src/mobius/integrations/onnx_genai/paged_hybrid_metadata_test.py @@ -0,0 +1,604 @@ +# 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): + 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() + ) + 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 + + +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() + with pytest.raises(ValueError, match="requires execution_provider"): + build_from_module( + Qwen35CausalLMModel(config), + config, + PagedHybridCausalLMTask(), + execution_provider=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 + + 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 938ae9f40..62e6970da 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -633,15 +633,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..cf650e156 --- /dev/null +++ b/src/mobius/models/paged_qwen35_export_test.py @@ -0,0 +1,133 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Structural tests for the packed Qwen3.5 hybrid serving graph.""" + +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, + Qwen35DecoderLayer, + Qwen35TextModel, +) +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_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 + 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..b717df733 100644 --- a/src/mobius/models/qwen35.py +++ b/src/mobius/models/qwen35.py @@ -22,6 +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._quantized_linear import make_quantized_linear_factory from mobius.components._rms_norm import OffsetRMSNorm from mobius.components._rotary_embedding import initialize_rope @@ -175,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) @@ -185,19 +187,39 @@ 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( - op, hidden_states, conv_state, recurrent_state + op, + hidden_states, + conv_state, + recurrent_state, + cumulative_sequence_lengths=( + 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) @@ -240,26 +262,37 @@ 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, ): - # Embed tokens: (batch, seq_len) → (batch, seq_len, hidden_size) + 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 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 = ( + None + 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 = [] @@ -273,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). @@ -306,6 +340,37 @@ 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 | 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 paged_context is None: + 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, + paged_context=paged_context, + ) + 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( self, state_dict: dict[str, torch.Tensor] ) -> dict[str, torch.Tensor]: 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/__init__.py b/src/mobius/tasks/__init__.py index 1677cd01e..4c1c557e8 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -61,6 +61,7 @@ "GlmMoeDsaTask", "GlmOcrVLTask", "HybridCausalLMTask", + "PagedHybridCausalLMTask", "HyV3MtpTask", "FalconH1CausalLMTask", "Cosmos3EdgeVLTask", @@ -133,6 +134,7 @@ from mobius.tasks._causal_lm import ( CausalLMTask, HybridCausalLMTask, + PagedHybridCausalLMTask, SmallThinkerGGUFCausalLMTask, ) from mobius.tasks._codec import CodecTask @@ -266,6 +268,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..5eded9e34 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,204 @@ 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: + 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 " + 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}" + ) + 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") + 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=None, + position_ids=position_ids, + past_key_values=states, + paged_context=context, + ) + 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..b28e6699f --- /dev/null +++ b/tests/integration/paged_hybrid_test.py @@ -0,0 +1,401 @@ +# 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._testing import make_config +from mobius.components._paged_attention import GENAI_REVISION, ORT_REVISION +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 _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) + 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", 2)], + [("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=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 + 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