diff --git a/src/mobius/_flags.py b/src/mobius/_flags.py index 12a9e7fb4..fe4a694a5 100644 --- a/src/mobius/_flags.py +++ b/src/mobius/_flags.py @@ -84,11 +84,6 @@ class _Flags: - ``False`` - Lower the ONNX opset declaration to 23 for non-CPU EPs (ORT ≤1.24.x workaround). Disabled by default. - * - ``use_gqa_for_kv_shared`` - - ``MOBIUS_USE_GQA_FOR_KV_SHARED`` - - ``False`` - - Use GQA for KV-shared layers on CUDA EP. Requires ORT GQA - new_kv_length=0 support. """ suppress_dedup_warning: bool = dataclasses.field( @@ -108,22 +103,6 @@ class _Flags: Set ``MOBIUS_ORT_CUDA_GROUPED_RMSNORM_WORKAROUND=1`` when targeting CUDA. """ - use_gqa_for_kv_shared: bool = dataclasses.field( - default_factory=lambda: _env_bool("MOBIUS_USE_GQA_FOR_KV_SHARED", False) - ) - """Use GroupQueryAttention for KV-shared layers on CUDA EP. - - When ``False`` (default), KV-shared layers that borrow K/V from a - source layer use standard ONNX Attention (which falls back to unfused - attention on CUDA EP). This avoids a CUTLASS MEA crash with the - aligned kernel for certain sequence lengths when ``past_key`` is - ``nullptr``. - - Set ``MOBIUS_USE_GQA_FOR_KV_SHARED=1`` when the ORT GQA kernel - supports ``new_kv_length=0`` (KV-shared pattern), which will enable - fused attention for these layers. - """ - ort_lower_opset_for_ep: bool = dataclasses.field( default_factory=lambda: _env_bool("MOBIUS_ORT_LOWER_OPSET_FOR_EP", False) ) diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 3dd6ac511..3236d2d97 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -655,14 +655,6 @@ def _write_genai_config( **audio_kwargs, ) - # Disable past_present_share_buffer for models with dual head_dim - # (e.g., different head_dim for sliding vs full-attention layers). - # Shared buffers require uniform head_dim across all KV cache layers. - global_head_dim = getattr(config, "global_head_dim", None) - head_dim = getattr(config, "head_dim", None) - if global_head_dim and head_dim and global_head_dim != head_dim: - generator._search_overrides["past_present_share_buffer"] = False - return generator.write(output_dir) diff --git a/src/mobius/models/gemma4.py b/src/mobius/models/gemma4.py index 09a92cccb..ba7caa858 100644 --- a/src/mobius/models/gemma4.py +++ b/src/mobius/models/gemma4.py @@ -35,7 +35,6 @@ from mobius._build_context import ep_capabilities from mobius._configs import ArchitectureConfig, Gemma4Config -from mobius._flags import flags from mobius._weight_utils import vlm_decoder_weights, vlm_embedding_weights from mobius.components import ( MLP, @@ -736,41 +735,79 @@ def forward( ) if self.is_kv_shared_layer: - # KV-shared layers always use standard Attention path because - # they borrow K,V from a source layer (no own KV cache). - if use_gqa: - raise ValueError( - "KV-shared GQA path not yet implemented. " - "Set MOBIUS_USE_GQA_FOR_KV_SHARED=0 or wait for " - "ORT GQA new_kv_length=0 support." - ) - # Borrow full-history K,V from source layer. - # present_key/value from the ONNX Attention op is 4D: - # [batch, kv_heads, total_seq, head_dim] - # The Attention op expects key/value as 3D: - # [batch, total_seq, kv_heads * head_dim] - # Transpose and reshape to match. + # KV-shared layers borrow K,V from a source layer (no own KV cache). src_key, src_value = shared_kv_states[self.kv_shared_layer_index] - # [B, kv_heads, total_seq, head_dim] → [B, total_seq, kv_heads*head_dim] - src_key = op.Transpose(src_key, perm=[0, 2, 1, 3]) - src_key = op.Reshape(src_key, [0, 0, -1]) - src_value = op.Transpose(src_value, perm=[0, 2, 1, 3]) - src_value = op.Reshape(src_value, [0, 0, -1]) - - attn_output, present_key, present_value = _apply_attention( - op, - query_states, - src_key, - src_value, - attention_bias, - past_key=None, - past_value=None, - num_attention_heads=self.num_attention_heads, - num_key_value_heads=self.num_key_value_heads, - scale=self.scaling, - softcap=self.softcap, - ) + if use_gqa: + # GQA path for shared KV: pass empty K/V tensors and wire the + # source layer's present_key/value as past_key/past_value. + # The shared buffer is already in BNSH format, so GQA reads it + # directly — no Transpose/Reshape needed. + gqa_ctx = attention_bias + + # Create empty K/V tensors with kv_sequence_length=0. + # Shape: [batch, 0, kv_heads * head_dim] + batch_dim = op.Shape(query_states, start=0, end=1) + kv_hidden = self.num_key_value_heads * self.head_dim + empty_shape = op.Concat( + batch_dim, + op.Constant(value_ints=[0, kv_hidden]), + axis=0, + ) + empty_kv = op.CastLike(op.ConstantOfShape(empty_shape), query_states) + + gqa_attrs: dict = { + "num_heads": self.num_attention_heads, + "kv_num_heads": self.num_key_value_heads, + "scale": self.scaling, + "do_rotary": 1, + "rotary_interleaved": int(self._rope_interleave), + } + if self.softcap: + gqa_attrs["softcap"] = self.softcap + if self.rotary_embedding_dim: + gqa_attrs["rotary_embedding_dim"] = self.rotary_embedding_dim + if gqa_ctx.local_window_size > 0: + gqa_attrs["local_window_size"] = gqa_ctx.local_window_size + + attn_output, present_key, present_value = op.GroupQueryAttention( + query_states, + empty_kv, # key: empty (kv_sequence_length=0) + empty_kv, # value: empty (kv_sequence_length=0) + src_key, # past_key: shared KV in BNSH + src_value, # past_value: shared KV in BNSH + gqa_ctx.seqlens_k, + gqa_ctx.total_seq_len, + gqa_ctx.cos_cache, + gqa_ctx.sin_cache, + _domain="com.microsoft", + _outputs=3, + **gqa_attrs, + ) + else: + # Fallback Attention path: transpose shared KV from BNSH to 3D. + # present_key/value from the ONNX Attention op is 4D: + # [batch, kv_heads, total_seq, head_dim] + # The Attention op expects key/value as 3D: + # [batch, total_seq, kv_heads * head_dim] + src_key = op.Transpose(src_key, perm=[0, 2, 1, 3]) + src_key = op.Reshape(src_key, [0, 0, -1]) + src_value = op.Transpose(src_value, perm=[0, 2, 1, 3]) + src_value = op.Reshape(src_value, [0, 0, -1]) + + attn_output, present_key, present_value = _apply_attention( + op, + query_states, + src_key, + src_value, + attention_bias, + past_key=None, + past_value=None, + num_attention_heads=self.num_attention_heads, + num_key_value_heads=self.num_key_value_heads, + scale=self.scaling, + softcap=self.softcap, + ) elif use_gqa: # GQA path: emit com.microsoft.GroupQueryAttention directly. # The op fuses RoPE + attention + KV cache into a single op, @@ -1477,11 +1514,12 @@ def forward( ) if use_gqa: - # Realize cos/sin caches as ONNX graph initializers. - # The returned gathered embeddings are saved for potential reuse - # by KV-shared layers that fall back to standard Attention. - local_pos_emb = self.rotary_emb_local(op, position_ids) - global_pos_emb = self.rotary_emb_global(op, position_ids) + # Calling forward() on the RoPE modules materializes their + # cos_cache / sin_cache nn.Parameters as ONNX graph initializers. + # GQA references these caches directly; without the call the + # parameters are never emitted into the graph. + _ = self.rotary_emb_local(op, position_ids) + _ = self.rotary_emb_global(op, position_ids) # seqlens_k[b] = sum(attention_mask[b]) - 1 (last valid KV idx) # total_seq_len = attention_mask.shape[1] (past + current) @@ -1525,18 +1563,11 @@ def forward( "full_attention": self.rotary_emb_global(op, position_ids), } - # Fallback attention bias for non-GQA layers (KV-shared layers use - # this when use_gqa_for_kv_shared is False, plus all layers when - # use_gqa is False). + # Fallback attention bias for non-GQA layers (used when use_gqa is False). query_input = input_ids if input_ids is not None else hidden_states fallback_bias_dict: dict[str, ir.Value | None] = {} - has_kv_shared = any(layer.self_attn.is_kv_shared_layer for layer in self.layers) - need_fallback = not use_gqa or (has_kv_shared and not flags.use_gqa_for_kv_shared) + need_fallback = not use_gqa if need_fallback: - # All fallback layers use float additive bias masks encoding - # causal + sliding window + padding constraints. Float bias - # works with both unfused and MEA kernel paths on CUDA EP. - # Padding mask is required for batch > 1 correctness. fallback_bias_dict = { "sliding_attention": create_attention_bias( op, @@ -1552,16 +1583,7 @@ def forward( dtype=self._dtype, ), } - # KV-shared layers also need position embeddings for the - # standard Attention path (manual RoPE). Reuse the embeddings - # already gathered when realizing cos/sin caches above. - if use_gqa: - fallback_pos_dict = { - "sliding_attention": local_pos_emb, - "full_attention": global_pos_emb, - } - else: - fallback_pos_dict = position_embeddings_dict + fallback_pos_dict = position_embeddings_dict else: fallback_pos_dict = {} @@ -1587,12 +1609,9 @@ def forward( ): per_layer_input = per_layer_inputs[i] if per_layer_inputs is not None else None - # Per-layer decision: use GQA for non-shared layers when - # available, fall back to standard Attention for KV-shared - # layers (unless the use_gqa_for_kv_shared flag is set). - is_shared = layer.self_attn.is_kv_shared_layer - use_gqa_this_layer = use_gqa and (not is_shared or flags.use_gqa_for_kv_shared) - if use_gqa_this_layer: + # Per-layer decision: use GQA when available. KV-shared layers + # also use GQA (with empty K/V and shared past buffer). + if use_gqa: attn_bias = gqa_ctx_dict[layer_type] pos_emb = None else: diff --git a/src/mobius/rewrite_rules/_skip_norm.py b/src/mobius/rewrite_rules/_skip_norm.py index e83567cfa..03e4a901f 100644 --- a/src/mobius/rewrite_rules/_skip_norm.py +++ b/src/mobius/rewrite_rules/_skip_norm.py @@ -64,6 +64,21 @@ def check(self, context, add_out, norm_out, **_): if producer is None or producer.op_type != "Add": return result.fail("Input to RMSNorm is not from an Add node") + # Both Add inputs must have the same rank (SkipSimplifiedLayerNormalization + # requires input and skip to have the same shape). A rank mismatch + # indicates broadcasting (e.g. Add(MatMul, bias) where bias is 1D). + input_a = producer.inputs[0] + input_b = producer.inputs[1] + shape_a = input_a.shape + shape_b = input_b.shape + rank_a = len(shape_a) if shape_a is not None else None + rank_b = len(shape_b) if shape_b is not None else None + if rank_a is not None and rank_b is not None and rank_a != rank_b: + return result.fail( + f"Add inputs have different ranks ({rank_a} vs {rank_b}); " + "SkipSimplifiedLayerNormalization requires same-shape inputs" + ) + # Don't fuse if add_out is itself a graph output — that indicates we're inside # an ONNX function body (e.g. SkipSimplifiedLayerNormalization_body) where # replace_all_uses_with would fail, or produce nested fusion.