Build RoPE cos, sin tensors on demand

PiperOrigin-RevId: 712561555

Author: talumbau

ai_edge_torch/generative/examples/gemma/gemma2.py

@@ -15,14 +15,13 @@
 
 """Example of building a Gemma2 model."""
 
-from typing import List, Optional, Tuple
+from typing import Optional, Tuple
 
 from ai_edge_torch.generative.layers import attention
 from ai_edge_torch.generative.layers import builder
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
 import ai_edge_torch.generative.layers.model_config as cfg
-import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
 from ai_edge_torch.generative.utilities import model_builder
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
@@ -104,12 +103,17 @@ def __init__(self, config: cfg.ModelConfig):
         config.embedding_dim,
         config.final_norm_config,
     )
-    self.mask_cache = attn_utils.build_causal_mask_cache(
-        size=config.kv_cache_max,
-    )
     # Gemma2 has same hyper parameters for each layer except for attention
     # types. Use the first layer.
     attn_config = config.block_config(0).attn_config
+    self.rope_cache = attn_utils.build_rope_cache(
+        size=config.kv_cache_max,
+        dim=int(attn_config.rotary_percentage * attn_config.head_dim),
+        base=attn_config.rotary_base,
+    )
+    self.mask_cache = attn_utils.build_causal_mask_cache(
+        size=config.kv_cache_max,
+    )
     self.sliding_window_mask_cache = attn_utils.build_sliding_window_mask_cache(
         size=config.kv_cache_max,
         window_size=attn_config.sliding_window_size,
@@ -136,48 +140,29 @@ def forward(
         f"Cannot forward sequence of length {seq_len}, max seq length is only"
         f" {self.config.max_seq_len}"
     )
-
-    # token embeddings of shape (b, t, n_embd)
-    input_embeds = self.tok_embedding(tokens)
-    # RoPE parameters are the same for all blocks. Use the first layer.
-    attn_config = self.config.block_config(0).attn_config
-    n_elem = int(attn_config.rotary_percentage * attn_config.head_dim)
-    rope = rotary_pos_emb.build_rope(
-        input_pos, n_elem, attn_config.head_dim, attn_config.rotary_base
-    )
-    mask = [self.get_attention_mask(
-        self.config.block_config(i).attn_config.attn_type, input_pos
-    ) for i in range(self.config.num_layers)]
-
-    return self._forward_with_embeds(
-        input_embeds, rope, mask, input_pos, kv_cache, export_config
-    )
-
-  def _forward_with_embeds(
-      self,
-      input_embeds: torch.Tensor,
-      rope: Tuple[torch.Tensor, torch.Tensor],
-      mask: List[torch.Tensor],
-      input_pos: torch.Tensor,
-      kv_cache: kv_utils.KVCache,
-      export_config: Optional[model_builder.ExportConfig] = None,
-  ) -> dict[torch.Tensor, kv_utils.KVCache]:
-    """Forwards the model with input embeddings."""
     assert len(self.transformer_blocks) == len(kv_cache.caches), (
         "The number of transformer blocks and the number of KV cache entries"
         " must be the same."
     )
 
-    if self.config.embedding_scale is not None:
-      input_embeds = input_embeds * self.config.embedding_scale
-    x = input_embeds
-    updated_kv_entries = []
+    cos, sin = self.rope_cache
+    cos = cos.index_select(0, input_pos)
+    sin = sin.index_select(0, input_pos)
+
+    # token embeddings of shape (b, t, n_embd)
+    x = self.tok_embedding(tokens)
+    x = x * (self.config.embedding_dim**0.5)
+
+    updated_kv_entires = []
     for i, block in enumerate(self.transformer_blocks):
+      mask = self.get_attention_mask(
+          block.config.attn_config.attn_type, input_pos
+      )
       kv_entry = kv_cache.caches[i] if kv_cache else None
-      x, kv_entry = block(x, rope, mask[i], input_pos, kv_entry)
+      x, kv_entry = block(x, (cos, sin), mask, input_pos, kv_entry)
       if kv_entry:
-        updated_kv_entries.append(kv_entry)
-    updated_kv_cache = kv_utils.KVCache(tuple(updated_kv_entries))
+        updated_kv_entires.append(kv_entry)
+    updated_kv_cache = kv_utils.KVCache(tuple(updated_kv_entires))
 
     if export_config is not None:
       if (
@@ -243,13 +228,11 @@ def get_block_config(idx: int) -> cfg.TransformerBlockConfig:
     )
 
   num_layers = 26
-  embedding_dim = 2304
   config = cfg.ModelConfig(
       vocab_size=256000,
       num_layers=num_layers,
       max_seq_len=8192,
-      embedding_dim=embedding_dim,
-      embedding_scale=embedding_dim**0.5,
+      embedding_dim=2304,
       kv_cache_max_len=kv_cache_max_len,
       block_configs=[get_block_config(i) for i in range(num_layers)],
       final_norm_config=norm_config,
@@ -266,7 +249,6 @@ def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
   config.num_layers = 2
   config.max_seq_len = 2 * kv_cache_max_len
   config.embedding_dim = 128
-  config.embedding_scale = config.embedding_dim**0.5
   config.block_configs = config.block_configs[: config.num_layers]
   for block_config in config.block_configs:
     block_config.attn_config.num_heads = 4

ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py

@@ -72,14 +72,14 @@ def forward(
     mask = self.mask_cache.index_select(2, input_pos)
     mask = mask[:, :, :, : self.config.max_seq_len]
 
-    updated_kv_entries = []
+    updated_kv_entires = []
     for i, block in enumerate(self.transformer_blocks):
       kv_entry = kv_cache.caches[i] if kv_cache else None
       x, kv_entry = block(x, (cos, sin), mask, input_pos, kv_entry)
       if kv_entry:
-        updated_kv_entries.append(kv_entry)
+        updated_kv_entires.append(kv_entry)
 
-    updated_kv_cache = kv_utils.KVCache(tuple(updated_kv_entries))
+    updated_kv_cache = kv_utils.KVCache(tuple(updated_kv_entires))
 
     if export_config is not None:
       if (

ai_edge_torch/generative/layers/attention.py

@@ -26,6 +26,33 @@
 from torch import nn
 
 
+def _embed_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    n_elem: int,
+    rope: Tuple[torch.Tensor, torch.Tensor],
+) -> Tuple[torch.Tensor, torch.Tensor]:
+  """Embed rotary positional embedding for query and key.
+
+  Args:
+    q (torch.Tensor): query tensor.
+    k (torch.Tensor): key tensor.
+    n_elem (int): number of elements to embed rotarty positional embedding.
+    rope (Tuple[torch.Tensor, torch.Tensor]): the input rope tensor.
+  """
+  if n_elem > 0:
+    cos, sin = rope
+    q_roped = rotary_pos_emb.apply_rope(
+        q[..., :n_elem], cos.repeat(1, 2), sin.repeat(1, 2)
+    )
+    k_roped = rotary_pos_emb.apply_rope(
+        k[..., :n_elem], cos.repeat(1, 2), sin.repeat(1, 2)
+    )
+    q = torch.cat((q_roped, q[..., n_elem:]), dim=-1)
+    k = torch.cat((k_roped, k[..., n_elem:]), dim=-1)
+  return q, k
+
+
 class TransformerBlock(nn.Module):
 
   def __init__(
@@ -211,8 +238,7 @@ def forward(
     if rope is not None:
       # Compute rotary positional embedding for query and key.
       n_elem = int(self.config.rotary_percentage * self.config.head_dim)
-      cos, sin = rope
-      q, k = rotary_pos_emb.apply_rope_inline(q, k, cos, sin)
+      q, k = _embed_rope(q, k, n_elem, rope)
 
     if kv_cache is not None:
       kv_cache = kv_utils.update(kv_cache, input_pos, k, v)
@@ -348,8 +374,7 @@ def forward(
     if rope is not None:
       # Compute rotary positional embedding for query and key.
       n_elem = int(self.config.rotary_percentage * self.config.head_dim)
-      cos, sin = rope
-      q, k = rotary_pos_emb.apply_rope_inline(q, k, cos, sin)
+      q, k = _embed_rope(q, k, n_elem, rope)
 
     if kv_cache is not None:
       kv_cache = kv_utils.update(kv_cache, input_pos, k, v)

ai_edge_torch/generative/layers/rotary_position_embedding.py

@@ -32,64 +32,57 @@ def apply_rope(
   """
   x = x.transpose(1, 2)
   head_size = x.size(-1)
-  x1, x2 = torch.split(x, head_size // 2, dim=-1)
-  left = x1 * cos - x2 * sin
-  right = x2 * cos + x1 * sin
-  roped = torch.cat([left, right], dim=-1)
+  x1 = x[..., : head_size // 2]  # (B, nh, T, hs/2)
+  x2 = x[..., head_size // 2 :]  # (B, nh, T, hs/2)
+  rotated = torch.cat((-x2, x1), dim=-1)  # (B, nh, T, hs)
+  roped = (x * cos) + (rotated * sin)
   return roped.transpose(1, 2).type_as(x)
 
 
-def build_rope(
+def apply_rope_inline(
+    q: torch.Tensor,
+    k: torch.Tensor,
     input_pos: torch.Tensor,
     n_elem: int,
-    head_dim: int,
     base: int = 10_000,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-  """Computes rotary positional embedding cosine and sine tensors.
+  """Computes rotary positional embedding inline for a query and key.
 
   Args:
+    q: the query tensor.
+    k: the key tensor.
     input_pos: the sequence indices for the query and key
     n_elem: number of elements of the head dimension for RoPE computation
-    base: the base of the exponentiated value for RoPE.
 
   Returns:
-    cos, sin tensors
+    output the RoPE'd query and key.
   """
 
   if n_elem <= 0:
-    return None, None
+    return q, k
 
   theta = 1.0 / (base ** (torch.arange(0, n_elem, 2).float() / n_elem))
   freq_exponents = (2.0 / n_elem) * torch.arange(
-      head_dim // 2, dtype=torch.float32
+      q.shape[-1] // 2, dtype=torch.float32
   )
   timescale = float(base) ** freq_exponents
   radians = input_pos.clone().unsqueeze(0).unsqueeze(-1) / timescale.unsqueeze(
       0
   ).unsqueeze(0)
-  cos = torch.cos(radians)
-  sin = torch.sin(radians)
-  return cos, sin
-
+  cos = torch.cos(radians).type_as(q)
+  sin = torch.sin(radians).type_as(q)
 
-def apply_rope_inline(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    cos: torch.Tensor,
-    sin: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-  """Computes rotary positional embedding inline for a query and key.
-
-  Args:
-    q: the query tensor.
-    k: the key tensor.
-    cos: the cosine tensor.
-    sin: the sine tensor.
-
-  Returns:
-    output the RoPE'd query and key.
-  """
+  def apply(x, sin, cos):
+    x = x.transpose(1, 2)
+    b, h, s, d = x.shape
+    ans = torch.split(x, d // 2, dim=-1)
+    x1, x2 = ans
+    left = x1 * cos - x2 * sin
+    right = x2 * cos + x1 * sin
+    res = torch.cat([left, right], dim=-1)
+    res = res.transpose(1, 2)
+    return res
 
-  q_roped = apply_rope(q, cos, sin)
-  k_roped = apply_rope(k, cos, sin)
+  q_roped = apply(q, sin, cos)
+  k_roped = apply(k, sin, cos)
   return q_roped, k_roped

ai_edge_torch/generative/utilities/model_builder.py

@@ -24,7 +24,6 @@
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
 import ai_edge_torch.generative.layers.model_config as cfg
-import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
 from torch import nn
@@ -86,6 +85,13 @@ def __init__(self, config: cfg.ModelConfig):
         config.embedding_dim,
         config.final_norm_config,
     )
+    # ROPE parameters for all attn_configs are the same. Take the first one.
+    attn_config = config.block_config(0).attn_config
+    self.rope_cache = attn_utils.build_rope_cache(
+        size=config.kv_cache_max,
+        dim=int(attn_config.rotary_percentage * attn_config.head_dim),
+        base=attn_config.rotary_base,
+    )
     self.mask_cache = attn_utils.build_causal_mask_cache(
         size=config.kv_cache_max,
     )
@@ -107,22 +113,16 @@ def forward(
 
     # token embeddings of shape (b, t, n_embd)
     input_embeds = self.tok_embedding(tokens)
-
-    # ROPE parameters for all attn_configs are the same. Take the first one.
-    attn_config = self.config.block_config(0).attn_config
-    n_elem = int(attn_config.rotary_percentage * attn_config.head_dim)
-    rope = rotary_pos_emb.build_rope(
-        input_pos, n_elem, attn_config.head_dim, attn_config.rotary_base
-    )
-
+    cos, sin = self.rope_cache
+    rope = (cos.index_select(0, input_pos), sin.index_select(0, input_pos))
     mask = self.mask_cache.index_select(2, input_pos)
     mask = mask[:, :, :, : self.config.kv_cache_max]
 
-    return self._forward_with_embeds(
+    return self.forward_with_embeds(
         input_embeds, rope, mask, input_pos, kv_cache, export_config
     )
 
-  def _forward_with_embeds(
+  def forward_with_embeds(
       self,
       input_embeds: torch.Tensor,
       rope: Tuple[torch.Tensor, torch.Tensor],
@@ -141,13 +141,13 @@ def _forward_with_embeds(
     if self.config.embedding_scale is not None:
       x = x * self.config.embedding_scale
 
-    updated_kv_entries = []
+    updated_kv_entires = []
     for i, block in enumerate(self.transformer_blocks):
       kv_entry = kv_cache.caches[i] if kv_cache else None
       x, kv_entry = block(x, rope, mask, input_pos, kv_entry)
       if kv_entry:
-        updated_kv_entries.append(kv_entry)
-    updated_kv_cache = kv_utils.KVCache(tuple(updated_kv_entries))
+        updated_kv_entires.append(kv_entry)
+    updated_kv_cache = kv_utils.KVCache(tuple(updated_kv_entires))
 
     if export_config is not None:
       if (