Attention mask for GQA fusion (#2452)

Expand the GQA fusion rule to handle attention mask better. 

* The patterns are extended to handle variations found in the
attention-mask logic for various models.
* It incorporates some optimizations of ModelBuilder that are arguably
not general-purpose, but make assumptions about the intended use-case
(which is the GenAI usage pattern).

---------

Signed-off-by: Ganesan Ramalingam <[email protected]>
Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: gramalingam <[email protected]>
Co-authored-by: Justin Chu <[email protected]>

Author: 3 people

onnxscript/rewriter/ort_fusions/gqa.py

@@ -8,7 +8,7 @@
 import onnx_ir as ir
 
 import onnxscript.rewriter._fusion_utils as _fusion_utils
-from onnxscript.rewriter import _ir_utils, pattern
+from onnxscript.rewriter import _basics, _ir_utils, pattern
 
 """
 GroupQueryAttention: This generalizes MHA by allowing the number of heads to be different
@@ -32,7 +32,20 @@
 Dim = Union[int, ir.SymbolicDim]
 
 
-def causal_mask_pattern(op, input_ids, past_kv_cache, shape_B111):
+def _is_model_input(value: ir.Value, name: str, model: ir.Model) -> bool:
+    return value in model.graph.inputs and value.name == name
+
+
+def _causal_mask(
+    op,
+    input_ids,
+    past_kv_cache,
+    shape_B111,
+    min_val,
+    window_size,
+    dtype,
+):
+    """Defines a pattern for a pure causal mask, with optional sliding window support."""
     seq_len = op.Shape(input_ids, end=2, start=1)
     seq_len_0D = op.Squeeze(seq_len)
 
@@ -42,28 +55,93 @@ def causal_mask_pattern(op, input_ids, past_kv_cache, shape_B111):
     total_seq_len_0D = op.Add(past_seq_len_0D, seq_len_0D)
     total_seq_len = op.Reshape(total_seq_len_0D, [-1])
 
-    # The Phi modeling code generates the following +1 as the target-length, which seems
-    # unnecessary in this context. But using it for pattern-matching against
-    # generated onnx model.
-    total_seq_len_plus_1_0D = op.Add(total_seq_len_0D, 1)
-    total_seq_len_plus_1 = op.Reshape(total_seq_len_plus_1_0D, [-1])
-
     current_range = op.Range(past_seq_len_0D, total_seq_len_0D, 1)
-    mask_shape = op.Concat(seq_len, total_seq_len_plus_1, axis=0)
-    min_float32 = float(np.finfo(np.float32).min)
-    mask_all_min = op.Expand(min_float32, mask_shape)
-    total_range_as_row = op.Range(0, total_seq_len_plus_1_0D, 1)
+    mask_shape = op.Concat(seq_len, total_seq_len, axis=0)
+    mask_all_min_expand = op.Expand(min_val, mask_shape)
+    # The following Trilu is optional: not used in Phi models, but used in LLama.
+    mask_all_min_trilu = op.Trilu(mask_all_min_expand, 1, upper=1)
+    mask_all_min = pattern.OrValue([mask_all_min_expand, mask_all_min_trilu])
+    total_range_as_row = op.Range(0, total_seq_len_0D, 1)
     current_range_as_column = op.Reshape(current_range, [-1, 1])
-    boolean_mask = op.Greater(total_range_as_row, current_range_as_column)
-    float_0_1_mask = op.Cast(boolean_mask, to=1)
+
+    non_causal = op.Greater(total_range_as_row, current_range_as_column)
+
+    # sliding window support:
+    current_range_minus_window = op.Sub(current_range_as_column, window_size)
+    out_of_sliding_window = op.LessOrEqual(total_range_as_row, current_range_minus_window)
+    non_causal_sliding_window = op.Or(non_causal, out_of_sliding_window)
+
+    boolean_mask = pattern.OrValue([non_causal, non_causal_sliding_window])
+
+    float_0_1_mask = op.Cast(boolean_mask, to=dtype)
     float_0_min_mask = op.Mul(mask_all_min, float_0_1_mask)
-    mask_4d = op.Unsqueeze(float_0_min_mask, [0, 1])
-    mask_B1ST_plus = op.Expand(mask_4d, shape_B111)
+    mask_4d_11ST = op.Unsqueeze(float_0_min_mask, [0, 1])
+    mask_4d_B1ST = op.Expand(mask_4d_11ST, shape_B111)
+
+    return mask_4d_B1ST
+
+
+class _CausalMaskPattern(pattern.PatternBase):
+    def pattern(
+        self,
+        op,
+        input_ids,
+        past_kv_cache,
+        shape_B111,
+        min_val,
+        window_size,
+        dtype1,
+        attn_mask_2d,
+        dtype2,
+    ):
+        causal_mask = _causal_mask(
+            op,
+            input_ids,
+            past_kv_cache,
+            shape_B111,
+            min_val,
+            window_size,
+            dtype1,
+        )
+
+        attn_mask_4d = op.Unsqueeze(attn_mask_2d, [1, 2])
+        attn_mask_4d_cast = op.Cast(attn_mask_4d, to=dtype2)
+
+        sum = op.Add(causal_mask, attn_mask_4d_cast)
+        sum_fp32 = op.Cast(sum, to=ir.DataType.FLOAT)
+        # The cast is optional, and may be absent if the sum is already in float32.
+        sum_fp32 = pattern.OrValue([sum_fp32, sum])
+        is_zero = op.Equal(sum_fp32, 0.0)
+        result = op.Where(is_zero, min_val, causal_mask)
+        return result
+
+    def check(self, context, dtype1, dtype2, min_val, attn_mask_2d, sliding_window=None, **_):
+        # Check that attn_mask_2d is the model input "attention_mask"
+        if not _is_model_input(attn_mask_2d, "attention_mask", context.model):
+            return pattern.MatchResult().fail("Invalid attention_mask input", attn_mask_2d)
+
+        if dtype1.as_int() != dtype2.as_int():
+            return pattern.MatchResult().fail("Dtype mismatch", [dtype1, dtype2])
+
+        # Check that min_val is a constant and matches the expected minimum value for the dtype.
+        min_value = _ir_utils.get_singleton_value(min_val)
+        if min_value is None:
+            return pattern.MatchResult().fail("Minval is not a constant.", min_val)
+        expected_min_value = np.finfo(min_val.dtype.numpy()).min
+        if min_value != expected_min_value:
+            return pattern.MatchResult().fail(
+                f"Expected min value {expected_min_value}, got {min_value}", min_val
+            )
+
+        # TODO(rama) Sliding window: not yet supported.
+        if sliding_window:
+            return pattern.MatchResult().fail(
+                "Sliding window not yet supported", sliding_window
+            )
+        return True
 
-    # Get rid of the extra +1 added above: total_seq_len is enough, no
-    # need for total_seq_len+1.
-    mask_B1ST = op.Slice(mask_B1ST_plus, [0], total_seq_len, [3], [1])
-    return mask_B1ST
+
+_causal_mask_pattern = _CausalMaskPattern()
 
 
 class GroupQueryAttention(pattern.RewriteRuleClassBase):
@@ -78,8 +156,7 @@ def pattern(
         value_BSDkv,
         past_key,
         past_value,
-        position_ids_q,
-        position_ids_k,
+        position_ids,
         cos,
         sin,
         mask,
@@ -101,15 +178,15 @@ def pattern(
 
         query_BHSDh_rope = op.RotaryEmbedding(
             query_BHSDh,
-            position_ids_q,
+            position_ids,
             cos,
             sin,
             _domain="com.microsoft",
             _outputs=["query_BHSDh_rope"],
         )
         key_BHkvSDh_rope = op.RotaryEmbedding(
             key_BHkvSDh,
-            position_ids_k,
+            position_ids,
             cos,
             sin,
             _domain="com.microsoft",
@@ -154,7 +231,7 @@ def pattern(
 
     def check(
         self,
-        op,
+        context: _basics.MatchContext,
         query_BSD,
         key_BSDkv,
         value_BSDkv,
@@ -164,6 +241,7 @@ def check(
         key_BHkvSDh_rope,
         query_BSHDh,
         key_BSHkvDh,
+        mask,
         **_,
     ):
         bindings: dict[str, Dim] = {}
@@ -210,6 +288,20 @@ def no_match(val: ir.Value, dims: Sequence[str]) -> bool:
             )
         self._interleaved = query_interleaved
 
+        # Check mask:
+        mask_node = mask.producer()
+        if mask_node is None:
+            return pattern.MatchResult().fail("Unhandled mask pattern", mask)
+        mask_match_result = _causal_mask_pattern.match(
+            context.model,
+            context.graph_or_function,
+            mask_node,
+            check_nodes_are_removable=False,
+        )
+        if mask_match_result is None:
+            return pattern.MatchResult().fail("Mask does not match causal mask pattern", mask)
+        # TODO: handle sliding window support in mask
+
         return True
 
     def rewrite(
@@ -220,104 +312,51 @@ def rewrite(
         value_BSDkv,
         past_key,
         past_value,
-        position_ids_q,
-        position_ids_k,
+        position_ids,
         cos,
         sin,
         mask,
         **_,
     ):
-        return op.GQA(
-            mask,
-            position_ids_k,
-            position_ids_q,
+        # Note that the following optimization is specific to current ORT GenAI attention-mask
+        # usage. Specifically, it assumes that the model-input "attention_mask" is a 2D
+        # mask with shape [batch_size, sequence_length], and that the mask is a 0/1 mask
+        # that is used only to indicate the current tokens. Hence, the input attention_mask
+        # is redundant as long as past-sequence-length and current-sequence-length can be
+        # computed.
+
+        # Construct seqlens_k and total_seq_length_int32 from position_ids
+        # seqlens_k : int32[batch_size] indicates total_sequence-length-1 for each batch
+        # position_ids: int64[batch_size, sequence_length] indicates the position of each token
+        one_int32_0d = op.Constant(value=ir.tensor(1, dtype=ir.DataType.INT32))
+        one_int64_1d = op.Constant(value=ir.tensor([1], dtype=ir.DataType.INT64))
+        zero_int64_1d = op.Constant(value=ir.tensor([0], dtype=ir.DataType.INT64))
+        seqlens_k_int64 = op.ReduceMax(position_ids, one_int64_1d, keepdims=0)
+        seqlens_k = op.Cast(seqlens_k_int64, to=ir.DataType.INT32)
+        max_seq_length = op.ReduceMax(seqlens_k, zero_int64_1d, keepdims=0)
+        total_seq_length_int32 = op.Add(max_seq_length, one_int32_0d)
+        return op.GroupQueryAttention(
             query_BSD,
             key_BSDkv,
             value_BSDkv,
             past_key,
             past_value,
-            None,  # seqlens_k,
-            None,  # total_seq_length_int32,
+            seqlens_k,
+            total_seq_length_int32,
             cos,
             sin,
             num_heads=self.num_heads,
             kv_num_heads=self.kv_num_heads,
             do_rotary=1,
             rotary_interleaved=self._interleaved,
             # skipped optional attributes: local_window_size, scale, smooth_softmax, softcap
-            _domain="ai.onnxruntime._fusion",
+            _domain="com.microsoft",
             _outputs=3,
         )
 
 
-class GQACausalMask(pattern.RewriteRuleClassBase):
-    def __init__(self):
-        super().__init__("GQACausalMask", remove_nodes=False)
-
-    def pattern(
-        self,
-        op,
-        mask,
-        input_ids,
-        some_kv_cache,
-        shape_B111,
-        past_seq_length,
-        total_seq_length,
-    ):
-        mask = causal_mask_pattern(op, input_ids, some_kv_cache, shape_B111)
-        position_ids = op.Range(past_seq_length, total_seq_length, 1)
-        position_ids_q = op.Unsqueeze(position_ids, [0])
-        position_ids_k = op.Unsqueeze(position_ids, [0])
-        return op.GQA(
-            mask,
-            position_ids_k,
-            position_ids_q,
-            _allow_other_inputs=True,
-            _domain="ai.onnxruntime._fusion",
-            _outputs=["attn_output", "key_seq", "value_seq"],
-        )
-
-    def rewrite(
-        self,
-        op,
-        total_seq_length,
-        attn_output,
-        **_,
-    ):
-        # Construct total_seq_length_int32 and seqlens_k
-        total_seq_length_int32 = op.Cast(total_seq_length, to=ir.DataType.INT32)
-        one_0D = op.Constant(value_int=1)
-        one_0D_int32 = op.Cast(one_0D, to=ir.DataType.INT32)
-        seqlens_k_0D = op.Sub(total_seq_length_int32, one_0D_int32)
-        zero_1D = op.Constant(value_int=0, dtype=ir.DataType.INT64, shape=[1])
-        seqlens_k = op.Unsqueeze(seqlens_k_0D, zero_1D)
-
-        gqa_node = attn_output.producer()
-        assert len(gqa_node.inputs) == 12, (
-            f"Expected 12 inputs for GQA node, got {len(gqa_node.inputs)}"
-        )
-        query, key, value, past_key, past_value = gqa_node.inputs[3:8]
-        cos, sin = gqa_node.inputs[10:12]
-        updated_inputs = [
-            query,
-            key,
-            value,
-            past_key,
-            past_value,
-            seqlens_k,
-            total_seq_length_int32,
-            cos,
-            sin,
-        ]
-        attributes = gqa_node.attributes
-        return op.GroupQueryAttention(
-            *updated_inputs, **attributes, _domain="com.microsoft", _outputs=3
-        )
-
-
 _basic_gqa_rule = GroupQueryAttention.rule()
-_gqa_causal_mask_rule = GQACausalMask.rule()
 
-gqa_rules = pattern.RewriteRuleSet([_basic_gqa_rule, _gqa_causal_mask_rule])
+gqa_rules = pattern.RewriteRuleSet([_basic_gqa_rule])
 
 fuse_gqa = _fusion_utils.apply_fusion_rules(gqa_rules)