Convert LlamaAttention as 1 op (attention)

[glistening]

onert would like to run llama model.

• [onert/llm] Support tinyllama model ONE#15627

Llama has attention block, which is defined as LlamaAttention in modeling_llama.py from HF.

I would like to merge the all opcodes (includign RoPE) from LlamaAttention as 1 opcode ( attention ).
(All I need is attention.)
(I am going to add attention op in circle schema.)

It looks similar to #90.

But it is different:

• llama model does not use torch.nn.functional.scaled_dot_product_attention to implement LlamaAttention.
• scaled_dot_product_attention is based on standard transformer attention. (It does not have RoPE).

What would be the best way to do this?

[mhs4670go]

Fusing attention patterns into a single node could be a one solution. But, the pattern can change among torch versions. We need to consider this change.

Maybe this graph matcher can be used for this case.

[glistening]

Fusing attention patterns into a single node could be a one solution. But, the pattern can change among torch versions. We need to consider this change.

Sure, recently I got the problem from version mismatch.
We may specify the required torch version like transformers version.

Maybe this graph matcher can be used for this case.

Looks interesting.

[dayo09]

@glistening Could you specify the version of transformer and the function that you want to convert as a custom operation? I am interested in this task, so it'd be helpful if you share the details :-D

[glistening]

I am currently using transformers 4.52.4 after #148 (comment).

But I am okay to change version if you have specific version preference.

[dayo09]

About using fx's graph matcher

I tried out capturing the attention graph in Llama model with a simple LlamaAttention-only module.
It failed at first attempt, and I learned several points.

1. fx's graph matcher assumes all the torch nn modules are exported; The example hints us to use to_edge but I tested that exported graph is okay to use the matcher.
2. attention_mask, past_key_value are optional in the function and it changes the pattern. The series of optional arguments must match to create the same pattern. (NOTE: Their dimensions don't care)

Therefore, graph matcher can be used inside TICO as a helper, as it simply matches the pattern.

[dayo09]

Export LlamaAttention with KV Cache

They are not exported as expected. In the circle model, we can see that K-V value filling-in is done well but K-V pulling-out is not done.

Details

# from torch.fx.passes.utils.matcher_utils import SubgraphMatcher
import torch
from torch.export import export

from transformers.cache_utils import DynamicCache
from transformers.models.llama.modeling_llama import LlamaAttention, LlamaConfig
# from test.modules.model.Llama.model import Llama
large_inputs = (torch.randn(3, 3),)
pattern_inputs = (torch.randn(5, 5),)

class LargeModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self._weight = torch.nn.Parameter(torch.ones(3, 3))
        self._bias = torch.nn.Parameter(torch.ones(3, 3))

    def forward(self, x):
        return torch.ops.aten.addmm.default(self._bias, x, self._weight)

mod = LlamaAttention(LlamaConfig(), layer_idx = 0)
hidden_size = mod.config.hidden_size
head_dim = mod.config.head_dim
num_heads = mod.config.num_attention_heads

seq_len = 3 # temporal
hidden_states = torch.randn(1, seq_len, hidden_size)
position_embeddings = (torch.randn(1, seq_len, head_dim), torch.randn(1, seq_len, head_dim))
# attention_mask의 causal_mask 버젼: (batch_size, 1, query_length, key_value_length)
attention_mask = torch.Tensor([[[[0.]] * seq_len]]) # shape: 1, 1, 3, 1
prev_seq_len = 4
past_key_values = DynamicCache()
past_key_values.update(torch.randn(1, prev_seq_len, num_heads, head_dim), torch.randn(1, prev_seq_len, num_heads, head_dim), 1)
cache_position = torch.tensor([prev_seq_len], dtype=torch.long)

res = mod.forward(hidden_states, position_embeddings, attention_mask, past_key_values, cache_position,)
llama_attention_ep = export(mod, (hidden_states, position_embeddings, attention_mask, past_key_values, cache_position,))
llama_attetnion_graph = llama_attention_ep.graph

from tico import convert_from_exported_program
convert_from_exported_program(llama_attention_ep).save('llama_attention_with_kvcache.circle')

(ADDED) @glistening hinted me that KV cache may have been failed tracking because it has only one layer. Let me try this with 2 layers so that KV cache updated output could be used and tracked.

[dayo09]

LlamaAttention with KV Cache (1-layer)

This has kv cache I/O, too.

To keep key_states and value_states outputs, I edited the modelling_llama.py code as below.

# LlamaAttention.forward()
        # return attn_output, attn_weights
        return attn_output, attn_weights, key_states, value_states

To use graph matcher with above pattern, the pattern model must be kept in the project.
I suggest introducing patterns directory.

tico/
  passes/
  patterns/ 
    <some model fragments to be custom operation>

[dayo09]

Hi. To match several versions of attention, let's clarify their variants that I found 'till now.

(Analysis based on transformer==4.51.3)

Attention Variant (1): config._attn_implementation

1. config._attn_implementation == "eager"

• It's default to LlamaConfig()
• It uses standard softmax

2. config._attn_implementation == "sdpa"

• It uses _safe_softmax

NOTE: Above doesn't change softmax implementation because, both maps to the same circle::softmax.

However, they differ in scaling location;

In the eager implementation, the scaling factor is applied as a single multiplication of 1/8 after the matmul operation, effectively merging the two factors into one. In contrast, the SDPA implementation applies the scaling separately, using the value (1/8)^0.5 for each relevant component before the matmul. This results in the scaling being distributed across the computation rather than combined into a single operation.

Attention Variant (2): KV Cache

1. No caching: Not a target
2. DynamicCache: IS A TARGET
3. StaticCache: To be configured

Attention Variant (3): Grouped-Query Attention

# Uses repeat_kv()
key_states = repeat_kv(key_states, self.num_key_value_groups)
value_states = repeat_kv(value_states, self.num_key_value_groups)

• It's activated when num_key_value_groups > 1
• Let's assume it's always ON

(etc)

(TO BE ADDED)

[dayo09]

Note: Matching is sensitive in arguments and attn_implementation.

LlamaAttention has 5 arguments.

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor],
        past_key_value: Optional[Cache] = None,
        cache_position: Optional[torch.LongTensor] = None,
        # **kwargs: Unpack[FlashAttentionKwargs], <-- can be ignored
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:

As SubgraphMatcher finds exact match, these cases won't match.

(1) when attention_mask is not given

Details

(2) when config._attn_implemtation differs by 'eager' or 'sdpa

Details

[dayo09]

Trying out torch.ops.library.custom_op -> FAILED...

@glistening, I have tried out the approach but it seems only functions are allowed to be registered as custom op.

We have tico/utils/register_custom_ops.py which registers temporal custom operations at initilization.
So I added the attention Module and decorated its forward function with custon_op.

It failed with this error. AttributeError: type object 'Module' has no attribute '__origin__'

Full errors

Traceback (most recent call last):
  File "/home/dayoung.lee/miniconda3/envs/py310/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/home/dayoung.lee/miniconda3/envs/py310/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/home/dayoung.lee/miniconda3/envs/py310/lib/python3.10/unittest/__main__.py", line 18, in <module>
    main(module=None)
  File "/home/dayoung.lee/miniconda3/envs/py310/lib/python3.10/unittest/main.py", line 100, in __init__
    self.parseArgs(argv)
  File "/home/dayoung.lee/miniconda3/envs/py310/lib/python3.10/unittest/main.py", line 147, in parseArgs
    self.createTests()
  File "/home/dayoung.lee/miniconda3/envs/py310/lib/python3.10/unittest/main.py", line 158, in createTests
    self.test = self.testLoader.loadTestsFromNames(self.testNames,
  File "/home/dayoung.lee/miniconda3/envs/py310/lib/python3.10/unittest/loader.py", line 220, in loadTestsFromNames
    suites = [self.loadTestsFromName(name, module) for name in names]
  File "/home/dayoung.lee/miniconda3/envs/py310/lib/python3.10/unittest/loader.py", line 220, in <listcomp>
    suites = [self.loadTestsFromName(name, module) for name in names]
  File "/home/dayoung.lee/miniconda3/envs/py310/lib/python3.10/unittest/loader.py", line 154, in loadTestsFromName
    module = __import__(module_name)
  File "/home/dayoung.lee/git/TICO/test/pt2_to_circle_test/test_model.py", line 21, in <module>
    from test.pt2_to_circle_test.builder import NormalTestDictBuilder
  File "/home/dayoung.lee/git/TICO/test/pt2_to_circle_test/builder.py", line 24, in <module>
    from test.pt2_to_circle_test.test_pt2_to_circle import (
  File "/home/dayoung.lee/git/TICO/test/pt2_to_circle_test/test_pt2_to_circle.py", line 24, in <module>
    import tico.pt2_to_circle
  File "/home/dayoung.lee/git/TICO/tico/__init__.py", line 21, in <module>
    from tico.utils.convert import convert, convert_from_exported_program, convert_from_pt2
  File "/home/dayoung.lee/git/TICO/tico/utils/convert.py", line 37, in <module>
    from tico.serialize.circle_serializer import build_circle
  File "/home/dayoung.lee/git/TICO/tico/serialize/circle_serializer.py", line 29, in <module>
    from tico.serialize.operators import *
  File "/home/dayoung.lee/git/TICO/tico/serialize/operators/__init__.py", line 18, in <module>
    from tico.utils.register_custom_op import RegisterOps
  File "/home/dayoung.lee/git/TICO/tico/utils/register_custom_op.py", line 29, in <module>
    class LlamaAttention(nn.Module):
  File "/home/dayoung.lee/git/TICO/tico/utils/register_custom_op.py", line 56, in LlamaAttention
    def forward(
  File "/home/dayoung.lee/miniconda3/envs/py310/lib/python3.10/site-packages/torch/_library/custom_ops.py", line 149, in inner
    schema_str = torch.library.infer_schema(fn, mutates_args=mutates_args)
  File "/home/dayoung.lee/miniconda3/envs/py310/lib/python3.10/site-packages/torch/_library/infer_schema.py", line 127, in infer_schema
    if annotation_type.__origin__ is tuple:
AttributeError: type object 'Module' has no attribute '__origin__'
Traceback (most recent call last):
  File "/home/dayoung.lee/git/TICO/infra/scripts/test_run.py", line 107, in <module>
    main()
  File "/home/dayoung.lee/git/TICO/infra/scripts/test_run.py", line 101, in main
    subprocess.run(cmd, check=True)
  File "/home/dayoung.lee/miniconda3/envs/py310/lib/python3.10/subprocess.py", line 526, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['python3', '-m', 'unittest', 'test.pt2_to_circle_test.test_model', '-k', 'LlamaWithKVCache', '-v']' returned non-zero exit status 1.

I also attache the file diff.

Git diff

diff --git a/tico/utils/register_custom_op.py b/tico/utils/register_custom_op.py
index 2015922..72d0292 100644
--- a/tico/utils/register_custom_op.py
+++ b/tico/utils/register_custom_op.py
@@ -16,10 +16,123 @@ from typing import List, Optional
 
 import torch
 from torch._subclasses.fake_tensor import FakeTensor
-from torch.library import custom_op, register_fake
 
 from tico.utils.mx.mx_ops import _quantize_mx
 
+from transformers.models.llama.modeling_llama import LlamaConfig, eager_attention_forward, apply_rotary_pos_emb
+import torch
+from torch import nn
+from typing import Callable, Optional, Tuple
+from transformers.cache_utils import Cache, DynamicCache
+from torch.library import custom_op, register_fake
+
+class LlamaAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: LlamaConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+
+        self.q_proj = nn.Linear(
+            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+        )
+
+    @custom_op("circle_custom::llama_attention", mutates_args=('past_key_value'))
+    def forward(
+        self: nn.Module,
+        hidden_states: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_value: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        # **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_value is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        # attention_interface: Callable = eager_attention_forward
+        # if self.config._attn_implementation != "eager":
+        #     if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False):
+        #         logger.warning_once(
+        #             "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to "
+        #             'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+        #         )
+        #     else:
+        #         attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+        
+        attention_interface: Callable = eager_attention_forward #sdpa_attention_forward
+        
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            # **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        # Without this, the matcher cannot find multiple attention layers.
+        #
+        # return attn_output, attn_weights
+        return attn_output, attn_weights, key_states, value_states
+    
+    def get_example_inputs(self):
+        seq_len = 5# 1  # Assume token generation
+        hidden_size = self.config.hidden_size
+        head_dim = self.config.head_dim
+        num_heads = self.config.num_attention_heads
+
+        hidden_states = torch.randn(1, seq_len, hidden_size)
+        position_embeddings = (
+            torch.randn(1, seq_len, head_dim),
+            torch.randn(1, seq_len, head_dim),
+        )
+        attention_mask = torch.Tensor([[[[0.0]] * seq_len]])  # shape: 1, 1, seq_len, 1
+        # This attention_mask will become a causal_mask of shape: (batch_size, 1, query_length, key_value_length)
+        prev_seq_len = 4
+        past_key_values = DynamicCache()
+
+        past_key_values.update(
+            torch.randn(1, num_heads, prev_seq_len, head_dim),
+            torch.randn(1, num_heads, prev_seq_len, head_dim),
+            0,
+        )
+        cache_position = torch.tensor([[prev_seq_len]])
+        return (hidden_states, position_embeddings, attention_mask, past_key_values,)
+    
+
+
 # Note that an operator assumes input tensor has NHWC format.
 def CircleResizeNearestNeighbor():
     @custom_op("circle_custom::resize_nearest_neighbor", mutates_args=())

[dayo09]

@glistening

Hi. I got the circle file replacing the llama_attention. Plz check the draft #190 :-D
With the draft, run ./ccex test -m LlamaWithKVCache. The test fails now for not able to run the circle interpreter, but you can find the generated LlamaWithKVCache model in test/pt2_to_circle_test/artifacts/model/LlamaWithKVCache/model/LlamaWithKVCache.circle.
Let me work on, to support LlamaConfig which is hard-coded and need manual job for fix.

# ver
torch 2.7.0
transformers 4.51.3

---

transformers compatiblity check ✔️

transformers	compatible	reason
4.51.3	O
4.50.0	O
4.49.0	X	<class 'transformers.cache_utils.DynamicCache'>` is not supported or pytree-flattenable.
4.48.0	X	<class 'transformers.cache_utils.DynamicCache'>` is not supported or pytree-flattenable.
4.44.0	X	<class 'transformers.cache_utils.DynamicCache'>` is not supported or pytree-flattenable.