[examples] Add torch ingress for llama example. (#41)

Author: kurapov-peter

examples/llama/ref_model.py

@@ -116,22 +116,24 @@ def __init__(self, args: ModelArgs):
             bias=False,
         )
 
-        self.cache_k = torch.zeros(
+        cache_k = torch.zeros(
             (
                 args.max_batch_size,
                 args.max_seq_len,
                 self.n_kv_heads,
                 self.head_dim,
             )
         )
-        self.cache_v = torch.zeros(
+        cache_v = torch.zeros(
             (
                 args.max_batch_size,
                 args.max_seq_len,
                 self.n_kv_heads,
                 self.head_dim,
             )
         )
+        self.register_buffer("cache_k", cache_k, persistent=False)
+        self.register_buffer("cache_v", cache_v, persistent=False)
 
     def forward(
         self,
@@ -149,14 +151,17 @@ def forward(
 
         xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
 
-        self.cache_k = self.cache_k.to(xq)
-        self.cache_v = self.cache_v.to(xq)
-
-        self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk
-        self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv
+        # TODO: the original implementation doesn't work with export.
+        # Local tensors instead of in-place buffer updates to please it.
+        cache_k_updated = self.cache_k.index_copy(
+            1, torch.arange(start_pos, start_pos + seqlen, device=xk.device), xk
+        )
+        cache_v_updated = self.cache_v.index_copy(
+            1, torch.arange(start_pos, start_pos + seqlen, device=xv.device), xv
+        )
 
-        keys = self.cache_k[:bsz, : start_pos + seqlen]
-        values = self.cache_v[:bsz, : start_pos + seqlen]
+        keys = cache_k_updated[:bsz, : start_pos + seqlen]
+        values = cache_v_updated[:bsz, : start_pos + seqlen]
 
         # repeat k/v heads if n_kv_heads < n_heads
         keys = repeat_kv(
@@ -246,17 +251,17 @@ def __init__(self, params: ModelArgs):
         self.norm = RMSNorm(params.dim, eps=params.norm_eps)
         self.output = nn.Linear(params.dim, params.vocab_size, bias=False)
 
-        self.freqs_cis = precompute_freqs_cis(
+        freqs_cis = precompute_freqs_cis(
             params.dim // params.n_heads,
             params.max_seq_len * 2,
             params.rope_theta,
         )
+        self.register_buffer("freqs_cis", freqs_cis, persistent=False)
 
     @torch.inference_mode()
     def forward(self, tokens: torch.Tensor, start_pos: int):
         _bsz, seqlen = tokens.shape
         h = self.tok_embeddings(tokens)
-        self.freqs_cis = self.freqs_cis.to(h.device)
         freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]
 
         mask = None

examples/llama/torch_ingress.py

@@ -0,0 +1,34 @@
+# RUN: %PYTHON %s
+# REQUIRES: torch
+
+import os
+from pathlib import Path
+
+import torch
+
+from lighthouse.ingress.torch import import_from_model
+from ref_model import ModelArgs, Transformer
+
+script_dir = Path(os.path.dirname(os.path.abspath(__file__)))
+model_path = script_dir / "ref_model.py"
+
+model_args = ModelArgs(
+    dim=512,
+    n_layers=2,
+    n_heads=8,
+    vocab_size=10000,
+    max_batch_size=1,
+    max_seq_len=128,
+)
+
+model = Transformer(model_args)
+sample_input = (torch.randint(0, model_args.vocab_size, (1, model_args.max_seq_len)), 0)
+
+mlir_module_str = import_from_model(
+    model, sample_args=sample_input, dialect="linalg-on-tensors"
+)
+
+dense_resource_idx = mlir_module_str.find("\n{-#\n  dialect_resources: {")
+assert dense_resource_idx != -1
+
+print(mlir_module_str[:dense_resource_idx])